ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Improve MoE performance

Browse source
This commit is contained in:
Datta Nimmaturi 2025-12-24 12:01:32 +00:00
commit 8678f5de43
7 changed files with 901 additions and 78 deletions
Download patch file Download diff file Expand all files Collapse all files

390
unsloth/kernels/moe/autotune_cache.py Normal file
View file

@ -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

149
unsloth/kernels/moe/grouped_gemm/interface.py
View file

@ -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

40
unsloth/kernels/moe/grouped_gemm/kernels/autotuning.py
View file

@ -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,

37
unsloth/kernels/moe/grouped_gemm/kernels/backward.py
View file

@ -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)

16
unsloth/kernels/moe/grouped_gemm/kernels/forward.py
View file

@ -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",

8
unsloth/models/loader.py
View file

@ -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:

339
unsloth/models/qwen3_moe_triton.py Normal file
View file

@ -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