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

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-01-01 05:25:16 +00:00
commit 4beba840ed
4 changed files with 243 additions and 177 deletions
Download patch file Download diff file Expand all files Collapse all files

221
unsloth/kernels/moe/autotune_cache.py
View file

@ -20,6 +20,7 @@ logger = logging.getLogger(__name__)
_kernel_config_cache: Dict[str, Any] = {}
_autotune_completed: Dict[str, bool] = {}
def _get_cache_key(
num_experts: int,
hidden_dim: int,
@ -39,15 +40,17 @@ def _get_cache_key(
"device_capability": device_capability,
"seq_len": seq_len,
}
key_str = json.dumps(key_data, sort_keys=True)
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)
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)
@ -55,7 +58,7 @@ def load_cached_config(cache_key: str) -> Optional[Dict[str, Any]]:
return None
try:
with open(cache_file, 'r') as f:
with open(cache_file, "r") as f:
cached_data = json.load(f)
# Verify cache is still valid (same device, etc.)
@ -75,12 +78,13 @@ def load_cached_config(cache_key: str) -> Optional[Dict[str, Any]]:
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
metadata: Dict[str, Any] = None,
) -> None:
"""Save kernel configuration to disk cache."""
cache_file = _get_cache_file_path(cache_key)
@ -88,19 +92,26 @@ def save_cached_config(
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),
"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)
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,
@ -127,7 +138,13 @@ def get_or_autotune_moe_kernels(
"""
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
num_experts,
hidden_dim,
intermediate_dim,
top_k,
dtype,
device_capability,
seq_len,
)
# Check if we already have cached configs
@ -163,7 +180,9 @@ def get_or_autotune_moe_kernels(
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}")
logger.info(
f"Configuration: {num_experts} experts, {hidden_dim} hidden, {intermediate_dim} intermediate, top_k={top_k}"
)
try:
configs = _run_moe_autotuning(
@ -181,7 +200,11 @@ def get_or_autotune_moe_kernels(
config_fwd,
config_bwd_dx,
config_bwd_dw,
{"num_experts": num_experts, "hidden_dim": hidden_dim, "intermediate_dim": intermediate_dim}
{
"num_experts": num_experts,
"hidden_dim": hidden_dim,
"intermediate_dim": intermediate_dim,
},
)
logger.info(f"MoE kernel auto-tuning completed: {cache_key}")
@ -189,11 +212,16 @@ def get_or_autotune_moe_kernels(
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.")
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,
@ -213,26 +241,28 @@ def _run_moe_autotuning(
total_tokens = num_tokens * top_k
# Create dummy tensors
hidden_states = torch.randn(num_tokens, hidden_dim, device=device, dtype=dtype)
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
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
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 = 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)
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
@ -255,86 +285,87 @@ def _run_moe_autotuning(
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,
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),
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)
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,
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),
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,
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),
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 (
@ -346,40 +377,41 @@ def _get_default_configs() -> Tuple[Any, Any, Any]:
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,
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,
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,
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
@ -387,6 +419,7 @@ def clear_cache() -> None:
_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

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

@ -42,24 +42,29 @@ logger.addHandler(ch)
# 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')
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():
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):
"""
@ -77,6 +82,7 @@ def _is_tracing(*tensors):
return True
return False
_per_device_alloc_fns = {}
@ -192,6 +198,7 @@ def grouped_gemm_forward(
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)
@ -398,6 +405,7 @@ def grouped_gemm_dX(
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)
@ -565,6 +573,7 @@ def grouped_gemm_dW(
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)

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

@ -39,8 +39,12 @@ def convert_args_to_list(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')
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

180
unsloth/models/qwen3_moe_triton.py
View file

@ -51,7 +51,10 @@ try:
permute,
unpermute,
)
from unsloth.kernels.moe.grouped_gemm.reference.moe_block import Qwen3MoeFusedGroupedGEMMBlock
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}")
@ -85,12 +88,12 @@ def get_moe_kernel_configs(
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,
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
@ -103,10 +106,7 @@ def get_moe_kernel_configs(
def Qwen3MoeSparseMoeBlock_triton_forward(
self,
hidden_states: torch.Tensor,
temp_gate = None,
temp_up = None
self, hidden_states: torch.Tensor, temp_gate = None, temp_up = None
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Fast forward implementation using Triton grouped GEMM kernels.
@ -116,7 +116,9 @@ def Qwen3MoeSparseMoeBlock_triton_forward(
"""
if not TRITON_MOE_AVAILABLE:
# Fallback to original implementation
return Qwen3MoeSparseMoeBlock_fast_forward(self, hidden_states, temp_gate, temp_up)
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
@ -124,14 +126,14 @@ def Qwen3MoeSparseMoeBlock_triton_forward(
# Router computation
router_logits = fast_linear_forward(
self.gate_proj, hidden_states_flat, out=temp_gate
self.gate_proj, hidden_states_flat, out = temp_gate
)
routing_weights = torch_nn_functional_softmax(
router_logits, dim=-1, dtype=torch.float32
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, 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
@ -140,16 +142,18 @@ def Qwen3MoeSparseMoeBlock_triton_forward(
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,
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)
return Qwen3MoeSparseMoeBlock_fast_forward(
self, hidden_states, temp_gate, temp_up
)
try:
# Prepare expert weights for grouped GEMM
@ -160,74 +164,82 @@ def Qwen3MoeSparseMoeBlock_triton_forward(
# 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_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]
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
selected_experts, num_classes = num_experts
).permute(2, 1, 0)
token_counts_by_expert = expert_mask.sum(dim=1).int()
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)
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()
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
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,
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)
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,
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.sum(dim = 1)
final_states = final_states.view(batch_size, seq_len, hidden_dim)
return final_states, router_logits
@ -235,7 +247,9 @@ def Qwen3MoeSparseMoeBlock_triton_forward(
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)
return Qwen3MoeSparseMoeBlock_fast_forward(
self, hidden_states, temp_gate, temp_up
)
def Qwen3MoeMLP_triton_forward(self, x: torch.Tensor) -> torch.Tensor:
@ -261,7 +275,9 @@ class FastTritonQwen3MoeModel(FastQwen3Model):
Qwen3MoeSparseMoeBlock.forward = Qwen3MoeSparseMoeBlock_triton_forward
Qwen3MoeMLP.forward = Qwen3MoeMLP_triton_forward
else:
logger.warning("Triton MoE kernels not available, using original implementation")
logger.warning(
"Triton MoE kernels not available, using original implementation"
)
@staticmethod
def from_pretrained(
@ -292,41 +308,45 @@ class FastTritonQwen3MoeModel(FastQwen3Model):
# 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,
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'):
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'):
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)
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}")
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,
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: