From 4beba840ed4e13ad5e0e795c57106e959a69d5b6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 1 Jan 2026 05:25:16 +0000 Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- unsloth/kernels/moe/autotune_cache.py | 221 ++++++++++-------- unsloth/kernels/moe/grouped_gemm/interface.py | 13 +- .../moe/grouped_gemm/kernels/autotuning.py | 6 +- unsloth/models/qwen3_moe_triton.py | 180 +++++++------- 4 files changed, 243 insertions(+), 177 deletions(-) diff --git a/unsloth/kernels/moe/autotune_cache.py b/unsloth/kernels/moe/autotune_cache.py index 4e93730974..0122165f62 100644 --- a/unsloth/kernels/moe/autotune_cache.py +++ b/unsloth/kernels/moe/autotune_cache.py @@ -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 diff --git a/unsloth/kernels/moe/grouped_gemm/interface.py b/unsloth/kernels/moe/grouped_gemm/interface.py index dd5307d601..ba7a3fa7b3 100644 --- a/unsloth/kernels/moe/grouped_gemm/interface.py +++ b/unsloth/kernels/moe/grouped_gemm/interface.py @@ -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) diff --git a/unsloth/kernels/moe/grouped_gemm/kernels/autotuning.py b/unsloth/kernels/moe/grouped_gemm/kernels/autotuning.py index 6c12be5b74..fb8e400991 100644 --- a/unsloth/kernels/moe/grouped_gemm/kernels/autotuning.py +++ b/unsloth/kernels/moe/grouped_gemm/kernels/autotuning.py @@ -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 diff --git a/unsloth/models/qwen3_moe_triton.py b/unsloth/models/qwen3_moe_triton.py index 805379a009..68b0ba522d 100644 --- a/unsloth/models/qwen3_moe_triton.py +++ b/unsloth/models/qwen3_moe_triton.py @@ -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: