Upload folder using huggingface_hub

flash_attn/cells/benchmark.py

@@ -0,0 +1,341 @@
+# /// script
+# dependencies = [
+#   "numpy",
+#   "torch",
+#   "kernels",
+#   "pandas",
+#   "matplotlib"
+# ]
+# ///
+# Benchmarking common shapes for Flux 1024x1024px image + varying text sequence lengths
+
+import functools
+import os
+import pathlib
+
+import matplotlib.pyplot as plt
+import torch
+import torch._dynamo.config
+import triton
+import triton.language as tl
+
+try:
+    from flash_attn import flash_attn_func
+except:
+    flash_attn_func = None
+    print("Flash Attention 2 not found.")
+
+try:
+    from flash_attn_interface import flash_attn_func as flash_attn_3_func
+except:
+    flash_attn_3_func = None
+    print("Flash Attention 3 not found.")
+
+try:
+    from kernels import get_kernel
+    hf_kernels_flash_attn = get_kernel("kernels-community/flash-attn")
+    hf_kernels_flash_attn_3 = get_kernel("kernels-community/flash-attn3")
+except:
+    hf_kernels_flash_attn = None
+    hf_kernels_flash_attn_3 = None
+    print("HF Kernels not found.")
+
+try:
+    from sageattention import sageattn_qk_int8_pv_fp16_cuda, sageattn_qk_int8_pv_fp16_triton, sageattn_qk_int8_pv_fp8_cuda_sm90
+except:
+    sageattn_qk_int8_pv_fp16_cuda = None
+    sageattn_qk_int8_pv_fp16_triton = None
+    sageattn_qk_int8_pv_fp8_cuda_sm90 = None
+    print("SageAttention not found.")
+
+try:
+    from transformer_engine.pytorch.attention import DotProductAttention
+except:
+    DotProductAttention = None
+    print("Transformer Engine not found.")
+
+try:
+    import xformers.ops as xops
+except:
+    xops = None
+    print("xFormers not found.")
+
+
+plt.rcParams.update({
+    "figure.figsize": (12, 10),
+    "figure.dpi": 120,
+    "font.size": 10,
+    "axes.titlesize": 12,
+    "axes.labelsize": 14,
+    "xtick.labelsize": 10,
+    "ytick.labelsize": 10,
+    "legend.fontsize": 8,
+    "axes.grid": True,
+    "grid.alpha": 0.3,
+    "grid.linestyle": "--",
+    "lines.linewidth": 2.0,
+    "lines.markersize": 6,
+    "legend.frameon": True,
+    "legend.framealpha": 0.9,
+    "legend.loc": "best",
+    "axes.spines.top": False,
+    "axes.spines.right": False,
+})
+
+
+# We want to compare the best compiled version for each specific shape (dynamic=False)
+torch._dynamo.config.cache_size_limit = 10000
+
+# We need to suppress_errors for FA3 to work. It makes it run in eager mode.
+# I can't seem to get it to work any other way under torch.compile, so any suggestions are welcome!
+torch._dynamo.config.suppress_errors = True
+
+output_dir = pathlib.Path("dump_attention_benchmark")
+output_dir.mkdir(parents=True, exist_ok=True)
+
+batch_size = 1
+num_attention_heads = 24
+attention_head_dim = 128
+image_sequence_length = 4096  # 1024x1024px
+text_sequence_lengths = [128, 256, 320, 384, 448, 512]
+sequence_lengths = [image_sequence_length + i for i in text_sequence_lengths]
+
+
+def _attention_torch(query, key, value, *, backend):
+    query, key, value = (x.transpose(1, 2).contiguous() for x in (query, key, value))
+    with torch.nn.attention.sdpa_kernel(backend):
+        out = torch.nn.functional.scaled_dot_product_attention(query, key, value)
+    out = out.transpose(1, 2).contiguous()
+    return out
+
+
+_compiled_attention_torch_default = torch.compile(_attention_torch, mode="default", fullgraph=True, dynamic=False)
+def _attention_torch_compile_default(query, key, value, *, backend):
+    return _compiled_attention_torch_default(query, key, value, backend=backend)
+
+
+_compiled_attention_torch_max_autotune = torch.compile(_attention_torch, mode="max-autotune", fullgraph=True, dynamic=False)
+def _attention_torch_compile_max_autotune(query, key, value, *, backend):
+    return _compiled_attention_torch_max_autotune(query, key, value, backend=backend)
+
+
+def _attention_flash_attn_2(query, key, value):
+    return flash_attn_func(query, key, value)
+
+
+_compiled_flash_attn_2_default = torch.compile(_attention_flash_attn_2, mode="default", fullgraph=True, dynamic=False)
+def _attention_flash_attn_2_compile_default(query, key, value):
+    return _compiled_flash_attn_2_default(query, key, value)
+
+
+_compiled_flash_attn_2_max_autotune = torch.compile(_attention_flash_attn_2, mode="max-autotune", fullgraph=True, dynamic=False)
+def _attention_flash_attn_2_compile_max_autotune(query, key, value):
+    return _compiled_flash_attn_2_max_autotune(query, key, value)
+
+
+# For fullgraph=True tracing to be compatible
+@torch.library.custom_op("flash_attn_3::_flash_attn_forward", mutates_args=(), device_types="cuda")
+def _wrapped_flash_attn_3(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+    out, lse = flash_attn_3_func(query, key, value)
+    return out
+
+
+@torch.library.register_fake("flash_attn_3::_flash_attn_forward")
+def _(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+    return torch.empty_like(query)
+
+
+def _attention_flash_attn_3(query, key, value):
+    out = _wrapped_flash_attn_3(query, key, value)
+    return out
+
+
+_compiled_flash_attn_3_default = torch.compile(_attention_flash_attn_3, mode="default", fullgraph=True, dynamic=False)
+def _attention_flash_attn_3_compile_default(query, key, value):
+    return _compiled_flash_attn_3_default(query, key, value)
+
+
+_compiled_flash_attn_3_max_autotune = torch.compile(_attention_flash_attn_3, mode="max-autotune", fullgraph=True, dynamic=False)
+def _attention_flash_attn_3_compile_max_autotune(query, key, value):
+    return _compiled_flash_attn_3_max_autotune(query, key, value)
+
+
+def _attention_hf_kernels_flash_attn(query, key, value):
+    return hf_kernels_flash_attn.fwd(query, key, value, is_causal=False)[0]
+
+
+def _attention_hf_kernels_flash_attn3(query, key, value):
+    return hf_kernels_flash_attn_3.flash_attn_func(query, key, value, causal=False)[0]
+
+
+def _attention_sageattn_qk_int8_pv_fp16_cuda(query, key, value):
+    return sageattn_qk_int8_pv_fp16_cuda(query, key, value, tensor_layout="NHD")
+
+
+def _attention_sageattn_qk_int8_pv_fp16_triton(query, key, value):
+    return sageattn_qk_int8_pv_fp16_triton(query, key, value, tensor_layout="NHD")
+
+
+def _attention_sageattn_qk_int8_pv_fp8_cuda_sm90(query, key, value):
+    return sageattn_qk_int8_pv_fp8_cuda_sm90(query, key, value, tensor_layout="NHD")
+
+
+if DotProductAttention is not None:
+    def set_te_backend(backend):
+        # must be applied before first use of
+        # transformer_engine.pytorch.attention
+        os.environ["NVTE_FLASH_ATTN"] = '0'
+        os.environ["NVTE_FUSED_ATTN"] = '0'
+        os.environ["NVTE_UNFUSED_ATTN"] = '0'
+        if backend == 'flash':
+            os.environ["NVTE_FLASH_ATTN"] = '1'
+        if backend == 'fused':
+            os.environ["NVTE_FUSED_ATTN"] = '1'
+        if backend == 'unfused':
+            os.environ["NVTE_UNFUSED_ATTN"] = '1'
+    
+    set_te_backend("fused")
+    te_attn_fn = DotProductAttention(
+        num_attention_heads=num_attention_heads,
+        kv_channels=attention_head_dim,
+        qkv_format="bshd",
+        attn_mask_type="no_mask",
+    )
+else:
+    def te_attn_fn(query, key, value):
+        raise RuntimeError("Transformer Engine is not available. Please install it for TE-based attention.")
+
+def _attention_te(query, key, value):
+    out = te_attn_fn(query, key, value)
+    out = out.unflatten(2, (num_attention_heads, attention_head_dim))
+    return out
+
+
+# Cannot fullgraph compile TE
+_compiled_te_attn_fn_default = torch.compile(_attention_te, mode="default", fullgraph=False, dynamic=False)
+def _attention_te_compile_default(query, key, value):
+    return _compiled_te_attn_fn_default(query, key, value)
+
+
+# Cannot fullgraph compile TE
+_compiled_te_attn_fn_max_autotune = torch.compile(_attention_te, mode="max-autotune", fullgraph=False, dynamic=False)
+def _attention_te_compile_max_autotune(query, key, value):
+    return _compiled_te_attn_fn_max_autotune(query, key, value)
+
+
+def _attention_xformers(query, key, value):
+    return xops.memory_efficient_attention(query, key, value)
+
+
+_compiled_xformers_default = torch.compile(_attention_xformers, mode="default", fullgraph=True, dynamic=False)
+def _attention_xformers_compile_default(query, key, value):
+    return _compiled_xformers_default(query, key, value)
+
+
+_compiled_xformers_max_autotune = torch.compile(_attention_xformers, mode="max-autotune", fullgraph=True, dynamic=False)
+def _attention_xformers_compile_max_autotune(query, key, value):
+    return _compiled_xformers_max_autotune(query, key, value)
+
+
+attention_ops = {}
+attention_ops["torch_cudnn"] = functools.partial(_attention_torch, backend=torch.nn.attention.SDPBackend.CUDNN_ATTENTION)
+attention_ops["torch_cudnn_compile_d"] = functools.partial(_attention_torch_compile_default, backend=torch.nn.attention.SDPBackend.CUDNN_ATTENTION)
+attention_ops["torch_cudnn_compile_ma"] = functools.partial(_attention_torch_compile_max_autotune, backend=torch.nn.attention.SDPBackend.CUDNN_ATTENTION)
+attention_ops["torch_flash"] = functools.partial(_attention_torch, backend=torch.nn.attention.SDPBackend.FLASH_ATTENTION)
+attention_ops["torch_flash_compile_d"] = functools.partial(_attention_torch_compile_default, backend=torch.nn.attention.SDPBackend.FLASH_ATTENTION)
+attention_ops["torch_flash_compile_ma"] = functools.partial(_attention_torch_compile_max_autotune, backend=torch.nn.attention.SDPBackend.FLASH_ATTENTION)
+if hf_kernels_flash_attn is not None:
+    attention_ops["hf_flash_attn"] = _attention_hf_kernels_flash_attn
+    attention_ops["hf_flash_attn3"] = _attention_hf_kernels_flash_attn3
+if flash_attn_func is not None:
+    attention_ops["flash_attn_2"] = _attention_flash_attn_2
+    attention_ops["flash_attn_2_compile_d"] = _attention_flash_attn_2_compile_default
+    attention_ops["flash_attn_2_compile_ma"] = _attention_flash_attn_2_compile_max_autotune
+if flash_attn_3_func is not None:
+    attention_ops["flash_attn_3"] = _attention_flash_attn_3
+    attention_ops["flash_attn_3_compile_d"] = _attention_flash_attn_3_compile_default
+    attention_ops["flash_attn_3_compile_ma"] = _attention_flash_attn_3_compile_max_autotune
+if sageattn_qk_int8_pv_fp16_cuda is not None:
+    attention_ops["sageattn_qk_int8_pv_fp16_cuda"] = _attention_sageattn_qk_int8_pv_fp16_cuda
+    attention_ops["sageattn_qk_int8_pv_fp16_triton"] = _attention_sageattn_qk_int8_pv_fp16_triton
+    if torch.cuda.get_device_capability()[0] >= 9:
+        attention_ops["sageattn_qk_int8_pv_fp8_cuda_sm90"] = _attention_sageattn_qk_int8_pv_fp8_cuda_sm90
+if DotProductAttention is not None:
+    attention_ops["te_fused"] = _attention_te
+    attention_ops["te_fused_compile_d"] = _attention_te_compile_default
+    attention_ops["te_fused_compile_ma"] = _attention_te_compile_max_autotune
+if xops is not None:
+    attention_ops["xformers"] = _attention_xformers
+    attention_ops["xformers_compile_d"] = _attention_xformers_compile_default
+    attention_ops["xformers_compile_ma"] = _attention_xformers_compile_max_autotune
+
+
+def get_color_and_linestyle(n: int) -> tuple[str, str]:
+    colors = ["#e41a1c", "#377eb8", "#4daf4a", "#984ea3", "#ff7f00", "#a65628", "#f781bf", "#999999"]
+    line_styles = ["-", ":", "-.", "--"]
+    if n > len(colors) * len(line_styles):
+        raise ValueError(f"Required {n=} styles but maximum is {len(colors) * len(line_styles)}")
+    styles = []
+    for i in range(n):
+        color = colors[i % len(colors)]
+        linestyle = line_styles[i // len(colors)]
+        styles.append((color, linestyle))
+    return styles
+
+
+def correctness():
+    for seq_len in sequence_lengths:
+        shape = (batch_size, seq_len, num_attention_heads, attention_head_dim)
+        print(f"\n\n===== Testing shape: {shape} =====")
+        
+        query = torch.randn(shape, device="cuda", dtype=torch.float32)
+        key = torch.randn(shape, device="cuda", dtype=torch.float32)
+        value = torch.randn(shape, device="cuda", dtype=torch.float32)
+
+        golden_truth = _attention_torch(query, key, value, backend=torch.nn.attention.SDPBackend.MATH)
+        query, key, value = (x.bfloat16() for x in (query, key, value))
+
+        for name, fn in attention_ops.items():
+            out = fn(query, key, value)
+            absdiff = (out - golden_truth).abs()
+            absmax = torch.max(absdiff)
+            mae = torch.mean(absdiff)
+            mse = torch.mean((golden_truth - out) ** 2)
+            print(f"{name:<30}: absmax={absmax:.6f}, mae={mae:.6f}, mse={mse:.6f}")
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["seq_len"],
+        x_vals=sequence_lengths,
+        x_log=False,
+        line_arg="provider",
+        line_vals=list(attention_ops.keys()),
+        line_names=[x.removeprefix("solution_") for x in attention_ops.keys()],
+        ylabel="Time (ms)",
+        styles=get_color_and_linestyle(len(attention_ops)),
+        plot_name="Attention Benchmark",
+        args={},
+    )
+)
+def benchmark_fn(seq_len: int, provider: str):
+    torch.manual_seed(0)
+    
+    shape = (batch_size, seq_len, num_attention_heads, attention_head_dim)
+    query = torch.randn(shape, device="cuda", dtype=torch.bfloat16) * torch.randint(1, 5, shape, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn(shape, device="cuda", dtype=torch.bfloat16) * torch.randint(1, 5, shape, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn(shape, device="cuda", dtype=torch.bfloat16) * torch.randint(1, 5, shape, device="cuda", dtype=torch.bfloat16)
+    
+    fn = attention_ops[provider]
+    ms, min_ms, max_ms = triton.testing.do_bench(
+        lambda: fn(query, key, value),
+        warmup=3,
+        rep=10,
+        quantiles=[0.5, 0.2, 0.8],
+    )
+    return ms, max_ms, min_ms
+
+
+with torch.inference_mode():
+    correctness()
+    benchmark_fn.run(print_data=True, save_path=output_dir.as_posix())

flash_attn/index.html

@@ -0,0 +1,24 @@
+<!DOCTYPE html>
+<html>
+<head>
+  <meta charset='UTF-8'>
+  <title>Directory Index</title>
+  <style>
+    body { font-family: monospace; margin: 20px; }
+    h1 { font-size: 1.5em; }
+    ul { list-style-type: none; padding-left: 20px; }
+    li { margin: 5px 0; }
+    .dir { font-weight: bold; }
+    .file { color: #0066cc; }
+    a { text-decoration: none; }
+    a:hover { text-decoration: underline; }
+  </style>
+</head>
+<body>
+  <h1>Index of /flash_attn</h1>
+  <ul>
+    <li><a href='../index.html' class='dir'>../</a></li>
+    <li><a href='benchmark.html' class='file'>benchmark.html</a></li>
+  </ul>
+</body>
+</html>

index.html

@@ -17,8 +17,8 @@
 <body>
   <h1>Index of /</h1>
   <ul>
-    <li><a href='megablocks/index.html' class='dir'>megablocks/</a></li>
-    <li><a href='megablocks_yamoe/index.html' class='dir'>megablocks_yamoe/</a></li>
+    <li><a href='flash_attn/index.html' class='dir'>flash_attn/</a></li>
+    <li><a href='moe_benchmarks/index.html' class='dir'>moe_benchmarks/</a></li>
   </ul>
 </body>
 </html>

moe_benchmarks/megablocks/cells/forward_and_backward.py

@@ -0,0 +1,196 @@
+# /// script
+# requires-python = ">=3.12"
+# dependencies = [
+#     "accelerate>=1.10.1",
+#     "torch>=2.7.0",
+#     "kernels==0.10.0",
+#     "transformers@https://github.com/huggingface/transformers.git",
+#     "ipdb>=0.13.13",
+#     "matplotlib>=3.7.2",
+#     "numpy>=1.24.3",
+# ]
+# ///
+
+import torch
+from transformers import GptOssForCausalLM, PreTrainedTokenizerFast, Mxfp4Config
+import time
+import torch.nn as nn
+from kernels import register_kernel_mapping, Mode, LayerRepository, replace_kernel_forward_from_hub
+import sys
+import torch.profiler
+import gc
+import logging
+from transformers.models.gpt_oss.modeling_gpt_oss import GptOssRMSNorm
+
+# remove liger kernel for testing 
+replace_kernel_forward_from_hub(GptOssRMSNorm, None)
+
+# set to debug logging
+logging.basicConfig(level=logging.INFO)
+
+def reset_peak_memory_stats():
+    """Clear CUDA cache and reset memory allocation counters."""
+    torch.cuda.empty_cache()
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats()
+    gc.collect()
+
+def get_memory_stats():
+    """Get current and peak CUDA memory usage."""
+    if not torch.cuda.is_available():
+        return {"allocated_gb": 0, "peak_gb": 0, "reserved_gb": 0}
+    return {
+        "allocated_gb": torch.cuda.memory_allocated() / 1e9,
+        "peak_gb": torch.cuda.max_memory_allocated() / 1e9,
+        "reserved_gb": torch.cuda.memory_reserved() / 1e9,
+    }
+
+def override_kernel_layer_name(cls_name: str, value) -> bool:
+    """Helper to dynamically override the kernel_layer_name in a model class."""
+    for mod in sys.modules.values():
+        if mod is None:
+            continue
+        obj = getattr(mod, cls_name, None)
+        if isinstance(obj, type) and issubclass(obj, nn.Module):
+            setattr(obj, "kernel_layer_name", value)
+            print(f"Overrode {cls_name}.kernel_layer_name to {value}")
+            return True
+    return False
+
+
+# Init the model the normal way
+model_id = "openai/gpt-oss-20b"
+tokenizer = PreTrainedTokenizerFast.from_pretrained(model_id)
+quantization_config = Mxfp4Config(dequantize=True)
+
+model = GptOssForCausalLM.from_pretrained(
+    model_id,
+    dtype="bfloat16",
+    device_map="auto",
+    use_kernels=True,
+    quantization_config=quantization_config,
+).eval()
+
+messages = [
+    {"role": "system", "content": "What is Tensor Parallelism?"},
+]
+
+inputs = tokenizer.apply_chat_template(
+    messages,
+    add_generation_prompt=True,
+    return_tensors="pt",
+    return_dict=True,
+    reasoning_effort="low",
+).to("cuda")
+
+max_tokens = 128  # Reduced to help with memory usage
+
+# Clear memory before backward pass
+reset_peak_memory_stats()
+print(f"Pre-generation memory: {get_memory_stats()}")
+
+# forward and backward pass
+with torch.autograd.set_grad_enabled(True):
+    start_time = time.perf_counter()
+    generated = model.generate(
+        **inputs,
+        max_new_tokens=max_tokens,
+        do_sample=False,
+        temperature=None,
+    )
+    end_time = time.perf_counter()
+    print(tokenizer.decode(generated[0], skip_special_tokens=False))
+    print(f"Generation took {end_time - start_time:.2f} seconds")
+    print(f"Post-generation memory: {get_memory_stats()}")
+
+    # Use gradient checkpointing to reduce memory usage
+    if hasattr(model, 'gradient_checkpointing_enable'):
+        model.gradient_checkpointing_enable()
+        print("Enabled gradient checkpointing")
+
+    # Reduce sequence length if needed for memory
+    max_seq_len = 512  # Limit sequence length for backward pass
+    if generated.size(1) > max_seq_len:
+        print(f"Truncating sequence from {generated.size(1)} to {max_seq_len} tokens")
+        full_sequence = generated[:, -max_seq_len:]
+    else:
+        full_sequence = generated
+
+    # Get model outputs for the full sequence
+    model.train()  # Enable dropout and other training behaviors
+
+    try:
+        outputs = model(
+            input_ids=full_sequence,
+            labels=full_sequence,  # This will compute loss internally
+            return_dict=True
+        )
+        print(f"Post-forward memory: {get_memory_stats()}")
+
+        # If model doesn't compute loss, compute it manually
+        if outputs.loss is None:
+            shift_logits = outputs.logits[..., :-1, :].contiguous()
+            shift_labels = full_sequence[..., 1:].contiguous()
+
+            # Use CrossEntropyLoss with ignore_index for padding tokens
+            loss_fct = torch.nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else -100)
+            loss = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1)
+            )
+        else:
+            loss = outputs.loss
+
+        print(f"Loss: {loss.item():.4f}")
+
+        # Clear intermediate tensors to save memory
+        del outputs
+        torch.cuda.empty_cache()
+
+        # Perform backward pass with memory management
+        print("Running backward pass...")
+        print(f"Pre-backward memory: {get_memory_stats()}")
+
+        loss.backward()
+        print(f"Post-backward memory: {get_memory_stats()}")
+
+    except torch.cuda.OutOfMemoryError as e:
+        print(f"OOM during forward/backward pass: {e}")
+        print("Try reducing max_tokens or max_seq_len")
+        raise
+
+    # Calculate gradient statistics and print sample gradients
+    total_norm = 0.0
+    param_count = 0
+    grad_samples = {}
+
+    for name, p in model.named_parameters():
+        if p.grad is not None:
+            param_count += 1
+            grad_norm = p.grad.data.norm(2).item()
+            total_norm += grad_norm ** 2
+
+            # Collect gradient statistics for key layers
+            if any(key in name for key in ['embed', 'lm_head', 'mlp.up', 'mlp.down', 'self_attn.q_proj', 'norm']):
+                grad_samples[name] = {
+                    'norm': grad_norm,
+                    'mean': p.grad.data.mean().item(),
+                    'std': p.grad.data.std().item(),
+                    'max': p.grad.data.max().item(),
+                    'min': p.grad.data.min().item(),
+                }
+
+    total_norm = total_norm ** 0.5
+
+    print(f"\nGradient norm: {total_norm:.4f}")
+    print(f"Parameters with gradients: {param_count}")
+
+    # Print sample gradients from important layers
+    print("\nSample gradient statistics:")
+    for i, (name, stats) in enumerate(list(grad_samples.items())[:10]):
+        print(f"  {name[:60]:<60} | norm: {stats['norm']:.4e} | mean: {stats['mean']:.4e} | std: {stats['std']:.4e}")
+
+    # Optional: zero gradients for next iteration
+    model.zero_grad()
+    model.eval()  # Switch back to eval mode
+

moe_benchmarks/megablocks/megablocks_only.html

@@ -3710,7 +3710,7 @@ span.linenos.special { color: #000000; background-color: #ffffc0; padding-left:
     <div class="system-info">
         <div class="system-info-header">Generated on:</div>
         <div class="system-info-content">
-            Linux x86_64 | Linux-6.11.0-1018-azure-x86_64-with-glibc2.39
+            Linux x86_64 | Linux-6.12.40-64.114.amzn2023.x86_64-x86_64-with-glibc2.36
         </div>
     </div>
     
@@ -3724,122 +3724,219 @@ span.linenos.special { color: #000000; background-color: #ffffc0; padding-left:
 <p>Next we can run with Megablocks kernels enabled.</p>
 <h3>Forward</h3>
 <p>First, we run a forward pass with Megablocks kernels.</p>
-<div class="cell cell-failed" id="cell-forward_only">
+<h2>Forward and Backward</h2>
+<p>Next, we run a forward and backward pass with Megablocks kernels enabled. This should be more memory efficient and allow us to complete the backward pass without running out of memory.</p>
+<div class="cell cell-failed" id="cell-forward_and_backward">
 <div class="cell-header">
 <span class="collapse-indicators">
-<span onclick="toggleCode('forward_only')" style="cursor: pointer;">▼ code</span> 
-<span onclick="toggleOutput('forward_only')" style="cursor: pointer;">▼ output</span>
- <span id="uv-indicator-forward_only" onclick="toggleUvLogsFromHeader('forward_only')" style="cursor: pointer;">▶ uv-logs</span>
+<span onclick="toggleCode('forward_and_backward')" style="cursor: pointer;">▼ code</span> 
+<span onclick="toggleOutput('forward_and_backward')" style="cursor: pointer;">▼ output</span>
+ <span id="uv-indicator-forward_and_backward" style="cursor: default; opacity: 0.3;">▶ uv-logs</span>
 </span> | 
-Cell: forward_only | 118.48s | FAILED
- | <button class="run-btn" onclick="runCell('forward_only')">▶ run</button>
-<button class="copy-btn" onclick="copyCell('forward_only')">Copy</button>
-<a href="cells/forward_only.py" target="_blank" class="raw-btn">Raw</a>
+Cell: forward_and_backward | 19.43s | FAILED
+ | <button class="run-btn" onclick="runCell('forward_and_backward')">▶ run</button>
+<button class="copy-btn" onclick="copyCell('forward_and_backward')">Copy</button>
+<a href="cells/forward_and_backward.py" target="_blank" class="raw-btn">Raw</a>
 </div>
-<div id="code-forward_only" class="cell-code" data-lines="101">
+<div id="code-forward_and_backward" class="cell-code" data-lines="196">
 <div class="highlight-with-lines">
-<div class="line-numbers" id="lines-forward_only">
-<a class="line-number" data-cell="forward_only" data-line="1" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 1, true);">1</a>
-<a class="line-number" data-cell="forward_only" data-line="2" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 2, true);">2</a>
-<a class="line-number" data-cell="forward_only" data-line="3" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 3, true);">3</a>
-<a class="line-number" data-cell="forward_only" data-line="4" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 4, true);">4</a>
-<a class="line-number" data-cell="forward_only" data-line="5" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 5, true);">5</a>
-<a class="line-number" data-cell="forward_only" data-line="6" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 6, true);">6</a>
-<a class="line-number" data-cell="forward_only" data-line="7" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 7, true);">7</a>
-<a class="line-number" data-cell="forward_only" data-line="8" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 8, true);">8</a>
-<a class="line-number" data-cell="forward_only" data-line="9" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 9, true);">9</a>
-<a class="line-number" data-cell="forward_only" data-line="10" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 10, true);">10</a>
-<a class="line-number" data-cell="forward_only" data-line="11" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 11, true);">11</a>
-<a class="line-number" data-cell="forward_only" data-line="12" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 12, true);">12</a>
-<a class="line-number" data-cell="forward_only" data-line="13" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 13, true);">13</a>
-<a class="line-number" data-cell="forward_only" data-line="14" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 14, true);">14</a>
-<a class="line-number" data-cell="forward_only" data-line="15" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 15, true);">15</a>
-<a class="line-number" data-cell="forward_only" data-line="16" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 16, true);">16</a>
-<a class="line-number" data-cell="forward_only" data-line="17" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 17, true);">17</a>
-<a class="line-number" data-cell="forward_only" data-line="18" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 18, true);">18</a>
-<a class="line-number" data-cell="forward_only" data-line="19" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 19, true);">19</a>
-<a class="line-number" data-cell="forward_only" data-line="20" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 20, true);">20</a>
-<a class="line-number" data-cell="forward_only" data-line="21" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 21, true);">21</a>
-<a class="line-number" data-cell="forward_only" data-line="22" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 22, true);">22</a>
-<a class="line-number" data-cell="forward_only" data-line="23" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 23, true);">23</a>
-<a class="line-number" data-cell="forward_only" data-line="24" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 24, true);">24</a>
-<a class="line-number" data-cell="forward_only" data-line="25" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 25, true);">25</a>
-<a class="line-number" data-cell="forward_only" data-line="26" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 26, true);">26</a>
-<a class="line-number" data-cell="forward_only" data-line="27" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 27, true);">27</a>
-<a class="line-number" data-cell="forward_only" data-line="28" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 28, true);">28</a>
-<a class="line-number" data-cell="forward_only" data-line="29" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 29, true);">29</a>
-<a class="line-number" data-cell="forward_only" data-line="30" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 30, true);">30</a>
-<a class="line-number" data-cell="forward_only" data-line="31" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 31, true);">31</a>
-<a class="line-number" data-cell="forward_only" data-line="32" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 32, true);">32</a>
-<a class="line-number" data-cell="forward_only" data-line="33" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 33, true);">33</a>
-<a class="line-number" data-cell="forward_only" data-line="34" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 34, true);">34</a>
-<a class="line-number" data-cell="forward_only" data-line="35" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 35, true);">35</a>
-<a class="line-number" data-cell="forward_only" data-line="36" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 36, true);">36</a>
-<a class="line-number" data-cell="forward_only" data-line="37" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 37, true);">37</a>
-<a class="line-number" data-cell="forward_only" data-line="38" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 38, true);">38</a>
-<a class="line-number" data-cell="forward_only" data-line="39" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 39, true);">39</a>
-<a class="line-number" data-cell="forward_only" data-line="40" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 40, true);">40</a>
-<a class="line-number" data-cell="forward_only" data-line="41" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 41, true);">41</a>
-<a class="line-number" data-cell="forward_only" data-line="42" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 42, true);">42</a>
-<a class="line-number" data-cell="forward_only" data-line="43" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 43, true);">43</a>
-<a class="line-number" data-cell="forward_only" data-line="44" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 44, true);">44</a>
-<a class="line-number" data-cell="forward_only" data-line="45" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 45, true);">45</a>
-<a class="line-number" data-cell="forward_only" data-line="46" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 46, true);">46</a>
-<a class="line-number" data-cell="forward_only" data-line="47" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 47, true);">47</a>
-<a class="line-number" data-cell="forward_only" data-line="48" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 48, true);">48</a>
-<a class="line-number" data-cell="forward_only" data-line="49" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 49, true);">49</a>
-<a class="line-number" data-cell="forward_only" data-line="50" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 50, true);">50</a>
-<a class="line-number" data-cell="forward_only" data-line="51" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 51, true);">51</a>
-<a class="line-number" data-cell="forward_only" data-line="52" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 52, true);">52</a>
-<a class="line-number" data-cell="forward_only" data-line="53" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 53, true);">53</a>
-<a class="line-number" data-cell="forward_only" data-line="54" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 54, true);">54</a>
-<a class="line-number" data-cell="forward_only" data-line="55" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 55, true);">55</a>
-<a class="line-number" data-cell="forward_only" data-line="56" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 56, true);">56</a>
-<a class="line-number" data-cell="forward_only" data-line="57" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 57, true);">57</a>
-<a class="line-number" data-cell="forward_only" data-line="58" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 58, true);">58</a>
-<a class="line-number" data-cell="forward_only" data-line="59" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 59, true);">59</a>
-<a class="line-number" data-cell="forward_only" data-line="60" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 60, true);">60</a>
-<a class="line-number" data-cell="forward_only" data-line="61" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 61, true);">61</a>
-<a class="line-number" data-cell="forward_only" data-line="62" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 62, true);">62</a>
-<a class="line-number" data-cell="forward_only" data-line="63" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 63, true);">63</a>
-<a class="line-number" data-cell="forward_only" data-line="64" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 64, true);">64</a>
-<a class="line-number" data-cell="forward_only" data-line="65" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 65, true);">65</a>
-<a class="line-number" data-cell="forward_only" data-line="66" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 66, true);">66</a>
-<a class="line-number" data-cell="forward_only" data-line="67" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 67, true);">67</a>
-<a class="line-number" data-cell="forward_only" data-line="68" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 68, true);">68</a>
-<a class="line-number" data-cell="forward_only" data-line="69" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 69, true);">69</a>
-<a class="line-number" data-cell="forward_only" data-line="70" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 70, true);">70</a>
-<a class="line-number" data-cell="forward_only" data-line="71" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 71, true);">71</a>
-<a class="line-number" data-cell="forward_only" data-line="72" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 72, true);">72</a>
-<a class="line-number" data-cell="forward_only" data-line="73" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 73, true);">73</a>
-<a class="line-number" data-cell="forward_only" data-line="74" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 74, true);">74</a>
-<a class="line-number" data-cell="forward_only" data-line="75" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 75, true);">75</a>
-<a class="line-number" data-cell="forward_only" data-line="76" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 76, true);">76</a>
-<a class="line-number" data-cell="forward_only" data-line="77" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 77, true);">77</a>
-<a class="line-number" data-cell="forward_only" data-line="78" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 78, true);">78</a>
-<a class="line-number" data-cell="forward_only" data-line="79" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 79, true);">79</a>
-<a class="line-number" data-cell="forward_only" data-line="80" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 80, true);">80</a>
-<a class="line-number" data-cell="forward_only" data-line="81" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 81, true);">81</a>
-<a class="line-number" data-cell="forward_only" data-line="82" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 82, true);">82</a>
-<a class="line-number" data-cell="forward_only" data-line="83" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 83, true);">83</a>
-<a class="line-number" data-cell="forward_only" data-line="84" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 84, true);">84</a>
-<a class="line-number" data-cell="forward_only" data-line="85" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 85, true);">85</a>
-<a class="line-number" data-cell="forward_only" data-line="86" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 86, true);">86</a>
-<a class="line-number" data-cell="forward_only" data-line="87" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 87, true);">87</a>
-<a class="line-number" data-cell="forward_only" data-line="88" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 88, true);">88</a>
-<a class="line-number" data-cell="forward_only" data-line="89" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 89, true);">89</a>
-<a class="line-number" data-cell="forward_only" data-line="90" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 90, true);">90</a>
-<a class="line-number" data-cell="forward_only" data-line="91" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 91, true);">91</a>
-<a class="line-number" data-cell="forward_only" data-line="92" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 92, true);">92</a>
-<a class="line-number" data-cell="forward_only" data-line="93" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 93, true);">93</a>
-<a class="line-number" data-cell="forward_only" data-line="94" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 94, true);">94</a>
-<a class="line-number" data-cell="forward_only" data-line="95" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 95, true);">95</a>
-<a class="line-number" data-cell="forward_only" data-line="96" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 96, true);">96</a>
-<a class="line-number" data-cell="forward_only" data-line="97" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 97, true);">97</a>
-<a class="line-number" data-cell="forward_only" data-line="98" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 98, true);">98</a>
-<a class="line-number" data-cell="forward_only" data-line="99" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 99, true);">99</a>
-<a class="line-number" data-cell="forward_only" data-line="100" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 100, true);">100</a>
-<a class="line-number" data-cell="forward_only" data-line="101" href="#cell-forward_only" onclick="event.preventDefault(); selectCellLine('forward_only', 101, true);">101</a>
+<div class="line-numbers" id="lines-forward_and_backward">
+<a class="line-number" data-cell="forward_and_backward" data-line="1" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 1, true);">1</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="2" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 2, true);">2</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="3" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 3, true);">3</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="4" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 4, true);">4</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="5" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 5, true);">5</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="6" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 6, true);">6</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="7" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 7, true);">7</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="8" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 8, true);">8</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="9" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 9, true);">9</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="10" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 10, true);">10</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="11" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 11, true);">11</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="12" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 12, true);">12</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="13" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 13, true);">13</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="14" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 14, true);">14</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="15" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 15, true);">15</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="16" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 16, true);">16</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="17" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 17, true);">17</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="18" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 18, true);">18</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="19" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 19, true);">19</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="20" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 20, true);">20</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="21" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 21, true);">21</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="22" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 22, true);">22</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="23" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 23, true);">23</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="24" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 24, true);">24</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="25" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 25, true);">25</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="26" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 26, true);">26</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="27" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 27, true);">27</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="28" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 28, true);">28</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="29" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 29, true);">29</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="30" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 30, true);">30</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="31" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 31, true);">31</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="32" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 32, true);">32</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="33" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 33, true);">33</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="34" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 34, true);">34</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="35" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 35, true);">35</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="36" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 36, true);">36</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="37" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 37, true);">37</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="38" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 38, true);">38</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="39" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 39, true);">39</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="40" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 40, true);">40</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="41" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 41, true);">41</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="42" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 42, true);">42</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="43" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 43, true);">43</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="44" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 44, true);">44</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="45" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 45, true);">45</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="46" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 46, true);">46</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="47" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 47, true);">47</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="48" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 48, true);">48</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="49" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 49, true);">49</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="50" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 50, true);">50</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="51" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 51, true);">51</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="52" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 52, true);">52</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="53" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 53, true);">53</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="54" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 54, true);">54</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="55" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 55, true);">55</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="56" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 56, true);">56</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="57" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 57, true);">57</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="58" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 58, true);">58</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="59" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 59, true);">59</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="60" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 60, true);">60</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="61" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 61, true);">61</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="62" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 62, true);">62</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="63" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 63, true);">63</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="64" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 64, true);">64</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="65" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 65, true);">65</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="66" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 66, true);">66</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="67" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 67, true);">67</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="68" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 68, true);">68</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="69" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 69, true);">69</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="70" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 70, true);">70</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="71" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 71, true);">71</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="72" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 72, true);">72</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="73" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 73, true);">73</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="74" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 74, true);">74</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="75" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 75, true);">75</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="76" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 76, true);">76</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="77" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 77, true);">77</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="78" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 78, true);">78</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="79" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 79, true);">79</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="80" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 80, true);">80</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="81" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 81, true);">81</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="82" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 82, true);">82</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="83" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 83, true);">83</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="84" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 84, true);">84</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="85" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 85, true);">85</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="86" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 86, true);">86</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="87" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 87, true);">87</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="88" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 88, true);">88</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="89" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 89, true);">89</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="90" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 90, true);">90</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="91" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 91, true);">91</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="92" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 92, true);">92</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="93" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 93, true);">93</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="94" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 94, true);">94</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="95" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 95, true);">95</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="96" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 96, true);">96</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="97" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 97, true);">97</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="98" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 98, true);">98</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="99" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 99, true);">99</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="100" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 100, true);">100</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="101" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 101, true);">101</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="102" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 102, true);">102</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="103" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 103, true);">103</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="104" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 104, true);">104</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="105" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 105, true);">105</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="106" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 106, true);">106</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="107" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 107, true);">107</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="108" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 108, true);">108</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="109" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 109, true);">109</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="110" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 110, true);">110</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="111" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 111, true);">111</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="112" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 112, true);">112</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="113" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 113, true);">113</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="114" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 114, true);">114</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="115" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 115, true);">115</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="116" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 116, true);">116</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="117" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 117, true);">117</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="118" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 118, true);">118</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="119" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 119, true);">119</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="120" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 120, true);">120</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="121" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 121, true);">121</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="122" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 122, true);">122</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="123" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 123, true);">123</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="124" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 124, true);">124</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="125" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 125, true);">125</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="126" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 126, true);">126</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="127" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 127, true);">127</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="128" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 128, true);">128</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="129" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 129, true);">129</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="130" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 130, true);">130</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="131" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 131, true);">131</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="132" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 132, true);">132</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="133" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 133, true);">133</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="134" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 134, true);">134</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="135" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 135, true);">135</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="136" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 136, true);">136</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="137" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 137, true);">137</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="138" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 138, true);">138</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="139" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 139, true);">139</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="140" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 140, true);">140</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="141" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 141, true);">141</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="142" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 142, true);">142</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="143" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 143, true);">143</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="144" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 144, true);">144</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="145" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 145, true);">145</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="146" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 146, true);">146</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="147" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 147, true);">147</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="148" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 148, true);">148</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="149" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 149, true);">149</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="150" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 150, true);">150</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="151" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 151, true);">151</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="152" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 152, true);">152</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="153" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 153, true);">153</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="154" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 154, true);">154</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="155" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 155, true);">155</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="156" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 156, true);">156</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="157" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 157, true);">157</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="158" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 158, true);">158</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="159" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 159, true);">159</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="160" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 160, true);">160</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="161" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 161, true);">161</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="162" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 162, true);">162</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="163" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 163, true);">163</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="164" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 164, true);">164</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="165" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 165, true);">165</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="166" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 166, true);">166</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="167" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 167, true);">167</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="168" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 168, true);">168</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="169" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 169, true);">169</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="170" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 170, true);">170</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="171" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 171, true);">171</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="172" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 172, true);">172</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="173" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 173, true);">173</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="174" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 174, true);">174</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="175" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 175, true);">175</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="176" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 176, true);">176</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="177" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 177, true);">177</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="178" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 178, true);">178</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="179" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 179, true);">179</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="180" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 180, true);">180</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="181" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 181, true);">181</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="182" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 182, true);">182</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="183" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 183, true);">183</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="184" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 184, true);">184</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="185" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 185, true);">185</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="186" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 186, true);">186</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="187" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 187, true);">187</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="188" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 188, true);">188</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="189" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 189, true);">189</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="190" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 190, true);">190</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="191" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 191, true);">191</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="192" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 192, true);">192</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="193" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 193, true);">193</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="194" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 194, true);">194</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="195" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 195, true);">195</a>
+<a class="line-number" data-cell="forward_and_backward" data-line="196" href="#cell-forward_and_backward" onclick="event.preventDefault(); selectCellLine('forward_and_backward', 196, true);">196</a>
 </div>
 <div class="code-wrap">
 <div class="highlight"><pre><span></span><span class="c1"># /// script</span>
