Add ROCm build debugging utilities

.gitignore

@@ -5,3 +5,5 @@ megablocks-moe/.bak
 .pytest_cache
 .readme_example.py.swp
 .torch_extensions/
+.torch_extensions_debug/
+strace.log

_dev/debug-build-1-env.sh

@@ -0,0 +1,79 @@
+#!/usr/bin/env bash
+
+# Debug script 1: Basic ROCm environment and tool availability check
+
+set -euo pipefail
+
+echo "=== ROCm Environment Debug Script 1 ==="
+echo "Testing basic ROCm/HIP environment setup"
+echo
+
+# Set ROCm environment variables
+export ROCM_PATH="${ROCM_PATH:-/opt/rocm-7.0.1}"
+export ROCM_HOME="${ROCM_HOME:-$ROCM_PATH}"
+export HIP_PATH="${HIP_PATH:-$ROCM_PATH}"
+export HIP_HOME="${HIP_HOME:-$ROCM_PATH}"
+export PATH="$ROCM_HOME/bin:$PATH"
+export TORCH_HIP_ARCH_LIST="${TORCH_HIP_ARCH_LIST:-gfx942}"
+export HSA_OVERRIDE_GFX_VERSION="${HSA_OVERRIDE_GFX_VERSION:-gfx942}"
+
+echo "Environment Variables:"
+echo "ROCM_PATH=$ROCM_PATH"
+echo "ROCM_HOME=$ROCM_HOME"
+echo "HIP_PATH=$HIP_PATH"
+echo "HIP_HOME=$HIP_HOME"
+echo "TORCH_HIP_ARCH_LIST=$TORCH_HIP_ARCH_LIST"
+echo "HSA_OVERRIDE_GFX_VERSION=$HSA_OVERRIDE_GFX_VERSION"
+echo "PATH (ROCm portion): $(echo $PATH | tr ':' '\n' | grep rocm || echo 'No ROCm in PATH')"
+echo
+
+echo "=== Directory Checks ==="
+echo "ROCm installation directory exists: $(test -d "$ROCM_PATH" && echo 'YES' || echo 'NO')"
+echo "ROCm bin directory exists: $(test -d "$ROCM_PATH/bin" && echo 'YES' || echo 'NO')"
+echo "ROCm include directory exists: $(test -d "$ROCM_PATH/include" && echo 'YES' || echo 'NO')"
+echo "ROCm lib directory exists: $(test -d "$ROCM_PATH/lib" && echo 'YES' || echo 'NO')"
+echo
+
+echo "=== Tool Availability ==="
+echo "hipcc available: $(which hipcc >/dev/null 2>&1 && echo 'YES' || echo 'NO')"
+echo "hip-clang available: $(which hip-clang >/dev/null 2>&1 && echo 'YES' || echo 'NO')"
+echo "rocm-smi available: $(which rocm-smi >/dev/null 2>&1 && echo 'YES' || echo 'NO')"
+echo "hipconfig available: $(which hipconfig >/dev/null 2>&1 && echo 'YES' || echo 'NO')"
+echo
+
+echo "=== Tool Versions ==="
+if which hipcc >/dev/null 2>&1; then
+    echo "hipcc version:"
+    hipcc --version || echo "Failed to get hipcc version"
+    echo
+fi
+
+if which hipconfig >/dev/null 2>&1; then
+    echo "HIP config:"
+    hipconfig --full || echo "Failed to get hipconfig"
+    echo
+fi
+
+if which rocm-smi >/dev/null 2>&1; then
+    echo "ROCm SMI:"
+    rocm-smi --showproductname || echo "Failed to get ROCm SMI info"
+    echo
+fi
+
+echo "=== Python Environment ==="
+python3 --version || echo "Python3 not available"
+python3 -c "import torch; print(f'PyTorch version: {torch.__version__}')" || echo "PyTorch not available"
+python3 -c "import torch; print(f'CUDA available: {torch.cuda.is_available()}')" || echo "Failed to check CUDA availability"
+python3 -c "import torch; print(f'HIP available: {hasattr(torch.version, \"hip\") and torch.version.hip is not None}')" || echo "Failed to check HIP availability"
+
+echo
+echo "=== Basic HIP Device Check ==="
+if which hipinfo >/dev/null 2>&1; then
+    echo "HIP devices:"
+    hipinfo || echo "hipinfo failed"
+else
+    echo "hipinfo not available"
+fi
+
+echo
+echo "=== Debug Script 1 Complete ==="

_dev/debug-build-2-hipcc.sh

@@ -0,0 +1,154 @@
+#!/usr/bin/env bash
+
+# Debug script 2: HIP compiler basic compilation test
+
+set -euo pipefail
+
+echo "=== HIP Compiler Debug Script 2 ==="
+echo "Testing basic HIP compilation capabilities"
+echo
+
+# Set ROCm environment variables
+export ROCM_PATH="${ROCM_PATH:-/opt/rocm-7.0.1}"
+export ROCM_HOME="${ROCM_HOME:-$ROCM_PATH}"
+export HIP_PATH="${HIP_PATH:-$ROCM_PATH}"
+export HIP_HOME="${HIP_HOME:-$ROCM_PATH}"
+export PATH="$ROCM_HOME/bin:$PATH"
+export TORCH_HIP_ARCH_LIST="${TORCH_HIP_ARCH_LIST:-gfx942}"
+export HSA_OVERRIDE_GFX_VERSION="${HSA_OVERRIDE_GFX_VERSION:-gfx942}"
+
+# Create a simple test directory
+mkdir -p /tmp/hip_test
+cd /tmp/hip_test
+
+echo "=== Creating Simple HIP Test Program ==="
+
+# Create a minimal HIP program
+cat > test_simple.hip << 'EOF'
+#include <hip/hip_runtime.h>
+#include <iostream>
+
+__global__ void simple_kernel(float* data, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) {
+        data[idx] = idx * 2.0f;
+    }
+}
+
+int main() {
+    const int n = 1024;
+    const size_t size = n * sizeof(float);
+
+    float* h_data = new float[n];
+    float* d_data;
+
+    // Allocate device memory
+    hipError_t err = hipMalloc(&d_data, size);
+    if (err != hipSuccess) {
+        std::cout << "Failed to allocate device memory: " << hipGetErrorString(err) << std::endl;
+        return 1;
+    }
+
+    // Launch kernel
+    dim3 block(256);
+    dim3 grid((n + block.x - 1) / block.x);
+    hipLaunchKernelGGL(simple_kernel, grid, block, 0, 0, d_data, n);
+
+    // Copy back
+    err = hipMemcpy(h_data, d_data, size, hipMemcpyDeviceToHost);
+    if (err != hipSuccess) {
+        std::cout << "Failed to copy from device: " << hipGetErrorString(err) << std::endl;
+        return 1;
+    }
+
+    // Check results
+    bool success = true;
+    for (int i = 0; i < 10; i++) {
+        if (h_data[i] != i * 2.0f) {
+            success = false;
+            break;
+        }
+    }
+
+    std::cout << "Simple HIP test: " << (success ? "PASSED" : "FAILED") << std::endl;
+
+    hipFree(d_data);
+    delete[] h_data;
+
+    return success ? 0 : 1;
+}
+EOF
+
+echo "=== Testing HIP Compilation ==="
+echo "Compiling simple HIP program..."
+
+if hipcc -o test_simple test_simple.hip --amdgpu-target=gfx942 -v; then
+    echo "✓ HIP compilation successful"
+    echo
+    echo "=== Testing HIP Execution ==="
+    if ./test_simple; then
+        echo "✓ HIP execution successful"
+    else
+        echo "✗ HIP execution failed"
+    fi
+else
+    echo "✗ HIP compilation failed"
+fi
+
+echo
+echo "=== Testing HIP with C++ Extensions Flags ==="
+
+# Test with flags similar to what torch uses
+echo "Compiling with torch-like flags..."
+if hipcc -o test_simple_torch test_simple.hip \
+    --amdgpu-target=gfx942 \
+    -O3 \
+    -std=c++17 \
+    -fPIC \
+    -shared \
+    -v; then
+    echo "✓ HIP compilation with torch flags successful"
+else
+    echo "✗ HIP compilation with torch flags failed"
+fi
+
+echo
+echo "=== Testing hipblaslt Library ==="
+echo "Checking if hipblaslt is available for linking..."
+
+cat > test_hipblaslt.cpp << 'EOF'
+#include <iostream>
+
+// Just test if we can link against hipblaslt
+int main() {
+    std::cout << "hipblaslt linkage test" << std::endl;
+    return 0;
+}
+EOF
+
+if hipcc -o test_hipblaslt test_hipblaslt.cpp -lhipblaslt -v; then
+    echo "✓ hipblaslt linkage successful"
+    if ./test_hipblaslt; then
+        echo "✓ hipblaslt execution successful"
+    else
+        echo "✗ hipblaslt execution failed"
+    fi
+else
+    echo "✗ hipblaslt linkage failed"
+fi
+
+echo
+echo "=== Library Search Paths ==="
+echo "ROCm lib directory contents:"
+ls -la "$ROCM_PATH/lib" | head -20 || echo "Could not list ROCm lib directory"
+
+echo
+echo "hipblaslt library files:"
+find "$ROCM_PATH" -name "*hipblaslt*" 2>/dev/null || echo "No hipblaslt files found"
+
+echo
+echo "=== Debug Script 2 Complete ==="
+
+# Cleanup
+cd /
+rm -rf /tmp/hip_test

_dev/debug-build-3-torch-ext.sh

@@ -0,0 +1,187 @@
+#!/usr/bin/env bash
+
+# Debug script 3: PyTorch C++ extension compilation test
+
+set -euo pipefail
+
+echo "=== PyTorch C++ Extension Debug Script 3 ==="
+echo "Testing PyTorch C++ extension compilation with HIP"
+echo
+
+# Set ROCm environment variables
+export ROCM_PATH="${ROCM_PATH:-/opt/rocm-7.0.1}"
+export ROCM_HOME="${ROCM_HOME:-$ROCM_PATH}"
+export HIP_PATH="${HIP_PATH:-$ROCM_PATH}"
+export HIP_HOME="${HIP_HOME:-$ROCM_PATH}"
+export PATH="$ROCM_HOME/bin:$PATH"
+export TORCH_HIP_ARCH_LIST="${TORCH_HIP_ARCH_LIST:-gfx942}"
+export HSA_OVERRIDE_GFX_VERSION="${HSA_OVERRIDE_GFX_VERSION:-gfx942}"
+export TORCH_EXTENSIONS_DIR="${TORCH_EXTENSIONS_DIR:-$PWD/.torch_extensions_debug}"
+
+# Create a test directory
+mkdir -p /tmp/torch_ext_test
+cd /tmp/torch_ext_test
+
+echo "=== Creating Simple PyTorch Extension ==="
+
+# Create a minimal CUDA/HIP kernel similar to megablocks
+cat > simple_kernel.cu << 'EOF'
+#include <torch/extension.h>
+#include <vector>
+
+#ifdef __HIP_PLATFORM_AMD__
+#include <hip/hip_runtime.h>
+#define CUDA_LAUNCH_KERNEL(kernel, grid, block, smem, stream, ...) \
+    hipLaunchKernelGGL(kernel, grid, block, smem, stream, __VA_ARGS__)
+#else
+#include <cuda_runtime.h>
+#define CUDA_LAUNCH_KERNEL(kernel, grid, block, smem, stream, ...) \
+    kernel<<<grid, block, smem, stream>>>(__VA_ARGS__)
+#endif
+
+__global__ void add_kernel(const float* a, const float* b, float* c, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) {
+        c[idx] = a[idx] + b[idx];
+    }
+}
+
+torch::Tensor add_tensors_cuda(torch::Tensor a, torch::Tensor b) {
+    auto c = torch::zeros_like(a);
+    int n = a.numel();
+
+    const int block_size = 256;
+    const int grid_size = (n + block_size - 1) / block_size;
+
+    CUDA_LAUNCH_KERNEL(
+        add_kernel,
+        dim3(grid_size),
+        dim3(block_size),
+        0,
+        0,
+        a.data_ptr<float>(),
+        b.data_ptr<float>(),
+        c.data_ptr<float>(),
+        n
+    );
+
+    return c;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("add_tensors", &add_tensors_cuda, "Add two tensors (CUDA/HIP)");
+}
+EOF
+
+# Create Python test script
+cat > test_extension.py << 'EOF'
+import os
+import sys
+import torch
+from torch.utils.cpp_extension import load
+
+print("=== PyTorch Extension Load Test ===")
+print(f"PyTorch version: {torch.__version__}")
+print(f"CUDA available: {torch.cuda.is_available()}")
+print(f"Device count: {torch.cuda.device_count()}")
+
+if hasattr(torch.version, 'hip') and torch.version.hip:
+    print(f"HIP version: {torch.version.hip}")
+
+print("\n=== Loading Extension ===")
+print("This may take a while and will show compilation output...")
+print("If this hangs, it indicates the same issue as build.py")
+
+try:
+    # Mimic the same load call as build.py
+    simple_ext = load(
+        name="simple_test_ext",
+        sources=["simple_kernel.cu"],
+        extra_cflags=["-O3", "-std=c++17"],
+        extra_cuda_cflags=["-O3"],  # torch switches this to hipcc on ROCm
+        verbose=True,
+        is_python_module=False
+    )
+    print("✓ Extension compilation successful!")
+
+    # Test the extension
+    print("\n=== Testing Extension ===")
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    a = torch.randn(1000, device=device)
+    b = torch.randn(1000, device=device)
+
+    if device == 'cuda':
+        result = simple_ext.add_tensors(a, b)
+        expected = a + b
+        if torch.allclose(result, expected):
+            print("✓ Extension execution successful!")
+        else:
+            print("✗ Extension execution failed - results don't match")
+    else:
+        print("⚠ No CUDA device, skipping execution test")
+
+except Exception as e:
+    print(f"✗ Extension compilation/loading failed: {e}")
+    import traceback
+    traceback.print_exc()
+EOF
+
+echo "=== Running PyTorch Extension Test ==="
+echo "This test mimics the same compilation process as build.py"
+echo "If this hangs, it shows the same issue as the main build"
+echo
+
+# Set a timeout to prevent infinite hang
+timeout 300 python3 test_extension.py || {
+    exit_code=$?
+    if [ $exit_code -eq 124 ]; then
+        echo "✗ Extension compilation timed out after 5 minutes (same as build.py hang)"
+    else
+        echo "✗ Extension compilation failed with exit code $exit_code"
+    fi
+}
+
+echo
+echo "=== Testing with Minimal Sources ==="
+
+# Create an even simpler version
+cat > minimal_kernel.cu << 'EOF'
+#include <torch/extension.h>
+
+torch::Tensor dummy_function(torch::Tensor input) {
+    return input.clone();
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("dummy", &dummy_function, "Dummy function");
+}
+EOF
+
+cat > test_minimal.py << 'EOF'
+import torch
+from torch.utils.cpp_extension import load
+
+print("=== Minimal Extension Test ===")
+
+try:
+    minimal_ext = load(
+        name="minimal_test_ext",
+        sources=["minimal_kernel.cu"],
+        extra_cflags=["-O3"],
+        verbose=True,
+        with_cuda=False  # Skip CUDA/HIP compilation
+    )
+    print("✓ Minimal extension (CPU only) successful!")
+except Exception as e:
+    print(f"✗ Even minimal extension failed: {e}")
+EOF
+
+echo "Testing minimal CPU-only extension..."
+timeout 120 python3 test_minimal.py || echo "Minimal extension also failed/timed out"
+
+echo
+echo "=== Debug Script 3 Complete ==="
+
+# Cleanup
+cd /
+rm -rf /tmp/torch_ext_test

_dev/debug-build-4-megablocks.sh

@@ -0,0 +1,257 @@
+#!/usr/bin/env bash
+
+# Debug script 4: MegaBlocks-specific build debugging
+
+set -euo pipefail
+
+echo "=== MegaBlocks Build Debug Script 4 ==="
+echo "Testing MegaBlocks-specific compilation components"
+echo
+
+# Set ROCm environment variables
+export ROCM_PATH="${ROCM_PATH:-/opt/rocm-7.0.1}"
+export ROCM_HOME="${ROCM_HOME:-$ROCM_PATH}"
+export HIP_PATH="${HIP_PATH:-$ROCM_PATH}"
+export HIP_HOME="${HIP_HOME:-$ROCM_PATH}"
+export PATH="$ROCM_HOME/bin:$PATH"
+export TORCH_HIP_ARCH_LIST="${TORCH_HIP_ARCH_LIST:-gfx942}"
+export HSA_OVERRIDE_GFX_VERSION="${HSA_OVERRIDE_GFX_VERSION:-gfx942}"
+export TORCH_EXTENSIONS_DIR="${TORCH_EXTENSIONS_DIR:-$PWD/.torch_extensions_debug}"
+
+echo "Working directory: $(pwd)"
+echo
+
+echo "=== Checking MegaBlocks Source Files ==="
+echo "Verifying all source files exist:"
+
+sources=(
+    "torch-ext/torch_binding.cpp"
+    "csrc/new_cumsum.cu"
+    "csrc/new_histogram.cu"
+    "csrc/new_indices.cu"
+    "csrc/new_replicate.cu"
+    "csrc/new_sort.cu"
+    "csrc/grouped_gemm/grouped_gemm.cu"
+)
+
+all_exist=true
+for src in "${sources[@]}"; do
+    if [ -f "$src" ]; then
+        echo "✓ $src exists ($(wc -l < "$src") lines)"
+    else
+        echo "✗ $src missing"
+        all_exist=false
+    fi
+done
+
+if [ "$all_exist" = false ]; then
+    echo "Cannot proceed - missing source files"
+    exit 1
+fi
+
+echo
+echo "=== Checking Include Directories ==="
+if [ -d "csrc" ]; then
+    echo "✓ csrc include directory exists"
+    echo "Headers in csrc/:"
+    find csrc -name "*.h" -o -name "*.hpp" | head -10
+else
+    echo "✗ csrc include directory missing"
+fi
+
+echo
+echo "=== Testing Individual Source Compilation ==="
+
+# Test compiling each .cu file individually
+for src in csrc/*.cu; do
+    if [ -f "$src" ]; then
+        echo "Testing compilation of $(basename "$src")..."
+        if timeout 60 hipcc -c "$src" -o "/tmp/$(basename "$src" .cu).o" \
+            --amdgpu-target=gfx942 \
+            -I./csrc \
+            -I"$(python3 -c 'import torch; print(torch.utils.cpp_extension.include_paths()[0])')" \
+            -std=c++17 \
+            -O3 \
+            -fPIC; then
+            echo "✓ $(basename "$src") compiled successfully"
+        else
+            echo "✗ $(basename "$src") compilation failed"
+        fi
+    fi
+done
+
+echo
+echo "=== Testing grouped_gemm.cu Specifically ==="
+echo "This is often the most complex kernel..."
+
+if timeout 120 hipcc -c csrc/grouped_gemm/grouped_gemm.cu -o /tmp/grouped_gemm.o \
+    --amdgpu-target=gfx942 \
+    -I./csrc \
+    -I"$(python3 -c 'import torch; print(torch.utils.cpp_extension.include_paths()[0])')" \
+    -std=c++17 \
+    -O3 \
+    -fPIC \
+    -lhipblaslt \
+    -v; then
+    echo "✓ grouped_gemm.cu compiled successfully"
+else
+    echo "✗ grouped_gemm.cu compilation failed"
+fi
+
+echo
+echo "=== Testing torch_binding.cpp ==="
+if timeout 60 hipcc -c torch-ext/torch_binding.cpp -o /tmp/torch_binding.o \
+    -I./csrc \
+    -I"$(python3 -c 'import torch; print(torch.utils.cpp_extension.include_paths()[0])')" \
+    -std=c++17 \
+    -O3 \
+    -fPIC; then
+    echo "✓ torch_binding.cpp compiled successfully"
+else
+    echo "✗ torch_binding.cpp compilation failed"
+fi
+
+echo
+echo "=== Testing Incremental PyTorch Extension Build ==="
+
+cat > debug_build.py << 'EOF'
+import os
+import pathlib
+import sys
+import signal
+import time
+from torch.utils.cpp_extension import load
+
+def timeout_handler(signum, frame):
+    print("Build timed out - this indicates a hanging issue")
+    sys.exit(1)
+
+# Set up timeout
+signal.signal(signal.SIGALRM, timeout_handler)
+signal.alarm(180)  # 3 minute timeout
+
+repo = pathlib.Path(".").resolve()
+os.environ.setdefault("TORCH_EXTENSIONS_DIR", str(repo / ".torch_extensions_debug"))
+
+print("=== Testing with Single Source File ===")
+try:
+    print("Building with just new_cumsum.cu...")
+    mod = load(
+        name="_megablocks_debug_single",
+        sources=["csrc/new_cumsum.cu"],
+        extra_include_paths=["csrc"],
+        extra_cflags=["-O3", "-std=c++17"],
+        extra_cuda_cflags=["-O3"],
+        verbose=True,
+        is_python_module=False,
+    )
+    print("✓ Single source build successful")
+except Exception as e:
+    print(f"✗ Single source build failed: {e}")
+
+print("\n=== Testing with Two Source Files ===")
+try:
+    print("Building with new_cumsum.cu and new_histogram.cu...")
+    mod = load(
+        name="_megablocks_debug_double",
+        sources=["csrc/new_cumsum.cu", "csrc/new_histogram.cu"],
+        extra_include_paths=["csrc"],
+        extra_cflags=["-O3", "-std=c++17"],
+        extra_cuda_cflags=["-O3"],
+        verbose=True,
+        is_python_module=False,
+    )
+    print("✓ Double source build successful")
+except Exception as e:
+    print(f"✗ Double source build failed: {e}")
+
+print("\n=== Testing with grouped_gemm.cu Only ===")
+try:
+    print("Building with just grouped_gemm.cu (most complex)...")
+    mod = load(
+        name="_megablocks_debug_gemm",
+        sources=["csrc/grouped_gemm/grouped_gemm.cu"],
+        extra_include_paths=["csrc"],
+        extra_cflags=["-O3", "-std=c++17"],
+        extra_cuda_cflags=["-O3"],
+        extra_ldflags=["-lhipblaslt"],
+        verbose=True,
+        is_python_module=False,
+    )
+    print("✓ grouped_gemm build successful")
+except Exception as e:
+    print(f"✗ grouped_gemm build failed: {e}")
+
+signal.alarm(0)  # Cancel timeout
+EOF
+
+echo "Running incremental build test..."
+python3 debug_build.py
+
+echo
+echo "=== Testing Full Build with Timeout ==="
+
+cat > debug_full_build.py << 'EOF'
+import os
+import pathlib
+import sys
+import signal
+from torch.utils.cpp_extension import load
+
+def timeout_handler(signum, frame):
+    print("Full build timed out - this confirms the hanging issue")
+    sys.exit(124)  # timeout exit code
+
+# Set up 5 minute timeout
+signal.signal(signal.SIGALRM, timeout_handler)
+signal.alarm(300)
+
+repo = pathlib.Path(".").resolve()
+os.environ.setdefault("TORCH_EXTENSIONS_DIR", str(repo / ".torch_extensions_debug"))
+
+sources = [
+    "torch-ext/torch_binding.cpp",
+    "csrc/new_cumsum.cu",
+    "csrc/new_histogram.cu",
+    "csrc/new_indices.cu",
+    "csrc/new_replicate.cu",
+    "csrc/new_sort.cu",
+    "csrc/grouped_gemm/grouped_gemm.cu",
+]
+
+print("=== Attempting Full MegaBlocks Build ===")
+print("This mimics the exact build.py process...")
+print("Sources:", sources)
+
+try:
+    mod = load(
+        name="_megablocks_debug_full",
+        sources=sources,
+        extra_include_paths=["csrc"],
+        extra_cflags=["-O3", "-std=c++17"],
+        extra_cuda_cflags=["-O3"],
+        extra_ldflags=["-lhipblaslt"],
+        verbose=True,
+        is_python_module=False,
+    )
+    print("✓ Full build successful!")
+    print("Built:", mod)
+
+except Exception as e:
+    print(f"✗ Full build failed: {e}")
+    import traceback
+    traceback.print_exc()
+
+signal.alarm(0)
+EOF
+
+echo "Running full build test (with timeout)..."
+python3 debug_full_build.py
+
+echo
+echo "=== Cleanup ==="
+rm -f /tmp/*.o
+rm -f debug_build.py debug_full_build.py
+
+echo
+echo "=== Debug Script 4 Complete ==="

_dev/debug-build-all.sh

@@ -0,0 +1,91 @@
+#!/usr/bin/env bash
+
+# Debug script: Run all debug tests in sequence
+
+set -euo pipefail
+
+echo "=== MegaBlocks Build Debugging Suite ==="
+echo "Running progressive debug tests to identify build hang issue"
+echo "ROCm installation: /opt/rocm-7.0.1"
+echo
+
+# Make all scripts executable
+chmod +x debug-build-1-env.sh debug-build-2-hipcc.sh debug-build-3-torch-ext.sh debug-build-4-megablocks.sh
+
+scripts=(
+    "debug-build-1-env.sh"
+    "debug-build-2-hipcc.sh"
+    "debug-build-3-torch-ext.sh"
+    "debug-build-4-megablocks.sh"
+)
+
+results=()
+start_time=$(date +%s)
+
+for script in "${scripts[@]}"; do
+    echo
+    echo "========================================"
+    echo "Running $script"
+    echo "========================================"
+
+    script_start=$(date +%s)
+
+    if ./"$script"; then
+        script_end=$(date +%s)
+        duration=$((script_end - script_start))
+        echo "✓ $script completed successfully in ${duration}s"
+        results+=("✓ $script: SUCCESS (${duration}s)")
+    else
+        script_end=$(date +%s)
+        duration=$((script_end - script_start))
+        echo "✗ $script failed in ${duration}s"
+        results+=("✗ $script: FAILED (${duration}s)")
+    fi
+
+    echo "----------------------------------------"
+done
+
+end_time=$(date +%s)
+total_duration=$((end_time - start_time))
+
+echo
+echo "========================================"
+echo "SUMMARY REPORT"
+echo "========================================"
+echo "Total runtime: ${total_duration}s"
+echo
+
+for result in "${results[@]}"; do
+    echo "$result"
+done
+
+echo
+echo "=== Analysis ==="
+echo "1. If debug-1-env.sh fails: ROCm installation/environment issue"
+echo "2. If debug-2-hipcc.sh fails: HIP compiler issue"
+echo "3. If debug-3-torch-ext.sh hangs: PyTorch extension compilation issue"
+echo "4. If debug-4-megablocks.sh hangs: MegaBlocks-specific compilation issue"
+echo
+echo "=== Next Steps Based on Results ==="
+echo "- If all pass: The issue may be intermittent or environment-specific"
+echo "- If script 3 or 4 hangs: Run with strace to see where it hangs:"
+echo "  strace -f -e trace=process,signal python3 build.py"
+echo "- Check compilation log files in .torch_extensions for more details"
+echo "- Consider using PYTORCH_JIT_LOG_LEVEL=1 for more verbose output"
+
+echo
+echo "=== Additional Debugging Commands ==="
+echo "# Check for stuck processes:"
+echo "ps aux | grep -E '(hipcc|hip-clang|python)'"
+echo
+echo "# Monitor system resources during build:"
+echo "htop"
+echo
+echo "# Check for device issues:"
+echo "dmesg | tail -20"
+echo
+echo "# Force clean rebuild:"
+echo "rm -rf .torch_extensions* && ./build.sh"
+
+echo
+echo "Debug suite complete."

_dev/debug_build.py

@@ -0,0 +1,68 @@
+import os
+import pathlib
+import sys
+import signal
+import time
+from torch.utils.cpp_extension import load
+
+def timeout_handler(signum, frame):
+    print("Build timed out - this indicates a hanging issue")
+    sys.exit(1)
+
+# Set up timeout
+signal.signal(signal.SIGALRM, timeout_handler)
+signal.alarm(180)  # 3 minute timeout
+
+repo = pathlib.Path(".").resolve()
+os.environ.setdefault("TORCH_EXTENSIONS_DIR", str(repo / ".torch_extensions_debug"))
+
+print("=== Testing with Single Source File ===")
+try:
+    print("Building with just new_cumsum.cu...")
+    mod = load(
+        name="_megablocks_debug_single",
+        sources=["csrc/new_cumsum.cu"],
+        extra_include_paths=["csrc"],
+        extra_cflags=["-O3", "-std=c++17"],
+        extra_cuda_cflags=["-O3"],
+        verbose=True,
+        is_python_module=False,
+    )
+    print("✓ Single source build successful")
+except Exception as e:
+    print(f"✗ Single source build failed: {e}")
+
+print("\n=== Testing with Two Source Files ===")
+try:
+    print("Building with new_cumsum.cu and new_histogram.cu...")
+    mod = load(
+        name="_megablocks_debug_double",
+        sources=["csrc/new_cumsum.cu", "csrc/new_histogram.cu"],
+        extra_include_paths=["csrc"],
+        extra_cflags=["-O3", "-std=c++17"],
+        extra_cuda_cflags=["-O3"],
+        verbose=True,
+        is_python_module=False,
+    )
+    print("✓ Double source build successful")
+except Exception as e:
+    print(f"✗ Double source build failed: {e}")
+
+print("\n=== Testing with grouped_gemm.cu Only ===")
+try:
+    print("Building with just grouped_gemm.cu (most complex)...")
+    mod = load(
+        name="_megablocks_debug_gemm",
+        sources=["csrc/grouped_gemm/grouped_gemm.cu"],
+        extra_include_paths=["csrc"],
+        extra_cflags=["-O3", "-std=c++17"],
+        extra_cuda_cflags=["-O3"],
+        extra_ldflags=["-lhipblaslt"],
+        verbose=True,
+        is_python_module=False,
+    )
+    print("✓ grouped_gemm build successful")
+except Exception as e:
+    print(f"✗ grouped_gemm build failed: {e}")
+
+signal.alarm(0)  # Cancel timeout

build-fixed.sh

@@ -0,0 +1,43 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+# Fixed build script with proper ROCm/HIP environment
+
+echo "=== Fixed Build Script ==="
+echo "Configuring environment for ROCm 7.0.1 with proper exports"
+
+# Default to the ROCm 7.0.1 install unless the caller overrides it.
+export ROCM_PATH="${ROCM_PATH:-/opt/rocm-7.0.1}"
+export ROCM_HOME="${ROCM_HOME:-$ROCM_PATH}"
+export HIP_PATH="${HIP_PATH:-$ROCM_PATH}"
+export HIP_HOME="${HIP_HOME:-$ROCM_PATH}"
+
+export PATH="$ROCM_HOME/bin:$PATH"
+
+# Fix architecture specifications - use gfx942 consistently
+export TORCH_HIP_ARCH_LIST="gfx942"
+export PYTORCH_ROCM_ARCH="gfx942"
+
+# Remove HSA_OVERRIDE_GFX_VERSION - not needed since MI300X is already gfx942
+unset HSA_OVERRIDE_GFX_VERSION
+
+# Force single-threaded compilation to avoid ninja hanging
+export MAX_JOBS=1
+
+# Enable PyTorch JIT logging for debugging
+export PYTORCH_JIT_LOG_LEVEL=1
+
+export TORCH_EXTENSIONS_DIR="${TORCH_EXTENSIONS_DIR:-$PWD/.torch_extensions}"
+
+echo "Environment configured:"
+echo "ROCM_PATH=$ROCM_PATH"
+echo "TORCH_HIP_ARCH_LIST=$TORCH_HIP_ARCH_LIST"
+echo "PYTORCH_ROCM_ARCH=$PYTORCH_ROCM_ARCH"
+echo "MAX_JOBS=$MAX_JOBS"
+echo "PYTORCH_JIT_LOG_LEVEL=$PYTORCH_JIT_LOG_LEVEL"
+echo "HSA_OVERRIDE_GFX_VERSION=${HSA_OVERRIDE_GFX_VERSION:-unset}"
+echo
+
+echo "Starting build..."
+python -u build.py

build-strace.sh

@@ -0,0 +1,54 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+# Build script with strace to debug hanging
+
+echo "=== Build with strace debugging ==="
+echo "This will trace system calls to identify where the build hangs"
+
+# Same environment as build-fixed.sh but without MAX_JOBS limit
+export ROCM_PATH="${ROCM_PATH:-/opt/rocm-7.0.1}"
+export ROCM_HOME="${ROCM_HOME:-$ROCM_PATH}"
+export HIP_PATH="${HIP_PATH:-$ROCM_PATH}"
+export HIP_HOME="${HIP_HOME:-$ROCM_PATH}"
+
+export PATH="$ROCM_HOME/bin:$PATH"
+
+# Fix architecture specifications
+export TORCH_HIP_ARCH_LIST="gfx942"
+export PYTORCH_ROCM_ARCH="gfx942"
+
+# Remove HSA_OVERRIDE_GFX_VERSION
+unset HSA_OVERRIDE_GFX_VERSION
+
+# Remove MAX_JOBS limit to see parallel compilation hang
+unset MAX_JOBS
+
+# Enable PyTorch JIT logging
+export PYTORCH_JIT_LOG_LEVEL=1
+
+export TORCH_EXTENSIONS_DIR="${TORCH_EXTENSIONS_DIR:-$PWD/.torch_extensions}"
+
+echo "Environment configured for strace:"
+echo "ROCM_PATH=$ROCM_PATH"
+echo "TORCH_HIP_ARCH_LIST=$TORCH_HIP_ARCH_LIST"
+echo "PYTORCH_ROCM_ARCH=$PYTORCH_ROCM_ARCH"
+echo "MAX_JOBS=${MAX_JOBS:-unset}"
+echo "PYTORCH_JIT_LOG_LEVEL=$PYTORCH_JIT_LOG_LEVEL"
+echo
+
+echo "Starting build with strace..."
+echo "Tracing process creation, signals, and file operations..."
+echo "Output will be saved to strace.log"
+
+# Use strace to trace the build process
+# -f: follow child processes
+# -e trace=process,signal: trace process creation and signals
+# -e trace=file: trace file operations
+# -o strace.log: save output to file
+# -T: show time spent in each syscall
+strace -f -e trace=process,signal,file -o strace.log -T python -u build.py
+
+echo "Build completed or interrupted"
+echo "Check strace.log for detailed system call trace"

build.sh

@@ -1,9 +1,39 @@
-# export TORCH_EXTENSIONS_DIR=/root/shisa-v2/train/v2.1/megablocks.kernels-community/.torch_extensions; export ROCM_HOME=/opt/rocm-6.4.1; export HIP_HOME=$ROCM_HOME; export TORCH_HIP_ARCH_LIST=gfx942; export HSA_OVERRIDE_GFX_VERSION=gfx942; python megablocks.kernels-community/build.py
-
-# 3-4min build
-export ROCM_HOME=/opt/rocm-6.4.1
-export HIP_HOME=$ROCM_HOME
-export TORCH_HIP_ARCH_LIST=gfx942
-export HSA_OVERRIDE_GFX_VERSION=gfx942
-export TORCH_EXTENSIONS_DIR="$PWD/megablocks.kernels-community/.torch_extensions"
-python megablocks.kernels-community/build.py
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+# 3-4min build with cleanup to prevent lock file hangs
+
+echo "=== MegaBlocks Build Script ==="
+echo "Cleaning up any previous build processes and lock files..."
+
+# Kill any hanging build processes
+echo "Killing any running build.py processes..."
+pkill -f "python.*build\.py" 2>/dev/null || true
+pkill -f "ninja" 2>/dev/null || true
+pkill -f "hipcc" 2>/dev/null || true
+
+# Wait a moment for processes to terminate
+sleep 2
+
+# Clean up lock files that cause infinite loops
+echo "Removing stale lock files..."
+if [ -d ".torch_extensions" ]; then
+  find .torch_extensions -name "lock" -delete 2>/dev/null || true
+  find .torch_extensions -name ".ninja_lock" -delete 2>/dev/null || true
+fi
+
+# Default to the ROCm 7.0.1 install unless the caller overrides it.
+export ROCM_PATH="${ROCM_PATH:-/opt/rocm-7.0.1}"
+export ROCM_HOME="${ROCM_HOME:-$ROCM_PATH}"
+export HIP_PATH="${HIP_PATH:-$ROCM_PATH}"
+export HIP_HOME="${HIP_HOME:-$ROCM_PATH}"
+
+export PATH="$ROCM_HOME/bin:$PATH"
+export LD_LIBRARY_PATH="$ROCM_HOME/lib:$ROCM_HOME/lib64:${LD_LIBRARY_PATH:-}"
+export TORCH_HIP_ARCH_LIST="${TORCH_HIP_ARCH_LIST:-gfx942}"
+export HSA_OVERRIDE_GFX_VERSION="${HSA_OVERRIDE_GFX_VERSION:-gfx942}"
+export TORCH_EXTENSIONS_DIR="${TORCH_EXTENSIONS_DIR:-$PWD/.torch_extensions}"
+
+echo "Environment configured. Starting build..."
+python -u build.py

run-tests.sh

@@ -0,0 +1,138 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+KERNEL_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+cd "$KERNEL_DIR"
+
+export KERNEL_DIR
+
+detect_variant() {
+python - <<'PY'
+import os
+import pathlib
+
+root = pathlib.Path(os.environ["KERNEL_DIR"])
+build_dir = root / "build"
+variant = None
+
+try:
+    from kernels.utils import build_variant as _build_variant
+except Exception:
+    _build_variant = None
+
+if _build_variant is not None:
+    try:
+        variant = _build_variant()
+    except Exception:
+        variant = None
+
+if variant is None:
+    candidates = sorted(build_dir.glob("torch*-rocm64-*") or build_dir.glob("torch*-cu*"))
+    if candidates:
+        variant = candidates[0].name
+
+if variant is None:
+    raise SystemExit("Could not determine MegaBlocks build variant. Run build.py first.")
+
+print(variant)
+PY
+}
+
+VARIANT=$(detect_variant)
+
+STAGED_DIR="$KERNEL_DIR/build/$VARIANT"
+find_staged_lib() {
+  local base="$1"
+  local candidates=(
+    "$base/_megablocks_rocm.so"
+    "$base/megablocks/_megablocks_rocm.so"
+  )
+  for path in "${candidates[@]}"; do
+    if [[ -f "$path" ]]; then
+      echo "$path"
+      return 0
+    fi
+  done
+  return 1
+}
+
+STAGED_LIB=$(find_staged_lib "$STAGED_DIR") || true
+
+if [[ -z "${STAGED_LIB:-}" ]]; then
+  echo "Staged ROCm extension not found under $STAGED_DIR; rebuilding kernels..."
+  python build.py
+  VARIANT=$(detect_variant)
+  STAGED_DIR="$KERNEL_DIR/build/$VARIANT"
+  STAGED_LIB=$(find_staged_lib "$STAGED_DIR") || true
+  if [[ -z "${STAGED_LIB:-}" ]]; then
+    echo "ERROR: build.py completed but no extension was found under $STAGED_DIR" >&2
+    exit 1
+  fi
+fi
+
+export PYTHONPATH="$STAGED_DIR:${PYTHONPATH:-}"
+
+echo "Using MegaBlocks build variant: $VARIANT"
+
+declare -i GPU_COUNT
+GPU_COUNT=$(python - <<'PY'
+import torch
+print(torch.cuda.device_count() if torch.cuda.is_available() else 0)
+PY
+)
+
+if (( GPU_COUNT == 0 )); then
+  echo "ERROR: No HIP/CUDA GPUs detected. Tests require at least one visible accelerator." >&2
+  exit 1
+fi
+
+echo "Detected $GPU_COUNT visible GPU(s)."
+
+log() {
+  echo
+  echo "==> $1"
+}
+
+run_pytest() {
+  local label="$1"
+  shift
+  log "$label"
+  set -x
+  "$@"
+  { set +x; } 2>/dev/null || true
+}
+
+SINGLE_GPU_ENV=(HIP_VISIBLE_DEVICES=0 CUDA_VISIBLE_DEVICES=0 WORLD_SIZE=1)
+MULTI2_GPU_ENV=(HIP_VISIBLE_DEVICES=0,1 CUDA_VISIBLE_DEVICES=0,1 WORLD_SIZE=2)
+MULTI8_GPU_ENV=(HIP_VISIBLE_DEVICES=$(seq -s, 0 7) CUDA_VISIBLE_DEVICES=$(seq -s, 0 7) WORLD_SIZE=8)
+
+SINGLE_TESTS=(
+  "test_mb_moe.py"
+  "test_mb_moe_shared_expert.py"
+  "layer_test.py"
+  "test_gg.py"
+  "ops_test.py"
+)
+
+for test in "${SINGLE_TESTS[@]}"; do
+  run_pytest "Single-GPU pytest ${test}" env "${SINGLE_GPU_ENV[@]}" python -m pytest "tests/${test}" -q
+done
+
+if (( GPU_COUNT >= 2 )); then
+  run_pytest "Distributed layer smoke (2 GPUs)" env "${MULTI2_GPU_ENV[@]}" python -m pytest "tests/parallel_layer_test.py::test_megablocks_moe_mlp_functionality" -q
+else
+  log "Skipping 2-GPU distributed layer test (requires >=2 GPUs, detected ${GPU_COUNT})."
+fi
+
+run_pytest "Shared expert functionality (world_size=1)" env "${SINGLE_GPU_ENV[@]}" python -m pytest 'tests/test_mb_moe_shared_expert_multi.py::test_shared_expert_distributed_functionality[1]' -q
+run_pytest "Shared expert weighted sum (world_size=1)" env "${SINGLE_GPU_ENV[@]}" python -m pytest 'tests/test_mb_moe_shared_expert_multi.py::test_shared_expert_distributed_weighted_sum[1]' -q
+
+if (( GPU_COUNT >= 8 )); then
+  run_pytest "Shared expert functionality (world_size=8)" env "${MULTI8_GPU_ENV[@]}" python -m pytest 'tests/test_mb_moe_shared_expert_multi.py::test_shared_expert_distributed_functionality[8]' -q
+  run_pytest "Shared expert weighted sum (world_size=8)" env "${MULTI8_GPU_ENV[@]}" python -m pytest 'tests/test_mb_moe_shared_expert_multi.py::test_shared_expert_distributed_weighted_sum[8]' -q
+else
+  log "Skipping 8-GPU shared expert tests (requires >=8 GPUs, detected ${GPU_COUNT})."
+fi
+
+echo
+echo "All requested tests completed."