# 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 ( KERNELS_MAX_VERSION, KERNELS_MIN_VERSION, is_kernels_available, is_torchdynamo_compiling, resolve_internal_import, ) from .hub_kernels import 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 @functools.cache def _load_deepgemm_kernel(requires_sm100: bool = False) -> DeepGEMM | str: """Load DeepGEMM once or returns an error message if env or any required symbol is missing. This is wrapped in a function that will raise an `ImportError` with the error message. The reason we raise in the wrapper rather than here is that @functools.cache will only cache a return value, not an exception. `requires_sm100` raises a Blackwell-specific error for callers (FP4 / Mega MoE) that won't work on Hopper, instead of the generic SM90+ message. """ if not is_torchdynamo_compiling(): if not is_kernels_available(): return ( "DeepGEMM kernel requires the `kernels` package. Please install a compatible version (" f"{KERNELS_MIN_VERSION} <= version < {KERNELS_MAX_VERSION}), e.g. `pip install kernels==" f"{KERNELS_MIN_VERSION}`" ) if not torch.cuda.is_available(): return "DeepGEMM kernel requires CUDA, but CUDA is not available." 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. allowed = (10,) if requires_sm100 else (9, 10) if major not in allowed: arch = "Blackwell (SM100)" if requires_sm100 else "Hopper (SM90) or Blackwell (SM100)" return f"DeepGEMM requires {arch}; current device is SM{major}{minor}." # 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 ( 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." ) # 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 ( 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)." ) # 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 ( 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." ) if nvcc_version < min_cuda: return ( 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." ) kernel = lazy_load_kernel("deep-gemm") if kernel is None: return "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: return ( f"DeepGEMM kernel is missing required symbols: {', '.join(missing)}. " f"Please install a compatible version ({KERNELS_MIN_VERSION} <= version < {KERNELS_MAX_VERSION}), " f"e.g. `pip install kernels=={KERNELS_MIN_VERSION}`" ) return 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(), ) @torch._dynamo.allow_in_graph def _populate_deepgemm_kernel(requires_sm100: bool = False) -> None: """Warm the `_load_deepgemm_kernel` cache from an opaque graph node, so Dynamo never traces the loader. Under `torch.compile`, Dynamo ignores `@functools.cache` and traces into `_load_deepgemm_kernel`, whose cold path (hub download + dynamic import via `lazy_load_kernel`) is untraceable and errors under `fullgraph`. `@allow_in_graph` turns the call into an opaque fx node instead — but an fx node's return must be proxyable, and the `DeepGEMM` bundle of Python callables isn't (`Unsupported: torch.* op returned non-Tensor`), so we can't just decorate the real loader. Hence two loaders: this one is opaque, returns `None`, and only warms the cache; the real `_load_deepgemm_kernel` right after is then a plain cache lookup. """ _load_deepgemm_kernel(requires_sm100=requires_sm100) def load_deepgemm_kernel(requires_sm100: bool = False) -> DeepGEMM: _populate_deepgemm_kernel(requires_sm100=requires_sm100) deepgemm_or_error = _load_deepgemm_kernel(requires_sm100=requires_sm100) if isinstance(deepgemm_or_error, str): raise ImportError(deepgemm_or_error) return deepgemm_or_error # ── Scale-factor helpers ─────────────────────────────────────────────────────── @functools.cache def _is_sm100(device: torch.device) -> bool: """``True`` for Blackwell (SM100+). Cached: device capability is fixed for the process lifetime and this gets hit on every linear/expert forward. """ return torch.cuda.get_device_capability(device)[0] >= 10 def _assert_sm100_scales_are_ue8m0(scale: torch.Tensor) -> None: """On B200 (SM100) DeepGEMM only supports UE8M0 (power-of-two) scales; the float32 scales that work on H100 (SM90) have no SM100 path. UE8M0 scales load as ``float8_e8m0fnu`` (the loader normalizes even float32-container checkpoints like dsv4-flash-base), so a plain ``float32`` scale here means a genuine non-UE8M0 checkpoint — fail loud rather than let ``_coerce_sf_for_kernel`` silently round it and corrupt the output. """ if not _is_sm100(scale.device): return # SM90 consumes float32 SFs directly (no UE8M0 round). if scale.dtype != torch.float32: return # already UE8M0 (`float8_e8m0fnu`) — kernel-ready as-is. raise ValueError( "DeepGEMM's Blackwell (SM100) experts kernel requires power-of-two (UE8M0) scale " "factors, but this checkpoint's expert scales are plain float32 " "(quantization_config.scale_fmt='float'). Rounding them to UE8M0 would scale the " "dequantized expert weights incorrectly and silently corrupt the output. Use a " "checkpoint quantized with scale_fmt='ue8m0', or an experts implementation that " "consumes float32 block scales directly, e.g. " "`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, 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. """ is_sm100 = _is_sm100(sf.device) 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. tokens_per_expert = torch.histc(expert_ids_sorted.int(), 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}") deepgemm = load_deepgemm_kernel(requires_sm100=weight.dtype == torch.int8) cast_kwargs = _select_fp8_cast_kwargs(weight, weight_scale_inv, block_size, _is_sm100(input.device)) 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, expected_mn=qinput_2d.size(0))), (weight, _coerce_sf_for_kernel(weight_scale_inv, 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(device) ) 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(device)) 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(device)) 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)." ) _assert_sm100_scales_are_ue8m0(self.down_proj_scale_inv) 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}") deepgemm = load_deepgemm_kernel(requires_sm100=self.down_proj.dtype == torch.int8) 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(device)) ( 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(device) ) 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, expected_mn=total_padded_rows)), (weight_up, _coerce_sf_for_kernel(weight_scale_up, expected_mn=weight_up.size(-2))), proj_out, grouped_layout, recipe=sf_recipe, use_psum_layout=_is_sm100(device), ) 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, expected_mn=total_padded_rows)), (weight_down, _coerce_sf_for_kernel(weight_scale_down, expected_mn=weight_down.size(-2))), out, grouped_layout, recipe=sf_recipe, use_psum_layout=_is_sm100(device), ) 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(requires_sm100=True) 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. """ _assert_sm100_scales_are_ue8m0(self.down_proj_scale_inv) if self.gate_up_proj.dtype != torch.int8: raise RuntimeError( 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(requires_sm100=True) # 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)