sail_scripts/model/custom_ops.py

import torch
import os
from torch.utils.cpp_extension import load
import time

# Disable JIT compilation by default - it's very slow
# Set SAIL_CUDA_KERNELS=1 to enable custom fused CUDA kernels
_CUSTOM_KERNELS_ENABLED = os.environ.get("SAIL_CUDA_KERNELS", "0") == "1"

_rmsnorm_ext = None
_swiglu_ext = None
_rmsnorm_failed = not _CUSTOM_KERNELS_ENABLED
_swiglu_failed = not _CUSTOM_KERNELS_ENABLED

def get_rmsnorm_ext():
    global _rmsnorm_ext, _rmsnorm_failed
    if _rmsnorm_ext is not None:
        return _rmsnorm_ext
    
    if _rmsnorm_failed:
        return None
        
    if not torch.cuda.is_available():
        print("CUDA not available. Falling back to native PyTorch RMSNorm.")
        _rmsnorm_failed = True
        return None

    try:
        current_dir = os.path.dirname(os.path.abspath(__file__))
        kernel_dir = os.path.join(current_dir, "kernels")
        
        sources = [
            os.path.join(kernel_dir, "rmsnorm.cpp"),
            os.path.join(kernel_dir, "rmsnorm_kernel.cu")
        ]
        
        print("JIT Compiling Custom Fused CUDA RMSNorm Extension... (This may take a minute)")
        _rmsnorm_ext = load(
            name="fused_rmsnorm",
            sources=sources,
            extra_cuda_cflags=['-O3', '--use_fast_math'],
            verbose=False
        )
        print("Custom Fused RMSNorm Loaded Successfully.")
        return _rmsnorm_ext
    except Exception as e:
        print(f"Warning: Failed to compile CUDA RMSNorm using standard JIT fallback: {e}. "
              "Check MSVC/NVCC installations. Falling back to native PyTorch.")
        _rmsnorm_failed = True
        return None

class FusedRMSNormFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, weight, eps):
        ext = get_rmsnorm_ext()
        if ext is None:
             raise RuntimeError("CUDA RMSNorm extension not loaded")
        
        output = ext.forward(input, weight, eps)
        ctx.save_for_backward(input, weight)
        ctx.eps = eps
        return output

    @staticmethod
    def backward(ctx, grad_output):
        # Fallback to PyTorch autograd backward pass for simplicity since our
        # forward pass does the heavy lifting of reducing memory overhead
        # In a full deployment, we would write a custom raw CUDA backward kernel too.
        input, weight = ctx.saved_tensors
        eps = ctx.eps
        
        with torch.enable_grad():
            input = input.detach().requires_grad_(True)
            weight = weight.detach().requires_grad_(True)
            
            # Recalculate forward purely for autograd graph mapping
            norm = input * torch.rsqrt(input.pow(2).mean(-1, keepdim=True) + eps)
            out = norm.float().type_as(input) * weight
            
            out.backward(grad_output)
            
        return input.grad, weight.grad, None

@torch.compiler.disable
def fused_rmsnorm(input, weight, eps):
    ext = get_rmsnorm_ext()
    if ext is not None and input.is_cuda and weight.is_cuda:
         return FusedRMSNormFunction.apply(input, weight, eps)
    else:
         # Fallback - MUST be out-of-place for CUDA Graphs
         variance = input.pow(2).mean(-1, keepdim=True)
         norm = input * torch.rsqrt(variance + eps)
         return (norm.to(input.dtype)) * weight

def get_swiglu_ext():
    global _swiglu_ext, _swiglu_failed
    if _swiglu_ext is not None:
        return _swiglu_ext
        
    if _swiglu_failed:
        return None

    if not torch.cuda.is_available():
        _swiglu_failed = True
        return None

    try:
        current_dir = os.path.dirname(os.path.abspath(__file__))
        kernel_dir = os.path.join(current_dir, "kernels")
        
        sources = [
            os.path.join(kernel_dir, "swiglu.cpp"),
            os.path.join(kernel_dir, "swiglu_kernel.cu")
        ]
        
        print("JIT Compiling Custom Fused CUDA SwiGLU Extension... (This may take a minute)")
        _swiglu_ext = load(
            name="fused_swiglu",
            sources=sources,
            extra_cuda_cflags=['-O3', '--use_fast_math'],
            verbose=False
        )
        print("Custom Fused SwiGLU Loaded Successfully.")
        return _swiglu_ext
    except Exception as e:
        print(f"Warning: Failed to compile CUDA SwiGLU using standard JIT fallback: {e}. ")
        _swiglu_failed = True
        return None

class FusedSwiGLUFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, y):
        ext = get_swiglu_ext()
        if ext is None:
             raise RuntimeError("CUDA SwiGLU extension not loaded")
        
        output = ext.forward(x, y)
        ctx.save_for_backward(x, y)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        x, y = ctx.saved_tensors
        
        with torch.enable_grad():
            x = x.detach().requires_grad_(True)
            y = y.detach().requires_grad_(True)
            
            # Recalculate forward purely for autograd graph
            out = torch.nn.functional.silu(x) * y
            out.backward(grad_output)
            
        return x.grad, y.grad

@torch.compiler.disable
def fused_swiglu(x, y):
    ext = get_swiglu_ext()
    if ext is not None and x.is_cuda and y.is_cuda:
         return FusedSwiGLUFunction.apply(x, y)
    else:
         # Fallback - out-of-place for CUDA graphs
         act = torch.nn.functional.silu(x)
         return act * y