from __future__ import annotations import importlib import os import atexit import traceback from types import SimpleNamespace from typing import Optional, Tuple import torch try: from torch._subclasses.fake_tensor import is_fake as _torch_is_fake_tensor except Exception: # pragma: no cover _torch_is_fake_tensor = None # Env toggles _ENV_ENABLE = "WAN2GP_QUANTO_INT8_KERNEL" _ENV_DEBUG = "WAN2GP_QUANTO_INT8_DEBUG" _ENV_ALLOW_RUNTIME_FALLBACK = "WAN2GP_QUANTO_INT8_ALLOW_RUNTIME_FALLBACK" _ENV_NATIVE_FALLBACK_MAX_M = "WAN2GP_QUANTO_INT8_NATIVE_FALLBACK_MAX_M" _ENV_PROFILE_SHAPES = "WAN2GP_QUANTO_INT8_PROFILE_SHAPES" _ENV_PROFILE_TIME = "WAN2GP_QUANTO_INT8_PROFILE_TIME" _STARTUP_PRINTED = False _RUNTIME_DISABLED = False _RUNTIME_DISABLE_REASON = "" _RUNTIME_DISABLE_PRINTED = False _TRITON_MODULE = None _TRITON_DIRECT_FUSED_READY = False _TRITON_DIRECT_SCALED_READY = False _KERNEL_USED_PRINTED = False _SHAPE_PROFILE_ON = False _SHAPE_COUNTS_FUSED = {} _SHAPE_COUNTS_SCALED = {} _TIME_PROFILE_ON = False _TIME_PROFILE_EVENTS = [] _TIME_PROFILE_CPU_MS = 0.0 _TIME_PROFILE_CALLS = 0 _DEBUG_OVERRIDE: Optional[bool] = None _PATCH_STATE = SimpleNamespace(enabled=False, orig_forward=None, orig_embedding_forward=None) _BASE_PATCH_STATE = SimpleNamespace(enabled=False, orig_forward=None) _OPS_REGISTERED = False _OPS_NAMESPACE = "wan2gp_int8" _OPS_LIBS = [] _FUSED_LAUNCH_CACHE_MAX = 4096 _FUSED_LAUNCH_CACHE = {} _FUSED_LAUNCH_CACHE_FIFO = [] _SCALED_LAUNCH_CACHE_MAX = 4096 _SCALED_LAUNCH_CACHE = {} _SCALED_LAUNCH_CACHE_FIFO = [] _QBYTES_TENSOR_CLS = None _WEIGHT_QBYTES_CLS = None _NATIVE_FALLBACK_MAX_M = 0 def _encode_dtype(dtype: torch.dtype) -> int: if dtype == torch.float16: return 1 if dtype == torch.float32: return 2 return 0 def _decode_dtype(code: int, fallback: torch.dtype = torch.bfloat16) -> torch.dtype: if int(code) == 1: return torch.float16 if int(code) == 2: return torch.float32 return torch.bfloat16 if fallback not in (torch.bfloat16, torch.float16, torch.float32) else fallback def _env_flag(name: str, default: str = "1") -> bool: val = os.environ.get(name, default) return str(val).strip().lower() in ("1", "true", "yes", "on") def _env_int(name: str, default: int) -> int: try: return int(os.environ.get(name, str(default))) except Exception: return default def _log(msg: str) -> None: print(f"[Quanto][INT8] {msg}") def _debug(msg: str) -> None: if _DEBUG_OVERRIDE is None: debug_on = _env_flag(_ENV_DEBUG, "0") else: debug_on = bool(_DEBUG_OVERRIDE) if debug_on: _log(msg) def _format_exception_detail(exc: Exception) -> str: try: return "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)).strip() except Exception: return str(exc) def _summarize_kernel_error(exc_or_text: Exception | str, max_chars: int = 480) -> str: text = str(exc_or_text) lines = [ln.strip() for ln in text.replace("\r", "\n").split("\n") if ln.strip()] if len(lines) == 0: return "Unknown Triton kernel failure" keywords = ( "CompilationError", "shape mismatch", "tl.dot", "K >=", "M >=", "N >=", "Triton", "unsupported", "invalid", "at ", ) picked = [ln for ln in lines if any(kw in ln for kw in keywords)] if len(picked) == 0: picked = [lines[-1]] unique: list[str] = [] seen = set() for ln in picked: if ln in seen: continue seen.add(ln) unique.append(ln) summary = " | ".join(unique[-4:]) if len(summary) > max_chars: summary = summary[: max_chars - 3] + "..." return summary def set_kernel_debug(enabled: Optional[bool] = None) -> None: global _DEBUG_OVERRIDE _DEBUG_OVERRIDE = None if enabled is None else bool(enabled) def _allow_runtime_fallback() -> bool: return _env_flag(_ENV_ALLOW_RUNTIME_FALLBACK, "1") def _startup_status(enabled: bool, detail: str) -> None: global _STARTUP_PRINTED if _STARTUP_PRINTED: return _STARTUP_PRINTED = True if enabled: _log(f"Injected int8 kernels ACTIVE (backend=triton).") else: _log(f"Injected int8 kernels INACTIVE. {detail}") def _disable_runtime(reason: str) -> None: global _RUNTIME_DISABLED, _RUNTIME_DISABLE_REASON, _RUNTIME_DISABLE_PRINTED _RUNTIME_DISABLED = True _RUNTIME_DISABLE_REASON = _summarize_kernel_error(reason) if not _RUNTIME_DISABLE_PRINTED: _RUNTIME_DISABLE_PRINTED = True _log( "Runtime fallback to non-injected Quanto path is now active. Reason: " f"{_RUNTIME_DISABLE_REASON}" ) def _reset_runtime_state(reset_triton_module: bool = True) -> None: global _STARTUP_PRINTED, _RUNTIME_DISABLED, _RUNTIME_DISABLE_REASON, _RUNTIME_DISABLE_PRINTED global _TRITON_MODULE, _TRITON_DIRECT_FUSED_READY, _TRITON_DIRECT_SCALED_READY, _KERNEL_USED_PRINTED global _FUSED_LAUNCH_CACHE, _FUSED_LAUNCH_CACHE_FIFO, _SCALED_LAUNCH_CACHE, _SCALED_LAUNCH_CACHE_FIFO global _SHAPE_COUNTS_FUSED, _SHAPE_COUNTS_SCALED, _TIME_PROFILE_EVENTS, _TIME_PROFILE_CPU_MS, _TIME_PROFILE_CALLS global _NATIVE_FALLBACK_MAX_M _STARTUP_PRINTED = False _RUNTIME_DISABLED = False _RUNTIME_DISABLE_REASON = "" _RUNTIME_DISABLE_PRINTED = False if reset_triton_module: _TRITON_MODULE = None _TRITON_DIRECT_FUSED_READY = False _TRITON_DIRECT_SCALED_READY = False _KERNEL_USED_PRINTED = False _FUSED_LAUNCH_CACHE = {} _FUSED_LAUNCH_CACHE_FIFO = [] _SCALED_LAUNCH_CACHE = {} _SCALED_LAUNCH_CACHE_FIFO = [] _SHAPE_COUNTS_FUSED = {} _SHAPE_COUNTS_SCALED = {} _TIME_PROFILE_EVENTS = [] _TIME_PROFILE_CPU_MS = 0.0 _TIME_PROFILE_CALLS = 0 _NATIVE_FALLBACK_MAX_M = 0 def _add_bias_in_place_or_fallback(output: torch.Tensor, bias: Optional[torch.Tensor]) -> torch.Tensor: if bias is None: return output if bias.device != output.device or bias.dtype != output.dtype: return output + bias output.add_(bias) return output def _default_quanto_qbytes_linear_forward(ctx, input, other, bias=None): ctx.save_for_backward(input, other) if _is_qbytes_tensor(input): # MPS: torch.ops.quanto.qbytes_mm has no MPS kernel → CPU fallback # → CPU/MPS op interleaving corrupts Metal command buffer. # Use native MPS dequant+matmul instead. if input.device.type == "mps": act = input._data.to(input._scale.dtype) * input._scale wgt = other._data.to(input._scale.dtype) * other._scale output = act @ wgt.t() return _add_bias_in_place_or_fallback(output, bias) output = torch.ops.quanto.qbytes_mm(input._data, other._data, input._scale * other._scale) else: in_features = input.shape[-1] out_features = other.shape[0] output_shape = input.shape[:-1] + (out_features,) # MPS: same reason — qbytes_mm falls back to CPU on MPS. if input.device.type == "mps": wgt = other._data.to(input.dtype) * other._scale output = input.reshape(-1, in_features) @ wgt.t() output = output.reshape(output_shape) return _add_bias_in_place_or_fallback(output, bias) output = torch.ops.quanto.qbytes_mm(input.reshape(-1, in_features), other._data, other._scale) output = output.reshape(output_shape) return _add_bias_in_place_or_fallback(output, bias) def _ensure_default_quanto_linear_patch() -> bool: try: from optimum.quanto.tensor.weights import qbytes as _qbytes except Exception: return False _init_quanto_tensor_types() current_forward = _qbytes.WeightQBytesLinearFunction.forward if not _BASE_PATCH_STATE.enabled: _BASE_PATCH_STATE.orig_forward = current_forward _BASE_PATCH_STATE.enabled = True _qbytes.WeightQBytesLinearFunction.forward = staticmethod(_default_quanto_qbytes_linear_forward) return True def _init_quanto_tensor_types() -> bool: global _QBYTES_TENSOR_CLS, _WEIGHT_QBYTES_CLS if _QBYTES_TENSOR_CLS is not None and _WEIGHT_QBYTES_CLS is not None: return True try: from optimum.quanto.tensor.qbytes import QBytesTensor from optimum.quanto.tensor.weights.qbytes import WeightQBytesTensor except Exception: return False _QBYTES_TENSOR_CLS = QBytesTensor _WEIGHT_QBYTES_CLS = WeightQBytesTensor return True def _refresh_triton_direct_kernel_flags() -> None: global _TRITON_DIRECT_FUSED_READY, _TRITON_DIRECT_SCALED_READY mod = _TRITON_MODULE triton_ns = getattr(mod, "triton", None) if mod is not None else None has_common = bool(mod is not None and triton_ns is not None and hasattr(triton_ns, "cdiv") and hasattr(mod, "_select_triton_int8_config")) _TRITON_DIRECT_FUSED_READY = bool(has_common and hasattr(mod, "_fused_dynamic_int8_blockscale_gemm_kernel")) _TRITON_DIRECT_SCALED_READY = bool(has_common and hasattr(mod, "_scaled_int8_gemm_kernel")) def _is_qbytes_tensor(t: torch.Tensor) -> bool: if not _init_quanto_tensor_types(): return False return isinstance(t, _QBYTES_TENSOR_CLS) def _is_weight_qbytes(t: torch.Tensor) -> bool: if not _init_quanto_tensor_types(): return False return isinstance(t, _WEIGHT_QBYTES_CLS) def _flatten_scale(scale: torch.Tensor) -> torch.Tensor: if scale.ndim == 2 and scale.shape[1] == 1: return scale.view(-1) if scale.ndim == 1: return scale return scale.reshape(-1) def _expand_scale_to_rows(scale: torch.Tensor, rows: int, dtype: torch.dtype, device: Optional[torch.device] = None) -> torch.Tensor: scale = _flatten_scale(scale) if scale.numel() == 1: scale = scale.reshape(1).expand(rows) elif scale.numel() != rows: raise RuntimeError(f"Activation scale length mismatch: expected {rows}, got {scale.numel()}") if device is None: return scale.contiguous().to(dtype=dtype) return scale.contiguous().to(device=device, dtype=dtype, non_blocking=True) def _prepare_weight_scale(scale: torch.Tensor, out_features: int, device: torch.device) -> torch.Tensor: flat_scale = _flatten_scale(scale) if flat_scale.numel() != out_features: raise RuntimeError("Weight scale length does not match output features") if flat_scale.device != device: flat_scale = flat_scale.to(device=device, non_blocking=True) if flat_scale.dtype != torch.float32: flat_scale = flat_scale.to(torch.float32) if not flat_scale.is_contiguous(): flat_scale = flat_scale.contiguous() return flat_scale def _cache_launch_params(cache: dict, fifo: list, max_size: int, key: tuple[int, int, int, int], params: tuple[int, int, int, int, int, int, int]) -> tuple[int, int, int, int, int, int, int]: if key in cache: return cache[key] cache[key] = params fifo.append(key) if len(fifo) > max_size: stale_key = fifo.pop(0) cache.pop(stale_key, None) return params def _replace_launch_params(cache: dict, fifo: list, max_size: int, key: tuple[int, int, int, int], params: tuple[int, int, int, int, int, int, int]) -> None: cache[key] = params if key not in fifo: fifo.append(key) while len(fifo) > max_size: stale_key = fifo.pop(0) cache.pop(stale_key, None) def _cache_recovered_triton_config(kind: str, device_index: int, m: int, k: int, n: int, cfg: tuple[int, int, int, int, int]) -> None: mod = _TRITON_MODULE if mod is None: return try: slot_id, _ = mod._resolve_autotune_slot(m, k, n) mod._set_cached_config(device_index, kind, slot_id, cfg) except Exception: pass def _compile_recovery_candidates(kind: str, preferred: tuple[int, int, int, int, int], m: int, k: int, n: int) -> list[tuple[int, int, int, int, int]]: mod = _TRITON_MODULE if mod is None: return [] try: baseline = mod._select_static_triton_int8_config(m, k, n) return list(mod._compile_recovery_candidates(kind, baseline, preferred, m, k, n)) except Exception: return [] def _fused_launch_params(m: int, k: int, n: int, device: torch.device) -> tuple[int, int, int, int, int, int, int]: device_index = int(device.index if device.type == "cuda" else -1) key = (device_index, m, k, n) cached = _FUSED_LAUNCH_CACHE.get(key) if cached is not None: return cached mod = _TRITON_MODULE if mod is None: raise RuntimeError("Triton backend not initialized") block_m, block_n, block_k, num_warps, num_stages = mod._select_triton_int8_config(m, k, n, device=device, kernel_kind="fused") grid_m = mod.triton.cdiv(m, block_m) grid_n = mod.triton.cdiv(n, block_n) params = (block_m, block_n, block_k, num_warps, num_stages, grid_m, grid_n) return _cache_launch_params(_FUSED_LAUNCH_CACHE, _FUSED_LAUNCH_CACHE_FIFO, _FUSED_LAUNCH_CACHE_MAX, key, params) def _scaled_launch_params(m: int, k: int, n: int, device: torch.device) -> tuple[int, int, int, int, int, int, int]: device_index = int(device.index if device.type == "cuda" else -1) key = (device_index, m, k, n) cached = _SCALED_LAUNCH_CACHE.get(key) if cached is not None: return cached mod = _TRITON_MODULE if mod is None: raise RuntimeError("Triton backend not initialized") block_m, block_n, block_k, num_warps, num_stages = mod._select_triton_int8_config(m, k, n, device=device, kernel_kind="scaled") grid_m = mod.triton.cdiv(m, block_m) grid_n = mod.triton.cdiv(n, block_n) params = (block_m, block_n, block_k, num_warps, num_stages, grid_m, grid_n) return _cache_launch_params(_SCALED_LAUNCH_CACHE, _SCALED_LAUNCH_CACHE_FIFO, _SCALED_LAUNCH_CACHE_MAX, key, params) def _is_compiling_graph() -> bool: try: if bool(torch.compiler.is_compiling()): return True except Exception: pass try: import torch._dynamo as _dynamo if bool(_dynamo.is_compiling()): return True except Exception: pass return False def _is_fake_tensor(t: object) -> bool: if not torch.is_tensor(t): return False if _torch_is_fake_tensor is not None: return bool(_torch_is_fake_tensor(t)) return False def _resolve_output_dtype(input: torch.Tensor, other: torch.Tensor) -> torch.dtype: other_scale = getattr(other, "_scale", None) if torch.is_tensor(other_scale) and other_scale.dtype in (torch.bfloat16, torch.float16, torch.float32): return other_scale.dtype if _is_qbytes_tensor(input): input_scale = getattr(input, "_scale", None) if torch.is_tensor(input_scale) and input_scale.dtype in (torch.bfloat16, torch.float16, torch.float32): return input_scale.dtype if isinstance(input, torch.Tensor) and input.dtype in (torch.bfloat16, torch.float16, torch.float32): return input.dtype return torch.bfloat16 def _probe_triton_backend() -> Tuple[Optional[object], str]: try: mod = importlib.import_module("shared.kernels.quanto_int8_triton") except Exception as exc: return None, f"failed to import shared.kernels.quanto_int8_triton ({exc})" if not hasattr(mod, "is_available"): return None, "shared.kernels.quanto_int8_triton.is_available() missing" try: if not bool(mod.is_available()): return None, "Triton backend unavailable on this runtime/GPU" except Exception as exc: return None, f"Triton availability check failed ({exc})" return mod, "ok" def _register_int8_ops_for_namespace(ns: str, lib: torch.library.Library) -> None: lib.define("fused_quant_scaled_mm(Tensor x2d, Tensor qweight, Tensor qweight_scale, int out_dtype_code=0) -> Tensor") lib.define("scaled_int8_mm(Tensor a_int8, Tensor b_int8, Tensor a_scale, Tensor b_scale, int out_dtype_code=0) -> Tensor") @torch.library.impl(f"{ns}::fused_quant_scaled_mm", "CUDA") def _fused_quant_scaled_mm_cuda(x2d: torch.Tensor, qweight: torch.Tensor, qweight_scale: torch.Tensor, out_dtype_code: int = 0): if _TRITON_MODULE is None: raise RuntimeError("Triton backend not initialized") out_dtype = _decode_dtype(out_dtype_code, x2d.dtype) return _TRITON_MODULE.fused_quant_scaled_mm(x2d, qweight, qweight_scale, out_dtype=out_dtype) @torch.library.impl(f"{ns}::scaled_int8_mm", "CUDA") def _scaled_int8_mm_cuda(a_int8: torch.Tensor, b_int8: torch.Tensor, a_scale: torch.Tensor, b_scale: torch.Tensor, out_dtype_code: int = 0): if _TRITON_MODULE is None: raise RuntimeError("Triton backend not initialized") out_dtype = _decode_dtype(out_dtype_code, torch.bfloat16) return _TRITON_MODULE.scaled_int8_mm(a_int8, b_int8, a_scale, b_scale, out_dtype=out_dtype) @torch.library.register_fake(f"{ns}::fused_quant_scaled_mm") def _fused_quant_scaled_mm_fake(x2d: torch.Tensor, qweight: torch.Tensor, qweight_scale: torch.Tensor, out_dtype_code: int = 0): if x2d.ndim != 2 or qweight.ndim != 2: raise RuntimeError("fused_quant_scaled_mm expects 2D tensors") out_dtype = _decode_dtype(out_dtype_code, x2d.dtype) return x2d.new_empty((x2d.shape[0], qweight.shape[0]), dtype=out_dtype) @torch.library.register_fake(f"{ns}::scaled_int8_mm") def _scaled_int8_mm_fake(a_int8: torch.Tensor, b_int8: torch.Tensor, a_scale: torch.Tensor, b_scale: torch.Tensor, out_dtype_code: int = 0): if a_int8.ndim != 2 or b_int8.ndim != 2: raise RuntimeError("scaled_int8_mm expects 2D tensors") out_dtype = _decode_dtype(out_dtype_code, torch.bfloat16) return a_int8.new_empty((a_int8.shape[0], b_int8.shape[0]), dtype=out_dtype) def _ensure_compile_safe_ops() -> None: global _OPS_REGISTERED, _OPS_LIBS if _OPS_REGISTERED: return libs = [] try: lib = torch.library.Library(_OPS_NAMESPACE, "DEF") libs.append(lib) _register_int8_ops_for_namespace(_OPS_NAMESPACE, lib) except Exception: # Namespace/op may already exist in long-lived processes. op_ns = getattr(torch.ops, _OPS_NAMESPACE, None) has_ops = bool( op_ns is not None and hasattr(op_ns, "fused_quant_scaled_mm") and hasattr(op_ns, "scaled_int8_mm") ) if not has_ops: raise _OPS_LIBS = libs _OPS_REGISTERED = True def _fused_quant_scaled_mm_direct_call(x2d: torch.Tensor, qweight: torch.Tensor, qweight_scale: torch.Tensor, output_dtype: torch.dtype) -> torch.Tensor: mod = _TRITON_MODULE if mod is None: raise RuntimeError("Triton backend not initialized") if not _TRITON_DIRECT_FUSED_READY: return mod.fused_quant_scaled_mm(x2d, qweight, qweight_scale, out_dtype=output_dtype) m, k = x2d.shape n, k2 = qweight.shape if k != k2: raise RuntimeError(f"Triton int8 GEMM shape mismatch: x={x2d.shape}, w={qweight.shape}") block_m, block_n, block_k, num_warps, num_stages, grid_m, grid_n = _fused_launch_params(m, k, n, x2d.device) selected_cfg = (block_m, block_n, block_k, num_warps, num_stages) out = torch.empty((m, n), device=x2d.device, dtype=output_dtype) try: mod._fused_dynamic_int8_blockscale_gemm_kernel[(grid_m, grid_n)]( x2d, qweight, qweight_scale, out, m, n, k, x2d.stride(0), x2d.stride(1), qweight.stride(0), qweight.stride(1), out.stride(0), out.stride(1), block_m=block_m, block_n=block_n, block_k=block_k, num_warps=num_warps, num_stages=num_stages, ) except Exception as exc: recovery_errors = [] device_index = int(x2d.device.index if x2d.device.type == "cuda" else -1) for candidate in _compile_recovery_candidates("fused", selected_cfg, m, k, n): if candidate == selected_cfg: continue block_m, block_n, block_k, num_warps, num_stages = candidate grid_m = mod.triton.cdiv(m, block_m) grid_n = mod.triton.cdiv(n, block_n) recovered_out = torch.empty((m, n), device=x2d.device, dtype=output_dtype) try: mod._fused_dynamic_int8_blockscale_gemm_kernel[(grid_m, grid_n)]( x2d, qweight, qweight_scale, recovered_out, m, n, k, x2d.stride(0), x2d.stride(1), qweight.stride(0), qweight.stride(1), recovered_out.stride(0), recovered_out.stride(1), block_m=block_m, block_n=block_n, block_k=block_k, num_warps=num_warps, num_stages=num_stages, ) params = (block_m, block_n, block_k, num_warps, num_stages, grid_m, grid_n) key = (device_index, m, k, n) _replace_launch_params(_FUSED_LAUNCH_CACHE, _FUSED_LAUNCH_CACHE_FIFO, _FUSED_LAUNCH_CACHE_MAX, key, params) _cache_recovered_triton_config("fused", device_index, m, k, n, candidate) _debug(f"Recovered fused int8 kernel config for shape=({m},{k},{n}): {selected_cfg} -> {candidate}") return recovered_out except Exception as recovery_exc: recovery_errors.append(f"{candidate}: {recovery_exc}") raise RuntimeError( "Triton fused int8 kernel launch failed " f"(shape m={m}, k={k}, n={n}; tile=({selected_cfg[0]},{selected_cfg[1]},{selected_cfg[2]}); " f"warps={selected_cfg[3]}, stages={selected_cfg[4]}). {exc}" + (f" Recovery candidates also failed: {' | '.join(recovery_errors[-4:])}" if recovery_errors else "") ) from exc return out def _scaled_int8_mm_direct_call( a_int8: torch.Tensor, b_int8: torch.Tensor, a_scale: torch.Tensor, b_scale: torch.Tensor, output_dtype: torch.dtype, ) -> torch.Tensor: mod = _TRITON_MODULE if mod is None: raise RuntimeError("Triton backend not initialized") if not _TRITON_DIRECT_SCALED_READY: return mod.scaled_int8_mm(a_int8, b_int8, a_scale, b_scale, out_dtype=output_dtype) m, k = a_int8.shape n, k2 = b_int8.shape if k != k2: raise RuntimeError(f"Triton int8 GEMM shape mismatch: a={a_int8.shape}, w={b_int8.shape}") block_m, block_n, block_k, num_warps, num_stages, grid_m, grid_n = _scaled_launch_params(m, k, n, a_int8.device) selected_cfg = (block_m, block_n, block_k, num_warps, num_stages) out = torch.empty((m, n), device=a_int8.device, dtype=output_dtype) try: mod._scaled_int8_gemm_kernel[(grid_m, grid_n)]( a_int8, b_int8, a_scale, b_scale, out, m, n, k, a_int8.stride(0), a_int8.stride(1), b_int8.stride(0), b_int8.stride(1), out.stride(0), out.stride(1), block_m=block_m, block_n=block_n, block_k=block_k, num_warps=num_warps, num_stages=num_stages, ) except Exception as exc: recovery_errors = [] device_index = int(a_int8.device.index if a_int8.device.type == "cuda" else -1) for candidate in _compile_recovery_candidates("scaled", selected_cfg, m, k, n): if candidate == selected_cfg: continue block_m, block_n, block_k, num_warps, num_stages = candidate grid_m = mod.triton.cdiv(m, block_m) grid_n = mod.triton.cdiv(n, block_n) recovered_out = torch.empty((m, n), device=a_int8.device, dtype=output_dtype) try: mod._scaled_int8_gemm_kernel[(grid_m, grid_n)]( a_int8, b_int8, a_scale, b_scale, recovered_out, m, n, k, a_int8.stride(0), a_int8.stride(1), b_int8.stride(0), b_int8.stride(1), recovered_out.stride(0), recovered_out.stride(1), block_m=block_m, block_n=block_n, block_k=block_k, num_warps=num_warps, num_stages=num_stages, ) params = (block_m, block_n, block_k, num_warps, num_stages, grid_m, grid_n) key = (device_index, m, k, n) _replace_launch_params(_SCALED_LAUNCH_CACHE, _SCALED_LAUNCH_CACHE_FIFO, _SCALED_LAUNCH_CACHE_MAX, key, params) _cache_recovered_triton_config("scaled", device_index, m, k, n, candidate) _debug(f"Recovered scaled int8 kernel config for shape=({m},{k},{n}): {selected_cfg} -> {candidate}") return recovered_out except Exception as recovery_exc: recovery_errors.append(f"{candidate}: {recovery_exc}") raise RuntimeError( "Triton scaled int8 kernel launch failed " f"(shape m={m}, k={k}, n={n}; tile=({selected_cfg[0]},{selected_cfg[1]},{selected_cfg[2]}); " f"warps={selected_cfg[3]}, stages={selected_cfg[4]}). {exc}" + (f" Recovery candidates also failed: {' | '.join(recovery_errors[-4:])}" if recovery_errors else "") ) from exc return out def _fused_quant_scaled_mm_call(x2d: torch.Tensor, qweight: torch.Tensor, qweight_scale: torch.Tensor, output_dtype: torch.dtype) -> torch.Tensor: if _TRITON_MODULE is not None and not _is_compiling_graph() and not (_is_fake_tensor(x2d) or _is_fake_tensor(qweight) or _is_fake_tensor(qweight_scale)): return _fused_quant_scaled_mm_direct_call(x2d, qweight, qweight_scale, output_dtype) return torch.ops.wan2gp_int8.fused_quant_scaled_mm(x2d, qweight, qweight_scale, _encode_dtype(output_dtype)) def _scaled_int8_mm_call( a_int8: torch.Tensor, b_int8: torch.Tensor, a_scale: torch.Tensor, b_scale: torch.Tensor, output_dtype: torch.dtype, ) -> torch.Tensor: if _TRITON_MODULE is not None and not _is_compiling_graph() and not ( _is_fake_tensor(a_int8) or _is_fake_tensor(b_int8) or _is_fake_tensor(a_scale) or _is_fake_tensor(b_scale)): return _scaled_int8_mm_direct_call(a_int8, b_int8, a_scale, b_scale, output_dtype) return torch.ops.wan2gp_int8.scaled_int8_mm(a_int8, b_int8, a_scale, b_scale, _encode_dtype(output_dtype)) def _use_int8_kernel(input: torch.Tensor, other: torch.Tensor) -> bool: if _RUNTIME_DISABLED: return False if _TRITON_MODULE is None: return False if not _is_weight_qbytes(other): return False if other._data.dtype != torch.int8: return False if not other._data.is_cuda: return False if _is_qbytes_tensor(input): return input._data.dtype == torch.int8 and input._data.is_cuda return input.is_cuda and input.dtype in (torch.bfloat16, torch.float16, torch.float32) def _use_int8_embedding_kernel(module, input: torch.Tensor) -> bool: if _RUNTIME_DISABLED: return False if _TRITON_MODULE is None: return False if not torch.is_tensor(input): return False qweight = getattr(module, "qweight", None) if not _is_weight_qbytes(qweight): return False if qweight._data.dtype != torch.int8: return False if getattr(module, "max_norm", None) is not None: return False if bool(getattr(module, "scale_grad_by_freq", False)) or bool(getattr(module, "sparse", False)): return False if input.device != qweight._data.device: return False scale = getattr(qweight, "_scale", None) if not torch.is_tensor(scale) or scale.device != input.device: return False return True def _gather_embedding_scale(qweight, input: torch.Tensor) -> torch.Tensor: scale = qweight._scale if scale.ndim == 0 or scale.numel() == 1: return scale if scale.ndim == 1: if scale.shape[0] != qweight._data.shape[0]: raise RuntimeError("Quanto embedding scale length mismatch.") scale = scale.unsqueeze(-1) elif scale.ndim == 2: if scale.shape[0] != qweight._data.shape[0]: raise RuntimeError("Quanto embedding scale row count mismatch.") if scale.shape[1] != 1: raise RuntimeError("Quanto embedding fast path only supports per-row scales.") else: raise RuntimeError("Quanto embedding fast path only supports scalar or per-row scales.") return torch.nn.functional.embedding(input, scale) def _int8_embedding_forward(module, input: torch.Tensor) -> torch.Tensor: qweight = module.qweight gathered = torch.nn.functional.embedding( input, qweight._data, module.padding_idx, None, module.norm_type, False, False, ) gathered = gathered.to(qweight._scale.dtype) scale = _gather_embedding_scale(qweight, input) if torch.is_tensor(scale): if scale.ndim == gathered.ndim - 1: scale = scale.unsqueeze(-1) return gathered * scale return gathered * scale def _activation_rows(input_shape: torch.Size) -> int: rows = 1 for dim in input_shape[:-1]: rows *= int(dim) return rows def _prefer_native_quanto_path(input: torch.Tensor) -> bool: if _NATIVE_FALLBACK_MAX_M < 0: return False return _activation_rows(input.shape) <= _NATIVE_FALLBACK_MAX_M def _mark_kernel_used() -> None: global _KERNEL_USED_PRINTED if _KERNEL_USED_PRINTED: return _KERNEL_USED_PRINTED = True _log("Injected Triton int8 kernels are being used.") def _int8_linear_forward_triton_dense_fast(ctx, input: torch.Tensor, other: torch.Tensor, bias: Optional[torch.Tensor]): ctx.save_for_backward(input, other) if _TRITON_MODULE is None: raise RuntimeError("Triton backend not initialized") _mark_kernel_used() input_shape = input.shape in_features = int(input_shape[-1]) out_features = int(other.shape[0]) a_2d = input.reshape(-1, in_features) if not a_2d.is_contiguous(): a_2d = a_2d.contiguous() b_int8 = other._data if not b_int8.is_contiguous(): b_int8 = b_int8.contiguous() b_scale = _prepare_weight_scale(other._scale, out_features, b_int8.device) if _SHAPE_PROFILE_ON: key = (int(a_2d.shape[0]), int(in_features), int(out_features)) _SHAPE_COUNTS_FUSED[key] = _SHAPE_COUNTS_FUSED.get(key, 0) + 1 if _TIME_PROFILE_ON and torch.cuda.is_available(): start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() out_2d = _fused_quant_scaled_mm_call(a_2d, b_int8, b_scale, input.dtype) end.record() _TIME_PROFILE_EVENTS.append((start, end)) else: out_2d = _fused_quant_scaled_mm_call(a_2d, b_int8, b_scale, input.dtype) out = out_2d.reshape(input_shape[:-1] + (out_features,)) if bias is not None: out += bias return out def _int8_linear_forward_triton(ctx, input: torch.Tensor, other: torch.Tensor, bias: Optional[torch.Tensor]): ctx.save_for_backward(input, other) if _TRITON_MODULE is None: raise RuntimeError("Triton backend not initialized") _mark_kernel_used() input_shape = input.shape in_features = int(input_shape[-1]) out_features = int(other.shape[0]) b_int8 = other._data if not b_int8.is_contiguous(): b_int8 = b_int8.contiguous() b_scale = _prepare_weight_scale(other._scale, out_features, b_int8.device) output_dtype = _resolve_output_dtype(input, other) input_is_qbytes = _is_qbytes_tensor(input) if input_is_qbytes: a_int8 = input._data.reshape(-1, in_features) if a_int8.dtype != torch.int8: raise RuntimeError("QBytes input must be int8 for injected path") if not a_int8.is_contiguous(): a_int8 = a_int8.contiguous() a_scale = _expand_scale_to_rows(input._scale, a_int8.shape[0], torch.float32, device=a_int8.device) if _SHAPE_PROFILE_ON: key = (int(a_int8.shape[0]), int(in_features), int(out_features)) _SHAPE_COUNTS_SCALED[key] = _SHAPE_COUNTS_SCALED.get(key, 0) + 1 if _TIME_PROFILE_ON and torch.cuda.is_available(): start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() out_2d = _scaled_int8_mm_call(a_int8, b_int8, a_scale, b_scale, output_dtype) end.record() _TIME_PROFILE_EVENTS.append((start, end)) else: out_2d = _scaled_int8_mm_call(a_int8, b_int8, a_scale, b_scale, output_dtype) else: a_2d = input.reshape(-1, in_features) if not a_2d.is_contiguous(): a_2d = a_2d.contiguous() if _SHAPE_PROFILE_ON: key = (int(a_2d.shape[0]), int(in_features), int(out_features)) _SHAPE_COUNTS_FUSED[key] = _SHAPE_COUNTS_FUSED.get(key, 0) + 1 if _TIME_PROFILE_ON and torch.cuda.is_available(): start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() out_2d = _fused_quant_scaled_mm_call(a_2d, b_int8, b_scale, output_dtype) end.record() _TIME_PROFILE_EVENTS.append((start, end)) else: out_2d = _fused_quant_scaled_mm_call(a_2d, b_int8, b_scale, output_dtype) out = out_2d.reshape(input_shape[:-1] + (out_features,)) return _add_bias_in_place_or_fallback(out, bias) def enable_quanto_int8_kernel(triton_mod=None) -> bool: global _TRITON_MODULE, _NATIVE_FALLBACK_MAX_M if _PATCH_STATE.enabled: _reset_runtime_state(reset_triton_module=False) if triton_mod is None: triton_mod = _TRITON_MODULE if triton_mod is None: triton_mod, _ = _probe_triton_backend() if triton_mod is None: return False _TRITON_MODULE = triton_mod _refresh_triton_direct_kernel_flags() _NATIVE_FALLBACK_MAX_M = _env_int(_ENV_NATIVE_FALLBACK_MAX_M, 0) _init_quanto_tensor_types() _ensure_compile_safe_ops() return True try: from optimum.quanto.tensor.weights import qbytes as _qbytes except Exception as exc: _debug(f"cannot import optimum.quanto qbytes ({exc})") return False try: from mmgp import offload as _mmgp_offload except Exception as exc: _debug(f"cannot import mmgp.offload ({exc})") return False if triton_mod is None: triton_mod, _ = _probe_triton_backend() if triton_mod is None: _ensure_default_quanto_linear_patch() return False _ensure_default_quanto_linear_patch() _reset_runtime_state() _TRITON_MODULE = triton_mod _refresh_triton_direct_kernel_flags() _NATIVE_FALLBACK_MAX_M = _env_int(_ENV_NATIVE_FALLBACK_MAX_M, 0) _init_quanto_tensor_types() _ensure_compile_safe_ops() orig_forward = _qbytes.WeightQBytesLinearFunction.forward orig_embedding_forward = _mmgp_offload.QEmbedding.forward def forward(ctx, input, other, bias=None): dense_hot_path = ( not _RUNTIME_DISABLED and type(input) is torch.Tensor and input.is_cuda and input.dtype in (torch.bfloat16, torch.float16, torch.float32) and _WEIGHT_QBYTES_CLS is not None and isinstance(other, _WEIGHT_QBYTES_CLS) and other._data.dtype == torch.int8 and other._data.is_cuda ) if dense_hot_path: if _prefer_native_quanto_path(input): return orig_forward(ctx, input, other, bias) try: return _int8_linear_forward_triton_dense_fast(ctx, input, other, bias) except Exception as exc: short_reason = _summarize_kernel_error(exc) if _allow_runtime_fallback(): _disable_runtime(short_reason) _debug(f"Full Triton failure detail:\n{_format_exception_detail(exc)}") return orig_forward(ctx, input, other, bias) full_detail = _format_exception_detail(exc) raise RuntimeError( "Injected Triton int8 kernel failed. " f"Set {_ENV_ALLOW_RUNTIME_FALLBACK}=1 to force fallback to non-injected Quanto path. " f"Reason: {short_reason}\n" f"Full Triton error details:\n{full_detail}" ) from exc if not _use_int8_kernel(input, other): return orig_forward(ctx, input, other, bias) if _prefer_native_quanto_path(input): return orig_forward(ctx, input, other, bias) try: return _int8_linear_forward_triton(ctx, input, other, bias) except Exception as exc: short_reason = _summarize_kernel_error(exc) if _allow_runtime_fallback(): _disable_runtime(short_reason) _debug(f"Full Triton failure detail:\n{_format_exception_detail(exc)}") return orig_forward(ctx, input, other, bias) full_detail = _format_exception_detail(exc) raise RuntimeError( "Injected Triton int8 kernel failed. " f"Set {_ENV_ALLOW_RUNTIME_FALLBACK}=1 to force fallback to non-injected Quanto path. " f"Reason: {short_reason}\n" f"Full Triton error details:\n{full_detail}" ) from exc def embedding_forward(self, input): if not _use_int8_embedding_kernel(self, input): return orig_embedding_forward(self, input) try: return _int8_embedding_forward(self, input) except Exception as exc: short_reason = _summarize_kernel_error(exc) if _allow_runtime_fallback(): _disable_runtime(short_reason) _debug(f"Full embedding fast-path failure detail:\n{_format_exception_detail(exc)}") return orig_embedding_forward(self, input) full_detail = _format_exception_detail(exc) raise RuntimeError( "Injected Quanto int8 embedding fast path failed. " f"Set {_ENV_ALLOW_RUNTIME_FALLBACK}=1 to force fallback to non-injected Quanto path. " f"Reason: {short_reason}\n" f"Full error details:\n{full_detail}" ) from exc _qbytes.WeightQBytesLinearFunction.forward = staticmethod(forward) _mmgp_offload.QEmbedding.forward = embedding_forward _PATCH_STATE.enabled = True _PATCH_STATE.orig_forward = orig_forward _PATCH_STATE.orig_embedding_forward = orig_embedding_forward return True def disable_quanto_int8_kernel(notify_disabled = False) -> bool: if not _PATCH_STATE.enabled: _ensure_default_quanto_linear_patch() _reset_runtime_state() return False from optimum.quanto.tensor.weights import qbytes as _qbytes _qbytes.WeightQBytesLinearFunction.forward = staticmethod(_PATCH_STATE.orig_forward) from mmgp import offload as _mmgp_offload if _PATCH_STATE.orig_embedding_forward is not None: _mmgp_offload.QEmbedding.forward = _PATCH_STATE.orig_embedding_forward _PATCH_STATE.enabled = False _PATCH_STATE.orig_forward = None _PATCH_STATE.orig_embedding_forward = None _reset_runtime_state() if notify_disabled: _startup_status(False, f"disabled by User.") return True def maybe_enable_quanto_int8_kernel(verbose_level: Optional[int] = None) -> bool: global _SHAPE_PROFILE_ON, _TIME_PROFILE_ON, _STARTUP_PRINTED _STARTUP_PRINTED = False verbose_debug: Optional[bool] = None if verbose_level is not None: try: verbose_debug = int(verbose_level) >= 2 except Exception: verbose_debug = False set_kernel_debug(verbose_debug) if not _env_flag(_ENV_ENABLE, "1"): _ensure_default_quanto_linear_patch() # _startup_status(False, f"disabled by {_ENV_ENABLE}=0; using non-injected Quanto path.") return False triton_mod, reason = _probe_triton_backend() if triton_mod is None: _ensure_default_quanto_linear_patch() # _startup_status(False, f"{reason}; using non-injected Quanto path.") return False set_triton_debug = getattr(triton_mod, "set_autotune_debug", None) if callable(set_triton_debug): set_triton_debug(verbose_debug) if not enable_quanto_int8_kernel(triton_mod=triton_mod): _ensure_default_quanto_linear_patch() _startup_status(False, "failed to patch Quanto linear forward; using non-injected Quanto path.") return False _SHAPE_PROFILE_ON = _env_flag(_ENV_PROFILE_SHAPES, "0") _TIME_PROFILE_ON = _env_flag(_ENV_PROFILE_TIME, "0") _startup_status( True, ( "Triton int8 kernels will be used for Quanto qint8 linear layers " "(QBytes int8 activations + fused dynamic int8 activation quantization)." ), ) return True def _print_shape_profile() -> None: if not _SHAPE_PROFILE_ON and not _TIME_PROFILE_ON: return if _SHAPE_PROFILE_ON and _SHAPE_COUNTS_FUSED: top_fused = sorted(_SHAPE_COUNTS_FUSED.items(), key=lambda kv: kv[1], reverse=True)[:10] _log(f"Fused shape profile (top {len(top_fused)}): {top_fused}") if _SHAPE_PROFILE_ON and _SHAPE_COUNTS_SCALED: top_scaled = sorted(_SHAPE_COUNTS_SCALED.items(), key=lambda kv: kv[1], reverse=True)[:10] _log(f"Scaled shape profile (top {len(top_scaled)}): {top_scaled}") if _TIME_PROFILE_ON: total_ms = 0.0 calls = 0 if _TIME_PROFILE_EVENTS: if torch.cuda.is_available(): torch.cuda.synchronize() for start, end in _TIME_PROFILE_EVENTS: total_ms += float(start.elapsed_time(end)) calls = len(_TIME_PROFILE_EVENTS) else: total_ms = _TIME_PROFILE_CPU_MS calls = _TIME_PROFILE_CALLS _log(f"Triton kernel time profile: {total_ms / 1000.0:.3f}s over {calls} calls") atexit.register(_print_shape_profile)