| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Compatibility wrapper for DeepGEMM API changes. |
| 4 | |
| 5 | Users of vLLM should always import **only** these wrappers. |
| 6 | """ |
| 7 | |
| 8 | import contextlib |
| 9 | import functools |
| 10 | import importlib |
| 11 | import os |
| 12 | from collections.abc import Callable |
| 13 | from enum import Enum |
| 14 | from typing import Any, NoReturn |
| 15 | |
| 16 | import torch |
| 17 | |
| 18 | import vllm.envs as envs |
| 19 | from vllm.logger import logger |
| 20 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 21 | get_fp8_min_max, |
| 22 | ) |
| 23 | from vllm.platforms import current_platform |
| 24 | from vllm.utils.import_utils import has_deep_gemm |
| 25 | from vllm.utils.math_utils import cdiv |
| 26 | |
| 27 | _DEEPGEMM_BLACKWELL_EXCLUDED_MODEL_TYPES: set[str] = { |
| 28 | "qwen3_5_text", |
| 29 | "qwen3_5_moe_text", |
| 30 | } |
| 31 | |
| 32 | # KV page sizes (in cache entries) supported by the paged-MQA logits kernels |
| 33 | # (fp8_fp4_paged_mqa_logits / get_paged_mqa_logits_metadata). Larger storage |
| 34 | # blocks must be virtually split into one of these page sizes. |
| 35 | PAGED_MQA_PAGE_SIZES = (32, 64) |
| 36 | |
| 37 | |
| 38 | def should_auto_disable_deep_gemm(model_type: str | None) -> bool: |
| 39 | """Check if DeepGemm should be auto-disabled for this model on Blackwell. |
| 40 | |
| 41 | Returns True if the model is known to have accuracy degradation with |
| 42 | DeepGemm's E8M0 scale format on Blackwell GPUs (SM100+). |
| 43 | """ |
| 44 | if model_type is None: |
| 45 | return False |
| 46 | if not ( |
| 47 | current_platform.is_device_capability_family(100) |
| 48 | or current_platform.is_device_capability_family(120) |
| 49 | ): |
| 50 | return False |
| 51 | return model_type in _DEEPGEMM_BLACKWELL_EXCLUDED_MODEL_TYPES |
| 52 | |
| 53 | |
| 54 | class DeepGemmQuantScaleFMT(Enum): |
| 55 | # Float32 scales in Float32 tensor |
| 56 | FLOAT32 = 0 |
| 57 | # Compute float32 scales and ceil the scales to UE8M0. |
| 58 | # Keep the scales in Float32 tensor. |
| 59 | FLOAT32_CEIL_UE8M0 = 1 |
| 60 | # Compute float32 scales and ceil the scales to UE8M0. |
| 61 | # Pack the scales into a int32 tensor where each int32 |
| 62 | # element contains 4 scale values. |
| 63 | UE8M0 = 2 |
| 64 | |
| 65 | @classmethod |
| 66 | def init_oracle_cache(cls) -> None: |
| 67 | """Initialize the oracle decision and store it in the class cache.""" |
| 68 | cached = getattr(cls, "_oracle_cache", None) |
| 69 | if cached is not None: |
| 70 | return |
| 71 | |
| 72 | use_e8m0 = ( |
| 73 | envs.VLLM_USE_DEEP_GEMM_E8M0 |
| 74 | and is_deep_gemm_supported() |
| 75 | and (_fp8_gemm_nt_impl is not None) |
| 76 | ) |
| 77 | if not use_e8m0: |
| 78 | cls._oracle_cache = cls.FLOAT32 # type: ignore |
| 79 | return |
| 80 | |
| 81 | cls._oracle_cache = ( # type: ignore |
| 82 | cls.UE8M0 |
| 83 | if ( |
| 84 | current_platform.is_device_capability_family(100) |
| 85 | or current_platform.is_device_capability_family(120) |
| 86 | ) |
| 87 | else cls.FLOAT32_CEIL_UE8M0 |
| 88 | ) |
| 89 | |
| 90 | @classmethod |
| 91 | def from_oracle(cls) -> "DeepGemmQuantScaleFMT": |
| 92 | """Return the oracle decision, initializing it on first use. |
| 93 | |
| 94 | The cache is normally populated by ``_lazy_init()`` (e.g. during |
| 95 | engine startup), but standalone consumers such as ``QuantFP8`` with an |
| 96 | explicit ``use_ue8m0=True`` can reach this before any DeepGEMM kernel |
| 97 | wrapper has run. Resolve the DeepGEMM symbols and initialize the |
| 98 | decision here instead of asserting; without DeepGEMM this yields |
| 99 | FLOAT32, matching ``is_deep_gemm_e8m0_used()``. |
| 100 | """ |
| 101 | cached = getattr(cls, "_oracle_cache", None) |
| 102 | if cached is None: |
| 103 | _lazy_init() |
| 104 | cls.init_oracle_cache() |
| 105 | cached = cls._oracle_cache # type: ignore[attr-defined] |
| 106 | return cached |
| 107 | |
| 108 | |
| 109 | @functools.cache |
| 110 | def is_deep_gemm_supported() -> bool: |
| 111 | """Return `True` if DeepGEMM is supported on the current platform. |
| 112 | Currently, only Hopper and Blackwell GPUs are supported. |
| 113 | """ |
| 114 | is_supported_arch = current_platform.support_deep_gemm() |
| 115 | return is_supported_arch and envs.VLLM_USE_DEEP_GEMM and has_deep_gemm() |
| 116 | |
| 117 | |
| 118 | @functools.cache |
| 119 | def is_deep_gemm_e8m0_used() -> bool: |
| 120 | """Return `True` if vLLM is configured to use DeepGEMM " |
| 121 | "E8M0 scale on a Hopper or Blackwell-class GPU. |
| 122 | """ |
| 123 | if not is_deep_gemm_supported(): |
| 124 | logger.debug_once( |
| 125 | "DeepGEMM E8M0 disabled: DeepGEMM not supported on this system." |
| 126 | ) |
| 127 | return False |
| 128 | |
| 129 | _lazy_init() |
| 130 | |
| 131 | if _fp8_gemm_nt_impl is None: |
| 132 | logger.info_once("DeepGEMM E8M0 disabled: _fp8_gemm_nt_impl not found") |
| 133 | return False |
| 134 | |
| 135 | if envs.VLLM_USE_DEEP_GEMM_E8M0: |
| 136 | logger.info_once("DeepGEMM E8M0 enabled on current platform.") |
| 137 | return True |
| 138 | |
| 139 | logger.info_once("DeepGEMM E8M0 disabled on current configuration.") |
| 140 | return False |
| 141 | |
| 142 | |
| 143 | def _missing(*_: Any, **__: Any) -> NoReturn: |
| 144 | """Placeholder for unavailable DeepGEMM backend.""" |
| 145 | raise RuntimeError( |
| 146 | "DeepGEMM backend is unavailable in the current vLLM environment, " |
| 147 | "or the available DeepGEMM package does not provide the required APIs " |
| 148 | "for these kernels." |
| 149 | ) |
| 150 | |
| 151 | |
| 152 | _cublaslt_gemm_nt_impl: Callable[..., Any] | None = None |
| 153 | _fp8_gemm_nt_impl: Callable[..., Any] | None = None |
| 154 | _fp8_einsum_impl: Callable[..., Any] | None = None |
| 155 | _grouped_impl: Callable[..., Any] | None = None |
| 156 | _grouped_masked_impl: Callable[..., Any] | None = None |
| 157 | _grouped_fp4_impl: Callable[..., Any] | None = None |
| 158 | _fp8_fp4_mqa_logits_impl: Callable[..., Any] | None = None |
| 159 | _fp8_fp4_paged_mqa_logits_impl: Callable[..., Any] | None = None |
| 160 | _get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None |
| 161 | _fp8_fp4_sparse_mqa_logits_impl: Callable[..., Any] | None = None |
| 162 | _fp8_fp4_paged_sparse_mqa_logits_impl: Callable[..., Any] | None = None |
| 163 | _get_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None |
| 164 | _get_paged_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None |
| 165 | _tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None |
| 166 | _mega_mhc_impl: Callable[..., Any] | None = None |
| 167 | _bf16_mega_gate_impl: Callable[..., Any] | None = None |
| 168 | _get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None |
| 169 | _get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None |
| 170 | _get_theoretical_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = ( |
| 171 | None |
| 172 | ) |
| 173 | _transform_sf_into_required_layout_impl: Callable[..., Any] | None = None |
| 174 | _pack_ue8m0_to_int_impl: Callable[..., Any] | None = None |
| 175 | _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl: Callable[..., Any] | None = None |
| 176 | _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl: ( |
| 177 | Callable[..., Any] | None |
| 178 | ) = None |
| 179 | |
| 180 | |
| 181 | @functools.cache |
| 182 | def _import_deep_gemm(): |
| 183 | """Import the deep_gemm module. |
| 184 | |
| 185 | Prefers an externally installed ``deep_gemm`` package (so users can |
| 186 | pin a specific version), then falls back to the vendored copy bundled |
| 187 | in the vLLM wheel. |
| 188 | |
| 189 | Returns ``None`` when neither source is usable. |
| 190 | """ |
| 191 | # 1. Try the external (pip-installed) package first. |
| 192 | try: |
| 193 | module = importlib.import_module("deep_gemm") |
| 194 | logger.debug_once("Imported deep_gemm module from site-packages") |
| 195 | return module |
| 196 | except ImportError: |
| 197 | logger.info_once( |
| 198 | "deep_gemm not found in site-packages, " |
| 199 | "trying vendored vllm.third_party.deep_gemm" |
| 200 | ) |
| 201 | |
| 202 | # 2. Fall back to the vendored copy bundled in the vLLM wheel. |
| 203 | try: |
| 204 | module = importlib.import_module("vllm.third_party.deep_gemm") |
| 205 | logger.debug_once("Imported deep_gemm module from vllm.third_party.deep_gemm") |
| 206 | return module |
| 207 | except ImportError: |
| 208 | logger.info_once("Vendored deep_gemm not found either") |
| 209 | except Exception as e: |
| 210 | # The vendored module may raise RuntimeError during _C.init() |
| 211 | # if JIT include files are missing (e.g. incomplete wheel). |
| 212 | logger.warning_once("Failed to import vendored deep_gemm: %s", e) |
| 213 | |
| 214 | return None |
| 215 | |
| 216 | |
| 217 | def _apply_pdl(mod, enable: bool = True) -> None: |
| 218 | mod_name = getattr(mod, "__name__", str(mod)) |
| 219 | try: |
| 220 | set_pdl_fn = getattr(mod, "set_pdl", None) |
| 221 | if set_pdl_fn is None: |
| 222 | return |
| 223 | set_pdl_fn(enable) |
| 224 | logger.info_once( |
| 225 | "DeepGEMM PDL %s on %s.", |
| 226 | "enabled" if enable else "disabled", |
| 227 | mod_name, |
| 228 | ) |
| 229 | except Exception as e: # noqa: BLE001 |
| 230 | logger.warning_once("Failed to set DeepGEMM PDL on %s: %s", mod_name, e) |
| 231 | |
| 232 | |
| 233 | def _lazy_init() -> None: |
| 234 | """Import deep_gemm and resolve symbols on first use.""" |
| 235 | global _cublaslt_gemm_nt_impl |
| 236 | global _fp8_gemm_nt_impl, _fp8_einsum_impl |
| 237 | global _grouped_impl, _grouped_masked_impl, _grouped_fp4_impl |
| 238 | global _fp8_fp4_mqa_logits_impl, _fp8_fp4_paged_mqa_logits_impl |
| 239 | global _get_paged_mqa_logits_metadata_impl |
| 240 | global _fp8_fp4_sparse_mqa_logits_impl, _fp8_fp4_paged_sparse_mqa_logits_impl |
| 241 | global _get_sparse_mqa_logits_metadata_impl |
| 242 | global _get_paged_sparse_mqa_logits_metadata_impl |
| 243 | global _tf32_hc_prenorm_gemm_impl, _mega_mhc_impl |
| 244 | global _bf16_mega_gate_impl |
| 245 | global _get_mn_major_tma_aligned_tensor_impl |
| 246 | global _get_mk_alignment_for_contiguous_layout_impl |
| 247 | global _get_theoretical_mk_alignment_for_contiguous_layout_impl |
| 248 | global _transform_sf_into_required_layout_impl |
| 249 | global _pack_ue8m0_to_int_impl |
| 250 | global _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl |
| 251 | global _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl |
| 252 | # fast path |
| 253 | if ( |
| 254 | _cublaslt_gemm_nt_impl is not None |
| 255 | or _fp8_gemm_nt_impl is not None |
| 256 | or _fp8_einsum_impl is not None |
| 257 | or _grouped_impl is not None |
| 258 | or _grouped_masked_impl is not None |
| 259 | or _grouped_fp4_impl is not None |
| 260 | or _fp8_fp4_mqa_logits_impl is not None |
| 261 | or _fp8_fp4_paged_mqa_logits_impl is not None |
| 262 | or _get_paged_mqa_logits_metadata_impl is not None |
| 263 | or _tf32_hc_prenorm_gemm_impl is not None |
| 264 | or _mega_mhc_impl is not None |
| 265 | or _bf16_mega_gate_impl is not None |
| 266 | or _get_mk_alignment_for_contiguous_layout_impl is not None |
| 267 | or _transform_sf_into_required_layout_impl is not None |
| 268 | or _pack_ue8m0_to_int_impl is not None |
| 269 | or _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl is not None |
| 270 | or _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl is not None |
| 271 | ): |
| 272 | return |
| 273 | |
| 274 | if not has_deep_gemm(): |
| 275 | return |
| 276 | |
| 277 | # Set up deep_gemm cache path |
| 278 | DEEP_GEMM_JIT_CACHE_ENV_NAME = "DG_JIT_CACHE_DIR" |
| 279 | if not os.environ.get(DEEP_GEMM_JIT_CACHE_ENV_NAME, None): |
| 280 | os.environ[DEEP_GEMM_JIT_CACHE_ENV_NAME] = os.path.join( |
| 281 | envs.VLLM_CACHE_ROOT, "deep_gemm" |
| 282 | ) |
| 283 | |
| 284 | _dg = _import_deep_gemm() |
| 285 | if _dg is None: |
| 286 | return |
| 287 | |
| 288 | # Enable PDL for DeepGEMM on architectures that support it (SM90+). |
| 289 | if current_platform.is_arch_support_pdl(): |
| 290 | _apply_pdl(_dg, True) |
| 291 | _cublaslt_gemm_nt_impl = getattr(_dg, "cublaslt_gemm_nt", None) |
| 292 | _fp8_gemm_nt_impl = getattr(_dg, "fp8_gemm_nt", None) |
| 293 | _fp8_einsum_impl = getattr(_dg, "fp8_einsum", None) |
| 294 | _grouped_impl = getattr(_dg, "m_grouped_fp8_gemm_nt_contiguous", None) |
| 295 | # DeepGEMM 2.8.0 dropped the legacy `fp8_*` aliases; keep both spellings so |
| 296 | # an externally pinned older DeepGEMM still resolves. |
| 297 | _grouped_masked_impl = getattr( |
| 298 | _dg, "m_grouped_fp8_gemm_nt_masked", None |
| 299 | ) or getattr(_dg, "fp8_m_grouped_gemm_nt_masked", None) |
| 300 | _grouped_fp4_impl = getattr(_dg, "m_grouped_fp8_fp4_gemm_nt_contiguous", None) |
| 301 | # DeepGEMM exposes fp8_fp4_*_mqa_logits as the canonical symbols that |
| 302 | # handle both the FP8 and FP4 Q/K paths via a tuple-typed `q`. |
| 303 | _fp8_fp4_mqa_logits_impl = getattr(_dg, "fp8_fp4_mqa_logits", None) |
| 304 | _fp8_fp4_paged_mqa_logits_impl = getattr(_dg, "fp8_fp4_paged_mqa_logits", None) |
| 305 | _get_paged_mqa_logits_metadata_impl = getattr( |
| 306 | _dg, "get_paged_mqa_logits_metadata", None |
| 307 | ) |
| 308 | # Sparse-indexer kernels (DeepGEMM >= 2.8, SM100 only). |
| 309 | _fp8_fp4_sparse_mqa_logits_impl = getattr(_dg, "fp8_fp4_sparse_mqa_logits", None) |
| 310 | _fp8_fp4_paged_sparse_mqa_logits_impl = getattr( |
| 311 | _dg, "fp8_fp4_paged_sparse_mqa_logits", None |
| 312 | ) |
| 313 | _get_sparse_mqa_logits_metadata_impl = getattr( |
| 314 | _dg, "get_sparse_mqa_logits_metadata", None |
| 315 | ) |
| 316 | _get_paged_sparse_mqa_logits_metadata_impl = getattr( |
| 317 | _dg, "get_paged_sparse_mqa_logits_metadata", None |
| 318 | ) |
| 319 | _tf32_hc_prenorm_gemm_impl = getattr(_dg, "tf32_hc_prenorm_gemm", None) |
| 320 | _mega_mhc_impl = getattr(_dg, "mega_mhc", None) |
| 321 | _bf16_mega_gate_impl = getattr(_dg, "bf16_mega_gate", None) |
| 322 | _get_mn_major_tma_aligned_tensor_impl = getattr( |
| 323 | _dg, "get_mn_major_tma_aligned_tensor", None |
| 324 | ) |
| 325 | _get_mk_alignment_for_contiguous_layout_impl = getattr( |
| 326 | _dg, "get_mk_alignment_for_contiguous_layout", None |
| 327 | ) |
| 328 | _get_theoretical_mk_alignment_for_contiguous_layout_impl = getattr( |
| 329 | _dg, "get_theoretical_mk_alignment_for_contiguous_layout", None |
| 330 | ) |
| 331 | _transform_sf_into_required_layout_impl = getattr( |
| 332 | _dg, "transform_sf_into_required_layout", None |
| 333 | ) |
| 334 | _pack_ue8m0_to_int_impl = getattr(_dg, "pack_ue8m0_to_int", None) |
| 335 | _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl = getattr( |
| 336 | _dg, "get_mn_major_tma_aligned_packed_ue8m0_tensor", None |
| 337 | ) |
| 338 | _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl = getattr( |
| 339 | _dg, "get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor", None |
| 340 | ) |
| 341 | DeepGemmQuantScaleFMT.init_oracle_cache() |
| 342 | |
| 343 | |
| 344 | def get_num_sms() -> int: |
| 345 | _lazy_init() |
| 346 | dg = _import_deep_gemm() |
| 347 | if dg is None: |
| 348 | raise RuntimeError("DeepGEMM is not available") |
| 349 | return int(dg.get_num_sms()) |
| 350 | |
| 351 | |
| 352 | def set_num_sms(num_sms: int) -> None: |
| 353 | _lazy_init() |
| 354 | dg = _import_deep_gemm() |
| 355 | if dg is None: |
| 356 | raise RuntimeError("DeepGEMM is not available") |
| 357 | dg.set_num_sms(num_sms) |
| 358 | |
| 359 | |
| 360 | def get_mk_alignment_for_contiguous_layout() -> list[int]: |
| 361 | _lazy_init() |
| 362 | if _get_mk_alignment_for_contiguous_layout_impl is None: |
| 363 | return _missing() |
| 364 | mk_align_size = _get_mk_alignment_for_contiguous_layout_impl() |
| 365 | return [mk_align_size, mk_align_size] |
| 366 | |
| 367 | |
| 368 | def get_theoretical_mk_alignment_for_contiguous_layout( |
| 369 | expected_m: int | None = None, |
| 370 | num_groups: int | None = None, |
| 371 | ) -> int: |
| 372 | """Per-call optimal M alignment for grouped contiguous GEMMs. |
| 373 | |
| 374 | `expected_m` is the TOTAL routed tokens (sum across experts, typically |
| 375 | M × num_topk). `num_groups` is the number of experts on this rank. |
| 376 | The helper divides to recover per-expert em and picks an alignment based |
| 377 | on data-driven thresholds (see deep_gemm runtime.hpp comments). |
| 378 | |
| 379 | Older callers that omit `num_groups` are interpreted as passing already |
| 380 | per-expert em (legacy behaviour preserved for backward compat). |
| 381 | """ |
| 382 | _lazy_init() |
| 383 | if _get_theoretical_mk_alignment_for_contiguous_layout_impl is None: |
| 384 | return _missing() |
| 385 | if num_groups is None: |
| 386 | return _get_theoretical_mk_alignment_for_contiguous_layout_impl(expected_m) |
| 387 | if num_groups <= 0: |
| 388 | raise ValueError(f"num_groups must be positive, got {num_groups}") |
| 389 | try: |
| 390 | return _get_theoretical_mk_alignment_for_contiguous_layout_impl( |
| 391 | expected_m, num_groups |
| 392 | ) |
| 393 | except TypeError: |
| 394 | per_group_m = None if expected_m is None else cdiv(expected_m, num_groups) |
| 395 | return _get_theoretical_mk_alignment_for_contiguous_layout_impl(per_group_m) |
| 396 | |
| 397 | |
| 398 | def set_mk_alignment_for_contiguous_layout(value: int) -> None: |
| 399 | """Set DeepGEMM's BLOCK_M cap for grouped contiguous GEMMs. |
| 400 | |
| 401 | The DG heuristic constrains BLOCK_M ≤ this value when picking a kernel |
| 402 | layout. Use this in concert with `compute_aligned_M_and_alignment`'s |
| 403 | per-call alignment so the workspace's per-expert padding matches the |
| 404 | kernel's BLOCK_M; a mismatch leads to the scheduler reading the wrong |
| 405 | expert_id from `m_indices` at `m_block_idx * BLOCK_M` stride and |
| 406 | OOB-indexing the B-weights tensor (manifests as IMA under CUDA-graph |
| 407 | replay). |
| 408 | """ |
| 409 | _lazy_init() |
| 410 | dg = _import_deep_gemm() |
| 411 | if dg is None: |
| 412 | raise RuntimeError("DeepGEMM is not available") |
| 413 | dg.set_mk_alignment_for_contiguous_layout(value) |
| 414 | |
| 415 | |
| 416 | @contextlib.contextmanager |
| 417 | def mk_alignment_scope(value: int): |
| 418 | """Temporarily set DeepGEMM's BLOCK_M cap, restoring on exit. |
| 419 | |
| 420 | Use around a sequence of grouped-contiguous GEMM calls whose workspace |
| 421 | is padded to `value` (typically the per_call_align returned by |
| 422 | `compute_aligned_M_and_alignment`). |
| 423 | """ |
| 424 | prev = get_mk_alignment_for_contiguous_layout()[0] |
| 425 | set_mk_alignment_for_contiguous_layout(value) |
| 426 | try: |
| 427 | yield |
| 428 | finally: |
| 429 | set_mk_alignment_for_contiguous_layout(prev) |
| 430 | |
| 431 | |
| 432 | def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor: |
| 433 | """Wrapper for DeepGEMM's get_mn_major_tma_aligned_tensor.""" |
| 434 | _lazy_init() |
| 435 | if _get_mn_major_tma_aligned_tensor_impl is None: |
| 436 | return _missing() |
| 437 | return _get_mn_major_tma_aligned_tensor_impl(x) |
| 438 | |
| 439 | |
| 440 | def pack_ue8m0_to_int(x: torch.Tensor) -> torch.Tensor: |
| 441 | """Pack 4 UE8M0 (uint8) scales into one int32. |
| 442 | |
| 443 | DeepGEMM's SM100/SM120 FP8/FP4 kernels accept either ``float32`` scales |
| 444 | (legacy format, 4 B/scale) or ``int32`` packed UE8M0 scales (1 B/scale |
| 445 | after 4:1 packing — 4× smaller than the legacy fp32 representation). |
| 446 | """ |
| 447 | _lazy_init() |
| 448 | if _pack_ue8m0_to_int_impl is None: |
| 449 | return _missing() |
| 450 | return _pack_ue8m0_to_int_impl(x) |
| 451 | |
| 452 | |
| 453 | def get_mn_major_tma_aligned_packed_ue8m0_tensor(x: torch.Tensor) -> torch.Tensor: |
| 454 | """Pack UE8M0 (uint8) → int32 with the MN-major TMA-aligned layout the |
| 455 | DeepGEMM kernels consume directly. 16× smaller than the fp32 legacy SF |
| 456 | format. Use for non-grouped 2D scale tensors. |
| 457 | """ |
| 458 | _lazy_init() |
| 459 | if _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl is None: |
| 460 | return _missing() |
| 461 | return _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl(x) |
| 462 | |
| 463 | |
| 464 | def get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor( |
| 465 | sf: torch.Tensor, |
| 466 | ks_tensor: torch.Tensor, |
| 467 | ks: list[int], |
| 468 | gran_k: int, |
| 469 | ) -> torch.Tensor: |
| 470 | """Grouped (3D, expert-batched) variant of |
| 471 | ``get_mn_major_tma_aligned_packed_ue8m0_tensor``. Use for MoE weight |
| 472 | scale tensors of shape ``(num_experts, mn, k_scale)``. |
| 473 | """ |
| 474 | _lazy_init() |
| 475 | if _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl is None: |
| 476 | return _missing() |
| 477 | return _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl( |
| 478 | sf, ks_tensor, ks, gran_k |
| 479 | ) |
| 480 | |
| 481 | |
| 482 | def cublaslt_gemm_nt(*args, **kwargs): |
| 483 | _lazy_init() |
| 484 | if _cublaslt_gemm_nt_impl is None: |
| 485 | return _missing(*args, **kwargs) |
| 486 | return _cublaslt_gemm_nt_impl(*args, **kwargs) |
| 487 | |
| 488 | |
| 489 | def bf16_mega_gate( |
| 490 | x: torch.Tensor, |
| 491 | weight: torch.Tensor, |
| 492 | num_topk: int, |
| 493 | *, |
| 494 | scoring_func: str, |
| 495 | routed_scaling_factor: float, |
| 496 | ep_rank: int, |
| 497 | bias: torch.Tensor | None = None, |
| 498 | image_bias: torch.Tensor | None = None, |
| 499 | image_token_mask: torch.Tensor | None = None, |
| 500 | fix_routing_mask: torch.Tensor | None = None, |
| 501 | unmapped_topk_idx: torch.Tensor | None = None, |
| 502 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 503 | """Run DeepGEMM Mega Gate. |
| 504 | |
| 505 | The output order follows vLLM's router contract: weights, then indices. |
| 506 | Passing PyTorch-owned outputs to DeepGEMM makes their storage part of CUDA |
| 507 | graph capture. |
| 508 | """ |
| 509 | _lazy_init() |
| 510 | if _bf16_mega_gate_impl is None: |
| 511 | raise RuntimeError( |
| 512 | "DeepGEMM MegaMoE requires a DeepGEMM build with bf16_mega_gate." |
| 513 | ) |
| 514 | |
| 515 | topk_idx = torch.empty((x.shape[0], num_topk), dtype=torch.int64, device=x.device) |
| 516 | topk_weights = torch.empty( |
| 517 | (x.shape[0], num_topk), dtype=torch.float32, device=x.device |
| 518 | ) |
| 519 | _bf16_mega_gate_impl( |
| 520 | x, |
| 521 | weight, |
| 522 | num_topk, |
| 523 | use_shared_as_routed=False, |
| 524 | num_shared_experts=0, |
| 525 | routed_scaling_factor=routed_scaling_factor, |
| 526 | ep_rank=ep_rank, |
| 527 | scoring_func=scoring_func, |
| 528 | bias=bias, |
| 529 | image_bias=image_bias, |
| 530 | image_token_mask=image_token_mask, |
| 531 | fix_routing_mask=fix_routing_mask, |
| 532 | unmapped_topk_idx=unmapped_topk_idx, |
| 533 | out=(topk_idx, topk_weights), |
| 534 | ) |
| 535 | return topk_weights, topk_idx |
| 536 | |
| 537 | |
| 538 | def fp8_gemm_nt(*args, **kwargs): |
| 539 | _lazy_init() |
| 540 | if _fp8_gemm_nt_impl is None: |
| 541 | return _missing(*args, **kwargs) |
| 542 | if "is_deep_gemm_e8m0_used" in kwargs: |
| 543 | use_ue8m0 = kwargs["is_deep_gemm_e8m0_used"] |
| 544 | del kwargs["is_deep_gemm_e8m0_used"] |
| 545 | else: |
| 546 | use_ue8m0 = is_deep_gemm_e8m0_used() |
| 547 | return _fp8_gemm_nt_impl(*args, disable_ue8m0_cast=not use_ue8m0, **kwargs) |
| 548 | |
| 549 | |
| 550 | def fp8_einsum(*args, **kwargs): |
| 551 | _lazy_init() |
| 552 | if _fp8_einsum_impl is None: |
| 553 | return _missing(*args, **kwargs) |
| 554 | return _fp8_einsum_impl(*args, **kwargs) |
| 555 | |
| 556 | |
| 557 | def m_grouped_fp8_gemm_nt_contiguous(*args, **kwargs): |
| 558 | _lazy_init() |
| 559 | if _grouped_impl is None: |
| 560 | return _missing(*args, **kwargs) |
| 561 | return _grouped_impl( |
| 562 | *args, disable_ue8m0_cast=not is_deep_gemm_e8m0_used(), **kwargs |
| 563 | ) |
| 564 | |
| 565 | |
| 566 | def m_grouped_fp8_fp4_gemm_nt_contiguous(*args, **kwargs): |
| 567 | _lazy_init() |
| 568 | if _grouped_fp4_impl is None: |
| 569 | return _missing(*args, **kwargs) |
| 570 | return _grouped_fp4_impl( |
| 571 | *args, disable_ue8m0_cast=not is_deep_gemm_e8m0_used(), **kwargs |
| 572 | ) |
| 573 | |
| 574 | |
| 575 | def fp8_m_grouped_gemm_nt_masked(*args, **kwargs): |
| 576 | _lazy_init() |
| 577 | if _grouped_masked_impl is None: |
| 578 | return _missing(*args, **kwargs) |
| 579 | return _grouped_masked_impl( |
| 580 | *args, disable_ue8m0_cast=not is_deep_gemm_e8m0_used(), **kwargs |
| 581 | ) |
| 582 | |
| 583 | |
| 584 | def transform_sf_into_required_layout(*args, **kwargs): |
| 585 | _lazy_init() |
| 586 | if _transform_sf_into_required_layout_impl is None: |
| 587 | return _missing(*args, **kwargs) |
| 588 | return _transform_sf_into_required_layout_impl( |
| 589 | *args, disable_ue8m0_cast=not is_deep_gemm_e8m0_used(), **kwargs |
| 590 | ) |
| 591 | |
| 592 | |
| 593 | def fp8_fp4_mqa_logits( |
| 594 | q: tuple[torch.Tensor, torch.Tensor | None], |
| 595 | kv: tuple[torch.Tensor, torch.Tensor], |
| 596 | weights: torch.Tensor, |
| 597 | cu_seqlen_ks: torch.Tensor, |
| 598 | cu_seqlen_ke: torch.Tensor, |
| 599 | clean_logits: bool, |
| 600 | ) -> torch.Tensor: |
| 601 | """Compute MQA logits for a single sequence without KV paging. |
| 602 | |
| 603 | Unified FP8/FP4 dispatch — the underlying DeepGEMM kernel takes |
| 604 | ``q = (values, scales_or_None)`` where ``scales`` is None for FP8 Q |
| 605 | (per-token scale is folded into ``weights``) and a packed block-scale |
| 606 | tensor for MXFP4 Q. |
| 607 | |
| 608 | Args: |
| 609 | q: Tuple ``(q_values, q_scale)``. FP8 path: q_values is [M, H, D] |
| 610 | float8_e4m3fn and q_scale is None (per-token scale is folded |
| 611 | into ``weights``). FP4 path: q_values is packed uint8 and |
| 612 | q_scale is the companion block-scale tensor. |
| 613 | kv: Tuple `(k_packed, k_scales)` — FP8 layout is [N, D] |
| 614 | float8_e4m3fn plus fp32 scales [N]; FP4 layout is packed uint8. |
| 615 | weights: weights of shape [M, H], dtype `torch.float32`. |
| 616 | cu_seqlen_ks: Start indices (inclusive) for valid K per query |
| 617 | position, shape [M], dtype int32. |
| 618 | cu_seqlen_ke: End indices (exclusive) for valid K per query |
| 619 | position, shape [M], dtype int32. |
| 620 | clean_logits: Whether to clean the unfilled logits into `-inf`. |
| 621 | |
| 622 | Returns: |
| 623 | Logits tensor of shape [M, N], dtype `torch.float32`. |
| 624 | |
| 625 | """ |
| 626 | _lazy_init() |
| 627 | if _fp8_fp4_mqa_logits_impl is None: |
| 628 | return _missing() |
| 629 | return _fp8_fp4_mqa_logits_impl( |
| 630 | q, |
| 631 | kv, |
| 632 | weights, |
| 633 | cu_seqlen_ks, |
| 634 | cu_seqlen_ke, |
| 635 | clean_logits=clean_logits, |
| 636 | ) |
| 637 | |
| 638 | |
| 639 | def native_next_n_supported(next_n: int) -> bool: |
| 640 | """Whether the paged MQA logits kernel takes `next_n` Q rows per request. |
| 641 | |
| 642 | SM90 implements only {1, 2, 4}; SM100 and SM120 schedule any `next_n` via |
| 643 | multi-atom tiles. Unsupported values must be flattened to one row per query. |
| 644 | """ |
| 645 | if current_platform.is_device_capability_family(90): |
| 646 | return next_n in (1, 2, 4) |
| 647 | return True |
| 648 | |
| 649 | |
| 650 | def _paged_mqa_logits_schedule_slots(num_sms: int, next_n: int) -> int: |
| 651 | """Scheduler tasks the paged MQA logits kernel launches. |
| 652 | |
| 653 | SM90 `next_n=4` runs one task per 2-CTA multicast cluster rather than per |
| 654 | SM, and `fp8_fp4_paged_mqa_logits` asserts its metadata is sized to match. |
| 655 | """ |
| 656 | num_kv_multicast = ( |
| 657 | 2 if next_n == 4 and current_platform.is_device_capability_family(90) else 1 |
| 658 | ) |
| 659 | return num_sms // num_kv_multicast |
| 660 | |
| 661 | |
| 662 | def get_paged_mqa_logits_metadata( |
| 663 | context_lens: torch.Tensor, |
| 664 | block_size: int, |
| 665 | num_sms: int, |
| 666 | indices: torch.Tensor | None = None, |
| 667 | ) -> torch.Tensor: |
| 668 | """Build scheduling metadata for paged MQA logits. |
| 669 | |
| 670 | Args: |
| 671 | context_lens: Tensor of shape [B, next_n], dtype int32; effective |
| 672 | context length per Q row. |
| 673 | block_size: KV-cache block size in tokens (e.g., 64). |
| 674 | num_sms: Number of SMs available. 132 for Hopper |
| 675 | indices: Optional request index for each varlen row. |
| 676 | |
| 677 | Returns: |
| 678 | Tensor of shape [slots + 1, 2] consumed by `fp8_fp4_paged_mqa_logits` |
| 679 | to schedule work across SMs. |
| 680 | |
| 681 | """ |
| 682 | _lazy_init() |
| 683 | if _get_paged_mqa_logits_metadata_impl is None: |
| 684 | return _missing() |
| 685 | next_n = context_lens.shape[1] if context_lens.dim() == 2 else 1 |
| 686 | num_slots = _paged_mqa_logits_schedule_slots(num_sms, next_n) |
| 687 | kwargs = {} if indices is None else {"indices": indices} |
| 688 | return _get_paged_mqa_logits_metadata_impl( |
| 689 | context_lens, block_size, num_slots, **kwargs |
| 690 | ) |
| 691 | |
| 692 | |
| 693 | def fp8_fp4_paged_mqa_logits( |
| 694 | q: tuple[torch.Tensor, torch.Tensor | None], |
| 695 | kv_cache: torch.Tensor, |
| 696 | weights: torch.Tensor, |
| 697 | context_lens: torch.Tensor, |
| 698 | block_tables: torch.Tensor, |
| 699 | schedule_metadata: torch.Tensor, |
| 700 | max_model_len: int, |
| 701 | clean_logits: bool, |
| 702 | indices: torch.Tensor | None = None, |
| 703 | ) -> torch.Tensor: |
| 704 | """Compute MQA logits using a paged KV-cache. |
| 705 | |
| 706 | Unified FP8/FP4 dispatch — the underlying DeepGEMM kernel takes |
| 707 | ``q = (values, scales_or_None)``; pass ``(q_tensor, None)`` for the FP8 |
| 708 | path and ``(q_values, q_scale)`` for MXFP4. |
| 709 | |
| 710 | Args: |
| 711 | q: Tuple ``(q_values, q_scale)``. FP8 path: q_values is |
| 712 | [B, next_n, H, D] float8_e4m3fn and q_scale is None. FP4 path: |
| 713 | q_values is packed uint8 and q_scale is the companion |
| 714 | block-scale tensor. |
| 715 | kv_cache: Paged KV-cache. FP8 layout is [num_blocks, block_size, 1, |
| 716 | D+4], dtype `torch.uint8`, with the last 4 bytes per (block, pos) |
| 717 | storing the float dequant scale. |
| 718 | weights: Tensor of shape [B * next_n, H], dtype `torch.float32`. |
| 719 | context_lens: Tensor of shape [B], dtype int32; effective context length |
| 720 | for each batch element. |
| 721 | block_tables: Tensor of shape [B, max_blocks], dtype int32; maps logical |
| 722 | block indices to physical blocks in the paged cache. |
| 723 | schedule_metadata: Returned by `get_paged_mqa_logits_metadata`; |
| 724 | used to distribute work across SMs. |
| 725 | max_model_len: Maximum sequence length used to size the logits output. |
| 726 | clean_logits: Whether to clean the unfilled logits into `-inf`. |
| 727 | indices: Optional request index for each varlen row. |
| 728 | |
| 729 | Returns: |
| 730 | Logits tensor of shape [B * next_n, max_model_len], dtype |
| 731 | `torch.float32`. |
| 732 | |
| 733 | """ |
| 734 | _lazy_init() |
| 735 | if _fp8_fp4_paged_mqa_logits_impl is None: |
| 736 | return _missing() |
| 737 | # DeepGEMM asserts block_tables.stride(-1)==1. A trailing size-1 dim |
| 738 | # (e.g. block_table shape [B,1] for short seqs under a large block_size) |
| 739 | # can be a transposed view where .contiguous() is a no-op (torch treats the |
| 740 | # size-1 dim's stride as irrelevant) yet stride(-1)!=1, failing the kernel. |
| 741 | # clone to contiguous format to force stride(-1)==1. |
| 742 | if block_tables.dim() >= 2 and block_tables.stride(-1) != 1: |
| 743 | block_tables = block_tables.clone(memory_format=torch.contiguous_format) |
| 744 | kwargs = {} if indices is None else {"indices": indices} |
| 745 | return _fp8_fp4_paged_mqa_logits_impl( |
| 746 | q, |
| 747 | kv_cache, |
| 748 | weights, |
| 749 | context_lens, |
| 750 | block_tables, |
| 751 | schedule_metadata, |
| 752 | max_model_len, |
| 753 | clean_logits=clean_logits, |
| 754 | **kwargs, |
| 755 | ) |
| 756 | |
| 757 | |
| 758 | def has_deep_gemm_sparse_mqa() -> bool: |
| 759 | """Whether the installed DeepGEMM provides the sparse-indexer kernels |
| 760 | (``fp8_fp4_(paged_)sparse_mqa_logits``, added in DeepGEMM 2.8, SM100-only).""" |
| 761 | _lazy_init() |
| 762 | return ( |
| 763 | _fp8_fp4_sparse_mqa_logits_impl is not None |
| 764 | and _fp8_fp4_paged_sparse_mqa_logits_impl is not None |
| 765 | and _get_sparse_mqa_logits_metadata_impl is not None |
| 766 | and _get_paged_sparse_mqa_logits_metadata_impl is not None |
| 767 | ) |
| 768 | |
| 769 | |
| 770 | def get_sparse_mqa_logits_metadata( |
| 771 | cu_seqlen_ks: torch.Tensor, |
| 772 | cu_seqlen_ke: torch.Tensor, |
| 773 | num_kv_tokens: int, |
| 774 | sparse_kv_block_indices: torch.Tensor, |
| 775 | qk_dtype: torch.dtype, |
| 776 | sparse_block_kv: int, |
| 777 | ) -> torch.Tensor: |
| 778 | """Build scheduling metadata for `fp8_fp4_sparse_mqa_logits`. |
| 779 | |
| 780 | Always uses the unaligned-ks variant, which degenerates to the aligned |
| 781 | variant when ``cu_seqlen_ks % sparse_block_kv == 0``, so callers never |
| 782 | need a host-side sync to pick one. |
| 783 | |
| 784 | Args: |
| 785 | cu_seqlen_ks: Per-row K range start bounds in the packed KV |
| 786 | workspace, shape [num_q_tokens], dtype int32. |
| 787 | cu_seqlen_ke: Per-row K range end bounds in the packed KV workspace, |
| 788 | shape [num_q_tokens], dtype int32. |
| 789 | num_kv_tokens: Total KV tokens in the packed workspace. |
| 790 | sparse_kv_block_indices: Per-row candidate block ids, shape |
| 791 | [num_q_tokens, num_max_sparse_blocks], dtype int32. Each row's |
| 792 | valid prefix must be sorted ascending (repeats are tolerated but |
| 793 | wasteful); pad by repeating the last valid block. Block ``i`` |
| 794 | covers tokens |
| 795 | ``[i * sparse_block_kv + ks % sparse_block_kv, ...)``. |
| 796 | qk_dtype: dtype of the packed Q values (``torch.int8`` for MXFP4, |
| 797 | ``torch.float8_e4m3fn`` for FP8). |
| 798 | sparse_block_kv: Tokens per sparse block, 8 or 16. |
| 799 | |
| 800 | """ |
| 801 | _lazy_init() |
| 802 | if _get_sparse_mqa_logits_metadata_impl is None: |
| 803 | return _missing() |
| 804 | return _get_sparse_mqa_logits_metadata_impl( |
| 805 | cu_seqlen_ks, |
| 806 | cu_seqlen_ke, |
| 807 | num_kv_tokens, |
| 808 | sparse_kv_block_indices, |
| 809 | qk_dtype, |
| 810 | sparse_block_kv, |
| 811 | use_unaligned_ks=True, |
| 812 | ) |
| 813 | |
| 814 | |
| 815 | def get_paged_sparse_mqa_logits_metadata( |
| 816 | context_lens: torch.Tensor, |
| 817 | block_table: torch.Tensor, |
| 818 | indices: torch.Tensor, |
| 819 | page_kv: int, |
| 820 | sparse_kv_block_indices: torch.Tensor, |
| 821 | qk_dtype: torch.dtype, |
| 822 | sparse_block_kv: int, |
| 823 | ) -> torch.Tensor: |
| 824 | """Build scheduling metadata for `fp8_fp4_paged_sparse_mqa_logits`. |
| 825 | |
| 826 | Args: |
| 827 | context_lens: Per-row context length, shape [num_q_tokens], int32. |
| 828 | block_table: Per-row page table, shape [num_q_tokens, max_pages], |
| 829 | int32; rows of queries paired with the same request must be |
| 830 | identical (expand it per query row before calling). |
| 831 | indices: Request index for each query row, shape [num_q_tokens]. |
| 832 | page_kv: Tokens per KV page; must be a multiple of sparse_block_kv. |
| 833 | sparse_kv_block_indices: Per-row logical (context-relative) candidate |
| 834 | block ids, [num_q_tokens, num_max_sparse_blocks], int32. |
| 835 | qk_dtype: dtype of the packed Q values. |
| 836 | sparse_block_kv: Tokens per sparse block, 8 or 16. |
| 837 | |
| 838 | """ |
| 839 | _lazy_init() |
| 840 | if _get_paged_sparse_mqa_logits_metadata_impl is None: |
| 841 | return _missing() |
| 842 | return _get_paged_sparse_mqa_logits_metadata_impl( |
| 843 | context_lens, |
| 844 | block_table, |
| 845 | indices, |
| 846 | page_kv, |
| 847 | sparse_kv_block_indices, |
| 848 | qk_dtype, |
| 849 | sparse_block_kv, |
| 850 | ) |
| 851 | |
| 852 | |
| 853 | def fp8_fp4_sparse_mqa_logits( |
| 854 | q: tuple[torch.Tensor, torch.Tensor], |
| 855 | kv: tuple[torch.Tensor, torch.Tensor], |
| 856 | weights: torch.Tensor, |
| 857 | metadata: torch.Tensor, |
| 858 | num_max_sparse_blocks: int, |
| 859 | sparse_block_kv: int, |
| 860 | ) -> torch.Tensor: |
| 861 | """Compute MQA logits only at the candidate blocks (prefill path). |
| 862 | |
| 863 | Args: |
| 864 | q: ``(q_values, q_scale)``. q_values is [M, H, D] (packed FP4 viewed |
| 865 | as int8, or FP8); q_scale is the packed UE8M0 scale tensor [M, H] |
| 866 | int32 (mandatory, unlike the dense kernels). |
| 867 | kv: ``(kv_values, kv_scale)`` — the packed KV workspace with UE8M0 |
| 868 | int32 scales. |
| 869 | weights: [M, H] ``torch.bfloat16`` (the Q scale is NOT folded in). |
| 870 | metadata: From `get_sparse_mqa_logits_metadata`. |
| 871 | num_max_sparse_blocks: Width of each sparse-index row. |
| 872 | sparse_block_kv: Tokens per sparse block, 8 or 16. |
| 873 | |
| 874 | Returns: |
| 875 | bf16 logits of shape [M, num_max_sparse_blocks * sparse_block_kv]; |
| 876 | column ``j * sparse_block_kv + o`` scores the token at |
| 877 | ``sparse_kv_block_indices[row, j] * sparse_block_kv + ks % sbk + o``. |
| 878 | |
| 879 | """ |
| 880 | _lazy_init() |
| 881 | if _fp8_fp4_sparse_mqa_logits_impl is None: |
| 882 | return _missing() |
| 883 | return _fp8_fp4_sparse_mqa_logits_impl( |
| 884 | q, |
| 885 | kv, |
| 886 | weights, |
| 887 | metadata, |
| 888 | num_max_sparse_blocks, |
| 889 | sparse_block_kv, |
| 890 | use_unaligned_ks=True, |
| 891 | ) |
| 892 | |
| 893 | |
| 894 | def fp8_fp4_paged_sparse_mqa_logits( |
| 895 | q: tuple[torch.Tensor, torch.Tensor], |
| 896 | kv_cache: torch.Tensor, |
| 897 | weights: torch.Tensor, |
| 898 | metadata: torch.Tensor, |
| 899 | num_max_sparse_blocks: int, |
| 900 | sparse_block_kv: int, |
| 901 | ) -> torch.Tensor: |
| 902 | """Compute MQA logits only at the candidate blocks (paged decode path). |
| 903 | |
| 904 | Args: |
| 905 | q: ``(q_values, q_scale)``; q_values is [num_q_tokens, 1, H, D]. |
| 906 | kv_cache: Fused paged cache [num_pages, page_kv, 1, D' + 4] uint8, |
| 907 | page stride 512B-aligned. |
| 908 | weights: [num_q_tokens, H] ``torch.bfloat16``. |
| 909 | metadata: From `get_paged_sparse_mqa_logits_metadata`. |
| 910 | num_max_sparse_blocks: Candidate blocks per row. |
| 911 | sparse_block_kv: Tokens per sparse block, 8 or 16. |
| 912 | |
| 913 | Returns: |
| 914 | bf16 logits of shape |
| 915 | [num_q_tokens, num_max_sparse_blocks * sparse_block_kv]. |
| 916 | |
| 917 | """ |
| 918 | _lazy_init() |
| 919 | if _fp8_fp4_paged_sparse_mqa_logits_impl is None: |
| 920 | return _missing() |
| 921 | return _fp8_fp4_paged_sparse_mqa_logits_impl( |
| 922 | q, |
| 923 | kv_cache, |
| 924 | weights, |
| 925 | metadata, |
| 926 | num_max_sparse_blocks, |
| 927 | sparse_block_kv, |
| 928 | ) |
| 929 | |
| 930 | |
| 931 | def tf32_hc_prenorm_gemm( |
| 932 | x: torch.Tensor, |
| 933 | fn: torch.Tensor, |
| 934 | out: torch.Tensor, |
| 935 | sqrsum: torch.Tensor, |
| 936 | num_split: int, |
| 937 | ) -> torch.Tensor: |
| 938 | """Perform the following computation: |
| 939 | out = x.float() @ fn.T |
| 940 | sqrsum = x.float().square().sum(-1) |
| 941 | |
| 942 | See the caller function for shape requirement |
| 943 | """ |
| 944 | _lazy_init() |
| 945 | if _tf32_hc_prenorm_gemm_impl is None: |
| 946 | return _missing() |
| 947 | return _tf32_hc_prenorm_gemm_impl( |
| 948 | x, |
| 949 | fn, |
| 950 | out, |
| 951 | sqrsum, |
| 952 | num_split, |
| 953 | ) |
| 954 | |
| 955 | |
| 956 | def mega_mhc(*args: Any, **kwargs: Any) -> None: |
| 957 | """Run DeepGEMM Mega mHC with caller-owned output tensors.""" |
| 958 | _lazy_init() |
| 959 | if _mega_mhc_impl is None: |
| 960 | _missing() |
| 961 | return |
| 962 | _mega_mhc_impl(*args, **kwargs) |
| 963 | |
| 964 | |
| 965 | def _ceil_to_ue8m0(x: torch.Tensor): |
| 966 | return torch.pow(2.0, torch.ceil(torch.log2(x.abs()))) |
| 967 | |
| 968 | |
| 969 | def _align(x: int, y: int) -> int: |
| 970 | return cdiv(x, y) * y |
| 971 | |
| 972 | |
| 973 | # Taken from https://github.com/deepseek-ai/DeepGEMM/blob/v2.1.1/csrc/utils/math.hpp#L19 |
| 974 | def get_tma_aligned_size(x: int, element_size: int) -> int: |
| 975 | return _align(x, 16 // element_size) |
| 976 | |
| 977 | |
| 978 | DEFAULT_BLOCK_SIZE = [128, 128] |
| 979 | |
| 980 | |
| 981 | # Taken from https://github.com/deepseek-ai/DeepGEMM/blob/dd6ed14acbc7445dcef224248a77ab4d22b5f240/deep_gemm/utils/math.py#L38 |
| 982 | @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) |
| 983 | def per_block_cast_to_fp8( |
| 984 | x: torch.Tensor, block_size: list[int] = DEFAULT_BLOCK_SIZE, use_ue8m0: bool = False |
| 985 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 986 | fp8_dtype = current_platform.fp8_dtype() |
| 987 | assert x.dim() == 2 |
| 988 | m, n = x.shape |
| 989 | block_m, block_n = block_size |
| 990 | x_padded = torch.zeros( |
| 991 | (_align(m, block_m), _align(n, block_n)), dtype=x.dtype, device=x.device |
| 992 | ) |
| 993 | x_padded[:m, :n] = x |
| 994 | x_view = x_padded.view(-1, block_m, x_padded.size(1) // block_n, block_n) |
| 995 | x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4) |
| 996 | _, fp8_max = get_fp8_min_max() |
| 997 | sf = x_amax / fp8_max |
| 998 | sf = _ceil_to_ue8m0(sf) if use_ue8m0 else sf |
| 999 | x_scaled = (x_view * (1.0 / sf)).to(fp8_dtype) |
| 1000 | return x_scaled.view_as(x_padded)[:m, :n].contiguous(), sf.view( |
| 1001 | x_view.size(0), x_view.size(2) |
| 1002 | ) |
| 1003 | |
| 1004 | |
| 1005 | def calc_diff(x: torch.Tensor, y: torch.Tensor): |
| 1006 | """Return a global difference metric for unit tests. |
| 1007 | |
| 1008 | DeepGEMM kernels on Blackwell/B200 currently exhibit noticeable per-element |
| 1009 | error, causing `torch.testing.assert_close` to fail. Instead of checking |
| 1010 | every element, we compute a cosine-style similarity over the whole tensor |
| 1011 | and report `1 - sim`. Once kernel accuracy improves this helper can be |
| 1012 | removed. |
| 1013 | """ |
| 1014 | x, y = x.double(), y.double() |
| 1015 | denominator = (x * x + y * y).sum() |
| 1016 | sim = 2 * (x * y).sum() / denominator |
| 1017 | return 1 - sim |
| 1018 | |
| 1019 | |
| 1020 | def should_use_deepgemm_for_fp8_linear( |
| 1021 | output_dtype: torch.dtype, |
| 1022 | weight_shape: tuple[int, int], |
| 1023 | supports_deep_gemm: bool | None = None, |
| 1024 | ): |
| 1025 | if supports_deep_gemm is None: |
| 1026 | supports_deep_gemm = is_deep_gemm_supported() |
| 1027 | |
| 1028 | # Verify DeepGEMM N/K dims requirements |
| 1029 | # NOTE: Also synchronized with test_w8a8_block_fp8_deep_gemm_matmul |
| 1030 | # test inside kernels/quantization/test_block_fp8.py |
| 1031 | N_MULTIPLE = 64 |
| 1032 | K_MULTIPLE = 128 |
| 1033 | |
| 1034 | return ( |
| 1035 | supports_deep_gemm |
| 1036 | and output_dtype == torch.bfloat16 |
| 1037 | and weight_shape[0] % N_MULTIPLE == 0 |
| 1038 | and weight_shape[1] % K_MULTIPLE == 0 |
| 1039 | ) |
| 1040 | |
| 1041 | |
| 1042 | __all__ = [ |
| 1043 | "calc_diff", |
| 1044 | "bf16_mega_gate", |
| 1045 | "DeepGemmQuantScaleFMT", |
| 1046 | "fp8_gemm_nt", |
| 1047 | "fp8_einsum", |
| 1048 | "m_grouped_fp8_gemm_nt_contiguous", |
| 1049 | "m_grouped_fp8_fp4_gemm_nt_contiguous", |
| 1050 | "fp8_m_grouped_gemm_nt_masked", |
| 1051 | "fp8_fp4_mqa_logits", |
| 1052 | "fp8_fp4_paged_mqa_logits", |
| 1053 | "get_paged_mqa_logits_metadata", |
| 1054 | "native_next_n_supported", |
| 1055 | "per_block_cast_to_fp8", |
| 1056 | "is_deep_gemm_e8m0_used", |
| 1057 | "is_deep_gemm_supported", |
| 1058 | "mega_mhc", |
| 1059 | "get_num_sms", |
| 1060 | "set_num_sms", |
| 1061 | "should_use_deepgemm_for_fp8_linear", |
| 1062 | "get_col_major_tma_aligned_tensor", |
| 1063 | "get_mk_alignment_for_contiguous_layout", |
| 1064 | "get_theoretical_mk_alignment_for_contiguous_layout", |
| 1065 | "pack_ue8m0_to_int", |
| 1066 | "get_mn_major_tma_aligned_packed_ue8m0_tensor", |
| 1067 | "get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor", |
| 1068 | ] |