@@ -3866,7 +3963,7 @@ Cell: forward_only | 118.48s | FAILED
 <span class="kn">import</span><span class="w"> </span><span class="nn">logging</span>
 <span class="kn">from</span><span class="w"> </span><span class="nn">transformers.models.gpt_oss.modeling_gpt_oss</span><span class="w"> </span><span class="kn">import</span> <span class="n">GptOssRMSNorm</span>
 
-
+<span class="c1"># remove liger kernel for testing </span>
 <span class="n">replace_kernel_forward_from_hub</span><span class="p">(</span><span class="n">GptOssRMSNorm</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
 
 <span class="c1"># set to debug logging</span>
@@ -3907,8 +4004,6 @@ Cell: forward_only | 118.48s | FAILED
 <span class="n">tokenizer</span> <span class="o">=</span> <span class="n">PreTrainedTokenizerFast</span><span class="o">.</span><span class="n">from_pretrained</span><span class="p">(</span><span class="n">model_id</span><span class="p">)</span>
 <span class="n">quantization_config</span> <span class="o">=</span> <span class="n">Mxfp4Config</span><span class="p">(</span><span class="n">dequantize</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
 
-
-
 <span class="n">model</span> <span class="o">=</span> <span class="n">GptOssForCausalLM</span><span class="o">.</span><span class="n">from_pretrained</span><span class="p">(</span>
     <span class="n">model_id</span><span class="p">,</span>
     <span class="n">dtype</span><span class="o">=</span><span class="s2">&quot;bfloat16&quot;</span><span class="p">,</span>
@@ -3929,9 +4024,14 @@ Cell: forward_only | 118.48s | FAILED
     <span class="n">reasoning_effort</span><span class="o">=</span><span class="s2">&quot;low&quot;</span><span class="p">,</span>
 <span class="p">)</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="s2">&quot;cuda&quot;</span><span class="p">)</span>
 
-<span class="n">max_tokens</span> <span class="o">=</span> <span class="mi">256</span>
+<span class="n">max_tokens</span> <span class="o">=</span> <span class="mi">128</span>  <span class="c1"># Reduced to help with memory usage</span>
+
+<span class="c1"># Clear memory before backward pass</span>
+<span class="n">reset_peak_memory_stats</span><span class="p">()</span>
+<span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Pre-generation memory: </span><span class="si">{</span><span class="n">get_memory_stats</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
 
-<span class="k">with</span> <span class="n">torch</span><span class="o">.</span><span class="n">inference_mode</span><span class="p">():</span>
+<span class="c1"># forward and backward pass</span>
+<span class="k">with</span> <span class="n">torch</span><span class="o">.</span><span class="n">autograd</span><span class="o">.</span><span class="n">set_grad_enabled</span><span class="p">(</span><span class="kc">True</span><span class="p">):</span>
     <span class="n">start_time</span> <span class="o">=</span> <span class="n">time</span><span class="o">.</span><span class="n">perf_counter</span><span class="p">()</span>
     <span class="n">generated</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">generate</span><span class="p">(</span>
         <span class="o">**</span><span class="n">inputs</span><span class="p">,</span>
@@ -3940,144 +4040,124 @@ Cell: forward_only | 118.48s | FAILED
         <span class="n">temperature</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
     <span class="p">)</span>
     <span class="n">end_time</span> <span class="o">=</span> <span class="n">time</span><span class="o">.</span><span class="n">perf_counter</span><span class="p">()</span>
-
-<span class="nb">print</span><span class="p">(</span><span class="n">tokenizer</span><span class="o">.</span><span class="n">decode</span><span class="p">(</span><span class="n">generated</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">skip_special_tokens</span><span class="o">=</span><span class="kc">False</span><span class="p">))</span>
-<span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Generation took </span><span class="si">{</span><span class="n">end_time</span><span class="w"> </span><span class="o">-</span><span class="w"> </span><span class="n">start_time</span><span class="si">:</span><span class="s2">.2f</span><span class="si">}</span><span class="s2"> seconds&quot;</span><span class="p">)</span>
+    <span class="nb">print</span><span class="p">(</span><span class="n">tokenizer</span><span class="o">.</span><span class="n">decode</span><span class="p">(</span><span class="n">generated</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">skip_special_tokens</span><span class="o">=</span><span class="kc">False</span><span class="p">))</span>
+    <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Generation took </span><span class="si">{</span><span class="n">end_time</span><span class="w"> </span><span class="o">-</span><span class="w"> </span><span class="n">start_time</span><span class="si">:</span><span class="s2">.2f</span><span class="si">}</span><span class="s2"> seconds&quot;</span><span class="p">)</span>
+    <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Post-generation memory: </span><span class="si">{</span><span class="n">get_memory_stats</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
+
+    <span class="c1"># Use gradient checkpointing to reduce memory usage</span>
+    <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="s1">&#39;gradient_checkpointing_enable&#39;</span><span class="p">):</span>
+        <span class="n">model</span><span class="o">.</span><span class="n">gradient_checkpointing_enable</span><span class="p">()</span>
+        <span class="nb">print</span><span class="p">(</span><span class="s2">&quot;Enabled gradient checkpointing&quot;</span><span class="p">)</span>
+
+    <span class="c1"># Reduce sequence length if needed for memory</span>
+    <span class="n">max_seq_len</span> <span class="o">=</span> <span class="mi">512</span>  <span class="c1"># Limit sequence length for backward pass</span>
+    <span class="k">if</span> <span class="n">generated</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span> <span class="o">&gt;</span> <span class="n">max_seq_len</span><span class="p">:</span>
+        <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Truncating sequence from </span><span class="si">{</span><span class="n">generated</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span><span class="si">}</span><span class="s2"> to </span><span class="si">{</span><span class="n">max_seq_len</span><span class="si">}</span><span class="s2"> tokens&quot;</span><span class="p">)</span>
+        <span class="n">full_sequence</span> <span class="o">=</span> <span class="n">generated</span><span class="p">[:,</span> <span class="o">-</span><span class="n">max_seq_len</span><span class="p">:]</span>
+    <span class="k">else</span><span class="p">:</span>
+        <span class="n">full_sequence</span> <span class="o">=</span> <span class="n">generated</span>
+
+    <span class="c1"># Get model outputs for the full sequence</span>
+    <span class="n">model</span><span class="o">.</span><span class="n">train</span><span class="p">()</span>  <span class="c1"># Enable dropout and other training behaviors</span>
+
+    <span class="k">try</span><span class="p">:</span>
+        <span class="n">outputs</span> <span class="o">=</span> <span class="n">model</span><span class="p">(</span>
+            <span class="n">input_ids</span><span class="o">=</span><span class="n">full_sequence</span><span class="p">,</span>
+            <span class="n">labels</span><span class="o">=</span><span class="n">full_sequence</span><span class="p">,</span>  <span class="c1"># This will compute loss internally</span>
+            <span class="n">return_dict</span><span class="o">=</span><span class="kc">True</span>
+        <span class="p">)</span>
+        <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Post-forward memory: </span><span class="si">{</span><span class="n">get_memory_stats</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
+
+        <span class="c1"># If model doesn&#39;t compute loss, compute it manually</span>
+        <span class="k">if</span> <span class="n">outputs</span><span class="o">.</span><span class="n">loss</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
+            <span class="n">shift_logits</span> <span class="o">=</span> <span class="n">outputs</span><span class="o">.</span><span class="n">logits</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="p">:</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="p">:]</span><span class="o">.</span><span class="n">contiguous</span><span class="p">()</span>
+            <span class="n">shift_labels</span> <span class="o">=</span> <span class="n">full_sequence</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">1</span><span class="p">:]</span><span class="o">.</span><span class="n">contiguous</span><span class="p">()</span>
+
+            <span class="c1"># Use CrossEntropyLoss with ignore_index for padding tokens</span>
+            <span class="n">loss_fct</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">CrossEntropyLoss</span><span class="p">(</span><span class="n">ignore_index</span><span class="o">=</span><span class="n">tokenizer</span><span class="o">.</span><span class="n">pad_token_id</span> <span class="k">if</span> <span class="n">tokenizer</span><span class="o">.</span><span class="n">pad_token_id</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="k">else</span> <span class="o">-</span><span class="mi">100</span><span class="p">)</span>
+            <span class="n">loss</span> <span class="o">=</span> <span class="n">loss_fct</span><span class="p">(</span>
+                <span class="n">shift_logits</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="n">shift_logits</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)),</span>
+                <span class="n">shift_labels</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span>
+            <span class="p">)</span>
+        <span class="k">else</span><span class="p">:</span>
+            <span class="n">loss</span> <span class="o">=</span> <span class="n">outputs</span><span class="o">.</span><span class="n">loss</span>
+
+        <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Loss: </span><span class="si">{</span><span class="n">loss</span><span class="o">.</span><span class="n">item</span><span class="p">()</span><span class="si">:</span><span class="s2">.4f</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
+
+        <span class="c1"># Clear intermediate tensors to save memory</span>
+        <span class="k">del</span> <span class="n">outputs</span>
+        <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">empty_cache</span><span class="p">()</span>
+
+        <span class="c1"># Perform backward pass with memory management</span>
+        <span class="nb">print</span><span class="p">(</span><span class="s2">&quot;Running backward pass...&quot;</span><span class="p">)</span>
+        <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Pre-backward memory: </span><span class="si">{</span><span class="n">get_memory_stats</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
+
+        <span class="n">loss</span><span class="o">.</span><span class="n">backward</span><span class="p">()</span>
+        <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Post-backward memory: </span><span class="si">{</span><span class="n">get_memory_stats</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
+
+    <span class="k">except</span> <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">OutOfMemoryError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
+        <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;OOM during forward/backward pass: </span><span class="si">{</span><span class="n">e</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
+        <span class="nb">print</span><span class="p">(</span><span class="s2">&quot;Try reducing max_tokens or max_seq_len&quot;</span><span class="p">)</span>
+        <span class="k">raise</span>
+
+    <span class="c1"># Calculate gradient statistics and print sample gradients</span>
+    <span class="n">total_norm</span> <span class="o">=</span> <span class="mf">0.0</span>
+    <span class="n">param_count</span> <span class="o">=</span> <span class="mi">0</span>
+    <span class="n">grad_samples</span> <span class="o">=</span> <span class="p">{}</span>
+
+    <span class="k">for</span> <span class="n">name</span><span class="p">,</span> <span class="n">p</span> <span class="ow">in</span> <span class="n">model</span><span class="o">.</span><span class="n">named_parameters</span><span class="p">():</span>
+        <span class="k">if</span> <span class="n">p</span><span class="o">.</span><span class="n">grad</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
+            <span class="n">param_count</span> <span class="o">+=</span> <span class="mi">1</span>
+            <span class="n">grad_norm</span> <span class="o">=</span> <span class="n">p</span><span class="o">.</span><span class="n">grad</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">norm</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">item</span><span class="p">()</span>
+            <span class="n">total_norm</span> <span class="o">+=</span> <span class="n">grad_norm</span> <span class="o">**</span> <span class="mi">2</span>
+
+            <span class="c1"># Collect gradient statistics for key layers</span>
+            <span class="k">if</span> <span class="nb">any</span><span class="p">(</span><span class="n">key</span> <span class="ow">in</span> <span class="n">name</span> <span class="k">for</span> <span class="n">key</span> <span class="ow">in</span> <span class="p">[</span><span class="s1">&#39;embed&#39;</span><span class="p">,</span> <span class="s1">&#39;lm_head&#39;</span><span class="p">,</span> <span class="s1">&#39;mlp.up&#39;</span><span class="p">,</span> <span class="s1">&#39;mlp.down&#39;</span><span class="p">,</span> <span class="s1">&#39;self_attn.q_proj&#39;</span><span class="p">,</span> <span class="s1">&#39;norm&#39;</span><span class="p">]):</span>
+                <span class="n">grad_samples</span><span class="p">[</span><span class="n">name</span><span class="p">]</span> <span class="o">=</span> <span class="p">{</span>
+                    <span class="s1">&#39;norm&#39;</span><span class="p">:</span> <span class="n">grad_norm</span><span class="p">,</span>
+                    <span class="s1">&#39;mean&#39;</span><span class="p">:</span> <span class="n">p</span><span class="o">.</span><span class="n">grad</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">mean</span><span class="p">()</span><span class="o">.</span><span class="n">item</span><span class="p">(),</span>
+                    <span class="s1">&#39;std&#39;</span><span class="p">:</span> <span class="n">p</span><span class="o">.</span><span class="n">grad</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">std</span><span class="p">()</span><span class="o">.</span><span class="n">item</span><span class="p">(),</span>
+                    <span class="s1">&#39;max&#39;</span><span class="p">:</span> <span class="n">p</span><span class="o">.</span><span class="n">grad</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">max</span><span class="p">()</span><span class="o">.</span><span class="n">item</span><span class="p">(),</span>
+                    <span class="s1">&#39;min&#39;</span><span class="p">:</span> <span class="n">p</span><span class="o">.</span><span class="n">grad</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">min</span><span class="p">()</span><span class="o">.</span><span class="n">item</span><span class="p">(),</span>
+                <span class="p">}</span>
+
+    <span class="n">total_norm</span> <span class="o">=</span> <span class="n">total_norm</span> <span class="o">**</span> <span class="mf">0.5</span>
+
+    <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="se">\n</span><span class="s2">Gradient norm: </span><span class="si">{</span><span class="n">total_norm</span><span class="si">:</span><span class="s2">.4f</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
+    <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Parameters with gradients: </span><span class="si">{</span><span class="n">param_count</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
+
+    <span class="c1"># Print sample gradients from important layers</span>
+    <span class="nb">print</span><span class="p">(</span><span class="s2">&quot;</span><span class="se">\n</span><span class="s2">Sample gradient statistics:&quot;</span><span class="p">)</span>
+    <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">stats</span><span class="p">)</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="nb">list</span><span class="p">(</span><span class="n">grad_samples</span><span class="o">.</span><span class="n">items</span><span class="p">())[:</span><span class="mi">10</span><span class="p">]):</span>
+        <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;  </span><span class="si">{</span><span class="n">name</span><span class="p">[:</span><span class="mi">60</span><span class="p">]</span><span class="si">:</span><span class="s2">&lt;60</span><span class="si">}</span><span class="s2"> | norm: </span><span class="si">{</span><span class="n">stats</span><span class="p">[</span><span class="s1">&#39;norm&#39;</span><span class="p">]</span><span class="si">:</span><span class="s2">.4e</span><span class="si">}</span><span class="s2"> | mean: </span><span class="si">{</span><span class="n">stats</span><span class="p">[</span><span class="s1">&#39;mean&#39;</span><span class="p">]</span><span class="si">:</span><span class="s2">.4e</span><span class="si">}</span><span class="s2"> | std: </span><span class="si">{</span><span class="n">stats</span><span class="p">[</span><span class="s1">&#39;std&#39;</span><span class="p">]</span><span class="si">:</span><span class="s2">.4e</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
+
+    <span class="c1"># Optional: zero gradients for next iteration</span>
+    <span class="n">model</span><span class="o">.</span><span class="n">zero_grad</span><span class="p">()</span>
+    <span class="n">model</span><span class="o">.</span><span class="n">eval</span><span class="p">()</span>  <span class="c1"># Switch back to eval mode</span>
 </pre></div>
 
-<div class="code-line-highlight" id="line-highlight-forward_only"></div>
+<div class="code-line-highlight" id="line-highlight-forward_and_backward"></div>
 </div>
 </div>
 </div>
-<div id="output-forward_only" class="cell-output">
-<div class="uv-install-logs" id="uv-logs-forward_only">
-<div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
-<div class="uv-logs-content" style="display: none;">
+<div id="output-forward_and_backward" class="cell-output">
+<div class="cell-stderr">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)
-   Building transformers @ git+https://github.com/huggingface/transformers.git@7258ea44bc0c0a425a468f66f8559d1de8c4126d
-Downloading triton (148.4MiB)
-Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
-Downloading hf-xet (3.0MiB)
-Downloading pillow (6.3MiB)
-Downloading tokenizers (3.1MiB)
-Downloading jedi (1.5MiB)
-Downloading nvidia-nvjitlink-cu12 (37.4MiB)
-Downloading nvidia-cufile-cu12 (1.1MiB)
-Downloading networkx (1.9MiB)
-Downloading nvidia-cusparselt-cu12 (273.9MiB)
-Downloading nvidia-cusolver-cu12 (255.1MiB)
-Downloading nvidia-cufft-cu12 (184.2MiB)
-Downloading nvidia-curand-cu12 (60.7MiB)
-Downloading nvidia-cublas-cu12 (566.8MiB)
-Downloading nvidia-nccl-cu12 (307.4MiB)
-Downloading nvidia-cudnn-cu12 (674.0MiB)
-Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
-Downloading sympy (6.0MiB)
-Downloading fonttools (4.7MiB)
-Downloading nvidia-cusparse-cu12 (274.9MiB)
-Downloading torch (846.8MiB)
-Downloading numpy (15.9MiB)
-Downloading matplotlib (8.3MiB)
-Downloading kiwisolver (1.4MiB)
- Downloading nvidia-cufile-cu12
- Downloading kiwisolver
- Downloading hf-xet
- Downloading tokenizers
- Downloading networkx
- Downloading fonttools
- Downloading pillow
- Downloading matplotlib
- Downloading nvidia-cuda-cupti-cu12
- Downloading sympy
- Downloading numpy
- Downloading jedi
-      Built transformers @ git+https://github.com/huggingface/transformers.git@7258ea44bc0c0a425a468f66f8559d1de8c4126d
- Downloading nvidia-nvjitlink-cu12
- Downloading nvidia-curand-cu12
- Downloading nvidia-cuda-nvrtc-cu12
- 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 321ms
-</div>
+    Updated https://github.com/huggingface/transformers.git (449533af73874470e914a203391635e04ac2ffc8)
+  × No solution found when resolving script dependencies:
+  ╰─▶ Because only transformers==4.57.0.dev0 is available and
+      transformers==4.57.0.dev0 depends on huggingface-hub==1.0.0rc1,
+      we can conclude that all versions of transformers depend on
+      huggingface-hub==1.0.0rc1.
+      And because kernels==0.10.0 depends on huggingface-hub&gt;=0.26.0,&lt;1.0,
+      we can conclude that kernels==0.10.0 and all versions of transformers
+      are incompatible.
+      And because you require kernels==0.10.0 and transformers, we can
+      conclude that your requirements are unsatisfiable.
 </div>
-<div class="cell-stderr">Fetching 3 files:   0%|          | 0/3 [00:00&lt;?, ?it/s]
-Fetching 3 files:   0%|          | 0/3 [00:50&lt;?, ?it/s]
-Traceback (most recent call last):
-  File &quot;/home/runner/work/kernels-uvnotes/kernels-uvnotes/moe_benchmarks/megablocks/.uvnote/cells/forward_only.py&quot;, line 68, in &lt;module&gt;
-    model = GptOssForCausalLM.from_pretrained(
-            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/home/runner/work/_temp/setup-uv-cache/environments-v2/forward-only-b65004b2d0cb4ca8/lib/python3.12/site-packages/transformers/modeling_utils.py&quot;, line 285, in _wrapper
-    return func(*args, **kwargs)
-           ^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/home/runner/work/_temp/setup-uv-cache/environments-v2/forward-only-b65004b2d0cb4ca8/lib/python3.12/site-packages/transformers/modeling_utils.py&quot;, line 4904, in from_pretrained
-    checkpoint_files, sharded_metadata = _get_resolved_checkpoint_files(
-                                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/home/runner/work/_temp/setup-uv-cache/environments-v2/forward-only-b65004b2d0cb4ca8/lib/python3.12/site-packages/transformers/modeling_utils.py&quot;, line 1239, in _get_resolved_checkpoint_files
-    checkpoint_files, sharded_metadata = get_checkpoint_shard_files(
-                                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/home/runner/work/_temp/setup-uv-cache/environments-v2/forward-only-b65004b2d0cb4ca8/lib/python3.12/site-packages/transformers/utils/hub.py&quot;, line 1116, in get_checkpoint_shard_files
-    cached_filenames = cached_files(
-                       ^^^^^^^^^^^^^
-  File &quot;/home/runner/work/_temp/setup-uv-cache/environments-v2/forward-only-b65004b2d0cb4ca8/lib/python3.12/site-packages/transformers/utils/hub.py&quot;, line 564, in cached_files
-    raise e
-  File &quot;/home/runner/work/_temp/setup-uv-cache/environments-v2/forward-only-b65004b2d0cb4ca8/lib/python3.12/site-packages/transformers/utils/hub.py&quot;, line 491, in cached_files
-    snapshot_download(
-  File &quot;/home/runner/work/_temp/setup-uv-cache/environments-v2/forward-only-b65004b2d0cb4ca8/lib/python3.12/site-packages/huggingface_hub/utils/_validators.py&quot;, line 114, in _inner_fn
-    return fn(*args, **kwargs)
-           ^^^^^^^^^^^^^^^^^^^
-  File &quot;/home/runner/work/_temp/setup-uv-cache/environments-v2/forward-only-b65004b2d0cb4ca8/lib/python3.12/site-packages/huggingface_hub/_snapshot_download.py&quot;, line 332, in snapshot_download
-    thread_map(
-  File &quot;/home/runner/work/_temp/setup-uv-cache/environments-v2/forward-only-b65004b2d0cb4ca8/lib/python3.12/site-packages/tqdm/contrib/concurrent.py&quot;, line 69, in thread_map
-    return _executor_map(ThreadPoolExecutor, fn, *iterables, **tqdm_kwargs)
-           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/home/runner/work/_temp/setup-uv-cache/environments-v2/forward-only-b65004b2d0cb4ca8/lib/python3.12/site-packages/tqdm/contrib/concurrent.py&quot;, line 51, in _executor_map
-    return list(tqdm_class(ex.map(fn, *iterables, chunksize=chunksize), **kwargs))
-           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/home/runner/work/_temp/setup-uv-cache/environments-v2/forward-only-b65004b2d0cb4ca8/lib/python3.12/site-packages/tqdm/std.py&quot;, line 1181, in __iter__
-    for obj in iterable:
-  File &quot;/usr/lib/python3.12/concurrent/futures/_base.py&quot;, line 619, in result_iterator
-    yield _result_or_cancel(fs.pop())
-          ^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/usr/lib/python3.12/concurrent/futures/_base.py&quot;, line 317, in _result_or_cancel
-    return fut.result(timeout)
-           ^^^^^^^^^^^^^^^^^^^
-  File &quot;/usr/lib/python3.12/concurrent/futures/_base.py&quot;, line 456, in result
-    return self.__get_result()
-           ^^^^^^^^^^^^^^^^^^^
-  File &quot;/usr/lib/python3.12/concurrent/futures/_base.py&quot;, line 401, in __get_result
-    raise self._exception
-  File &quot;/usr/lib/python3.12/concurrent/futures/thread.py&quot;, line 58, in run
-    result = self.fn(*self.args, **self.kwargs)
-             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/home/runner/work/_temp/setup-uv-cache/environments-v2/forward-only-b65004b2d0cb4ca8/lib/python3.12/site-packages/huggingface_hub/_snapshot_download.py&quot;, line 306, in _inner_hf_hub_download
-    return hf_hub_download(
-           ^^^^^^^^^^^^^^^^
-  File &quot;/home/runner/work/_temp/setup-uv-cache/environments-v2/forward-only-b65004b2d0cb4ca8/lib/python3.12/site-packages/huggingface_hub/utils/_validators.py&quot;, line 114, in _inner_fn
-    return fn(*args, **kwargs)
-           ^^^^^^^^^^^^^^^^^^^
-  File &quot;/home/runner/work/_temp/setup-uv-cache/environments-v2/forward-only-b65004b2d0cb4ca8/lib/python3.12/site-packages/huggingface_hub/file_download.py&quot;, line 1010, in hf_hub_download
-    return _hf_hub_download_to_cache_dir(
-           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/home/runner/work/_temp/setup-uv-cache/environments-v2/forward-only-b65004b2d0cb4ca8/lib/python3.12/site-packages/huggingface_hub/file_download.py&quot;, line 1171, in _hf_hub_download_to_cache_dir
-    _download_to_tmp_and_move(
-  File &quot;/home/runner/work/_temp/setup-uv-cache/environments-v2/forward-only-b65004b2d0cb4ca8/lib/python3.12/site-packages/huggingface_hub/file_download.py&quot;, line 1723, in _download_to_tmp_and_move
-    xet_get(
-  File &quot;/home/runner/work/_temp/setup-uv-cache/environments-v2/forward-only-b65004b2d0cb4ca8/lib/python3.12/site-packages/huggingface_hub/file_download.py&quot;, line 629, in xet_get
-    download_files(
-RuntimeError: Data processing error: CAS service error : IO Error: No space left on device (os error 28)</div>
 </div>
 </div>
-
-<h2>Forward and Backward</h2>
-<p>Next, we run a forward and backward pass with Megablocks kernels enabled. This should be more memory efficient and allow us to complete the backward pass without running out of memory.</p>
     </div>
     
 </body>

moe_benchmarks/megablocks_yamoe/artifacts/binned_run/binned_results.json

@@ -0,0 +1,24 @@
+{
+  "implementation": "binned_results",
+  "config": {
+    "warmup": 10,
+    "iters": 50,
+    "device": "cuda",
+    "dtype": "torch.float32",
+    "tokens": 100,
+    "vary_inputs": true
+  },
+  "stats": {
+    "avg_ms": 36.06324691992995,
+    "min_ms": 33.29206800026441,
+    "max_ms": 38.40615900026023,
+    "std_ms": 1.258567678508065,
+    "p50_ms": 36.21510599987232,
+    "p95_ms": 37.524451049966956,
+    "p99_ms": 38.03603995002959,
+    "num_iters": 50,
+    "tokens_per_s": 2772.906172925215,
+    "throughput_variance": 98.28636435515342
+  },
+  "output_sum": 3.97190523147583
+}

moe_benchmarks/megablocks_yamoe/artifacts/gptoss_run/gptoss_results.json

@@ -0,0 +1,24 @@
+{
+  "implementation": "gptoss_results",
+  "config": {
+    "warmup": 10,
+    "iters": 50,
+    "device": "cuda",
+    "dtype": "torch.float32",
+    "tokens": 100,
+    "vary_inputs": true
+  },
+  "stats": {
+    "avg_ms": 45.286630379978305,
+    "min_ms": 38.91367899996112,
+    "max_ms": 49.84392799997295,
+    "std_ms": 3.2326168009526866,
+    "p50_ms": 45.42240999990099,
+    "p95_ms": 49.729684149951936,
+    "p99_ms": 49.82545450991893,
+    "num_iters": 50,
+    "tokens_per_s": 2208.1572234663554,
+    "throughput_variance": 161.27578702324564
+  },
+  "output_sum": 11.53223705291748
+}

moe_benchmarks/megablocks_yamoe/artifacts/gptoss_training_run/gptoss_training_results.json

@@ -0,0 +1,24 @@
+{
+  "implementation": "gptoss_training_results",
+  "config": {
+    "warmup": 10,
+    "iters": 50,
+    "device": "cuda",
+    "dtype": "torch.float32",
+    "tokens": 100,
+    "vary_inputs": true
+  },
+  "stats": {
+    "avg_ms": 46.01034353989235,
+    "min_ms": 39.20698799993261,
+    "max_ms": 51.09754699969926,
+    "std_ms": 3.2594474712819497,
+    "p50_ms": 46.132551999562565,
+    "p95_ms": 50.721096600273086,
+    "p99_ms": 51.0080171399477,
+    "num_iters": 50,
+    "tokens_per_s": 2173.4243282338675,
+    "throughput_variance": 158.68467070353637
+  },
+  "output_sum": 11.53223705291748
+}

moe_benchmarks/megablocks_yamoe/artifacts/yamoe_run/yamoe_results.json

@@ -0,0 +1,24 @@
+{
+  "implementation": "yamoe_results",
+  "config": {
+    "warmup": 10,
+    "iters": 50,
+    "device": "cuda",
+    "dtype": "torch.float32",
+    "tokens": 100,
+    "vary_inputs": true
+  },
+  "stats": {
+    "avg_ms": 4.2510544400101935,
+    "min_ms": 4.144352999901457,
+    "max_ms": 4.320155999266717,
+    "std_ms": 0.02873328656403644,
+    "p50_ms": 4.2539659998510615,
+    "p95_ms": 4.2857709999225335,
+    "p99_ms": 4.306132199617423,
+    "num_iters": 50,
+    "tokens_per_s": 23523.575482547854,
+    "throughput_variance": 160.28680309512873
+  },
+  "output_sum": 3.97190523147583
+}

moe_benchmarks/megablocks_yamoe/cells/binned_run.py

@@ -0,0 +1,195 @@
+# /// script
+# dependencies = [
+#     "torch",
+#     "numpy",
+# ]
+# ///
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from bench_utils import to_dtype, tensor_stats, set_seed, bench_context
+from config import (
+    NUM_EXPERTS, HIDDEN_SIZE, TOP_K,
+    BATCH_SIZE, SEQ_LEN, DTYPE, DEVICE,
+    WEIGHT_SEED, EXPERT_SEED, INPUT_SEED, GENERAL_SEED
+)
+from pathlib import Path
+import os
+
+# Discover the upstream artifact directory from env
+data_dir = os.environ.get('UVNOTE_INPUT_SAVE_DATA', '.')
+
+router_weight = torch.load(Path(data_dir) / 'router_weight.pt')
+router_bias = torch.load(Path(data_dir) / 'router_bias.pt')
+gate_up_proj = torch.load(Path(data_dir) / 'gate_up_proj.pt')
+gate_up_proj_bias = torch.load(Path(data_dir) / 'gate_up_proj_bias.pt')
+down_proj = torch.load(Path(data_dir) / 'down_proj.pt')
+down_proj_bias = torch.load(Path(data_dir) / 'down_proj_bias.pt')
+
+print("Loaded shared weights from artifacts")
+print(f"Router weight sum: {router_weight.sum().item():.6f}")
+print(f"Gate/up sum: {gate_up_proj.sum().item():.6f}")
+print(f"Down sum: {down_proj.sum().item():.6f}")
+
+def binned_gather(x, indices, bins, expert_capacity, top_k):
+    E, H = bins.shape[0], x.shape[1]
+    out = torch.zeros((E, expert_capacity, H), device=x.device, dtype=x.dtype)
+    for e in range(E):
+        start = 0 if e == 0 else bins[e - 1]
+        end = bins[e]
+        n = min(end - start, expert_capacity)
+        for i in range(n):
+            flat_pos = indices[start + i]
+            tok = flat_pos // top_k
+            out[e, i] = x[tok]
+    return out
+
+def binned_scatter(x, indices, weights, bins, expert_capacity, top_k):
+    E, C, H = x.shape
+    N = indices.shape[0] // top_k
+    out = torch.zeros((N, top_k, H), dtype=x.dtype, device=x.device)
+    for e in range(E):
+        start = 0 if e == 0 else bins[e - 1]
+        end = bins[e]
+        n = end - start
+        if n == 0:
+            continue
+        take = min(n, expert_capacity)
+        for i in range(take):
+            flat_pos = indices[start + i]
+            tok = flat_pos // top_k
+            slot = flat_pos % top_k
+            scale = weights[flat_pos] if weights is not None else 1.0
+            out[tok, slot] = x[e, i] * scale
+    return out.sum(dim=1)
+
+def sort_tokens_by_expert(router_indices, num_experts):
+    flat_indices = router_indices.flatten()
+    sorted_values, sorted_indices = torch.sort(flat_indices)
+    tokens_per_expert = torch.bincount(sorted_values, minlength=num_experts)
+    bins = torch.cumsum(tokens_per_expert, dim=0)
+    return sorted_indices, sorted_values, bins, tokens_per_expert
+
+def binned_experts_ref(
+    hidden_states,
+    router_indices,
+    routing_weights,
+    gate_up_proj,
+    gate_up_proj_bias,
+    down_proj,
+    down_proj_bias,
+    expert_capacity,
+):
+    B, S, H = hidden_states.shape
+    E, K = routing_weights.shape[1], router_indices.shape[1]
+
+    indices, _, bins, _ = sort_tokens_by_expert(router_indices, E)
+    x = binned_gather(hidden_states.view(-1, H), indices, bins, expert_capacity, K)
+
+    gate_up = torch.bmm(x, gate_up_proj) 
+    gate_up += gate_up_proj_bias[..., None, :]
+
+    gate, up = gate_up[..., ::2], gate_up[..., 1::2]
+
+    # clamp to limit
+    limit = 7.0
+    gate = gate.clamp(min=None, max=limit)
+    up = up.clamp(min=-limit, max=limit)
+
+    glu = gate * torch.sigmoid(gate * 1.702)
+    x = (up + 1) * glu
+    x = torch.bmm(x, down_proj) + down_proj_bias[..., None, :]
+
+    # build routing weights aligned to (token, slot)
+    flat_dense = routing_weights.view(-1, E)
+    flat_router = router_indices.view(-1, K)
+    selected = torch.gather(flat_dense, 1, flat_router).reshape(-1)
+
+    # scatter back
+    y = binned_scatter(x, indices, selected, bins, expert_capacity, K)
+
+    return y.view(B, S, H)
+
+class BinnedRouter(nn.Module):
+    def __init__(self, router_weight, router_bias):
+        super().__init__()
+        self.top_k = TOP_K
+        self.num_experts = NUM_EXPERTS
+        self.hidden_dim = HIDDEN_SIZE
+        self.weight = nn.Parameter(router_weight.clone())
+        self.bias = nn.Parameter(router_bias.clone())
+
+    def forward(self, hidden_states):
+        hidden_states = hidden_states.reshape(-1, self.hidden_dim)
+        router_logits = F.linear(hidden_states, self.weight, self.bias)
+        router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1)
+        router_top_value = torch.nn.functional.softmax(router_top_value, dim=1, dtype=router_top_value.dtype)
+        router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
+        return router_scores, router_indices
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+class BinnedMoEMLP(nn.Module):
+    def __init__(self, router_weight, router_bias, gate_up_proj, gate_up_proj_bias, down_proj, down_proj_bias):
+        super().__init__()
+        self.router = BinnedRouter(router_weight, router_bias)
+        self.num_experts = NUM_EXPERTS
+        self.hidden_size = HIDDEN_SIZE
+        self.top_k = TOP_K
+        
+        # Expert weights - use the loaded weights
+        self.gate_up_proj = nn.Parameter(gate_up_proj.clone())
+        self.gate_up_proj_bias = nn.Parameter(gate_up_proj_bias.clone())
+        self.down_proj = nn.Parameter(down_proj.clone())
+        self.down_proj_bias = nn.Parameter(down_proj_bias.clone())
+
+    def forward(self, hidden_states):
+        router_scores, router_indices = self.router(hidden_states)
+        batch_size = hidden_states.shape[0]
+        expert_capacity = ceil_div(batch_size * self.top_k, self.num_experts)
+
+        output = binned_experts_ref(
+            hidden_states,
+            router_indices,
+            router_scores,
+            self.gate_up_proj,
+            self.gate_up_proj_bias,
+            self.down_proj,
+            self.down_proj_bias,
+            expert_capacity,
+        )
+        
+        return output, router_scores
+
+# Run the model
+set_seed(GENERAL_SEED)
+
+device = torch.device(DEVICE)
+dtype = to_dtype(DTYPE)
+
+print("\n=== Binned Implementation ===")
+# Initialize model with loaded weights
+model = BinnedMoEMLP(
+    router_weight.to(device),
+    router_bias.to(device),
+    gate_up_proj.to(device),
+    gate_up_proj_bias.to(device),
+    down_proj.to(device),
+    down_proj_bias.to(device)
+).to(device=device)
+
+print(f"Router weight sum: {model.router.weight.sum().item():.6f}")
+print(f"Gate/up proj sum: {model.gate_up_proj.sum().item():.6f}")
+print(f"Down proj sum: {model.down_proj.sum().item():.6f}")
+
+# Generate the same input as Yamoe
+set_seed(INPUT_SEED)
+x = torch.randn(BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE, device=device, dtype=dtype) * 0.1
+
+# Benchmark the model with varied inputs to prevent caching artifacts
+tokens = BATCH_SIZE * SEQ_LEN
+with bench_context(warmup=10, iters=50, device=device, dtype=dtype, tokens=tokens, save_json="binned_results.json", vary_inputs=True) as bench:
+    output, stats = bench(model, x)
+    print(f"\nOutput sum: {output[0].sum().item():.6f}")

moe_benchmarks/megablocks_yamoe/cells/gptoss_run.py

@@ -0,0 +1,147 @@
+# /// script
+# dependencies = [
+#     "torch",
+#     "numpy",
+# ]
+# ///
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from bench_utils import to_dtype, tensor_stats, set_seed, bench_context
+from config import (
+    NUM_EXPERTS, HIDDEN_SIZE, TOP_K,
+    BATCH_SIZE, SEQ_LEN, DTYPE, DEVICE,
+    WEIGHT_SEED, EXPERT_SEED, INPUT_SEED, GENERAL_SEED
+)
+from pathlib import Path
+import os
+
+# Discover the upstream artifact directory from env
+data_dir = os.environ.get('UVNOTE_INPUT_SAVE_DATA', '.')
+
+router_weight = torch.load(Path(data_dir) / 'router_weight.pt')
+router_bias = torch.load(Path(data_dir) / 'router_bias.pt')
+gate_up_proj = torch.load(Path(data_dir) / 'gate_up_proj.pt')
+gate_up_proj_bias = torch.load(Path(data_dir) / 'gate_up_proj_bias.pt')
+down_proj = torch.load(Path(data_dir) / 'down_proj.pt')
+down_proj_bias = torch.load(Path(data_dir) / 'down_proj_bias.pt')
+
+print("Loaded shared weights from artifacts")
+print(f"Router weight sum: {router_weight.sum().item():.6f}")
+print(f"Gate/up sum: {gate_up_proj.sum().item():.6f}")
+print(f"Down sum: {down_proj.sum().item():.6f}")
+
+class GptOssRouter(nn.Module):
+    def __init__(self, router_weight, router_bias):
+        super().__init__()
+        self.top_k = TOP_K
+        self.num_experts = NUM_EXPERTS
+        self.hidden_dim = HIDDEN_SIZE
+        self.weight = nn.Parameter(router_weight.clone())
+        self.bias = nn.Parameter(router_bias.clone())
+
+    def forward(self, hidden_states):
+        hidden_states = hidden_states.reshape(-1, self.hidden_dim)
+        router_logits = F.linear(hidden_states, self.weight, self.bias)
+        router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1)
+        router_top_value = torch.nn.functional.softmax(router_top_value, dim=1, dtype=router_top_value.dtype)
+        router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
+        return router_scores, router_indices
+
+class GptOssExperts(nn.Module):
+    def __init__(self, gate_up_proj, gate_up_proj_bias, down_proj, down_proj_bias):
+        super().__init__()
+        self.num_experts = NUM_EXPERTS
+        self.hidden_size = HIDDEN_SIZE
+        self.expert_dim = self.hidden_size
+        self.gate_up_proj = nn.Parameter(gate_up_proj.clone())
+        self.gate_up_proj_bias = nn.Parameter(gate_up_proj_bias.clone())
+        self.down_proj = nn.Parameter(down_proj.clone())
+        self.down_proj_bias = nn.Parameter(down_proj_bias.clone())
+        self.alpha = 1.702
+        self.limit = 7.0
+
+    def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None) -> torch.Tensor:
+        batch_size = hidden_states.shape[0]
+        hidden_states = hidden_states.reshape(-1, self.hidden_size)
+        num_experts = routing_weights.shape[1]
+        
+        if hidden_states.device.type == "cpu" or self.training:
+            next_states = torch.zeros_like(hidden_states, dtype=hidden_states.dtype, device=hidden_states.device)
+            with torch.no_grad():
+                expert_mask = torch.nn.functional.one_hot(router_indices, num_classes=num_experts)
+                expert_mask = expert_mask.permute(2, 1, 0)
+                expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+            
+            for expert_idx in expert_hit[:]:
+                expert_idx = expert_idx[0]
+                with torch.no_grad():
+                    _, token_idx = torch.where(expert_mask[expert_idx])
+                current_state = hidden_states[token_idx]
+                gate_up = current_state @ self.gate_up_proj[expert_idx] + self.gate_up_proj_bias[expert_idx]
+                gate, up = gate_up[..., ::2], gate_up[..., 1::2]
+                gate = gate.clamp(min=None, max=self.limit)
+                up = up.clamp(min=-self.limit, max=self.limit)
+                glu = gate * torch.sigmoid(gate * self.alpha)
+                gated_output = (up + 1) * glu
+                out = gated_output @ self.down_proj[expert_idx] + self.down_proj_bias[expert_idx]
+                weighted_output = out * routing_weights[token_idx, expert_idx, None]
+                next_states.index_add_(0, token_idx, weighted_output.to(hidden_states.dtype))
+            next_states = next_states.view(batch_size, -1, self.hidden_size)
+        else:
+            hidden_states = hidden_states.repeat(num_experts, 1)
+            hidden_states = hidden_states.view(num_experts, -1, self.hidden_size)
+            gate_up = torch.bmm(hidden_states, self.gate_up_proj) + self.gate_up_proj_bias[..., None, :]
+            gate, up = gate_up[..., ::2], gate_up[..., 1::2]
+            gate = gate.clamp(min=None, max=self.limit)
+            up = up.clamp(min=-self.limit, max=self.limit)
+            glu = gate * torch.sigmoid(gate * self.alpha)
+            next_states = torch.bmm(((up + 1) * glu), self.down_proj)
+            next_states = next_states + self.down_proj_bias[..., None, :]
+            next_states = next_states.view(num_experts, batch_size, -1, self.hidden_size)
+            next_states = next_states * routing_weights.transpose(0, 1).view(num_experts, batch_size, -1)[..., None]
+            next_states = next_states.sum(dim=0)
+        return next_states
+
+class GptOssMoEMLP(nn.Module):
+    def __init__(self, router_weight, router_bias, gate_up_proj, gate_up_proj_bias, down_proj, down_proj_bias):
+        super().__init__()
+        self.router = GptOssRouter(router_weight, router_bias)
+        self.experts = GptOssExperts(gate_up_proj, gate_up_proj_bias, down_proj, down_proj_bias)
+
+    def forward(self, hidden_states):
+        router_scores, router_indices = self.router(hidden_states)
+        routed_out = self.experts(hidden_states, router_indices=router_indices, routing_weights=router_scores)
+        return routed_out, router_scores
+
+# Run the model
+set_seed(GENERAL_SEED)
+
+device = torch.device(DEVICE)
+dtype = to_dtype(DTYPE)
+
+print("\n=== GPT-OSS Implementation ===")
+# Initialize model with loaded weights
+model = GptOssMoEMLP(
+    router_weight.to(device),
+    router_bias.to(device),
+    gate_up_proj.to(device),
+    gate_up_proj_bias.to(device),
+    down_proj.to(device),
+    down_proj_bias.to(device)
+).to(device=device)
+
+print(f"Router weight sum: {model.router.weight.sum().item():.6f}")
+print(f"Gate/up proj sum: {model.experts.gate_up_proj.sum().item():.6f}")
+print(f"Down proj sum: {model.experts.down_proj.sum().item():.6f}")
+
+# Generate the same input as other implementations
+set_seed(INPUT_SEED)
+x = torch.randn(BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE, device=device, dtype=dtype) * 0.1
+
+# Benchmark the model with varied inputs to prevent caching artifacts
+tokens = BATCH_SIZE * SEQ_LEN
+with bench_context(warmup=10, iters=50, device=device, dtype=dtype, tokens=tokens, save_json="gptoss_results.json", vary_inputs=True) as bench:
+    output, stats = bench(model, x)
+    print(f"\nOutput sum: {output[0].sum().item():.6f}")

moe_benchmarks/megablocks_yamoe/cells/gptoss_training_run.py

@@ -0,0 +1,138 @@
+# /// script
+# dependencies = [
+#     "torch",
+#     "numpy",
+# ]
+# ///
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from bench_utils import to_dtype, tensor_stats, set_seed, bench_context
+from config import (
+    NUM_EXPERTS, HIDDEN_SIZE, TOP_K,
+    BATCH_SIZE, SEQ_LEN, DTYPE, DEVICE,
+    WEIGHT_SEED, EXPERT_SEED, INPUT_SEED, GENERAL_SEED
+)
+from pathlib import Path
+import os
+
+# Discover the upstream artifact directory from env
+data_dir = os.environ.get('UVNOTE_INPUT_SAVE_DATA', '.')
+
+router_weight = torch.load(Path(data_dir) / 'router_weight.pt')
+router_bias = torch.load(Path(data_dir) / 'router_bias.pt')
+gate_up_proj = torch.load(Path(data_dir) / 'gate_up_proj.pt')
+gate_up_proj_bias = torch.load(Path(data_dir) / 'gate_up_proj_bias.pt')
+down_proj = torch.load(Path(data_dir) / 'down_proj.pt')
+down_proj_bias = torch.load(Path(data_dir) / 'down_proj_bias.pt')
+
+print("Loaded shared weights from artifacts")
+print(f"Router weight sum: {router_weight.sum().item():.6f}")
+print(f"Gate/up sum: {gate_up_proj.sum().item():.6f}")
+print(f"Down sum: {down_proj.sum().item():.6f}")
+
+class GptOssTrainingRouter(nn.Module):
+    def __init__(self, router_weight, router_bias):
+        super().__init__()
+        self.top_k = TOP_K
+        self.num_experts = NUM_EXPERTS
+        self.hidden_dim = HIDDEN_SIZE
+        self.weight = nn.Parameter(router_weight.clone())
+        self.bias = nn.Parameter(router_bias.clone())
+
+    def forward(self, hidden_states):
+        hidden_states = hidden_states.reshape(-1, self.hidden_dim)
+        router_logits = F.linear(hidden_states, self.weight, self.bias)
+        router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1)
+        router_top_value = torch.nn.functional.softmax(router_top_value, dim=1, dtype=router_top_value.dtype)
+        router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
+        return router_scores, router_indices
+
+class GptOssTrainingExperts(nn.Module):
+    def __init__(self, gate_up_proj, gate_up_proj_bias, down_proj, down_proj_bias):
+        super().__init__()
+        self.num_experts = NUM_EXPERTS
+        self.hidden_size = HIDDEN_SIZE
+        self.expert_dim = self.hidden_size
+        self.gate_up_proj = nn.Parameter(gate_up_proj.clone())
+        self.gate_up_proj_bias = nn.Parameter(gate_up_proj_bias.clone())
+        self.down_proj = nn.Parameter(down_proj.clone())
+        self.down_proj_bias = nn.Parameter(down_proj_bias.clone())
+        self.alpha = 1.702
+        self.limit = 7.0
+
+    def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None) -> torch.Tensor:
+        batch_size = hidden_states.shape[0]
+        hidden_states = hidden_states.reshape(-1, self.hidden_size)
+        num_experts = routing_weights.shape[1]
+        
+        # Force training mode path (expert loop instead of batched)
+        next_states = torch.zeros_like(hidden_states, dtype=hidden_states.dtype, device=hidden_states.device)
+        with torch.no_grad():
+            expert_mask = torch.nn.functional.one_hot(router_indices, num_classes=num_experts)
+            expert_mask = expert_mask.permute(2, 1, 0)
+            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+        
+        for expert_idx in expert_hit[:]:
+            expert_idx = expert_idx[0]
+            with torch.no_grad():
+                _, token_idx = torch.where(expert_mask[expert_idx])
+            current_state = hidden_states[token_idx]
+            gate_up = current_state @ self.gate_up_proj[expert_idx] + self.gate_up_proj_bias[expert_idx]
+            gate, up = gate_up[..., ::2], gate_up[..., 1::2]
+            gate = gate.clamp(min=None, max=self.limit)
+            up = up.clamp(min=-self.limit, max=self.limit)
+            glu = gate * torch.sigmoid(gate * self.alpha)
+            gated_output = (up + 1) * glu
+            out = gated_output @ self.down_proj[expert_idx] + self.down_proj_bias[expert_idx]
+            weighted_output = out * routing_weights[token_idx, expert_idx, None]
+            next_states.index_add_(0, token_idx, weighted_output.to(hidden_states.dtype))
+        next_states = next_states.view(batch_size, -1, self.hidden_size)
+        return next_states
+
+class GptOssTrainingMoEMLP(nn.Module):
+    def __init__(self, router_weight, router_bias, gate_up_proj, gate_up_proj_bias, down_proj, down_proj_bias):
+        super().__init__()
+        self.router = GptOssTrainingRouter(router_weight, router_bias)
+        self.experts = GptOssTrainingExperts(gate_up_proj, gate_up_proj_bias, down_proj, down_proj_bias)
+
+    def forward(self, hidden_states):
+        router_scores, router_indices = self.router(hidden_states)
+        routed_out = self.experts(hidden_states, router_indices=router_indices, routing_weights=router_scores)
+        return routed_out, router_scores
+
+# Run the model
+set_seed(GENERAL_SEED)
+
+device = torch.device(DEVICE)
+dtype = to_dtype(DTYPE)
+
+print("\n=== GPT-OSS Implementation (Training Mode - Expert Loop) ===")
+# Initialize model with loaded weights and force training mode
+model = GptOssTrainingMoEMLP(
+    router_weight.to(device),
+    router_bias.to(device),
+    gate_up_proj.to(device),
+    gate_up_proj_bias.to(device),
+    down_proj.to(device),
+    down_proj_bias.to(device)
+).to(device=device)
+
+# Set to training mode to force expert loop path
+model.train()
+
+print(f"Router weight sum: {model.router.weight.sum().item():.6f}")
+print(f"Gate/up proj sum: {model.experts.gate_up_proj.sum().item():.6f}")
+print(f"Down proj sum: {model.experts.down_proj.sum().item():.6f}")
+print(f"Model training mode: {model.training}")
+
+# Generate the same input as other implementations
+set_seed(INPUT_SEED)
+x = torch.randn(BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE, device=device, dtype=dtype) * 0.1
+
+# Benchmark the model with varied inputs to prevent caching artifacts
+tokens = BATCH_SIZE * SEQ_LEN
+with bench_context(warmup=10, iters=50, device=device, dtype=dtype, tokens=tokens, save_json="gptoss_training_results.json", vary_inputs=True) as bench:
+    output, stats = bench(model, x)
+    print(f"\nOutput sum: {output[0].sum().item():.6f}")

moe_benchmarks/megablocks_yamoe/cells/megablocks_run.py

@@ -0,0 +1,103 @@
+# /// script
+# dependencies = [
+#     "torch",
+#     "numpy",
+#     "kernels",
+# ]
+# ///
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from kernels import get_kernel, get_local_kernel
+from bench_utils import to_dtype, tensor_stats, set_seed, bench_context
+from config import (
+    NUM_EXPERTS, HIDDEN_SIZE, TOP_K,
+    BATCH_SIZE, SEQ_LEN, DTYPE, DEVICE,
+    WEIGHT_SEED, EXPERT_SEED, INPUT_SEED, GENERAL_SEED
+)
+from pathlib import Path
+from collections import namedtuple
+import os
+
+# Discover the upstream artifact directory from env
+data_dir = os.environ.get('UVNOTE_INPUT_SAVE_DATA', '.')
+
+print(f"Loading weights from: {data_dir}")
+
+router_weight = torch.load(Path(data_dir) / 'router_weight.pt')
+router_bias = torch.load(Path(data_dir) / 'router_bias.pt')
+gate_up_proj = torch.load(Path(data_dir) / 'gate_up_proj.pt')
+gate_up_proj_bias = torch.load(Path(data_dir) / 'gate_up_proj_bias.pt')
+down_proj = torch.load(Path(data_dir) / 'down_proj.pt')
+down_proj_bias = torch.load(Path(data_dir) / 'down_proj_bias.pt')
+
+print("Loaded shared weights from artifacts")
+print(f"Router weight sum: {router_weight.sum().item():.6f}")
+print(f"Gate/up sum: {gate_up_proj.sum().item():.6f}")
+print(f"Down sum: {down_proj.sum().item():.6f}")
+
+def build_megablocks_model(device: torch.device):
+    # Download optimized kernels from the Hugging Face hub
+    megablocks = get_kernel("kernels-community/megablocks", revision="v0.0.2")
+    model = megablocks.layers.MegaBlocksMoeMLP()
+
+    # Create attribute container for expert weights
+    model.experts = namedtuple(
+        "Experts", ["gate_up_proj", "gate_up_proj_bias", "down_proj", "down_proj_bias", "hidden_size"]
+    )
+
+    # Use loaded router weights for consistency
+    model.router = torch.nn.Linear(HIDDEN_SIZE, NUM_EXPERTS, device=device)
+    with torch.no_grad():
+        model.router.weight.copy_(router_weight)
+        model.router.bias.copy_(router_bias)
+
+    # 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))
+    e.down_proj_bias = torch.nn.Parameter(down_proj_bias.clone().to(device))
+    e.hidden_size = HIDDEN_SIZE
+    
+    # Log weight statistics for comparison
+    print(f"[MegaBlocks] Router weight sum: {model.router.weight.sum().item():.6f}")
+    print(f"[MegaBlocks] Gate/up projection shape: {tuple(e.gate_up_proj.shape)}, sum: {e.gate_up_proj.sum().item():.6f}")
+    print(f"[MegaBlocks] Down projection shape: {tuple(e.down_proj.shape)}, sum: {e.down_proj.sum().item():.6f}")
+
+    return model
+
+# Create a wrapper to match the interface of other implementations
+class MegaBlocksMoEWrapper(nn.Module):
+    def __init__(self, megablocks_model):
+        super().__init__()
+        self.model = megablocks_model
+        
+    def forward(self, hidden_states):
+        # MegaBlocks expects input in the format (batch, seq_len, hidden_dim)
+        output, dummy_routing_weights = self.model(hidden_states)
+        return output, dummy_routing_weights
+
+# Run the model
+set_seed(GENERAL_SEED)
+
+device = torch.device(DEVICE)
+dtype = to_dtype(DTYPE)
+
+print("\n=== MegaBlocks Implementation ===")
+# Build MegaBlocks model with loaded weights
+megablocks_model = build_megablocks_model(device)
+model = MegaBlocksMoEWrapper(megablocks_model).to(device=device)
+
+# Generate the same input as other implementations
+set_seed(INPUT_SEED)
+x = torch.randn(BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE, device=device, dtype=dtype) * 0.1
+
+# Benchmark the model with varied inputs to prevent caching artifacts
+tokens = BATCH_SIZE * SEQ_LEN
+with bench_context(warmup=10, iters=50, device=device, dtype=dtype, tokens=tokens, save_json="megablocks_results.json", vary_inputs=True) as bench:
+    output, stats = bench(model, x)
+    print(f"\nOutput sum: {output[0].sum().item():.6f}")

moe_benchmarks/megablocks_yamoe/cells/setup.py

@@ -0,0 +1,116 @@
+# /// script
+# requires-python = ">=3.12"
+# dependencies = [
+#     "accelerate>=1.10.1",
+#     "torch>=2.7.0",
+#     "kernels==0.10.0",
+#     "transformers@https://github.com/huggingface/transformers.git",
+#     "ipdb>=0.13.13",
+#     "matplotlib>=3.7.2",
+#     "numpy>=1.24.3",
+# ]
+# ///
+
+import torch
+from transformers import GptOssForCausalLM, PreTrainedTokenizerFast, Mxfp4Config
+import time
+import torch.nn as nn
+from kernels import register_kernel_mapping, Mode, LayerRepository
+import sys
+import torch.profiler
+import gc
+import logging
+
+# set to debug logging
+logging.basicConfig(level=logging.INFO)
+
+def reset_peak_memory_stats():
+    """Clear CUDA cache and reset memory allocation counters."""
+    torch.cuda.empty_cache()
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats()
+    gc.collect()
+
+def get_memory_stats():
+    """Get current and peak CUDA memory usage."""
+    if not torch.cuda.is_available():
+        return {"allocated_gb": 0, "peak_gb": 0, "reserved_gb": 0}
+    return {
+        "allocated_gb": torch.cuda.memory_allocated() / 1e9,
+        "peak_gb": torch.cuda.max_memory_allocated() / 1e9,
+        "reserved_gb": torch.cuda.memory_reserved() / 1e9,
+    }
+
+def override_kernel_layer_name(cls_name: str, value) -> bool:
+    """Helper to dynamically override the kernel_layer_name in a model class."""
+    for mod in sys.modules.values():
+        if mod is None:
+            continue
+        obj = getattr(mod, cls_name, None)
+        if isinstance(obj, type) and issubclass(obj, nn.Module):
+            setattr(obj, "kernel_layer_name", value)
+            print(f"Overrode {cls_name}.kernel_layer_name to {value}")
+            return True
+    return False
+
+
+# Init the model the normal way
+model_id = "openai/gpt-oss-20b"
+tokenizer = PreTrainedTokenizerFast.from_pretrained(model_id)
+quantization_config = Mxfp4Config(dequantize=True)
+
+
+from kernels import replace_kernel_forward_from_hub, register_kernel_mapping, LayerRepository, Mode
+
+from transformers.models.gpt_oss.modeling_gpt_oss import GptOssMLP, GptOssRMSNorm
+
+replace_kernel_forward_from_hub(GptOssMLP, "Yamoe")
+replace_kernel_forward_from_hub(GptOssRMSNorm, None)
+custom_mapping = {
+    "Yamoe": {
+        "cuda": {
+            Mode.INFERENCE: LayerRepository(
+                repo_id="drbh/yamoe",
+                layer_name="Yamoe",
+                revision="v0.3.0",
+            )
+        }
+    }
+}
+register_kernel_mapping(custom_mapping)
+
+
+model = GptOssForCausalLM.from_pretrained(
+    model_id,
+    dtype="bfloat16",
+    device_map="auto",
+    use_kernels=True,
+    quantization_config=quantization_config,
+).eval()
+
+messages = [
+    {"role": "system", "content": "What is Tensor Parallelism?"},
+]
+
+inputs = tokenizer.apply_chat_template(
+    messages,
+    add_generation_prompt=True,
+    return_tensors="pt",
+    return_dict=True,
+    reasoning_effort="low",
+).to("cuda")
+
+max_tokens = 256
+
+with torch.inference_mode():
+    start_time = time.perf_counter()
+    generated = model.generate(
+        **inputs,
+        max_new_tokens=max_tokens,
+        do_sample=False,
+        temperature=None,
+    )
+    end_time = time.perf_counter()
+
+print(tokenizer.decode(generated[0], skip_special_tokens=False))
+print(f"Generation took {end_time - start_time:.2f} seconds")

moe_benchmarks/megablocks_yamoe/megablocks_yamoe.html

@@ -3710,61 +3710,288 @@ span.linenos.special { color: #000000; background-color: #ffffc0; padding-left:
     <div class="system-info">
         <div class="system-info-header">Generated on:</div>
         <div class="system-info-content">
-            Linux x86_64 | Linux-6.11.0-1018-azure-x86_64-with-glibc2.39
+            Linux x86_64 | Linux-6.12.40-64.114.amzn2023.x86_64-x86_64-with-glibc2.36
         </div>
     </div>
     
     <div class="main-content">
-        <div class="cell cell-failed" id="cell-nv">
+        <h1>Comparison of Megablocks and Yamoe Kernels</h1>
+<p>This note compares the performance of the Megablocks and Yamoe kernels on the GPT-OSS-20B model.</p>
+<h2>Megablocks kernel</h2>
+<h2>Yamoe Kernel</h2>
+<div class="cell cell-failed" id="cell-setup">
 <div class="cell-header">
 <span class="collapse-indicators">
-<span onclick="toggleCode('nv')" style="cursor: pointer;">▼ code</span> 
-<span onclick="toggleOutput('nv')" style="cursor: pointer;">▼ output</span>
- <span id="uv-indicator-nv" style="cursor: default; opacity: 0.3;">▶ uv-logs</span>
+<span onclick="toggleCode('setup')" style="cursor: pointer;">▼ code</span> 
+<span onclick="toggleOutput('setup')" style="cursor: pointer;">▼ output</span>
+ <span id="uv-indicator-setup" style="cursor: default; opacity: 0.3;">▶ uv-logs</span>
 </span> | 
-Cell: nv | 0.07s | FAILED
- | <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>
+Cell: setup | 19.20s | FAILED
+ | <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>
 </div>
-<div id="code-nv" class="cell-code" data-lines="3">
+<div id="code-setup" class="cell-code" data-lines="116">
 <div class="highlight-with-lines">
-<div class="line-numbers" id="lines-nv">
-<a class="line-number" data-cell="nv" data-line="1" href="#cell-nv" onclick="event.preventDefault(); selectCellLine('nv', 1, true);">1</a>
-<a class="line-number" data-cell="nv" data-line="2" href="#cell-nv" onclick="event.preventDefault(); selectCellLine('nv', 2, true);">2</a>
-<a class="line-number" data-cell="nv" data-line="3" href="#cell-nv" onclick="event.preventDefault(); selectCellLine('nv', 3, true);">3</a>
+<div class="line-numbers" id="lines-setup">
+<a class="line-number" data-cell="setup" data-line="1" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 1, true);">1</a>
+<a class="line-number" data-cell="setup" data-line="2" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 2, true);">2</a>
+<a class="line-number" data-cell="setup" data-line="3" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 3, true);">3</a>
+<a class="line-number" data-cell="setup" data-line="4" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 4, true);">4</a>
+<a class="line-number" data-cell="setup" data-line="5" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 5, true);">5</a>
+<a class="line-number" data-cell="setup" data-line="6" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 6, true);">6</a>
+<a class="line-number" data-cell="setup" data-line="7" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 7, true);">7</a>
+<a class="line-number" data-cell="setup" data-line="8" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 8, true);">8</a>
+<a class="line-number" data-cell="setup" data-line="9" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 9, true);">9</a>
+<a class="line-number" data-cell="setup" data-line="10" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 10, true);">10</a>
+<a class="line-number" data-cell="setup" data-line="11" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 11, true);">11</a>
+<a class="line-number" data-cell="setup" data-line="12" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 12, true);">12</a>
+<a class="line-number" data-cell="setup" data-line="13" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 13, true);">13</a>
+<a class="line-number" data-cell="setup" data-line="14" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 14, true);">14</a>
+<a class="line-number" data-cell="setup" data-line="15" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 15, true);">15</a>
+<a class="line-number" data-cell="setup" data-line="16" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 16, true);">16</a>
+<a class="line-number" data-cell="setup" data-line="17" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 17, true);">17</a>
+<a class="line-number" data-cell="setup" data-line="18" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 18, true);">18</a>
+<a class="line-number" data-cell="setup" data-line="19" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 19, true);">19</a>
+<a class="line-number" data-cell="setup" data-line="20" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 20, true);">20</a>
+<a class="line-number" data-cell="setup" data-line="21" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 21, true);">21</a>
+<a class="line-number" data-cell="setup" data-line="22" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 22, true);">22</a>
+<a class="line-number" data-cell="setup" data-line="23" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 23, true);">23</a>
+<a class="line-number" data-cell="setup" data-line="24" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 24, true);">24</a>
+<a class="line-number" data-cell="setup" data-line="25" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 25, true);">25</a>
+<a class="line-number" data-cell="setup" data-line="26" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 26, true);">26</a>
+<a class="line-number" data-cell="setup" data-line="27" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 27, true);">27</a>
+<a class="line-number" data-cell="setup" data-line="28" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 28, true);">28</a>
+<a class="line-number" data-cell="setup" data-line="29" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 29, true);">29</a>
+<a class="line-number" data-cell="setup" data-line="30" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 30, true);">30</a>
+<a class="line-number" data-cell="setup" data-line="31" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 31, true);">31</a>
+<a class="line-number" data-cell="setup" data-line="32" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 32, true);">32</a>
+<a class="line-number" data-cell="setup" data-line="33" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 33, true);">33</a>
+<a class="line-number" data-cell="setup" data-line="34" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 34, true);">34</a>
+<a class="line-number" data-cell="setup" data-line="35" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 35, true);">35</a>
+<a class="line-number" data-cell="setup" data-line="36" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 36, true);">36</a>
+<a class="line-number" data-cell="setup" data-line="37" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 37, true);">37</a>
+<a class="line-number" data-cell="setup" data-line="38" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 38, true);">38</a>
+<a class="line-number" data-cell="setup" data-line="39" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 39, true);">39</a>
+<a class="line-number" data-cell="setup" data-line="40" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 40, true);">40</a>
+<a class="line-number" data-cell="setup" data-line="41" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 41, true);">41</a>
+<a class="line-number" data-cell="setup" data-line="42" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 42, true);">42</a>
+<a class="line-number" data-cell="setup" data-line="43" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 43, true);">43</a>
+<a class="line-number" data-cell="setup" data-line="44" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 44, true);">44</a>
+<a class="line-number" data-cell="setup" data-line="45" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 45, true);">45</a>
+<a class="line-number" data-cell="setup" data-line="46" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 46, true);">46</a>
+<a class="line-number" data-cell="setup" data-line="47" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 47, true);">47</a>
+<a class="line-number" data-cell="setup" data-line="48" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 48, true);">48</a>
+<a class="line-number" data-cell="setup" data-line="49" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 49, true);">49</a>
+<a class="line-number" data-cell="setup" data-line="50" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 50, true);">50</a>
+<a class="line-number" data-cell="setup" data-line="51" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 51, true);">51</a>
+<a class="line-number" data-cell="setup" data-line="52" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 52, true);">52</a>
+<a class="line-number" data-cell="setup" data-line="53" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 53, true);">53</a>
+<a class="line-number" data-cell="setup" data-line="54" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 54, true);">54</a>
+<a class="line-number" data-cell="setup" data-line="55" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 55, true);">55</a>
+<a class="line-number" data-cell="setup" data-line="56" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 56, true);">56</a>
+<a class="line-number" data-cell="setup" data-line="57" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 57, true);">57</a>
+<a class="line-number" data-cell="setup" data-line="58" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 58, true);">58</a>
+<a class="line-number" data-cell="setup" data-line="59" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 59, true);">59</a>
+<a class="line-number" data-cell="setup" data-line="60" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 60, true);">60</a>
+<a class="line-number" data-cell="setup" data-line="61" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 61, true);">61</a>
+<a class="line-number" data-cell="setup" data-line="62" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 62, true);">62</a>
+<a class="line-number" data-cell="setup" data-line="63" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 63, true);">63</a>
+<a class="line-number" data-cell="setup" data-line="64" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 64, true);">64</a>
+<a class="line-number" data-cell="setup" data-line="65" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 65, true);">65</a>
+<a class="line-number" data-cell="setup" data-line="66" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 66, true);">66</a>
+<a class="line-number" data-cell="setup" data-line="67" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 67, true);">67</a>
+<a class="line-number" data-cell="setup" data-line="68" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 68, true);">68</a>
+<a class="line-number" data-cell="setup" data-line="69" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 69, true);">69</a>
+<a class="line-number" data-cell="setup" data-line="70" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 70, true);">70</a>
+<a class="line-number" data-cell="setup" data-line="71" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 71, true);">71</a>
+<a class="line-number" data-cell="setup" data-line="72" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 72, true);">72</a>
+<a class="line-number" data-cell="setup" data-line="73" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 73, true);">73</a>
+<a class="line-number" data-cell="setup" data-line="74" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 74, true);">74</a>
+<a class="line-number" data-cell="setup" data-line="75" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 75, true);">75</a>
+<a class="line-number" data-cell="setup" data-line="76" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 76, true);">76</a>
+<a class="line-number" data-cell="setup" data-line="77" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 77, true);">77</a>
+<a class="line-number" data-cell="setup" data-line="78" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 78, true);">78</a>
+<a class="line-number" data-cell="setup" data-line="79" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 79, true);">79</a>
+<a class="line-number" data-cell="setup" data-line="80" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 80, true);">80</a>
+<a class="line-number" data-cell="setup" data-line="81" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 81, true);">81</a>
+<a class="line-number" data-cell="setup" data-line="82" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 82, true);">82</a>
+<a class="line-number" data-cell="setup" data-line="83" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 83, true);">83</a>
+<a class="line-number" data-cell="setup" data-line="84" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 84, true);">84</a>
+<a class="line-number" data-cell="setup" data-line="85" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 85, true);">85</a>
+<a class="line-number" data-cell="setup" data-line="86" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 86, true);">86</a>
+<a class="line-number" data-cell="setup" data-line="87" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 87, true);">87</a>
+<a class="line-number" data-cell="setup" data-line="88" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 88, true);">88</a>
+<a class="line-number" data-cell="setup" data-line="89" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 89, true);">89</a>
+<a class="line-number" data-cell="setup" data-line="90" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 90, true);">90</a>
+<a class="line-number" data-cell="setup" data-line="91" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 91, true);">91</a>
+<a class="line-number" data-cell="setup" data-line="92" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 92, true);">92</a>
+<a class="line-number" data-cell="setup" data-line="93" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 93, true);">93</a>
+<a class="line-number" data-cell="setup" data-line="94" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 94, true);">94</a>
+<a class="line-number" data-cell="setup" data-line="95" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 95, true);">95</a>
+<a class="line-number" data-cell="setup" data-line="96" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 96, true);">96</a>
+<a class="line-number" data-cell="setup" data-line="97" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 97, true);">97</a>
+<a class="line-number" data-cell="setup" data-line="98" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 98, true);">98</a>
+<a class="line-number" data-cell="setup" data-line="99" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 99, true);">99</a>
+<a class="line-number" data-cell="setup" data-line="100" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 100, true);">100</a>
+<a class="line-number" data-cell="setup" data-line="101" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 101, true);">101</a>
+<a class="line-number" data-cell="setup" data-line="102" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 102, true);">102</a>
+<a class="line-number" data-cell="setup" data-line="103" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 103, true);">103</a>
+<a class="line-number" data-cell="setup" data-line="104" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 104, true);">104</a>
+<a class="line-number" data-cell="setup" data-line="105" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 105, true);">105</a>
+<a class="line-number" data-cell="setup" data-line="106" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 106, true);">106</a>
+<a class="line-number" data-cell="setup" data-line="107" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 107, true);">107</a>
+<a class="line-number" data-cell="setup" data-line="108" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 108, true);">108</a>
+<a class="line-number" data-cell="setup" data-line="109" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 109, true);">109</a>
+<a class="line-number" data-cell="setup" data-line="110" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 110, true);">110</a>
+<a class="line-number" data-cell="setup" data-line="111" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 111, true);">111</a>
+<a class="line-number" data-cell="setup" data-line="112" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 112, true);">112</a>
+<a class="line-number" data-cell="setup" data-line="113" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 113, true);">113</a>
+<a class="line-number" data-cell="setup" data-line="114" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 114, true);">114</a>
+<a class="line-number" data-cell="setup" data-line="115" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 115, true);">115</a>
+<a class="line-number" data-cell="setup" data-line="116" href="#cell-setup" onclick="event.preventDefault(); selectCellLine('setup', 116, true);">116</a>
 </div>
 <div class="code-wrap">
