# Copyright 2026 The HuggingFace 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. """DeepGEMM integration: fused grouped GEMM kernels from `kernels-community/deep-gemm`. Provides: - `deepgemm_bf16_experts_forward`: BF16 M-grouped experts forward. - `deepgemm_fp8_fp4_linear`: end-to-end FP8/FP4 linear (output dtype follows the input). - `deepgemm_fp8_fp4_experts_forward`: FP8 (or FP4 on SM100+) M-grouped experts forward. - `deepgemm_fp8_fp4_megamoe_experts_forward`: FP8xFP4 Mega MoE forward (SM100+). Requirements: CUDA, Hopper (SM90+), CUDA runtime ≥ 12.3, kernels-community/deep-gemm ≥ 2.5 (Mega MoE symbols required). Mega MoE additionally needs SM100+ at call time. """ from __future__ import annotations import functools import json import os import re import shutil from collections.abc import Callable from dataclasses import dataclass import torch from ..utils import logging from ..utils.deprecation import deprecate_kwarg from ..utils.import_utils import ( is_kernels_available, maybe_import_error, resolve_internal_import, ) from .hub_kernels import _MISSING_KERNELS_MESSAGE, lazy_load_kernel from .tensor_parallel import to_local logger = logging.get_logger(__name__) # ── Kernel loading ───────────────────────────────────────────────────────────── @dataclass(frozen=True) class DeepGEMM: """Curated entry points exposed by `kernels-community/deep-gemm`.""" fp8_fp4_matmul: Callable grouped_fp8_fp4_matmul_nt: Callable grouped_fp8_fp4_matmul_nn: Callable grouped_bf16_matmul_nt: Callable grouped_bf16_matmul_nn: Callable per_token_cast_to_fp8: Callable transform_sf_into_required_layout: Callable transform_weights_for_mega_moe: Callable get_symm_buffer_for_mega_moe: Callable fp8_fp4_mega_moe: Callable # M/K-dimension alignment for TMA-based contiguous grouped GEMM. Sourced from # `get_mk_alignment_for_contiguous_layout()` at load time. The kernel exposes a # `set_mk_alignment_for_contiguous_layout` setter, but we don't call it: the # build-time default (128) was empirically the best across MoE workloads # (bench showed kernel-recommended 240 is slower and 256 doesn't even compile). # Same stance as vLLM, which caches and never sets it. m_alignment: int @functools.cache def _get_cuda_home() -> str | None: """Resolve the CUDA toolkit root the way DeepGEMM's JIT does: ``CUDA_HOME`` → ``CUDA_PATH`` → dir of ``which nvcc`` → ``/usr/local/cuda`` (``None`` if none found). Mirrors DeepGEMM's own ``_find_cuda_home`` so we agree on the path it will actually use, rather than reusing ``torch.utils.cpp_extension.CUDA_HOME`` whose resolution inits a CUDA context (fork-unsafe). """ cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") if cuda_home: return cuda_home nvcc = shutil.which("nvcc") if nvcc: return os.path.dirname(os.path.dirname(nvcc)) if os.path.isdir("/usr/local/cuda"): return "/usr/local/cuda" return None @functools.cache def _get_nvcc_version() -> tuple[int, int] | None: """Version of the CUDA toolkit nvcc will use, as ``(major, minor)``, read off disk without a subprocess from (in order) ``{CUDA_HOME}/version.json``, ``version.txt``, or the ``CUDA_VERSION`` define in ``include/cuda.h``. ``None`` if unreadable. This is the compiler that builds the kernels, unlike ``torch.version.cuda`` (torch's bundled runtime, which never drives a JIT compile). """ cuda_home = _get_cuda_home() if cuda_home is None: return None version_json = os.path.join(cuda_home, "version.json") if os.path.isfile(version_json): try: with open(version_json) as f: components = json.load(f) version = components.get("cuda_nvcc", components.get("cuda", {})).get("version", "") major, minor = version.split(".")[:2] return int(major), int(minor) except (OSError, ValueError, AttributeError): pass version_txt = os.path.join(cuda_home, "version.txt") if os.path.isfile(version_txt): try: with open(version_txt) as f: match = re.search(r"CUDA Version (\d+)\.(\d+)", f.read()) if match: return int(match.group(1)), int(match.group(2)) except (OSError, ValueError): # ValueError covers UnicodeDecodeError on a non-text file pass # `cuda.h` ships with every toolkit (incl. distro packages that have no version file). cuda_h = os.path.join(cuda_home, "include", "cuda.h") if os.path.isfile(cuda_h): try: with open(cuda_h) as f: match = re.search(r"#define CUDA_VERSION (\d+)", f.read()) if match: cuda_version = int(match.group(1)) return cuda_version // 1000, (cuda_version % 1000) // 10 except (OSError, ValueError): # ValueError covers UnicodeDecodeError on a non-text file pass return None @torch._dynamo.assume_constant_result def is_sm100() -> bool: """Whether the current CUDA device is Blackwell (SM100+). `@assume_constant_result` makes dynamo evaluate it once at trace time and inline the bool — keeping the untraceable `torch.cuda.get_device_capability()` pybind out of the graph — so the branches on it (SF layout, cast kwargs, psum layout, and the arch guards) stay compile-safe. Re-queries each eager call, so a faked capability in tests is honoured. """ return torch.cuda.get_device_capability()[0] >= 10 @torch._dynamo.assume_constant_result def is_deepgemm_loadable(raise_error: bool = False) -> bool: """Whether the DeepGEMM kernel can be loaded in this environment: `kernels` installed, a CUDA GPU on a supported arch (Hopper SM90 or Blackwell SM100), and a CUDA toolkit/nvcc new enough to JIT-compile it. A one-glance gate for callers — including external stacks — deciding whether to dispatch to DeepGEMM. With `raise_error=True` (used by the loader) it re-raises the specific `ImportError` explaining what's missing instead of returning `False`. `@assume_constant_result` makes dynamo evaluate this once at trace time and inline the bool (its probe is untraceable — an `lru_cache`'d import check + the `get_device_capability` pybind + nvcc lookup), so the check can sit in a compiled dispatch gate (e.g. `fp8_linear`) without breaking the graph. FP4 / Mega MoE additionally need Blackwell (SM100); the relevant forwards enforce that via `_assert_sm100_requirements` / `is_sm100()`, since this env check can't see which dtype path the caller will take. """ if not is_kernels_available(): return maybe_import_error(f"DeepGEMM kernel unavailable: {_MISSING_KERNELS_MESSAGE}", raise_error=raise_error) if not torch.cuda.is_available(): return maybe_import_error("DeepGEMM kernel requires CUDA, but CUDA is not available.", raise_error=raise_error) major, minor = torch.cuda.get_device_capability() # DeepGEMM ships kernels only for SM90 (Hopper) and SM100 (Blackwell); anything # else — Ada (SM89), Ampere (SM80), or future archs (SM110+) — has no build. if major not in (9, 10): return maybe_import_error( f"DeepGEMM requires Hopper (SM90) or Blackwell (SM100); current device is SM{major}{minor}.", raise_error=raise_error, ) # DeepGEMM JIT-compiles kernels with the system nvcc, so a resolvable CUDA toolkit is required. # Per the DeepGEMM README: SM90 needs CUDA 12.3+, SM100 needs CUDA 12.9+. min_cuda = (12, 9) if major == 10 else (12, 3) cuda_home = _get_cuda_home() if cuda_home is None: return maybe_import_error( f"DeepGEMM's JIT needs a CUDA toolkit ≥ {min_cuda[0]}.{min_cuda[1]}, but none was found. " "Set `CUDA_HOME` to a CUDA toolkit.", raise_error=raise_error, ) # The Kernel Hub `deep-gemm` build always uses nvcc and ignores `DG_JIT_USE_NVRTC` (there is no # NVRTC fallback), so `CUDA_HOME` must hold an nvcc of the required version. if not os.path.isfile(os.path.join(cuda_home, "bin", "nvcc")): return maybe_import_error( f"DeepGEMM's JIT compiles with nvcc, but none was found in `{cuda_home}/bin`. Point " f"`CUDA_HOME` at a full CUDA ≥ {min_cuda[0]}.{min_cuda[1]} toolkit (not a runtime-only install).", raise_error=raise_error, ) # Treat an unreadable version as unsupported: `cuda.h` (with `CUDA_VERSION`) ships with every real # toolkit, so `None` here means an incomplete install we can't vouch for — fail early to Triton. nvcc_version = _get_nvcc_version() if nvcc_version is None: return maybe_import_error( f"DeepGEMM found nvcc in `{cuda_home}/bin` but could not read its CUDA version " f"(no parseable `version.json`, `version.txt`, or `include/cuda.h`). Point `CUDA_HOME` at a " f"complete CUDA ≥ {min_cuda[0]}.{min_cuda[1]} toolkit.", raise_error=raise_error, ) if nvcc_version < min_cuda: return maybe_import_error( f"DeepGEMM on SM{major}{minor} needs a CUDA ≥ {min_cuda[0]}.{min_cuda[1]} toolkit, but nvcc " f"{nvcc_version[0]}.{nvcc_version[1]} in `{cuda_home}` is too old. Point `CUDA_HOME` at a " f"CUDA ≥ {min_cuda[0]}.{min_cuda[1]} toolkit.", raise_error=raise_error, ) return True # Cache the loaded bundle but not failures: re-checking each call is cheap and intended, since the env # can change between attempts. A module global (not `@functools.cache`) avoids Dynamo warning about # tracing a cache-wrapped function on every compile. _DEEPGEMM: DeepGEMM | None = None @torch._dynamo.allow_in_graph def _load_deepgemm_kernel() -> None: """Load DeepGEMM once into the `_DEEPGEMM` module global, raising `ImportError` if the env or any required symbol is missing. Under NO circumstances may this function return a value: it rides through `@allow_in_graph` as an opaque fx node, whose return must be proxyable — returning the bundle (e.g. from the warm-cache short-circuit) breaks torch.compile with `Unsupported: torch.* op returned non-Tensor`. Callers read `_DEEPGEMM` back from the global. `@allow_in_graph` makes `torch.compile` treat the untraceable cold path (hub download + dynamic import via `lazy_load_kernel`) as a single opaque node instead of tracing into it; it returns `None` (proxyable) and populates the global, which `load_deepgemm_kernel` then returns. """ global _DEEPGEMM if _DEEPGEMM is not None: return is_deepgemm_loadable(raise_error=True) kernel = lazy_load_kernel("deep-gemm") if kernel is None: raise ImportError( "Failed to load `kernels-community/deep-gemm` — check that a build matches the current torch/CUDA." ) fp8_fp4_matmul = getattr(kernel, "fp8_fp4_gemm_nt", None) grouped_fp8_fp4_matmul_nt = getattr(kernel, "m_grouped_fp8_fp4_gemm_nt_contiguous", None) grouped_fp8_fp4_matmul_nn = getattr(kernel, "m_grouped_fp8_fp4_gemm_nn_contiguous", None) grouped_bf16_matmul_nt = getattr(kernel, "m_grouped_bf16_gemm_nt_contiguous", None) grouped_bf16_matmul_nn = getattr(kernel, "m_grouped_bf16_gemm_nn_contiguous", None) per_token_cast_to_fp8 = resolve_internal_import(kernel, chained_path="utils.per_token_cast_to_fp8") transform_sf_into_required_layout = getattr(kernel, "transform_sf_into_required_layout", None) transform_weights_for_mega_moe = getattr(kernel, "transform_weights_for_mega_moe", None) get_symm_buffer_for_mega_moe = getattr(kernel, "get_symm_buffer_for_mega_moe", None) get_mk_alignment = getattr(kernel, "get_mk_alignment_for_contiguous_layout", None) fp8_fp4_mega_moe = getattr(kernel, "fp8_fp4_mega_moe", None) missing = [ name for name, attr in [ ("fp8_fp4_gemm_nt", fp8_fp4_matmul), ("m_grouped_fp8_fp4_gemm_nt_contiguous", grouped_fp8_fp4_matmul_nt), ("m_grouped_fp8_fp4_gemm_nn_contiguous", grouped_fp8_fp4_matmul_nn), ("m_grouped_bf16_gemm_nt_contiguous", grouped_bf16_matmul_nt), ("m_grouped_bf16_gemm_nn_contiguous", grouped_bf16_matmul_nn), ("utils.per_token_cast_to_fp8", per_token_cast_to_fp8), ("transform_sf_into_required_layout", transform_sf_into_required_layout), ("transform_weights_for_mega_moe", transform_weights_for_mega_moe), ("get_symm_buffer_for_mega_moe", get_symm_buffer_for_mega_moe), ("get_mk_alignment_for_contiguous_layout", get_mk_alignment), ("fp8_fp4_mega_moe", fp8_fp4_mega_moe), ] if attr is None ] if missing: raise ImportError( f"DeepGEMM kernel is missing required symbols: {', '.join(missing)}. {_MISSING_KERNELS_MESSAGE}" ) _DEEPGEMM = DeepGEMM( fp8_fp4_matmul=fp8_fp4_matmul, grouped_fp8_fp4_matmul_nt=grouped_fp8_fp4_matmul_nt, grouped_fp8_fp4_matmul_nn=grouped_fp8_fp4_matmul_nn, grouped_bf16_matmul_nt=grouped_bf16_matmul_nt, grouped_bf16_matmul_nn=grouped_bf16_matmul_nn, per_token_cast_to_fp8=per_token_cast_to_fp8, transform_sf_into_required_layout=transform_sf_into_required_layout, transform_weights_for_mega_moe=transform_weights_for_mega_moe, get_symm_buffer_for_mega_moe=get_symm_buffer_for_mega_moe, fp8_fp4_mega_moe=fp8_fp4_mega_moe, m_alignment=get_mk_alignment(), ) def load_deepgemm_kernel() -> DeepGEMM: _load_deepgemm_kernel() return _DEEPGEMM # ── Scale-factor helpers ─────────────────────────────────────────────────────── def _assert_sm100_requirements(weight: torch.Tensor, scale: torch.Tensor) -> None: """Before-load guard for DeepGEMM's FP8/FP4 arch constraints on the given weight/scale dtypes: - FP4 (``int8``-packed) weights have no Hopper (SM90) kernel — they need Blackwell (SM100+). - Blackwell has no float32 scale-factor path: ``_coerce_sf_for_kernel`` would silently round a ``float32`` scale to UE8M0 and corrupt the output. UE8M0 scales load as ``float8_e8m0fnu`` (the loader normalizes even float32-container checkpoints like dsv4-flash-base), so a ``float32`` scale on SM100 means a genuine non-UE8M0 checkpoint. Uses `is_sm100()` (compile-safe via `assume_constant_result`), so the whole guard folds to a constant under ``torch.compile``: the valid case compiles away to nothing, while an unsupported combo fails loud rather than letting the hot path silently corrupt (unlike an ``is_compiling`` skip, which would miss a model compiled from cold with no eager warmup). Both raise ``NotImplementedError``, which ``fp8_linear`` treats as "DeepGEMM declined" and falls back to Triton (SM90 consuming float32 SFs directly is fine, so those cases are no-ops). """ if not is_sm100(): # SM90: DeepGEMM has no FP4 (int8-packed) kernel, but consumes float32 SFs directly. if weight.dtype == torch.int8: raise NotImplementedError( "DeepGEMM's FP4 (int8-packed) path requires a Blackwell (SM100+) GPU; FP4 weights have no " "Hopper (SM90) kernel. Use an FP8 checkpoint, or run on a Blackwell GPU." ) return # SM100: DeepGEMM has no float32 scale-factor path. if scale.dtype == torch.float32: raise NotImplementedError( "DeepGEMM has no float32 scale-factor path on Blackwell (SM100): these scales are plain float32 " "(quantization_config.scale_fmt='float'), and rounding them to UE8M0 would silently corrupt the " "output. Use a checkpoint quantized with scale_fmt='ue8m0', or a path that consumes float32 block " "scales directly — the FP8 linear falls back to Triton automatically; for experts use " "`model.set_experts_implementation('grouped_mm')`." ) def _ceil_to_ue8m0(sf: torch.Tensor) -> torch.Tensor: """Round each fp32 SF up to the nearest power of 2 (zero mantissa). Mirrors `deep_gemm.utils.math.ceil_to_ue8m0`. On SM100 the kernel's `pack_fp32_into_ue8m0` cleanly extracts the biased exponent only when the mantissa is already zero — its inner shifts (`>> 15`, `>> 7`, `<< 1`) otherwise leak mantissa bits into adjacent UE8M0 byte slots and silently corrupt the SF. SM90 consumes raw fp32 SFs without going through this path. """ int_view = sf.view(torch.int32) return (int_view + ((1 << 23) - 1)).bitwise_and_(~((1 << 23) - 1)).view(torch.float) def _coerce_sf_for_kernel(sf: torch.Tensor, is_sm100: bool, expected_mn: int | None = None) -> torch.Tensor: """Lay out `sf` as DeepGEMM's dispatch expects, per arch. On SM100 the int-SF path only *checks* the SF (`tma_stride_check`) and never transforms it, so we hand it a TMA-aligned MN-major layout (`stride(-2) == 1`, `stride(-1) == align(mn, 16/esize)`). On SM90 DeepGEMM transforms SFA itself (`get_mn_major_tma_aligned_tensor`) and only *checks* SFB against `sm90_sfb_check`, which rejects TMA padding (`stride(-1)` must equal `size(-2)`, not `align(mn, …)`); a padded weight SF trips `layout.hpp` whenever `mn` isn't a multiple of `16/esize` (e.g. N=576 → mn=5). So on SM90 we return the raw row-major SF and let DeepGEMM lay it out. Inputs come in three flavors: - `float8_e8m0fnu` on SM100: raw UE8M0 bytes — pack 4 K-bytes → int32 (last dim /4) for the kernel's `(INT, 1, gran_k)` path. - `float8_e8m0fnu` on SM90: SM90 dispatch only accepts FP32 SFs, so cast UE8M0 → FP32 (exact upcast — UE8M0 is the biased-exponent half of a pow-of-2 FP32, so `.float()` rebuilds the original FP32 scale exactly). - `float32`: per-token / per-block SFs from `per_token_cast_to_fp8` or on-disk weights — round to UE8M0 on SM100 (see `_ceil_to_ue8m0`). - `int32`: already-packed UE8M0 — pass through. When `expected_mn` is set and the SF's M-dim is smaller (block-quantized UE8M0, e.g. DSv4-Flash compressor weights with `(N/128, K/128)` SFs), we repeat the SF on the M-axis to per-row before packing — the `(INT, 1, gran_k)` DeepGEMM kernel branch is the only UE8M0 path on SM100; for `gran_mn > 1` the kernel only handles FP32 SFs and would otherwise reject our INT SF here. """ if sf.dtype == torch.float8_e8m0fnu: if expected_mn is not None and sf.size(-2) < expected_mn: gran_mn = expected_mn // sf.size(-2) sf = sf.repeat_interleave(gran_mn, dim=-2) if is_sm100: sf = sf.contiguous().view(torch.int32) else: sf = sf.float() elif sf.dtype == torch.float32 and is_sm100: sf = _ceil_to_ue8m0(sf) if sf.dim() not in (2, 3): raise ValueError(f"DeepGEMM SF must be 2D or 3D, got {sf.dim()}D") # SM90 dispatch transforms SFA and only checks SFB (`sm90_sfb_check`), which needs # an unpadded contiguous layout — DeepGEMM does the MN-major alignment itself. if not is_sm100: return sf.contiguous() mn = sf.size(-2) kf = sf.size(-1) align_to = 16 // sf.element_size() # `get_tma_aligned_size`: align(mn, 16 / element_size) aligned_mn = -(-mn // align_to) * align_to target_strides = (1, aligned_mn) if sf.dim() == 2 else (kf * aligned_mn, 1, aligned_mn) if tuple(sf.stride()) == target_strides: return sf out = torch.empty_strided(sf.shape, target_strides, dtype=sf.dtype, device=sf.device) out.copy_(sf) return out def _select_fp8_cast_kwargs( weight: torch.Tensor, weight_scale_inv: torch.Tensor, block_size: tuple | None, is_sm100: bool ) -> dict: """Pick the `per_token_cast_to_fp8` kwargs from weight dtype + SF dtype + arch. Cases mirror the kernel's recipes: - FP4 weights (`int8`): gran_k=32 packed-UE8M0 SF. SM100+ only. - FP8 weights + UE8M0 SF on SM100: gran_k=128 packed-UE8M0 SF (DSv4). - FP8 weights + UE8M0 SF on SM90: gran_k=128 FP32 SF — the SM90 dispatch in `layout.hpp` only matches FP32 SFs, so we keep act SFs as FP32 (and float the weight SF in `_coerce_sf_for_kernel`; UE8M0 → FP32 is an exact upcast). - FP8 weights + float SF: gran_k=128 float SF (DSv3). """ if weight.dtype == torch.int8: # FP4 return {"use_ue8m0": True, "gran_k": 32, "use_packed_ue8m0": True} # FP8 weights: validate block_size (informational; kernel infers recipe from SF dtype/shape). if block_size is None: raise ValueError( "DeepGEMM requires block-wise quantized FP8 weights, but the experts have no `block_size` set." ) block_size = tuple(block_size) if block_size not in ((128, 128), (1, 128)): raise ValueError(f"DeepGEMM requires `block_size` ∈ {{(128, 128), (1, 128)}}, got {block_size}.") if weight_scale_inv.dtype == torch.float8_e8m0fnu and is_sm100: return {"use_ue8m0": True, "gran_k": 128, "use_packed_ue8m0": True} return {"use_ue8m0": False, "gran_k": 128} # ── Layout helpers (M-grouped contiguous, TMA-aligned) ───────────────────────── def _build_deepgemm_contiguous_layout( expert_ids_sorted: torch.Tensor, num_experts: int, alignment: int, use_psum_layout: bool ) -> tuple[torch.Tensor, torch.Tensor, int]: """Build the TMA-aligned grouped layout DeepGEMM expects. Returns `(sorted_to_padded, grouped_layout, total_padded_rows)`: - `grouped_layout` is per-row expert id (Hopper, with `-1` for padding / sentinels) or a cumsum of aligned per-expert counts (Blackwell). - EP sentinels (values == `num_experts`) are routed past the last expert block so DeepGEMM skips them. """ device = expert_ids_sorted.device num_tokens = expert_ids_sorted.size(0) # `histc` drops values > max, so EP sentinels (== num_experts) don't count. # CPU requires float input, CUDA requires int input (deterministic mode). histc_input = expert_ids_sorted.float() if device.type == "cpu" else expert_ids_sorted.int() tokens_per_expert = torch.histc(histc_input, bins=num_experts, min=0, max=num_experts - 1).long() aligned_tokens_per_expert = ((tokens_per_expert + alignment - 1) // alignment) * alignment # Upper bound — avoids GPU→CPU sync; padding rows are skipped. total_padded_rows = num_tokens + min(num_tokens, num_experts) * (alignment - 1) # Exclusive cumsum of per-expert padding (index `num_experts` = total padding, # which routes EP sentinels past all aligned blocks on Blackwell). padding_per_expert = aligned_tokens_per_expert - tokens_per_expert cumulative_padding = torch.nn.functional.pad(padding_per_expert.cumsum(0), (1, 0)) sorted_to_padded = torch.arange(num_tokens, device=device) + cumulative_padding[expert_ids_sorted] if use_psum_layout: # SM100+: kernel reads cumsum of aligned counts as expert boundaries. grouped_layout = aligned_tokens_per_expert.cumsum(0).int() else: # SM90: per-row expert id, -1 = skip (padding & sentinels). grouped_layout = torch.full((total_padded_rows,), -1, device=device, dtype=torch.int32) grouped_layout[sorted_to_padded] = torch.where(expert_ids_sorted < num_experts, expert_ids_sorted.int(), -1) return sorted_to_padded, grouped_layout, total_padded_rows def _pad_for_deepgemm(x: torch.Tensor, sorted_to_padded: torch.Tensor, total_padded_rows: int) -> torch.Tensor: """Pad a sorted tensor into the TMA-aligned contiguous layout.""" padded = torch.empty(total_padded_rows, *x.shape[1:], device=x.device, dtype=x.dtype) padded[sorted_to_padded] = x return padded def _unpad_from_deepgemm_contiguous_layout(x_padded: torch.Tensor, sorted_to_padded: torch.Tensor) -> torch.Tensor: return x_padded[sorted_to_padded] # ── Routing helpers (sort → matmul → restore) ───────────────────────────────── def _dispatch_routed_input( hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor, num_experts: int, m_alignment: int, use_psum_layout: bool, ) -> tuple: """Sort tokens by expert id and build the M-grouped padded layout. Returns `(sorted_hidden_states_g, sample_weights_g, expert_ids_g, sentinel_mask, perm, sorted_to_padded, grouped_layout, total_padded_rows)`. """ # S is the number of selected token-expert pairs (S = num_tokens * num_top_k) num_top_k = top_k_index.size(-1) expert_ids = top_k_index.reshape(-1) # (S,) sample_weights = top_k_weights.reshape(-1) # (S,) # Sort by expert for grouped processing expert_ids_g, perm = torch.sort(expert_ids) sorted_hidden_states_g = hidden_states[perm // num_top_k] sample_weights_g = sample_weights[perm] # Build the M-grouped padded layout (DeepGEMM contract: each expert's rows # start on the kernel's M-alignment boundary, sentinels routed past valid # expert blocks). sorted_to_padded, grouped_layout, total_padded_rows = _build_deepgemm_contiguous_layout( expert_ids_g, num_experts, m_alignment, use_psum_layout ) # EP sentinel mask is captured before the in-place clamp; used by the post-mask in # `_combine_routed_output` to zero sentinel rows before the per-token reduction. The clamp # keeps any per-row gather (e.g. bias) in-bounds — bias added at sentinel positions falls # in rows the kernel skips, so harmless. Safe to mutate now: the layout was built from the # unclamped tensor and nothing downstream needs the sentinel info from `expert_ids_g` itself. sentinel_mask = (expert_ids_g >= num_experts).unsqueeze(-1) expert_ids_g.clamp_(max=num_experts - 1) return ( sorted_hidden_states_g, sample_weights_g, expert_ids_g, sentinel_mask, perm, sorted_to_padded, grouped_layout, total_padded_rows, ) def _combine_routed_output( out_padded: torch.Tensor, sorted_weights: torch.Tensor, sentinel_mask: torch.Tensor, perm: torch.Tensor, sorted_to_padded: torch.Tensor, num_tokens: int, num_top_k: int, hidden_dim: int, out_dtype: torch.dtype, ) -> torch.Tensor: """Unpad → weighted multiply → mask sentinels → restore order → top-k reduce.""" out = _unpad_from_deepgemm_contiguous_layout(out_padded, sorted_to_padded) weighted = out * sorted_weights.to(out.dtype).unsqueeze(-1) # Sentinel rows past the valid expert blocks may carry NaN from allocator # reuse (`0 * NaN = NaN`); zero them so the top-k reduction stays finite. weighted.masked_fill_(sentinel_mask, 0.0) inv_perm = torch.empty_like(perm) inv_perm[perm] = torch.arange(perm.size(0), device=out.device) # Deterministic reshape+sum (index_add_ with duplicates is non-deterministic on CUDA). return weighted[inv_perm].view(num_tokens, num_top_k, hidden_dim).sum(dim=1).to(out_dtype) # ── Public dispatches ────────────────────────────────────────────────────────── @deprecate_kwarg("output_dtype", version="v5.16") def deepgemm_fp8_fp4_linear( input: torch.Tensor, weight: torch.Tensor, weight_scale_inv: torch.Tensor, bias: torch.Tensor | None = None, block_size: tuple[int, int] | None = None, output_dtype: torch.dtype | None = None, activation_scale: torch.Tensor | None = None, ) -> torch.Tensor: """End-to-end DeepGEMM linear: per-token activation quant + FP8/FP4 matmul. Static (per-tensor) activation quantization is rejected — DeepGEMM needs per-row SFs. Callers should route static activations through the Triton fallback. """ if activation_scale is not None: raise NotImplementedError("DeepGEMM linear does not support static activation quantization.") if input.dtype not in (torch.bfloat16, torch.float16): raise ValueError(f"DeepGEMM linear requires FP16 or BF16 activations, got {input.dtype}") # Fail before the (hub-download + JIT) load if this device can't serve these dtypes. _assert_sm100_requirements(weight, weight_scale_inv) deepgemm = load_deepgemm_kernel() cast_kwargs = _select_fp8_cast_kwargs(weight, weight_scale_inv, block_size, is_sm100()) input_2d = input.view(-1, input.shape[-1]) qinput_2d, scale_2d = deepgemm.per_token_cast_to_fp8(input_2d, **cast_kwargs) output = torch.empty(qinput_2d.shape[0], weight.shape[0], device=input.device, dtype=input.dtype) # Pass `(1, 1, gran_k)` for int-SF paths so the kernel uses the right K granularity # (the default `(1, 1, 128)` mismatches FP4's gran_k=32). Float-SF leaves it None. sf_recipe = (1, 1, cast_kwargs["gran_k"]) if cast_kwargs.get("use_packed_ue8m0") else None deepgemm.fp8_fp4_matmul( (qinput_2d, _coerce_sf_for_kernel(scale_2d, is_sm100(), expected_mn=qinput_2d.size(0))), (weight, _coerce_sf_for_kernel(weight_scale_inv, is_sm100(), expected_mn=weight.size(0))), output, recipe=sf_recipe, ) output = output.view(input.shape[:-1] + (weight.shape[0],)) if bias is not None: output.add_(bias) return output def deepgemm_bf16_experts_forward( self: torch.nn.Module, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor, ) -> torch.Tensor: if hidden_states.dtype != torch.bfloat16: raise ValueError(f"DeepGEMM experts path requires bfloat16 hidden states, got {hidden_states.dtype}") deepgemm = load_deepgemm_kernel() # Non-transposed weights (E, N, K) → NT kernel; transposed (E, K, N) → NN kernel. grouped_bf16_matmul = deepgemm.grouped_bf16_matmul_nn if self.is_transposed else deepgemm.grouped_bf16_matmul_nt device = hidden_states.device num_top_k = top_k_index.size(-1) num_tokens = hidden_states.size(0) hidden_dim = hidden_states.size(-1) ( sorted_hidden, sorted_weights, expert_ids_g, sentinel_mask, perm, sorted_to_padded, grouped_layout, total_padded_rows, ) = _dispatch_routed_input( hidden_states, top_k_index, top_k_weights, self.num_experts, deepgemm.m_alignment, is_sm100() ) weight_up = to_local(self.gate_up_proj if self.has_gate else self.up_proj) weight_down = to_local(self.down_proj) up_bias = to_local(self.gate_up_proj_bias if self.has_gate else self.up_proj_bias) if self.has_bias else None down_bias = to_local(self.down_proj_bias) if self.has_bias else None # Up projection. up_out_dim = weight_up.shape[-1] if self.is_transposed else weight_up.shape[1] act = _pad_for_deepgemm(sorted_hidden, sorted_to_padded, total_padded_rows) proj_out = torch.empty(total_padded_rows, up_out_dim, device=device, dtype=hidden_states.dtype) grouped_bf16_matmul(act, weight_up, proj_out, grouped_layout, use_psum_layout=is_sm100()) if self.has_bias: proj_out.index_add_(0, sorted_to_padded, up_bias[expert_ids_g]) proj_out = self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out) # Down projection. out = torch.empty(total_padded_rows, hidden_dim, device=device, dtype=hidden_states.dtype) grouped_bf16_matmul(proj_out, weight_down, out, grouped_layout, use_psum_layout=is_sm100()) if self.has_bias: out.index_add_(0, sorted_to_padded, down_bias[expert_ids_g]) return _combine_routed_output( out, sorted_weights, sentinel_mask, perm, sorted_to_padded, num_tokens, num_top_k, hidden_dim, hidden_states.dtype, ) def deepgemm_fp8_fp4_experts_forward( self: torch.nn.Module, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor, ) -> torch.Tensor: if self._deepgemm_disabled: # Set at load when the model spans >1 CUDA device in this process, where DeepGEMM's # context-bound kernels corrupt across devices (see `quantizer_finegrained_fp8.py`). raise RuntimeError( "DeepGEMM experts selected on a model spanning multiple CUDA devices in one process; " "its kernels are bound to a single CUDA context and corrupt across devices. Use " "`experts_implementation='grouped_mm'`, or run one device per process (TP/EP)." ) # Fail before the (hub-download + JIT) load if this device can't serve these dtypes. _assert_sm100_requirements(self.down_proj, self.down_proj_scale_inv) deepgemm = load_deepgemm_kernel() if self.activation_scheme == "static": raise NotImplementedError("DeepGEMM experts dispatch does not support static activation quantization.") if hidden_states.dtype != torch.bfloat16: raise ValueError(f"DeepGEMM experts path requires bfloat16 hidden states, got {hidden_states.dtype}") grouped_fp8_fp4_matmul = ( deepgemm.grouped_fp8_fp4_matmul_nn if self.is_transposed else deepgemm.grouped_fp8_fp4_matmul_nt ) device = hidden_states.device num_top_k = top_k_index.size(-1) num_tokens = hidden_states.size(0) hidden_dim = hidden_states.size(-1) weight_up = to_local(self.gate_up_proj if self.has_gate else self.up_proj) weight_scale_up = to_local(self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv) weight_down = to_local(self.down_proj) weight_scale_down = to_local(self.down_proj_scale_inv) cast_kwargs = _select_fp8_cast_kwargs(weight_up, weight_scale_up, self.block_size, is_sm100()) ( sorted_hidden, sorted_weights, _expert_ids_g, sentinel_mask, perm, sorted_to_padded, grouped_layout, total_padded_rows, ) = _dispatch_routed_input( hidden_states, top_k_index, top_k_weights, self.num_experts, deepgemm.m_alignment, is_sm100() ) sf_recipe = (1, 1, cast_kwargs["gran_k"]) if cast_kwargs.get("use_packed_ue8m0") else None # Up projection. act_fp8, act_scales = deepgemm.per_token_cast_to_fp8(sorted_hidden, **cast_kwargs) act_fp8 = _pad_for_deepgemm(act_fp8, sorted_to_padded, total_padded_rows) act_scales = _pad_for_deepgemm(act_scales, sorted_to_padded, total_padded_rows) proj_out = torch.empty(total_padded_rows, weight_up.shape[1], device=device, dtype=torch.bfloat16) grouped_fp8_fp4_matmul( (act_fp8, _coerce_sf_for_kernel(act_scales, is_sm100(), expected_mn=total_padded_rows)), (weight_up, _coerce_sf_for_kernel(weight_scale_up, is_sm100(), expected_mn=weight_up.size(-2))), proj_out, grouped_layout, recipe=sf_recipe, use_psum_layout=is_sm100(), ) proj_out = self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out) # Down projection. proj_fp8, proj_scales = deepgemm.per_token_cast_to_fp8(proj_out, **cast_kwargs) out = torch.empty(total_padded_rows, hidden_dim, device=device, dtype=torch.bfloat16) grouped_fp8_fp4_matmul( (proj_fp8, _coerce_sf_for_kernel(proj_scales, is_sm100(), expected_mn=total_padded_rows)), (weight_down, _coerce_sf_for_kernel(weight_scale_down, is_sm100(), expected_mn=weight_down.size(-2))), out, grouped_layout, recipe=sf_recipe, use_psum_layout=is_sm100(), ) return _combine_routed_output( out, sorted_weights, sentinel_mask, perm, sorted_to_padded, num_tokens, num_top_k, hidden_dim, hidden_states.dtype, ) def setup_megamoe_weights(module: torch.nn.Module) -> None: """One-shot pack + permute of an FP8Experts module's L1/L2 weights into the Mega MoE UTCCP layout. Called lazily on the first megamoe forward; idempotent via the caller's ``_megamoe_transformed`` flag. Steps: 1. Cast UE8M0 SF → FP32 and call ``transform_sf_into_required_layout`` → packed int32 in MN-major TMA-aligned layout. 2. Run ``transform_weights_for_mega_moe``: interleaves gate/up on L1 and transposes both SFs for UTCCP. 3. Overwrite the loader-side parameters in place; the interleave preserves the ``[E_local, 2*I, *]`` leading dims so downstream ``.size(...)`` reads stay valid. Unwraps any ``DTensor`` wrappers FSDP2/EP may have placed around the loader- side Parameters — the kernel takes raw pointers. """ deepgemm = load_deepgemm_kernel() gate_up_sf_raw = to_local(module.gate_up_proj_scale_inv.data) down_sf_raw = to_local(module.down_proj_scale_inv.data) # Force int8 view: the kernel's interleave reshape/empty_like/copy_ is bit-level. gate_up_w = to_local(module.gate_up_proj.data).view(torch.int8).contiguous() down_w = to_local(module.down_proj.data).view(torch.int8).contiguous() intermediate_hidden = module.intermediate_dim num_local_experts = module.num_experts hidden_dim = module.hidden_dim if hidden_dim % 32 != 0 or intermediate_hidden % 32 != 0: raise ValueError( f"DeepGEMM Mega MoE requires `hidden_dim` and `intermediate_hidden` divisible by 32 " f"(FP8 SF granularity); got hidden_dim={hidden_dim}, intermediate_hidden={intermediate_hidden}." ) gate_up_sf = deepgemm.transform_sf_into_required_layout( gate_up_sf_raw.float(), 2 * intermediate_hidden, hidden_dim, recipe=(1, 32), num_groups=num_local_experts, ) down_sf = deepgemm.transform_sf_into_required_layout( down_sf_raw.float(), hidden_dim, intermediate_hidden, recipe=(1, 32), num_groups=num_local_experts, ) (gate_up, gate_up_sf), (down, down_sf) = deepgemm.transform_weights_for_mega_moe( (gate_up_w, gate_up_sf), (down_w, down_sf), ) module.gate_up_proj = torch.nn.Parameter(gate_up, requires_grad=False) module.gate_up_proj_scale_inv = torch.nn.Parameter(gate_up_sf, requires_grad=False) module.down_proj = torch.nn.Parameter(down, requires_grad=False) module.down_proj_scale_inv = torch.nn.Parameter(down_sf, requires_grad=False) def deepgemm_fp8_fp4_megamoe_experts_forward( self: torch.nn.Module, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor, process_group: torch.distributed.ProcessGroup | None = None, ) -> torch.Tensor: """FP8 acts × FP4 weights Mega MoE forward (SM100+). Fuses EP dispatch + L1 + SwiGLU + L2 + EP combine into one kernel, overlapping NVLink with tensor-core compute. The kernel handles the full `(num_tokens, hidden) → (num_tokens, hidden)` MoE forward including the weighted top-k reduction; the caller must NOT all-reduce the output. `process_group` is supplied automatically by `MoeTensorParalellExperts._prepare_input_fn` when the module is wrapped for TP — it's required for the symm-buffer rendezvous on first forward. `top_k_index` is GLOBAL expert ids (`-1` marks skipped slots). Caller-managed `self` attributes: - `gate_up_proj`, `gate_up_proj_scale_inv`: L1 weight + UE8M0 SF. - `down_proj`, `down_proj_scale_inv`: L2 weight + UE8M0 SF. Both pairs must be transformed together via `transform_weights_for_mega_moe((gate_up, gate_up_sf), (down, down_sf))`. - `config.swiglu_limit` (optional): SwiGLU clamp; absent → unclamped. """ # Fail before the (hub-download + JIT) load if this device can't serve these dtypes. Mega MoE is # Blackwell-only, and its weights are always FP4 (int8) — so the FP4 arch check doubles as the # SM100 gate; the explicit `!= int8` check below covers a non-FP4 (misconfigured) checkpoint. _assert_sm100_requirements(self.gate_up_proj, self.down_proj_scale_inv) if self.gate_up_proj.dtype != torch.int8: raise NotImplementedError( f"DeepGEMM Mega MoE requires FP4-packed expert weights (dtype=`int8`), got " f"`{self.gate_up_proj.dtype}`. Use the 'deepgemm' dispatch for FP8 experts." ) if process_group is None: raise ValueError( "DeepGEMM Mega MoE requires a `process_group` for the EP group. The TP wrapping " "(MoeTensorParalellMegaMoeExperts) supplies it automatically; pass it explicitly otherwise." ) deepgemm = load_deepgemm_kernel() # First-forward one-shot: pack UE8M0 SFs and interleave the L1/L2 weights for UTCCP. # Kept lazy here (instead of in a quantizer load-time hook) so the megamoe-specific # setup lives alongside the megamoe forward — `set_experts_implementation` refuses # to flip in/out of `deepgemm_megamoe` at runtime, so the flag won't go stale. if not getattr(self, "_megamoe_transformed", False): setup_megamoe_weights(self) self._megamoe_transformed = True num_top_k = top_k_index.size(-1) num_tokens = hidden_states.size(0) hidden_dim = hidden_states.size(-1) num_local_experts = self.gate_up_proj.size(0) intermediate_hidden = self.gate_up_proj.size(1) // 2 num_global_experts = num_local_experts * process_group.size() # Lazily (re)allocate the symmetric buffer when the cached one is too small. if getattr(self, "symm_buffer", None) is None or self.symm_buffer.num_max_tokens_per_rank < num_tokens: self.symm_buffer = deepgemm.get_symm_buffer_for_mega_moe( process_group, hidden=hidden_dim, num_topk=num_top_k, num_experts=num_global_experts, num_max_tokens_per_rank=num_tokens, intermediate_hidden=intermediate_hidden, ) x_fp8, x_sf = deepgemm.per_token_cast_to_fp8(hidden_states, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) self.symm_buffer.x[:num_tokens].copy_(x_fp8) self.symm_buffer.x_sf[:num_tokens].copy_(x_sf) self.symm_buffer.topk_idx[:num_tokens].copy_(top_k_index) self.symm_buffer.topk_weights[:num_tokens].copy_(top_k_weights) # `activation_clamp` must match `_apply_gate`'s clamp on the regular path so the kernel's # fused SwiGLU sees the same value range the model was calibrated for. y = torch.empty((num_tokens, hidden_dim), dtype=torch.bfloat16, device=hidden_states.device) deepgemm.fp8_fp4_mega_moe( y, (self.gate_up_proj, self.gate_up_proj_scale_inv), (self.down_proj, self.down_proj_scale_inv), self.symm_buffer, activation_clamp=getattr(getattr(self, "config", None), "swiglu_limit", None), ) return y.to(hidden_states.dtype)