From 8678f5de435595aab901dd233be308057f0c4a02 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Wed, 24 Dec 2025 12:01:32 +0000 Subject: [PATCH] Improve MoE performance --- unsloth/kernels/moe/autotune_cache.py | 390 ++++++++++++++++++ unsloth/kernels/moe/grouped_gemm/interface.py | 149 ++++--- .../moe/grouped_gemm/kernels/autotuning.py | 40 +- .../moe/grouped_gemm/kernels/backward.py | 37 +- .../moe/grouped_gemm/kernels/forward.py | 16 +- unsloth/models/loader.py | 8 +- unsloth/models/qwen3_moe_triton.py | 339 +++++++++++++++ 7 files changed, 901 insertions(+), 78 deletions(-) create mode 100644 unsloth/kernels/moe/autotune_cache.py create mode 100644 unsloth/models/qwen3_moe_triton.py diff --git a/unsloth/kernels/moe/autotune_cache.py b/unsloth/kernels/moe/autotune_cache.py new file mode 100644 index 0000000000..ae05b6959d --- /dev/null +++ b/unsloth/kernels/moe/autotune_cache.py @@ -0,0 +1,390 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +""" +Auto-tuning cache system for MoE kernels to ensure tuning runs only once at training start. +""" + +import hashlib +import json +import logging +import os +import time +from typing import Dict, List, Optional, Tuple, Any +import torch +import triton + +logger = logging.getLogger(__name__) + +# Global cache for kernel configurations +_kernel_config_cache: Dict[str, Any] = {} +_autotune_completed: Dict[str, bool] = {} + +def _get_cache_key( + num_experts: int, + hidden_dim: int, + intermediate_dim: int, + top_k: int, + dtype: torch.dtype, + device_capability: Tuple[int, int], + seq_len: int = 8192, # Default sequence length for tuning +) -> str: + """Generate a unique cache key based on model configuration.""" + key_data = { + "num_experts": num_experts, + "hidden_dim": hidden_dim, + "intermediate_dim": intermediate_dim, + "top_k": top_k, + "dtype": str(dtype), + "device_capability": device_capability, + "seq_len": seq_len, + } + key_str = json.dumps(key_data, sort_keys=True) + return hashlib.md5(key_str.encode()).hexdigest() + +def _get_cache_file_path(cache_key: str) -> str: + """Get the file path for the cache file.""" + cache_dir = os.path.expanduser("~/.cache/unsloth/moe_autotune") + os.makedirs(cache_dir, exist_ok=True) + return os.path.join(cache_dir, f"{cache_key}.json") + +def load_cached_config(cache_key: str) -> Optional[Dict[str, Any]]: + """Load cached kernel configuration from disk.""" + cache_file = _get_cache_file_path(cache_key) + if not os.path.exists(cache_file): + return None + + try: + with open(cache_file, 'r') as f: + cached_data = json.load(f) + + # Verify cache is still valid (same device, etc.) + current_device_capability = torch.cuda.get_device_capability() + if cached_data.get("device_capability") != current_device_capability: + logger.info("Device capability changed, invalidating cache") + os.remove(cache_file) + return None + + logger.info(f"Loaded cached MoE kernel config: {cache_key}") + return cached_data + except Exception as e: + logger.warning(f"Failed to load cache file {cache_file}: {e}") + try: + os.remove(cache_file) + except: + pass + return None + +def save_cached_config( + cache_key: str, + config_fwd: Any, + config_bwd_dx: Any, + config_bwd_dw: Any, + metadata: Dict[str, Any] = None +) -> None: + """Save kernel configuration to disk cache.""" + cache_file = _get_cache_file_path(cache_key) + + cache_data = { + "timestamp": time.time(), + "device_capability": torch.cuda.get_device_capability(), + "config_fwd": config_fwd.__dict__ if hasattr(config_fwd, '__dict__') else str(config_fwd), + "config_bwd_dx": config_bwd_dx.__dict__ if hasattr(config_bwd_dx, '__dict__') else str(config_bwd_dx), + "config_bwd_dw": config_bwd_dw.__dict__ if hasattr(config_bwd_dw, '__dict__') else str(config_bwd_dw), + "metadata": metadata or {}, + } + + try: + with open(cache_file, 'w') as f: + json.dump(cache_data, f, indent=2) + logger.info(f"Saved MoE kernel config cache: {cache_key}") + except Exception as e: + logger.warning(f"Failed to save cache file {cache_file}: {e}") + +def get_or_autotune_moe_kernels( + num_experts: int, + hidden_dim: int, + intermediate_dim: int, + top_k: int, + dtype: torch.dtype, + force_autotune: bool = False, + seq_len: int = 8192, +) -> Tuple[Any, Any, Any]: + """ + Get cached kernel configurations or run auto-tuning. + + Args: + num_experts: Number of experts in the MoE layer + hidden_dim: Hidden dimension of the model + intermediate_dim: Intermediate dimension for MoE MLP + top_k: Number of experts to route to + dtype: Data type for computation + force_autotune: Force re-running autotuning even if cache exists + seq_len: Sequence length to use for tuning benchmarks + + Returns: + Tuple of (config_fwd, config_bwd_dx, config_bwd_dw) + """ + device_capability = torch.cuda.get_device_capability() + cache_key = _get_cache_key( + num_experts, hidden_dim, intermediate_dim, top_k, dtype, device_capability, seq_len + ) + + # Check if we already have cached configs + if not force_autotune and cache_key in _kernel_config_cache: + logger.info(f"Using in-memory cached MoE kernel configs: {cache_key}") + return _kernel_config_cache[cache_key] + + # Try to load from disk + if not force_autotune: + cached_data = load_cached_config(cache_key) + if cached_data is not None: + # Reconstruct config objects from cached data + try: + from grouped_gemm.kernels.tuning import ( + KernelConfigForward, + KernelConfigBackward_dX, + KernelConfigBackward_dW, + ) + + config_fwd = KernelConfigForward(**cached_data["config_fwd"]) + config_bwd_dx = KernelConfigBackward_dX(**cached_data["config_bwd_dx"]) + config_bwd_dw = KernelConfigBackward_dW(**cached_data["config_bwd_dw"]) + + configs = (config_fwd, config_bwd_dx, config_bwd_dw) + _kernel_config_cache[cache_key] = configs + return configs + except Exception as e: + logger.warning(f"Failed to reconstruct cached configs: {e}") + + # Run autotuning + if cache_key in _autotune_completed and not force_autotune: + logger.info(f"Autotuning already completed for: {cache_key}") + return _kernel_config_cache[cache_key] + + logger.info(f"Running MoE kernel auto-tuning for: {cache_key}") + logger.info(f"Configuration: {num_experts} experts, {hidden_dim} hidden, {intermediate_dim} intermediate, top_k={top_k}") + + try: + configs = _run_moe_autotuning( + num_experts, hidden_dim, intermediate_dim, top_k, dtype, seq_len + ) + + # Cache the results + _kernel_config_cache[cache_key] = configs + _autotune_completed[cache_key] = True + + # Save to disk + config_fwd, config_bwd_dx, config_bwd_dw = configs + save_cached_config( + cache_key, + config_fwd, + config_bwd_dx, + config_bwd_dw, + {"num_experts": num_experts, "hidden_dim": hidden_dim, "intermediate_dim": intermediate_dim} + ) + + logger.info(f"MoE kernel auto-tuning completed: {cache_key}") + return configs + + except Exception as e: + logger.error(f"MoE kernel auto-tuning failed: {e}") + if "AttributeError" in str(e) and "_experimental_make_tensor_descriptor" in str(e): + logger.warning("Unsloth: Your Triton version might be incompatible with TMA features. Falling back to default configs.") + logger.info("Falling back to default kernel configurations") + return _get_default_configs() + +def _run_moe_autotuning( + num_experts: int, + hidden_dim: int, + intermediate_dim: int, + top_k: int, + dtype: torch.dtype, + seq_len: int, +) -> Tuple[Any, Any, Any]: + """Run the actual auto-tuning for MoE kernels.""" + + # Create dummy inputs for tuning + device = "cuda" + batch_size = 2 # Small batch for tuning + num_tokens = batch_size * seq_len + total_tokens = num_tokens * top_k + + # Create dummy tensors + hidden_states = torch.randn(num_tokens, hidden_dim, device=device, dtype=dtype) + + # Create dummy weights + gate_up_weights = torch.randn( + num_experts, 2 * intermediate_dim, hidden_dim, device=device, dtype=dtype + ) + down_weights = torch.randn( + num_experts, hidden_dim, intermediate_dim, device=device, dtype=dtype + ) + + # Create dummy routing data + m_sizes = torch.randint(1, total_tokens // num_experts + 1, (num_experts,), device=device) + m_sizes = m_sizes * (total_tokens // m_sizes.sum().item()) + # Adjust to ensure exact total + diff = total_tokens - m_sizes.sum().item() + if diff != 0: + m_sizes[0] += diff + + gather_indices = torch.arange(total_tokens, device=device) + torch.randperm(total_tokens, out=gather_indices) + + # Autotune forward kernel - use the interface function with autotune=True + # This properly invokes the kernel and lets triton handle the autotuning + from grouped_gemm.interface import ( + grouped_gemm_forward, + grouped_gemm_dX, + grouped_gemm_dW, + ) + from grouped_gemm.kernels.forward import _autotuned_grouped_gemm_forward_kernel + from grouped_gemm.kernels.backward import ( + _autotuned_grouped_gemm_dX_kernel, + _autotuned_grouped_gemm_dW_kernel, + ) + from grouped_gemm.kernels.tuning import ( + KernelConfigForward, + KernelConfigBackward_dX, + KernelConfigBackward_dW, + ) + + logger.info("Autotuning forward kernel (first GEMM)...") + # Run with autotune=True to trigger autotuning + _ = grouped_gemm_forward( + X=hidden_states, + W=gate_up_weights, + topk=top_k, + m_sizes=m_sizes, + gather_indices=gather_indices, + permute_x=True, + permute_y=False, + autotune=True, + ) + triton_config_fwd = _autotuned_grouped_gemm_forward_kernel.best_config + + # Convert triton.Config to KernelConfigForward + config_fwd = KernelConfigForward( + BLOCK_SIZE_M=triton_config_fwd.kwargs["BLOCK_SIZE_M"], + BLOCK_SIZE_N=triton_config_fwd.kwargs["BLOCK_SIZE_N"], + BLOCK_SIZE_K=triton_config_fwd.kwargs["BLOCK_SIZE_K"], + num_warps=triton_config_fwd.num_warps, + num_stages=triton_config_fwd.num_stages, + use_tma_load_x=triton_config_fwd.kwargs.get("USE_TMA_LOAD_X", False), + use_tma_load_w=triton_config_fwd.kwargs.get("USE_TMA_LOAD_W", False), + use_tma_store=triton_config_fwd.kwargs.get("USE_TMA_STORE", False), + ) + + # Autotune backward dX kernel + logger.info("Autotuning backward dX kernel...") + dummy_grad = torch.randn(total_tokens, 2 * intermediate_dim, device=device, dtype=dtype) + _ = grouped_gemm_dX( + dY=dummy_grad, + W=gate_up_weights, + gather_indices=gather_indices, + m_sizes=m_sizes, + topk=top_k, + permute_x=True, + permute_y=False, + autotune=True, + ) + triton_config_bwd_dx = _autotuned_grouped_gemm_dX_kernel.best_config + + # Convert triton.Config to KernelConfigBackward_dX + config_bwd_dx = KernelConfigBackward_dX( + BLOCK_SIZE_M=triton_config_bwd_dx.kwargs["BLOCK_SIZE_M"], + BLOCK_SIZE_N=triton_config_bwd_dx.kwargs["BLOCK_SIZE_N"], + BLOCK_SIZE_K=triton_config_bwd_dx.kwargs["BLOCK_SIZE_K"], + num_warps=triton_config_bwd_dx.num_warps, + num_stages=triton_config_bwd_dx.num_stages, + use_tma_load_dy=triton_config_bwd_dx.kwargs.get("USE_TMA_LOAD_dY", False), + use_tma_load_w=triton_config_bwd_dx.kwargs.get("USE_TMA_LOAD_W", False), + use_tma_store=triton_config_bwd_dx.kwargs.get("USE_TMA_STORE", False), + ) + + # Autotune backward dW kernel + logger.info("Autotuning backward dW kernel...") + _ = grouped_gemm_dW( + X=hidden_states, + dY=dummy_grad, + m_sizes=m_sizes, + gather_indices=gather_indices, + topk=top_k, + permute_x=True, + permute_y=False, + autotune=True, + ) + triton_config_bwd_dw = _autotuned_grouped_gemm_dW_kernel.best_config + + # Convert triton.Config to KernelConfigBackward_dW + config_bwd_dw = KernelConfigBackward_dW( + BLOCK_SIZE_M=triton_config_bwd_dw.kwargs["BLOCK_SIZE_M"], + BLOCK_SIZE_N=triton_config_bwd_dw.kwargs["BLOCK_SIZE_N"], + BLOCK_SIZE_K=triton_config_bwd_dw.kwargs["BLOCK_SIZE_K"], + num_warps=triton_config_bwd_dw.num_warps, + num_stages=triton_config_bwd_dw.num_stages, + use_tma_load_dy=triton_config_bwd_dw.kwargs.get("USE_TMA_LOAD_dY", False), + use_tma_load_x=triton_config_bwd_dw.kwargs.get("USE_TMA_LOAD_X", False), + use_tma_store=triton_config_bwd_dw.kwargs.get("USE_TMA_STORE", False), + ) + + return config_fwd, config_bwd_dx, config_bwd_dw + + + +def _get_default_configs() -> Tuple[Any, Any, Any]: + """Get default kernel configurations as fallback.""" + from grouped_gemm.kernels.tuning import ( + KernelConfigForward, + KernelConfigBackward_dX, + KernelConfigBackward_dW, + ) + + logger.warning("Using default MoE kernel configurations (not optimal)") + + config_fwd = KernelConfigForward( + BLOCK_SIZE_M=128, + BLOCK_SIZE_N=128, + BLOCK_SIZE_K=64, + num_warps=8, + num_stages=3, + use_tma_load_x=False, + use_tma_load_w=False, + use_tma_store=False, + ) + + config_bwd_dx = KernelConfigBackward_dX( + BLOCK_SIZE_M=128, + BLOCK_SIZE_N=128, + BLOCK_SIZE_K=64, + num_warps=8, + num_stages=3, + use_tma_load_dy=False, + use_tma_load_w=False, + use_tma_store=False, + ) + + config_bwd_dw = KernelConfigBackward_dW( + BLOCK_SIZE_M=128, + BLOCK_SIZE_N=128, + BLOCK_SIZE_K=64, + num_warps=8, + num_stages=3, + use_tma_load_dy=False, + use_tma_load_x=False, + use_tma_store=False, + ) + + return config_fwd, config_bwd_dx, config_bwd_dw + +def clear_cache() -> None: + """Clear all cached kernel configurations.""" + global _kernel_config_cache, _autotune_completed + _kernel_config_cache.clear() + _autotune_completed.clear() + logger.info("Cleared MoE kernel cache") + +def is_autotuning_completed(cache_key: str) -> bool: + """Check if autotuning has been completed for a given cache key.""" + return cache_key in _autotune_completed diff --git a/unsloth/kernels/moe/grouped_gemm/interface.py b/unsloth/kernels/moe/grouped_gemm/interface.py index d278d89645..dd5307d601 100644 --- a/unsloth/kernels/moe/grouped_gemm/interface.py +++ b/unsloth/kernels/moe/grouped_gemm/interface.py @@ -35,17 +35,48 @@ ch = logging.StreamHandler() ch.setFormatter(formatter) logger.addHandler(ch) -_FUSED_MUL_WARN = False -_SUPPORTS_TMA = None +# Precompute TMA support to avoid graph breaks +# TMA requires both: +# 1. GPU capability >= 9 (Hopper+) +# 2. Triton version with TMA API (make_tensor_descriptor or _experimental_make_tensor_descriptor) +def _check_tma_support(): + import triton.language as tl + gpu_supports_tma = torch.cuda.get_device_capability()[0] >= 9 + # Check for both old experimental and new stable API names + triton_has_tma_api = hasattr(tl, 'make_tensor_descriptor') or hasattr(tl, '_experimental_make_tensor_descriptor') + return gpu_supports_tma and triton_has_tma_api + +_SUPPORTS_TMA = _check_tma_support() def supports_tma(): - global _SUPPORTS_TMA - if _SUPPORTS_TMA is None: - _SUPPORTS_TMA = torch.cuda.get_device_capability()[0] >= 9 return _SUPPORTS_TMA + +# Helper to support allow_in_graph +try: + from torch.compiler import allow_in_graph +except ImportError: + from torch._dynamo import allow_in_graph + +# Helper to detect if we're in tracing/compilation mode +def _is_tracing(*tensors): + """ + Check if tensors are fake tensors used during torch.compile tracing. + During tracing, tensors are FakeTensor/FunctionalTensor and we can't run Triton kernels. + During execution, tensors are real Tensors and we MUST run the kernels. + + NOTE: We do NOT use torch.compiler.is_compiling() because it returns True + during both tracing AND execution. We only want to skip kernels during tracing + when tensors are actually fake. + """ + for t in tensors: + name = type(t).__name__ + if name in ("FakeTensor", "FunctionalTensor", "FunctionalTensorWrapper"): + return True + return False + _per_device_alloc_fns = {} @@ -83,6 +114,7 @@ def log_kernel_info( logger.debug(f"{kernel_name} autotuned best_config: {best_config}") +@allow_in_graph def grouped_gemm_forward( X: torch.Tensor, W: torch.Tensor, @@ -158,11 +190,20 @@ def grouped_gemm_forward( use_tma_store = False if use_tma or autotune: + # Respect global persistent allocator if set + if not getattr(triton, "_unsloth_allocator_set", False): + def alloc_fn(size: int, alignment: int, stream: int): + return torch.empty(size, device = "cuda", dtype = torch.int8) - def alloc_fn(size: int, alignment: int, stream: int): - return torch.empty(size, device = "cuda", dtype = torch.int8) + triton.set_allocator(alloc_fn) - triton.set_allocator(alloc_fn) + if W.ndim == 3: + num_experts = W.shape[0] + N = W.shape[1] + # K = W.shape[2] + else: + num_experts = m_sizes.shape[0] + N = W.shape[0] // num_experts X = X.view(-1, X.shape[-1]) W = W.view(-1, W.shape[-1]) @@ -188,9 +229,7 @@ def grouped_gemm_forward( total_tokens = X.shape[0] num_tokens = total_tokens // topk - num_experts = m_sizes.shape[0] _, K = X.shape - N = W.shape[0] // num_experts assert K == W.shape[1], f"K ({K}) must match W.shape[1] ({W.shape[1]})" if fuse_mul_post: @@ -212,8 +251,8 @@ def grouped_gemm_forward( ) y = torch.empty((total_tokens, N), device = X.device, dtype = X.dtype) - if total_tokens == 0 or N == 0: - return y + # if total_tokens == 0 or N == 0: + # return y NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count @@ -221,9 +260,9 @@ def grouped_gemm_forward( return (NUM_SMS,) if not autotune: - BLOCK_SIZE_K = min(K, BLOCK_SIZE_K) - BLOCK_SIZE_N = min(N, BLOCK_SIZE_N) - BLOCK_SIZE_M = min(total_tokens, BLOCK_SIZE_M) + # BLOCK_SIZE_K = min(K, BLOCK_SIZE_K) + # BLOCK_SIZE_N = min(N, BLOCK_SIZE_N) + pass if debug: print( @@ -276,16 +315,19 @@ def grouped_gemm_forward( if autotune else _grouped_gemm_forward_kernel ) - compiled_kernel: triton.compiler.CompiledKernel = kernel[grid](**kernel_args) - if autotune: - log_kernel_info(compiled_kernel, kernel.best_config) - else: - log_kernel_info(compiled_kernel) + is_fake = _is_tracing(X, W) + if not is_fake: + compiled_kernel: triton.compiler.CompiledKernel = kernel[grid](**kernel_args) + if autotune: + log_kernel_info(compiled_kernel, kernel.best_config) + else: + log_kernel_info(compiled_kernel) return y +@allow_in_graph def grouped_gemm_dX( dY: torch.Tensor, W: torch.Tensor, @@ -354,20 +396,27 @@ def grouped_gemm_dX( use_tma_store = False if use_tma or autotune: + # Respect global persistent allocator if set + if not getattr(triton, "_unsloth_allocator_set", False): + def alloc_fn(size: int, alignment: int, stream: int): + # print(f"DEBUG::GROUPED_GEMM alloc_fn {size=} {alignment=} {stream=}") + return torch.empty(size, device = "cuda", dtype = torch.int8) - def alloc_fn(size: int, alignment: int, stream: int): - # print(f"DEBUG::GROUPED_GEMM alloc_fn {size=} {alignment=} {stream=}") - return torch.empty(size, device = "cuda", dtype = torch.int8) + triton.set_allocator(alloc_fn) - triton.set_allocator(alloc_fn) + if W.ndim == 3: + num_experts = W.shape[0] + N = W.shape[1] + else: + num_experts = m_sizes.shape[0] + N = W.shape[0] // num_experts - num_experts = m_sizes.shape[0] dY = dY.view(-1, dY.shape[-1]) W = W.view(-1, W.shape[-1]) M_total, N_grad = dY.shape N_total, K = W.shape - N = N_total // num_experts + # N = N_total // num_experts assert N_grad == N, f"Grad_output N ({N_grad}) must match weight N ({N})" assert ( @@ -393,9 +442,9 @@ def grouped_gemm_dX( return (NUM_SMS,) if not autotune: - BLOCK_SIZE_M = min(M_total, BLOCK_SIZE_M) - BLOCK_SIZE_N = min(N_grad, BLOCK_SIZE_N) - BLOCK_SIZE_K = min(K, BLOCK_SIZE_K) + # BLOCK_SIZE_N = min(N_grad, BLOCK_SIZE_N) + # BLOCK_SIZE_K = min(K, BLOCK_SIZE_K) + pass if debug: print( @@ -437,15 +486,19 @@ def grouped_gemm_dX( } ) kernel = _autotuned_grouped_gemm_dX_kernel if autotune else _grouped_gemm_dX_kernel - compiled_kernel: triton.compiler.CompiledKernel = kernel[grid](**kernel_args) - if autotune: - log_kernel_info(compiled_kernel, kernel.best_config) - else: - log_kernel_info(compiled_kernel) + is_fake = _is_tracing(dY, W) + if not is_fake: + compiled_kernel: triton.compiler.CompiledKernel = kernel[grid](**kernel_args) + + if autotune: + log_kernel_info(compiled_kernel, kernel.best_config) + else: + log_kernel_info(compiled_kernel) return dX +@allow_in_graph def grouped_gemm_dW( X: torch.Tensor, dY: torch.Tensor, @@ -510,11 +563,12 @@ def grouped_gemm_dW( use_tma_store = False if use_tma or autotune: + # Respect global persistent allocator if set + if not getattr(triton, "_unsloth_allocator_set", False): + def alloc_fn(size: int, alignment: int, stream: int): + return torch.empty(size, device = "cuda", dtype = torch.int8) - def alloc_fn(size: int, alignment: int, stream: int): - return torch.empty(size, device = "cuda", dtype = torch.int8) - - triton.set_allocator(alloc_fn) + triton.set_allocator(alloc_fn) if permute_x or permute_y: assert gather_indices is not None @@ -541,9 +595,9 @@ def grouped_gemm_dW( dW = torch.zeros((num_experts, N, K), device = X.device, dtype = X.dtype) if not autotune: - BLOCK_SIZE_M = min(total_tokens, BLOCK_SIZE_M) - BLOCK_SIZE_N = min(N, BLOCK_SIZE_N) - BLOCK_SIZE_K = min(K, BLOCK_SIZE_K) + # BLOCK_SIZE_N = min(N, BLOCK_SIZE_N) + # BLOCK_SIZE_K = min(K, BLOCK_SIZE_K) + pass def grid(META): return (NUM_SMS,) @@ -607,12 +661,15 @@ def grouped_gemm_dW( ) kernel = _autotuned_grouped_gemm_dW_kernel if autotune else _grouped_gemm_dW_kernel - compiled_kernel: triton.compiler.CompiledKernel = kernel[grid](**kernel_args) - if autotune: - log_kernel_info(compiled_kernel, kernel.best_config) - else: - log_kernel_info(compiled_kernel) + is_fake = _is_tracing(X, dY) + if not is_fake: + compiled_kernel: triton.compiler.CompiledKernel = kernel[grid](**kernel_args) + + if autotune: + log_kernel_info(compiled_kernel, kernel.best_config) + else: + log_kernel_info(compiled_kernel) return dW diff --git a/unsloth/kernels/moe/grouped_gemm/kernels/autotuning.py b/unsloth/kernels/moe/grouped_gemm/kernels/autotuning.py index a185b5fd37..6c12be5b74 100644 --- a/unsloth/kernels/moe/grouped_gemm/kernels/autotuning.py +++ b/unsloth/kernels/moe/grouped_gemm/kernels/autotuning.py @@ -36,17 +36,35 @@ def convert_args_to_list(args): return [val_to_list(arg) for arg in args] +def _triton_supports_tma(): + """Check if current Triton version supports TMA API.""" + import triton.language as tl + # Check for both old experimental and new stable API names + return hasattr(tl, 'make_tensor_descriptor') or hasattr(tl, '_experimental_make_tensor_descriptor') + +# Precompute at module import +# NOTE: TMA is disabled for now due to compatibility issues with permute_x/permute_y settings +# in the MoE grouped GEMM forward/backward passes. Re-enable once these are resolved. +_TRITON_HAS_TMA = False # _triton_supports_tma() + + def get_forward_configs( BLOCK_M = DEFAULT_M_BLOCK_SIZES, BLOCK_N = DEFAULT_N_BLOCK_SIZES, BLOCK_K = DEFAULT_K_BLOCK_SIZES, - TMA_LOAD_X = True, - TMA_LOAD_W = True, + TMA_LOAD_X = None, # Auto-detect if not specified + TMA_LOAD_W = None, # Auto-detect if not specified TMA_STORE = False, # NOTE: TMA_STORE is disabled for now num_warps = DEFAULT_NUM_WARPS, num_stages = DEFAULT_NUM_STAGES, num_ctas = DEFAULT_NUM_CTAS, ): + # Auto-detect TMA support + if TMA_LOAD_X is None: + TMA_LOAD_X = _TRITON_HAS_TMA + if TMA_LOAD_W is None: + TMA_LOAD_W = _TRITON_HAS_TMA + ( BLOCK_M, BLOCK_N, @@ -115,13 +133,18 @@ def get_dX_kernel_configs( BLOCK_M = DEFAULT_M_BLOCK_SIZES, BLOCK_N = DEFAULT_N_BLOCK_SIZES, BLOCK_K = DEFAULT_K_BLOCK_SIZES, - TMA_LOAD_dY = True, - TMA_LOAD_W = True, + TMA_LOAD_dY = None, # Auto-detect if not specified + TMA_LOAD_W = None, # Auto-detect if not specified TMA_STORE = False, # NOTE: TMA_STORE is disabled for now num_warps = DEFAULT_NUM_WARPS, num_stages = DEFAULT_NUM_STAGES, num_ctas = DEFAULT_NUM_CTAS, ): + # Auto-detect TMA support + if TMA_LOAD_dY is None: + TMA_LOAD_dY = _TRITON_HAS_TMA + if TMA_LOAD_W is None: + TMA_LOAD_W = _TRITON_HAS_TMA ( BLOCK_M, BLOCK_N, @@ -193,10 +216,15 @@ def get_dW_kernel_configs( num_warps = DEFAULT_NUM_WARPS, num_stages = DEFAULT_NUM_STAGES, num_ctas = DEFAULT_NUM_CTAS, - TMA_LOAD_dY = True, - TMA_LOAD_X = True, + TMA_LOAD_dY = None, # Auto-detect if not specified + TMA_LOAD_X = None, # Auto-detect if not specified TMA_STORE = False, ): + # Auto-detect TMA support + if TMA_LOAD_dY is None: + TMA_LOAD_dY = _TRITON_HAS_TMA + if TMA_LOAD_X is None: + TMA_LOAD_X = _TRITON_HAS_TMA ( BLOCK_M, BLOCK_N, diff --git a/unsloth/kernels/moe/grouped_gemm/kernels/backward.py b/unsloth/kernels/moe/grouped_gemm/kernels/backward.py index d8bdcb57e2..970e51b619 100644 --- a/unsloth/kernels/moe/grouped_gemm/kernels/backward.py +++ b/unsloth/kernels/moe/grouped_gemm/kernels/backward.py @@ -53,11 +53,11 @@ def _grouped_gemm_dX_kernel( m_sizes_ptr, # problem sizes NUM_EXPERTS: tl.constexpr, - NUM_TOKENS: tl.constexpr, + NUM_TOKENS, TOPK: tl.constexpr, N: tl.constexpr, K: tl.constexpr, - NUM_SMS: tl.constexpr, + NUM_SMS, # Tuning parameters BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, @@ -69,7 +69,7 @@ def _grouped_gemm_dX_kernel( USE_TMA_STORE: tl.constexpr = False, FLATTEN: tl.constexpr = True, ) -> None: - TOTAL_TOKENS: tl.constexpr = NUM_TOKENS * TOPK + TOTAL_TOKENS = NUM_TOKENS * TOPK output_dtype = dX_ptr.dtype.element_ty tidx = tl.program_id(0) @@ -82,7 +82,7 @@ def _grouped_gemm_dX_kernel( # Also, we are defining a single global descriptor with single block shape # Need to check that this does not result in errors when crossing expert boundaries if USE_TMA_LOAD_dY: - dY_desc = tl._experimental_make_tensor_descriptor( + dY_desc = tl.make_tensor_descriptor( dY_ptr, shape = [TOTAL_TOKENS, N], strides = [N, 1], @@ -91,7 +91,7 @@ def _grouped_gemm_dX_kernel( if USE_TMA_LOAD_W: expert_stride = N * K - w_desc = tl._experimental_make_tensor_descriptor( + w_desc = tl.make_tensor_descriptor( w_ptr, shape = [NUM_EXPERTS, N, K], strides = [expert_stride, K, 1], @@ -123,7 +123,7 @@ def _grouped_gemm_dX_kernel( tl.static_assert( K % BLOCK_SIZE_K == 0, "K must be divisible by BLOCK_SIZE_K" ) - dX_desc = tl._experimental_make_tensor_descriptor( + dX_desc = tl.make_tensor_descriptor( dX_ptr, shape = [m_end, K], strides = [K, 1], @@ -232,6 +232,7 @@ def _grouped_gemm_dX_kernel( # TODO: check if predication along K is needed since we checked that K is divisible by BLOCK_SIZE_K in the forward kernel # [M, N] @ [N, K] -> [M, K] + dY = dY.to(w.dtype) accumulator += tl.dot(dY, w) # NOTE: no transpose of b # Advance A along contiguous dimension @@ -266,7 +267,8 @@ def _grouped_gemm_dX_kernel( _autotuned_grouped_gemm_dX_kernel = triton.autotune( configs = get_dX_kernel_configs(), prune_configs_by = {"early_config_prune": prune_dX_configs}, - key = ["NUM_EXPERTS", "NUM_TOKENS", "N", "K", "PERMUTE_X", "PERMUTE_Y"], + # NOTE: NUM_TOKENS removed from key to avoid recompilation for every sequence length + key = ["NUM_EXPERTS", "N", "K", "PERMUTE_X", "PERMUTE_Y"], )(_grouped_gemm_dX_kernel) """ @@ -298,12 +300,12 @@ def _grouped_gemm_dW_kernel( m_sizes_ptr, gather_indices_ptr, # problem sizes - NUM_TOKENS: tl.constexpr, + NUM_TOKENS, TOPK: tl.constexpr, NUM_EXPERTS: tl.constexpr, N: tl.constexpr, K: tl.constexpr, - NUM_SMS: tl.constexpr, + NUM_SMS, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, @@ -315,14 +317,14 @@ def _grouped_gemm_dW_kernel( FLATTEN: tl.constexpr = True, acc_dtype: tl.constexpr = tl.float32, ) -> None: - TOTAL_TOKENS: tl.constexpr = NUM_TOKENS * TOPK + TOTAL_TOKENS = NUM_TOKENS * TOPK TMA_LOAD_BOTH: tl.constexpr = USE_TMA_LOAD_X and USE_TMA_LOAD_dY tidx = tl.program_id(0) output_dtype = dW_ptr.dtype.element_ty if USE_TMA_LOAD_dY and not TMA_LOAD_BOTH: - dY_desc = tl._experimental_make_tensor_descriptor( + dY_desc = tl.make_tensor_descriptor( dY_ptr, shape = [TOTAL_TOKENS, N], strides = [N, 1], @@ -330,7 +332,7 @@ def _grouped_gemm_dW_kernel( ) if USE_TMA_LOAD_X and not TMA_LOAD_BOTH: - x_desc = tl._experimental_make_tensor_descriptor( + x_desc = tl.make_tensor_descriptor( x_ptr, shape = [TOTAL_TOKENS, K], strides = [K, 1], @@ -349,7 +351,7 @@ def _grouped_gemm_dW_kernel( if USE_TMA_STORE: tl.static_assert(N % BLOCK_SIZE_N == 0, "N must be divisible by BLOCK_SIZE_N") tl.static_assert(K % BLOCK_SIZE_K == 0, "K must be divisible by BLOCK_SIZE_K") - dW_desc = tl._experimental_make_tensor_descriptor( + dW_desc = tl.make_tensor_descriptor( dW_ptr, shape = [NUM_EXPERTS, N, K], strides = [N * K, K, 1], @@ -390,14 +392,14 @@ def _grouped_gemm_dW_kernel( if m_size > 0: if TMA_LOAD_BOTH: - dY_desc = tl._experimental_make_tensor_descriptor( + dY_desc = tl.make_tensor_descriptor( dY_ptr, shape = [m_end, N], strides = [N, 1], block_shape = [BLOCK_SIZE_M, BLOCK_SIZE_N], ) - x_desc = tl._experimental_make_tensor_descriptor( + x_desc = tl.make_tensor_descriptor( x_ptr, shape = [m_end, K], strides = [K, 1], @@ -475,7 +477,7 @@ def _grouped_gemm_dW_kernel( ) accumulator += tl.dot( - dY.T, # [BLOCK_N, BLOCK_M] + dY.T.to(x.dtype), # [BLOCK_N, BLOCK_M] x, # [BLOCK_M, BLOCK_K] ) @@ -498,5 +500,6 @@ def _grouped_gemm_dW_kernel( _autotuned_grouped_gemm_dW_kernel = triton.autotune( configs = get_dW_kernel_configs(), prune_configs_by = {"early_config_prune": prune_kernel_configs_backward_dW}, - key = ["NUM_EXPERTS", "NUM_TOKENS", "N", "K", "PERMUTE_X", "PERMUTE_Y"], + # NOTE: NUM_TOKENS removed from key to avoid recompilation for every sequence length + key = ["NUM_EXPERTS", "N", "K", "PERMUTE_X", "PERMUTE_Y"], )(_grouped_gemm_dW_kernel) diff --git a/unsloth/kernels/moe/grouped_gemm/kernels/forward.py b/unsloth/kernels/moe/grouped_gemm/kernels/forward.py index f84694e91e..a9f79fac0d 100644 --- a/unsloth/kernels/moe/grouped_gemm/kernels/forward.py +++ b/unsloth/kernels/moe/grouped_gemm/kernels/forward.py @@ -31,11 +31,11 @@ def _grouped_gemm_forward_kernel( topk_weights_ptr, # Constant problem shapes NUM_EXPERTS: tl.constexpr, - NUM_TOKENS: tl.constexpr, + NUM_TOKENS, TOPK: tl.constexpr, N: tl.constexpr, K: tl.constexpr, - NUM_SMS: tl.constexpr, + NUM_SMS, # Tuning params BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, @@ -53,7 +53,7 @@ def _grouped_gemm_forward_kernel( ) -> None: tl.static_assert(K % BLOCK_SIZE_K == 0) - TOTAL_TOKENS: tl.constexpr = NUM_TOKENS * TOPK + TOTAL_TOKENS = NUM_TOKENS * TOPK SHOULD_PERMUTE: tl.constexpr = PERMUTE_X or PERMUTE_Y SHOULD_FUSE_MUL: tl.constexpr = FUSE_MUL_PRE or FUSE_MUL_POST SHOULD_PERMUTE_OR_FUSE: tl.constexpr = SHOULD_PERMUTE or SHOULD_FUSE_MUL @@ -66,7 +66,7 @@ def _grouped_gemm_forward_kernel( # Also, we are defining a single global descriptor with single block shape # Need to check that this does not result in errors when crossing expert boundaries if USE_TMA_LOAD_X: - x_desc = tl._experimental_make_tensor_descriptor( + x_desc = tl.make_tensor_descriptor( x_ptr, shape = [TOTAL_TOKENS, K], strides = [K, 1], @@ -75,7 +75,7 @@ def _grouped_gemm_forward_kernel( if USE_TMA_LOAD_W: expert_stride = N * K - w_desc = tl._experimental_make_tensor_descriptor( + w_desc = tl.make_tensor_descriptor( w_ptr, shape = [NUM_EXPERTS, N, K], strides = [expert_stride, K, 1], @@ -100,7 +100,7 @@ def _grouped_gemm_forward_kernel( # Need to create tma_store within loop since we need to predicate stores based on m_size if USE_TMA_STORE: - y_desc = tl._experimental_make_tensor_descriptor( + y_desc = tl.make_tensor_descriptor( y_ptr, # + m_start * N, shape = [m_end, N], strides = [N, 1], @@ -213,6 +213,7 @@ def _grouped_gemm_forward_kernel( ) w = tl.reshape(w, (BLOCK_SIZE_N, BLOCK_SIZE_K)) + x = x.to(w.dtype) accumulator += tl.dot(x, w.T) if not USE_TMA_LOAD_X: @@ -253,9 +254,10 @@ def _grouped_gemm_forward_kernel( _autotuned_grouped_gemm_forward_kernel = triton.autotune( configs = get_forward_configs(), prune_configs_by = {"early_config_prune": prune_kernel_configs_fwd}, + # NOTE: NUM_TOKENS removed from key to avoid recompilation for every sequence length + # The kernel handles variable token counts via m_sizes and tile-based processing key = [ "NUM_EXPERTS", - "NUM_TOKENS", "N", "K", "PERMUTE_X", diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 91016a13ba..8e7ed3cb1b 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -323,7 +323,7 @@ class FastLanguageModel(FastLlamaModel): ) model_types = get_transformers_model_type( peft_config if peft_config is not None else model_config, - trust_remote_code = trust_remote_code, + # trust_remote_code = trust_remote_code, ) if len(model_types) == 1: model_type = model_types[0] @@ -855,7 +855,7 @@ class FastModel(FastBaseModel): ) model_types = get_transformers_model_type( peft_config if peft_config is not None else model_config, - trust_remote_code = trust_remote_code, + # trust_remote_code = trust_remote_code, ) model_types_all = ",".join(model_types) + "," @@ -949,6 +949,10 @@ class FastModel(FastBaseModel): ";" "os.environ['TRITON_F32_DEFAULT'] = 'ieee'" ) + elif "qwen3_moe" in model_types_all: + # Qwen3 MoE uses Triton grouped GEMM kernels which don't work well + # with torch.compile due to autograd.Function backward pass issues + os.environ["UNSLOTH_COMPILE_DISABLE"] = "partial" elif "gpt_oss" in model_types_all: os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1" if not load_in_4bit: diff --git a/unsloth/models/qwen3_moe_triton.py b/unsloth/models/qwen3_moe_triton.py new file mode 100644 index 0000000000..805379a009 --- /dev/null +++ b/unsloth/models/qwen3_moe_triton.py @@ -0,0 +1,339 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Qwen3 MoE integration with Triton kernels for faster training. +""" + +import logging +import os +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F + +from .llama import * +from ._utils import __version__ +from .llama import ( + LlamaRotaryEmbedding, + LlamaLinearScalingRotaryEmbedding, +) +from .qwen3 import ( + Qwen3Attention_fast_forward, + FastQwen3Model, +) +from transformers.models.qwen3_moe.modeling_qwen3_moe import ( + Qwen3MoeAttention, + Qwen3MoeSparseMoeBlock, + Qwen3MoeMLP, + Qwen3MoeDecoderLayer, + Qwen3MoeModel, + Qwen3MoeForCausalLM, +) + +# Try to import Triton kernels +try: + from unsloth.kernels.moe.autotune_cache import get_or_autotune_moe_kernels + from unsloth.kernels.moe.grouped_gemm.interface import grouped_gemm + from unsloth.kernels.moe.grouped_gemm.reference.moe_ops import ( + Qwen3MoeGroupedGEMMBlock, + permute, + unpermute, + ) + from unsloth.kernels.moe.grouped_gemm.reference.moe_block import Qwen3MoeFusedGroupedGEMMBlock + TRITON_MOE_AVAILABLE = True +except ImportError as e: + logging.warning(f"Triton MoE kernels not available: {e}") + TRITON_MOE_AVAILABLE = False + +torch_nn_functional_softmax = torch.nn.functional.softmax + +logger = logging.getLogger(__name__) + +# Global variable to store kernel configs +_moe_kernel_configs = {} +_moe_autotuning_done = False + + +def get_moe_kernel_configs( + num_experts: int, + hidden_dim: int, + intermediate_dim: int, + top_k: int, + dtype: torch.dtype, + force_autotune: bool = False, +) -> Tuple[Optional[any], Optional[any], Optional[any]]: + """Get or create MoE kernel configurations.""" + global _moe_kernel_configs, _moe_autotuning_done + + if not TRITON_MOE_AVAILABLE: + return None, None, None + + config_key = (num_experts, hidden_dim, intermediate_dim, top_k, str(dtype)) + + if config_key not in _moe_kernel_configs or force_autotune: + try: + configs = get_or_autotune_moe_kernels( + num_experts=num_experts, + hidden_dim=hidden_dim, + intermediate_dim=intermediate_dim, + top_k=top_k, + dtype=dtype, + force_autotune=force_autotune, + ) + _moe_kernel_configs[config_key] = configs + _moe_autotuning_done = True + logger.info(f"MoE kernel configs ready for {config_key}") + except Exception as e: + logger.error(f"Failed to get MoE kernel configs: {e}") + return None, None, None + + return _moe_kernel_configs.get(config_key, (None, None, None)) + + +def Qwen3MoeSparseMoeBlock_triton_forward( + self, + hidden_states: torch.Tensor, + temp_gate = None, + temp_up = None +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Fast forward implementation using Triton grouped GEMM kernels. + + This replaces the original Qwen3MoeSparseMoeBlock_fast_forward with + Triton-optimized kernels when available. + """ + if not TRITON_MOE_AVAILABLE: + # Fallback to original implementation + return Qwen3MoeSparseMoeBlock_fast_forward(self, hidden_states, temp_gate, temp_up) + + batch_size, seq_len, hidden_dim = hidden_states.shape + num_tokens = batch_size * seq_len + hidden_states_flat = hidden_states.view(-1, hidden_dim) + + # Router computation + router_logits = fast_linear_forward( + self.gate_proj, hidden_states_flat, out=temp_gate + ) + + routing_weights = torch_nn_functional_softmax( + router_logits, dim=-1, dtype=torch.float32 + ) + routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + routing_weights = routing_weights.to(hidden_states.dtype) + + # Get kernel configs + num_experts = self.num_experts + top_k = self.top_k + intermediate_dim = self.experts[0].gate_proj.in_features // 2 # Assuming SwiGLU + + config_fwd, config_bwd_dx, config_bwd_dw = get_moe_kernel_configs( + num_experts=num_experts, + hidden_dim=hidden_dim, + intermediate_dim=intermediate_dim, + top_k=top_k, + dtype=hidden_states.dtype, + ) + + # If we don't have kernel configs, fallback to original + if config_fwd is None: + return Qwen3MoeSparseMoeBlock_fast_forward(self, hidden_states, temp_gate, temp_up) + + try: + # Prepare expert weights for grouped GEMM + gate_up_weights = [] + down_weights = [] + + for expert in self.experts: + # Combine gate and up projections for first GEMM + gate_weight = expert.gate_proj.weight + up_weight = expert.up_proj.weight + gate_up_weight = torch.cat([gate_weight, up_weight], dim=0) + gate_up_weights.append(gate_up_weight) + down_weights.append(expert.down_proj.weight) + + gate_up_weights = torch.stack(gate_up_weights) # [num_experts, 2*intermediate_dim, hidden_dim] + down_weights = torch.stack(down_weights) # [num_experts, hidden_dim, intermediate_dim] + + # Compute token counts and gather indices without array operations + expert_mask = torch.nn.functional.one_hot( + selected_experts, num_classes=num_experts + ).permute(2, 1, 0) + + token_counts_by_expert = expert_mask.sum(dim=1).int() + + # Create gather indices for routing - avoid complex array operations + total_tokens = num_tokens * top_k + gather_indices = torch.zeros(total_tokens, dtype=torch.long, device=hidden_states.device) + + # Simple sequential assignment for gather indices + current_idx = 0 + for expert_idx in range(num_experts): + expert_tokens = expert_mask[expert_idx].sum(dim=0).bool() + num_expert_tokens = expert_tokens.sum().item() + if num_expert_tokens > 0: + expert_indices = torch.where(expert_tokens)[0] + gather_indices[current_idx:current_idx + num_expert_tokens] = expert_indices + current_idx += num_expert_tokens + + # First grouped GEMM: gate_up projection + intermediate_states = grouped_gemm( + X=hidden_states_flat, + W=gate_up_weights, + m_sizes=token_counts_by_expert, + gather_indices=gather_indices, + topk=top_k, + permute_x=True, + permute_y=False, + autotune=False, # Use pre-tuned configs + kernel_config_fwd=config_fwd, + kernel_config_bwd_dX=config_bwd_dx, + kernel_config_bwd_dW=config_bwd_dw, + is_first_gemm=True, + ) + + # Apply activation and multiply + gate, up = intermediate_states.chunk(2, dim=-1) + intermediate_states = F.silu(gate) * up + + # Second grouped GEMM: down projection + final_states = grouped_gemm( + X=intermediate_states, + W=down_weights, + m_sizes=token_counts_by_expert, + gather_indices=gather_indices, + topk=top_k, + permute_x=False, + permute_y=True, + autotune=False, # Use pre-tuned configs + kernel_config_fwd=config_fwd, + kernel_config_bwd_dX=config_bwd_dx, + kernel_config_bwd_dW=config_bwd_dw, + is_first_gemm=False, + ) + + # Reshape and apply routing weights + final_states = final_states.view(num_tokens, top_k, hidden_dim) + final_states = final_states * routing_weights.unsqueeze(-1) + final_states = final_states.sum(dim=1) + final_states = final_states.view(batch_size, seq_len, hidden_dim) + + return final_states, router_logits + + except Exception as e: + logger.error(f"Triton MoE kernel failed: {e}") + # Fallback to original implementation + return Qwen3MoeSparseMoeBlock_fast_forward(self, hidden_states, temp_gate, temp_up) + + +def Qwen3MoeMLP_triton_forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Fast forward implementation for Qwen3MoeMLP using optimized kernels. + """ + # This is for individual expert MLPs, still use the original fast implementation + return fast_swiglu_inference(self, x) + + +class FastTritonQwen3MoeModel(FastQwen3Model): + """Fast Qwen3 MoE model with Triton kernel integration.""" + + @staticmethod + def pre_patch(): + """Patch Qwen3 MoE components with Triton optimizations.""" + # Apply original patches first + FastQwen3Model.pre_patch() + + # Override MoE-specific components with Triton versions + if TRITON_MOE_AVAILABLE: + logger.info("Patching Qwen3 MoE with Triton kernels") + Qwen3MoeSparseMoeBlock.forward = Qwen3MoeSparseMoeBlock_triton_forward + Qwen3MoeMLP.forward = Qwen3MoeMLP_triton_forward + else: + logger.warning("Triton MoE kernels not available, using original implementation") + + @staticmethod + def from_pretrained( + model_name = "Qwen/Qwen3-7B", + max_seq_length = 4096, + dtype = None, + load_in_4bit = True, + token = None, + device_map = "sequential", + rope_scaling = None, + fix_tokenizer = True, + model_patcher = None, + tokenizer_name = None, + trust_remote_code = False, + # MoE-specific parameters + use_triton_moe = True, + force_moe_autotune = False, + **kwargs, + ): + """Load Qwen3 MoE model with Triton optimizations.""" + + # Set environment variable for MoE kernel preference + if use_triton_moe and TRITON_MOE_AVAILABLE: + os.environ["UNSLOTH_TRITON_MOE"] = "1" + logger.info("Enabling Triton MoE kernels") + else: + os.environ["UNSLOTH_TRITON_MOE"] = "0" + + # Load the model using the original method + model = FastQwen3Model.from_pretrained( + model_name=model_name, + max_seq_length=max_seq_length, + dtype=dtype, + load_in_4bit=load_in_4bit, + token=token, + device_map=device_map, + rope_scaling=rope_scaling, + fix_tokenizer=fix_tokenizer, + model_patcher=FastTritonQwen3MoeModel, + tokenizer_name=tokenizer_name, + trust_remote_code=trust_remote_code, + **kwargs, + ) + + # Pre-autotune MoE kernels if requested and available + if use_triton_moe and TRITON_MOE_AVAILABLE and hasattr(model, 'model'): + try: + # Extract MoE configuration from the model + config = model.config + if hasattr(config, 'num_experts') and hasattr(config, 'hidden_size'): + num_experts = config.num_experts + hidden_dim = config.hidden_size + intermediate_dim = config.intermediate_size or (hidden_dim * 4) # Common ratio + top_k = getattr(config, 'num_experts_per_tok', 2) + + logger.info(f"Pre-autotuning MoE kernels: {num_experts} experts, hidden={hidden_dim}, intermediate={intermediate_dim}") + + # Trigger autotuning + get_moe_kernel_configs( + num_experts=num_experts, + hidden_dim=hidden_dim, + intermediate_dim=intermediate_dim, + top_k=top_k, + dtype=model.dtype, + force_autotune=force_moe_autotune, + ) + + except Exception as e: + logger.error(f"Failed to pre-autotune MoE kernels: {e}") + + return model + + +# Import the original fast forward function for fallback +from .qwen3_moe import Qwen3MoeSparseMoeBlock_fast_forward