-<div class="highlight"><pre><span></span><span class="kn">import</span><span class="w"> </span><span class="nn">subprocess</span>
-
-<span class="nb">print</span><span class="p">(</span><span class="n">subprocess</span><span class="o">.</span><span class="n">run</span><span class="p">([</span><span class="s2">&quot;nvidia-smi&quot;</span><span class="p">],</span> <span class="n">capture_output</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">text</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span><span class="o">.</span><span class="n">stdout</span><span class="p">)</span>
+<div class="highlight"><pre><span></span><span class="c1"># /// script</span>
+<span class="c1"># requires-python = &quot;&gt;=3.12&quot;</span>
+<span class="c1"># dependencies = [</span>
+<span class="c1">#     &quot;accelerate&gt;=1.10.1&quot;,</span>
+<span class="c1">#     &quot;torch&gt;=2.7.0&quot;,</span>
+<span class="c1">#     &quot;kernels==0.10.0&quot;,</span>
+<span class="c1">#     &quot;transformers@https://github.com/huggingface/transformers.git&quot;,</span>
+<span class="c1">#     &quot;ipdb&gt;=0.13.13&quot;,</span>
+<span class="c1">#     &quot;matplotlib&gt;=3.7.2&quot;,</span>
+<span class="c1">#     &quot;numpy&gt;=1.24.3&quot;,</span>
+<span class="c1"># ]</span>
+<span class="c1"># ///</span>
+
+<span class="kn">import</span><span class="w"> </span><span class="nn">torch</span>
+<span class="kn">from</span><span class="w"> </span><span class="nn">transformers</span><span class="w"> </span><span class="kn">import</span> <span class="n">GptOssForCausalLM</span><span class="p">,</span> <span class="n">PreTrainedTokenizerFast</span><span class="p">,</span> <span class="n">Mxfp4Config</span>
+<span class="kn">import</span><span class="w"> </span><span class="nn">time</span>
+<span class="kn">import</span><span class="w"> </span><span class="nn">torch.nn</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">nn</span>
+<span class="kn">from</span><span class="w"> </span><span class="nn">kernels</span><span class="w"> </span><span class="kn">import</span> <span class="n">register_kernel_mapping</span><span class="p">,</span> <span class="n">Mode</span><span class="p">,</span> <span class="n">LayerRepository</span>
+<span class="kn">import</span><span class="w"> </span><span class="nn">sys</span>
+<span class="kn">import</span><span class="w"> </span><span class="nn">torch.profiler</span>
+<span class="kn">import</span><span class="w"> </span><span class="nn">gc</span>
+<span class="kn">import</span><span class="w"> </span><span class="nn">logging</span>
+
+<span class="c1"># set to debug logging</span>
+<span class="n">logging</span><span class="o">.</span><span class="n">basicConfig</span><span class="p">(</span><span class="n">level</span><span class="o">=</span><span class="n">logging</span><span class="o">.</span><span class="n">INFO</span><span class="p">)</span>
+
+<span class="k">def</span><span class="w"> </span><span class="nf">reset_peak_memory_stats</span><span class="p">():</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;Clear CUDA cache and reset memory allocation counters.&quot;&quot;&quot;</span>
+    <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">empty_cache</span><span class="p">()</span>
+    <span class="k">if</span> <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">is_available</span><span class="p">():</span>
+        <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">reset_peak_memory_stats</span><span class="p">()</span>
+    <span class="n">gc</span><span class="o">.</span><span class="n">collect</span><span class="p">()</span>
+
+<span class="k">def</span><span class="w"> </span><span class="nf">get_memory_stats</span><span class="p">():</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;Get current and peak CUDA memory usage.&quot;&quot;&quot;</span>
+    <span class="k">if</span> <span class="ow">not</span> <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">is_available</span><span class="p">():</span>
+        <span class="k">return</span> <span class="p">{</span><span class="s2">&quot;allocated_gb&quot;</span><span class="p">:</span> <span class="mi">0</span><span class="p">,</span> <span class="s2">&quot;peak_gb&quot;</span><span class="p">:</span> <span class="mi">0</span><span class="p">,</span> <span class="s2">&quot;reserved_gb&quot;</span><span class="p">:</span> <span class="mi">0</span><span class="p">}</span>
+    <span class="k">return</span> <span class="p">{</span>
+        <span class="s2">&quot;allocated_gb&quot;</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">memory_allocated</span><span class="p">()</span> <span class="o">/</span> <span class="mf">1e9</span><span class="p">,</span>
+        <span class="s2">&quot;peak_gb&quot;</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">max_memory_allocated</span><span class="p">()</span> <span class="o">/</span> <span class="mf">1e9</span><span class="p">,</span>
+        <span class="s2">&quot;reserved_gb&quot;</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">memory_reserved</span><span class="p">()</span> <span class="o">/</span> <span class="mf">1e9</span><span class="p">,</span>
+    <span class="p">}</span>
+
+<span class="k">def</span><span class="w"> </span><span class="nf">override_kernel_layer_name</span><span class="p">(</span><span class="n">cls_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">bool</span><span class="p">:</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;Helper to dynamically override the kernel_layer_name in a model class.&quot;&quot;&quot;</span>
+    <span class="k">for</span> <span class="n">mod</span> <span class="ow">in</span> <span class="n">sys</span><span class="o">.</span><span class="n">modules</span><span class="o">.</span><span class="n">values</span><span class="p">():</span>
+        <span class="k">if</span> <span class="n">mod</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
+            <span class="k">continue</span>
+        <span class="n">obj</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">mod</span><span class="p">,</span> <span class="n">cls_name</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
+        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">obj</span><span class="p">,</span> <span class="nb">type</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">issubclass</span><span class="p">(</span><span class="n">obj</span><span class="p">,</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
+            <span class="nb">setattr</span><span class="p">(</span><span class="n">obj</span><span class="p">,</span> <span class="s2">&quot;kernel_layer_name&quot;</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span>
+            <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Overrode </span><span class="si">{</span><span class="n">cls_name</span><span class="si">}</span><span class="s2">.kernel_layer_name to </span><span class="si">{</span><span class="n">value</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
+            <span class="k">return</span> <span class="kc">True</span>
+    <span class="k">return</span> <span class="kc">False</span>
+
+
+<span class="c1"># Init the model the normal way</span>
+<span class="n">model_id</span> <span class="o">=</span> <span class="s2">&quot;openai/gpt-oss-20b&quot;</span>
+<span class="n">tokenizer</span> <span class="o">=</span> <span class="n">PreTrainedTokenizerFast</span><span class="o">.</span><span class="n">from_pretrained</span><span class="p">(</span><span class="n">model_id</span><span class="p">)</span>
+<span class="n">quantization_config</span> <span class="o">=</span> <span class="n">Mxfp4Config</span><span class="p">(</span><span class="n">dequantize</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
+
+
+<span class="kn">from</span><span class="w"> </span><span class="nn">kernels</span><span class="w"> </span><span class="kn">import</span> <span class="n">replace_kernel_forward_from_hub</span><span class="p">,</span> <span class="n">register_kernel_mapping</span><span class="p">,</span> <span class="n">LayerRepository</span><span class="p">,</span> <span class="n">Mode</span>
+
+<span class="kn">from</span><span class="w"> </span><span class="nn">transformers.models.gpt_oss.modeling_gpt_oss</span><span class="w"> </span><span class="kn">import</span> <span class="n">GptOssMLP</span><span class="p">,</span> <span class="n">GptOssRMSNorm</span>
+
+<span class="n">replace_kernel_forward_from_hub</span><span class="p">(</span><span class="n">GptOssMLP</span><span class="p">,</span> <span class="s2">&quot;Yamoe&quot;</span><span class="p">)</span>
+<span class="n">replace_kernel_forward_from_hub</span><span class="p">(</span><span class="n">GptOssRMSNorm</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
+<span class="n">custom_mapping</span> <span class="o">=</span> <span class="p">{</span>
+    <span class="s2">&quot;Yamoe&quot;</span><span class="p">:</span> <span class="p">{</span>
+        <span class="s2">&quot;cuda&quot;</span><span class="p">:</span> <span class="p">{</span>
+            <span class="n">Mode</span><span class="o">.</span><span class="n">INFERENCE</span><span class="p">:</span> <span class="n">LayerRepository</span><span class="p">(</span>
+                <span class="n">repo_id</span><span class="o">=</span><span class="s2">&quot;drbh/yamoe&quot;</span><span class="p">,</span>
+                <span class="n">layer_name</span><span class="o">=</span><span class="s2">&quot;Yamoe&quot;</span><span class="p">,</span>
+                <span class="n">revision</span><span class="o">=</span><span class="s2">&quot;v0.3.0&quot;</span><span class="p">,</span>
+            <span class="p">)</span>
+        <span class="p">}</span>
+    <span class="p">}</span>
+<span class="p">}</span>
+<span class="n">register_kernel_mapping</span><span class="p">(</span><span class="n">custom_mapping</span><span class="p">)</span>
+
+
+<span class="n">model</span> <span class="o">=</span> <span class="n">GptOssForCausalLM</span><span class="o">.</span><span class="n">from_pretrained</span><span class="p">(</span>
+    <span class="n">model_id</span><span class="p">,</span>
+    <span class="n">dtype</span><span class="o">=</span><span class="s2">&quot;bfloat16&quot;</span><span class="p">,</span>
+    <span class="n">device_map</span><span class="o">=</span><span class="s2">&quot;auto&quot;</span><span class="p">,</span>
+    <span class="n">use_kernels</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
+    <span class="n">quantization_config</span><span class="o">=</span><span class="n">quantization_config</span><span class="p">,</span>
+<span class="p">)</span><span class="o">.</span><span class="n">eval</span><span class="p">()</span>
+
+<span class="n">messages</span> <span class="o">=</span> <span class="p">[</span>
+    <span class="p">{</span><span class="s2">&quot;role&quot;</span><span class="p">:</span> <span class="s2">&quot;system&quot;</span><span class="p">,</span> <span class="s2">&quot;content&quot;</span><span class="p">:</span> <span class="s2">&quot;What is Tensor Parallelism?&quot;</span><span class="p">},</span>
+<span class="p">]</span>
+
+<span class="n">inputs</span> <span class="o">=</span> <span class="n">tokenizer</span><span class="o">.</span><span class="n">apply_chat_template</span><span class="p">(</span>
+    <span class="n">messages</span><span class="p">,</span>
+    <span class="n">add_generation_prompt</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
+    <span class="n">return_tensors</span><span class="o">=</span><span class="s2">&quot;pt&quot;</span><span class="p">,</span>
+    <span class="n">return_dict</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
+    <span class="n">reasoning_effort</span><span class="o">=</span><span class="s2">&quot;low&quot;</span><span class="p">,</span>
+<span class="p">)</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="s2">&quot;cuda&quot;</span><span class="p">)</span>
+
+<span class="n">max_tokens</span> <span class="o">=</span> <span class="mi">256</span>
+
+<span class="k">with</span> <span class="n">torch</span><span class="o">.</span><span class="n">inference_mode</span><span class="p">():</span>
+    <span class="n">start_time</span> <span class="o">=</span> <span class="n">time</span><span class="o">.</span><span class="n">perf_counter</span><span class="p">()</span>
+    <span class="n">generated</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">generate</span><span class="p">(</span>
+        <span class="o">**</span><span class="n">inputs</span><span class="p">,</span>
+        <span class="n">max_new_tokens</span><span class="o">=</span><span class="n">max_tokens</span><span class="p">,</span>
+        <span class="n">do_sample</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
+        <span class="n">temperature</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
+    <span class="p">)</span>
+    <span class="n">end_time</span> <span class="o">=</span> <span class="n">time</span><span class="o">.</span><span class="n">perf_counter</span><span class="p">()</span>
+
+<span class="nb">print</span><span class="p">(</span><span class="n">tokenizer</span><span class="o">.</span><span class="n">decode</span><span class="p">(</span><span class="n">generated</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">skip_special_tokens</span><span class="o">=</span><span class="kc">False</span><span class="p">))</span>
+<span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Generation took </span><span class="si">{</span><span class="n">end_time</span><span class="w"> </span><span class="o">-</span><span class="w"> </span><span class="n">start_time</span><span class="si">:</span><span class="s2">.2f</span><span class="si">}</span><span class="s2"> seconds&quot;</span><span class="p">)</span>
 </pre></div>
 
-<div class="code-line-highlight" id="line-highlight-nv"></div>
+<div class="code-line-highlight" id="line-highlight-setup"></div>
 </div>
 </div>
 </div>
-<div id="output-nv" class="cell-output">
-<div class="cell-stderr">Traceback (most recent call last):
-  File &quot;/home/runner/work/kernels-uvnotes/kernels-uvnotes/moe_benchmarks/megablocks_yamoe/.uvnote/cells/nv.py&quot;, line 3, in &lt;module&gt;
-    print(subprocess.run([&quot;nvidia-smi&quot;], capture_output=True, text=True).stdout)
-          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/opt/hostedtoolcache/Python/3.11.13/x64/lib/python3.11/subprocess.py&quot;, line 548, in run
-    with Popen(*popenargs, **kwargs) as process:
-         ^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/opt/hostedtoolcache/Python/3.11.13/x64/lib/python3.11/subprocess.py&quot;, line 1026, in __init__
-    self._execute_child(args, executable, preexec_fn, close_fds,
-  File &quot;/opt/hostedtoolcache/Python/3.11.13/x64/lib/python3.11/subprocess.py&quot;, line 1955, in _execute_child
-    raise child_exception_type(errno_num, err_msg, err_filename)
-FileNotFoundError: [Errno 2] No such file or directory: &#x27;nvidia-smi&#x27;
+<div id="output-setup" class="cell-output">
+<div class="cell-stderr">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 (449533af73874470e914a203391635e04ac2ffc8)
+  × No solution found when resolving script dependencies:
+  ╰─▶ Because only transformers==4.57.0.dev0 is available and
+      transformers==4.57.0.dev0 depends on huggingface-hub==1.0.0rc1,
+      we can conclude that all versions of transformers depend on
+      huggingface-hub==1.0.0rc1.
+      And because kernels==0.10.0 depends on huggingface-hub&gt;=0.26.0,&lt;1.0,
+      we can conclude that kernels==0.10.0 and all versions of transformers
+      are incompatible.
+      And because you require kernels==0.10.0 and transformers, we can
+      conclude that your requirements are unsatisfiable.
 </div>
 </div>
 </div>
-
-<h1>Comparison of Megablocks and Yamoe Kernels</h1>
-<p>This note compares the performance of the Megablocks and Yamoe kernels on the GPT-OSS-20B model.</p>
-<h2>Megablocks kernel</h2>
-<h2>Yamoe Kernel</h2>
     </div>
     
 </body>