diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 1dbc33e6b20..6d70548aee8 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -1,6 +1,20 @@ Changelog ========= +0.45 (2026-06-xx) +^^^^^^^^^^^^^^^^^ + +**Backward Breaking Changes** + +- Reorganize custom CUDA / Triton kernels under ``modelopt.torch.kernels`` into ``common/attention``, ``quantization/{conv,gemm}``, and ``sparsity/attention``. High-level APIs (``mtq.quantize``, ``mtsa.sparsify``, etc.) are unchanged, but **any code importing directly from the kernel subpackages must be updated**: there is no backwards-compatibility shim; the old import paths will raise ``ImportError`` / ``ModuleNotFoundError``. Migration table: + + - ``from modelopt.torch.kernels import IS_AVAILABLE, attention, attention_calibrate, register_triton_attention`` → ``from modelopt.torch.kernels.common.attention import ...`` + - ``from modelopt.torch.kernels.triton_fa import ...`` → ``from modelopt.torch.kernels.common.attention.triton_fa import ...`` + - ``from modelopt.torch.kernels.hf_triton_attention import ...`` → ``from modelopt.torch.kernels.common.attention.hf_triton_attention import ...`` + - ``from modelopt.torch.quantization.triton import ...`` → ``from modelopt.torch.kernels.quantization.gemm import ...`` + - ``from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import ...`` → ``from modelopt.torch.kernels.quantization.conv.implicit_gemm_cuda import ...`` + - ``from modelopt.torch.sparsity.attention_sparsity.kernels import ...`` → ``from modelopt.torch.kernels.sparsity.attention import ...`` + 0.44 (2026-05-xx) ^^^^^^^^^^^^^^^^^ @@ -9,14 +23,14 @@ Changelog - Support full Transformer Engine spec for Minitron pruning (``mcore_minitron``). Now we no longer need to use custom ModelOpt spec. Note that this does not affect the usage of the pruning workflow but makes pruning slightly faster and may result in slightly different pruned model because of different kernel and numerics. - Add Puzzletron - a new algorithm for heterogeneous pruning of LLM and VLM models. See `examples/puzzletron/README.md `_ for more details. - Added iterator interface using CalibrationDataReader in ONNX quantization workflow. -- Add N:M sparse softmax support to the Triton flash attention kernel (``modelopt.torch.kernels.triton_fa``). See `examples/llm_sparsity/attention_sparsity/README.md `_ for usage. -- Add skip-softmax skipping to the Triton flash attention kernel (``modelopt.torch.kernels.triton_fa``). See `examples/llm_sparsity/attention_sparsity/README.md `_ for usage. +- Add N:M sparse softmax support to the Triton flash attention kernel (``modelopt.torch.kernels.common.attention.triton_fa``). See `examples/llm_sparsity/attention_sparsity/README.md `_ for usage. +- Add skip-softmax skipping to the Triton flash attention kernel (``modelopt.torch.kernels.common.attention.triton_fa``). See `examples/llm_sparsity/attention_sparsity/README.md `_ for usage. - Add Video Sparse Attention (VSA) method for video diffusion models (``modelopt.torch.sparsity.attention_sparsity``). VSA uses 3D block tiling with a two-branch architecture for attention speedup. - Enable PTQ workflow for the Step3.5-Flash MoE model with NVFP4 W4A4 + FP8 KV cache quantization. See `modelopt_recipes/models/Step3.5-Flash/nvfp4-mlp-only.yaml `_ for more details. - Add support for vLLM fakequant reload using ModelOpt state for HF models. See `examples/vllm_serve/README.md `_ for more details. - [Early Testing] Add Claude Code PTQ skill (``.claude/skills/ptq/``) for agent-assisted post-training quantization. The skill guides the agent through environment detection, model support checking, format selection, and execution via the launcher or manual SLURM/Docker/bare GPU paths. Includes handling for unlisted models with custom module patching. This feature is in early testing — use with caution. - Add performant layerwise calibration for large models that don't fit on GPU (e.g. DeepSeek-R1, Kimi-K2). See `modelopt_recipes/general/ptq/nvfp4_experts_only-fp8_kv.yaml `_ for usage. Layerwise calibration also supports PTQ with intermediate progress saving — useful when long PTQ runs get hit with Slurm timeouts. See `modelopt_recipes/general/ptq/nvfp4_default-none_kv_gptq.yaml `_ for usage. -- Add implicit GEMM CUDA kernel for Conv3D with fused NVFP4 fake quantization (``modelopt.torch.quantization.src.conv``). When NVFP4 quantization is applied to an ``nn.Conv3d`` layer via ModelOpt PTQ, the implicit GEMM path is used automatically instead of cuDNN. Uses BF16 WMMA tensor cores (SM80+) with FP32 accumulation and in-kernel FP4 (E2M1) activation quantization. Grouped convolution (``groups > 1``) falls back to the default cuDNN path. Inference only — training mode falls back to cuDNN with a warning. +- Add implicit GEMM CUDA kernel for Conv3D with fused NVFP4 fake quantization (``modelopt.torch.kernels.quantization.conv``). When NVFP4 quantization is applied to an ``nn.Conv3d`` layer via ModelOpt PTQ, the implicit GEMM path is used automatically instead of cuDNN. Uses BF16 WMMA tensor cores (SM80+) with FP32 accumulation and in-kernel FP4 (E2M1) activation quantization. Grouped convolution (``groups > 1``) falls back to the default cuDNN path. Inference only — training mode falls back to cuDNN with a warning. **Backward Breaking Changes** diff --git a/CLAUDE.md b/CLAUDE.md index 4af38586788..d0b47148c38 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -65,6 +65,7 @@ ModelOpt code base is organized into four top-level namespaces: | `nas` | `modelopt/torch/nas/` | Neural architecture search | | `export` | `modelopt/torch/export/` | Checkpoint export for TRT-LLM / Megatron | | `peft` | `modelopt/torch/peft/` | QLoRA and PEFT integration | +| `kernels` | `modelopt/torch/kernels/` | Custom CUDA/Triton kernels grouped by role: `common/attention` (baseline Triton FA), `quantization/{conv,gemm}` (implicit-GEMM CUDA + tensor-quant C++/CUDA + fp4/fp8 Triton), `sparsity/attention` (skip-softmax / N:M / diffusers+LTX backends) | | `_deploy` | `modelopt/torch/_deploy/` | Internal deployment utilities | | `utils` | `modelopt/torch/utils/` | Shared utilities and plugin infrastructure | diff --git a/examples/deepseek/ptq.py b/examples/deepseek/ptq.py index 30574b3eee5..d60d011ed00 100644 --- a/examples/deepseek/ptq.py +++ b/examples/deepseek/ptq.py @@ -55,8 +55,8 @@ import modelopt.torch.quantization as mtq from modelopt.torch.export.model_config import KV_CACHE_FP8 from modelopt.torch.export.quant_utils import get_quant_config +from modelopt.torch.kernels.quantization.gemm import weight_dequant from modelopt.torch.quantization.nn import TensorQuantizer -from modelopt.torch.quantization.triton import weight_dequant from modelopt.torch.quantization.utils import ( is_quantized_column_parallel_linear, is_quantized_parallel_linear, diff --git a/examples/deepseek/quantize_to_nvfp4.py b/examples/deepseek/quantize_to_nvfp4.py index db6d5f6a24f..e54fdbebf46 100644 --- a/examples/deepseek/quantize_to_nvfp4.py +++ b/examples/deepseek/quantize_to_nvfp4.py @@ -47,8 +47,8 @@ from safetensors.torch import load_file, save_file from tqdm import tqdm +from modelopt.torch.kernels.quantization.gemm import weight_dequant from modelopt.torch.quantization.qtensor import NVFP4QTensor -from modelopt.torch.quantization.triton import weight_dequant def _remap_key(key_dict: dict[str, Any]): diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index ac14d982279..8e0f7cdef90 100644 --- a/examples/diffusers/README.md +++ b/examples/diffusers/README.md @@ -119,7 +119,7 @@ python quantize.py \ #### Wan 2.2 VAE NVFP4 (Conv3D Implicit GEMM) -The Wan 2.2 VAE (`AutoencoderKLWan`, shared between the 5B and 14B pipelines) is built from 3D convolutions. When quantizing the VAE with NVFP4, the `Conv3d` layers are automatically dispatched through a custom BF16 WMMA implicit-GEMM kernel with fused FP4 activation quantization. Requires SM80+ (Ampere or newer). See [`modelopt/torch/quantization/src/conv/README.md`](../../modelopt/torch/quantization/src/conv/README.md) for kernel details. +The Wan 2.2 VAE (`AutoencoderKLWan`, shared between the 5B and 14B pipelines) is built from 3D convolutions. When quantizing the VAE with NVFP4, the `Conv3d` layers are automatically dispatched through a custom BF16 WMMA implicit-GEMM kernel with fused FP4 activation quantization. Requires SM80+ (Ampere or newer). See [`modelopt/torch/kernels/quantization/conv/README.md`](../../modelopt/torch/kernels/quantization/conv/README.md) for kernel details. ```sh python quantize.py \ diff --git a/modelopt/torch/kernels/__init__.py b/modelopt/torch/kernels/__init__.py index fa07b06e20c..151b6e21d9e 100644 --- a/modelopt/torch/kernels/__init__.py +++ b/modelopt/torch/kernels/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -13,38 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Shared Triton kernels for modelopt (attention, quantization, etc.).""" - -import torch - -from modelopt.torch.utils import import_plugin - -IS_AVAILABLE = False -attention = None -attention_calibrate = None -register_triton_attention = None - -if torch.cuda.is_available(): - with import_plugin( - "triton", - msg_if_missing=( - "Your device is potentially capable of using the triton attention " - "kernel. Try to install triton with `pip install triton`." - ), - ): - from .triton_fa import attention as _attention - from .triton_fa import attention_calibrate as _attention_calibrate - - attention = _attention - attention_calibrate = _attention_calibrate - IS_AVAILABLE = True - from .hf_triton_attention import register_triton_attention as _register_triton_attention - - register_triton_attention = _register_triton_attention - -__all__ = [ - "IS_AVAILABLE", - "attention", - "attention_calibrate", - "register_triton_attention", -] +"""ModelOpt kernel library: common, quantization (conv, gemm), sparsity (attention, gemm).""" diff --git a/modelopt/torch/kernels/common/__init__.py b/modelopt/torch/kernels/common/__init__.py new file mode 100644 index 00000000000..f5c9e562d32 --- /dev/null +++ b/modelopt/torch/kernels/common/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Common (non-domain-specific) kernels. Base FA lives in ``common/attention``.""" diff --git a/modelopt/torch/kernels/common/attention/__init__.py b/modelopt/torch/kernels/common/attention/__init__.py new file mode 100644 index 00000000000..caf319a765e --- /dev/null +++ b/modelopt/torch/kernels/common/attention/__init__.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Shared Triton kernels for modelopt (attention, quantization, etc.).""" + +import torch + +from modelopt.torch.utils import import_plugin + +IS_AVAILABLE = False +attention = None +attention_calibrate = None +register_triton_attention = None + +if torch.cuda.is_available(): + with import_plugin( + "triton", + msg_if_missing=( + "Your device is potentially capable of using the triton attention " + "kernel. Try to install triton with `pip install triton`." + ), + ): + from .triton_fa import attention as _attention + + attention = _attention + IS_AVAILABLE = True + from .hf_triton_attention import register_triton_attention as _register_triton_attention + + register_triton_attention = _register_triton_attention + + # Calibration lives in the sparsity subpackage (skip-softmax specific). + # Imported here so ``from modelopt.torch.kernels.common.attention import + # attention_calibrate`` keeps working. + from modelopt.torch.kernels.sparsity.attention.calibrate import ( + attention_calibrate as _attention_calibrate, + ) + + attention_calibrate = _attention_calibrate + +__all__ = [ + "IS_AVAILABLE", + "attention", + "attention_calibrate", + "register_triton_attention", +] diff --git a/modelopt/torch/kernels/hf_triton_attention.py b/modelopt/torch/kernels/common/attention/hf_triton_attention.py similarity index 98% rename from modelopt/torch/kernels/hf_triton_attention.py rename to modelopt/torch/kernels/common/attention/hf_triton_attention.py index 5021d34e379..235487462f2 100644 --- a/modelopt/torch/kernels/hf_triton_attention.py +++ b/modelopt/torch/kernels/common/attention/hf_triton_attention.py @@ -25,7 +25,7 @@ import torch import torch.nn as nn -from modelopt.torch.kernels.triton_fa import attention +from modelopt.torch.kernels.common.attention.triton_fa import attention def _seq_lens_from_mask( diff --git a/modelopt/torch/kernels/triton_fa.py b/modelopt/torch/kernels/common/attention/triton_fa.py similarity index 66% rename from modelopt/torch/kernels/triton_fa.py rename to modelopt/torch/kernels/common/attention/triton_fa.py index 8044383889f..a4b3cc90e32 100644 --- a/modelopt/torch/kernels/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -24,11 +24,44 @@ """ import math +from typing import Any import torch import triton import triton.language as tl +# Helpers for optional N:M sparsity and sink/window-aware dense regions live +# in the sparsity package. The baseline forward kernel below calls them +# conditionally under constexpr guards, so the unified single-kernel design +# stays intact while keeping feature-specific logic in its own subpackage. +# +# Lazy import: Triton resolves @triton.jit names at kernel compile time (first +# call), not at definition time, so populating the module globals before the +# first ``attention()`` call is sufficient. Deferring avoids a circular import +# (common.attention/__init__.py ↔ sparsity.attention/__init__.py via this file). +_apply_sparse_nm_to_qk_tile: Any = None +_is_dense_region: Any = None +_skip_softmax_decision: Any = None + + +def _load_sparsity_helpers() -> None: + global _apply_sparse_nm_to_qk_tile, _is_dense_region, _skip_softmax_decision + if _apply_sparse_nm_to_qk_tile is None: + from modelopt.torch.kernels.sparsity.attention.skip_softmax_helpers import ( + _apply_sparse_nm_to_qk_tile as _nm, + ) + from modelopt.torch.kernels.sparsity.attention.skip_softmax_helpers import ( + _is_dense_region as _dense, + ) + from modelopt.torch.kernels.sparsity.attention.skip_softmax_helpers import ( + _skip_softmax_decision as _skip, + ) + + _apply_sparse_nm_to_qk_tile = _nm + _is_dense_region = _dense + _skip_softmax_decision = _skip + + LOG2E: float = 1.44269504088896 # --------------------------------------------------------------------------- @@ -47,145 +80,6 @@ _FWD_CONFIGS = [triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_stages=1, num_warps=4)] -# --------------------------------------------------------------------------- -# N:M sparse softmax helpers -# --------------------------------------------------------------------------- -@triton.jit -def _sparse_nm_masks_m4(x0, x1, x2, x3, N: tl.constexpr): - """Top-N of 4 selection via pure boolean logic (6 comparisons, no int casts). - - Uses ``>=`` so that ties are broken by index (lower index wins). - Guarantees exactly N masks are True for any input including all-equal. - - Boolean formulas for "at least K of 3 wins": - K=3 (N=1): AND of all — must beat all 3 others - K=2 (N=2): majority — must beat at least 2 (sorting network) - K=1 (N=3): OR of all — must beat at least 1 - """ - c01 = x0 >= x1 - c02 = x0 >= x2 - c03 = x0 >= x3 - c12 = x1 >= x2 - c13 = x1 >= x3 - c23 = x2 >= x3 - - nc01 = ~c01 - nc02 = ~c02 - nc03 = ~c03 - nc12 = ~c12 - nc13 = ~c13 - nc23 = ~c23 - - if N == 1: - # Keep max only: must beat all 3 - m0 = c01 & c02 & c03 - m1 = nc01 & c12 & c13 - m2 = nc02 & nc12 & c23 - m3 = nc03 & nc13 & nc23 - elif N == 2: - # Majority vote: must beat at least 2 of 3 - m0 = (c01 & c02) | (c01 & c03) | (c02 & c03) - m1 = (nc01 & c12) | (nc01 & c13) | (c12 & c13) - m2 = (nc02 & nc12) | (nc02 & c23) | (nc12 & c23) - m3 = (nc03 & nc13) | (nc03 & nc23) | (nc13 & nc23) - elif N == 3: - # Keep all but min: must beat at least 1 - m0 = c01 | c02 | c03 - m1 = nc01 | c12 | c13 - m2 = nc02 | nc12 | c23 - m3 = nc03 | nc13 | nc23 - else: - tl.static_assert(False, "N must be 1, 2, or 3 for M=4") - - return m0, m1, m2, m3 - - -@triton.jit -def _apply_sparse_nm_to_qk_tile( - qk, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, - SPARSITY_N: tl.constexpr, - SPARSITY_M: tl.constexpr, -): - """Apply N:M sparse softmax to a QK score tile. - - For every ``SPARSITY_M`` consecutive elements along the N (key) dimension, - keeps the top ``SPARSITY_N`` values and sets the rest to ``-inf``. - ``BLOCK_N`` must be divisible by ``SPARSITY_M``. - - For M=4, exactly N values are retained (ties broken by position). - For M=8, a threshold-based approach (``tl.sort``) may retain more - than N values when ties straddle the threshold boundary. - """ - tl.static_assert(SPARSITY_M == 4 or SPARSITY_M == 8, "SPARSITY_M must be 4 or 8") # noqa: PLR1714 - MASK_VAL: tl.constexpr = float("-inf") - - if SPARSITY_M == 4: - tl.static_assert(BLOCK_N % 4 == 0, "BLOCK_N must be divisible by 4") - reshaped = tl.reshape(qk, (BLOCK_M, BLOCK_N // 4, 4)) - cols = tl.arange(0, 4)[None, None, :] - x0 = tl.sum(tl.where(cols == 0, reshaped, 0.0), axis=2) - x1 = tl.sum(tl.where(cols == 1, reshaped, 0.0), axis=2) - x2 = tl.sum(tl.where(cols == 2, reshaped, 0.0), axis=2) - x3 = tl.sum(tl.where(cols == 3, reshaped, 0.0), axis=2) - - m0, m1, m2, m3 = _sparse_nm_masks_m4(x0, x1, x2, x3, SPARSITY_N) - - out = tl.full((BLOCK_M, BLOCK_N // 4, 4), 0.0, dtype=qk.dtype) - out = tl.where(cols == 0, tl.expand_dims(tl.where(m0, x0, MASK_VAL), 2), out) - out = tl.where(cols == 1, tl.expand_dims(tl.where(m1, x1, MASK_VAL), 2), out) - out = tl.where(cols == 2, tl.expand_dims(tl.where(m2, x2, MASK_VAL), 2), out) - out = tl.where(cols == 3, tl.expand_dims(tl.where(m3, x3, MASK_VAL), 2), out) - return tl.reshape(out, (BLOCK_M, BLOCK_N)) - - else: # SPARSITY_M == 8 - tl.static_assert(BLOCK_N % 8 == 0, "BLOCK_N must be divisible by 8") - reshaped = tl.reshape(qk, (BLOCK_M, BLOCK_N // 8, 8)) - - # Sort each group of 8 ascending; N-th largest is at index (8 - N) - sorted_vals = tl.sort(reshaped, dim=2) - KTH_IDX: tl.constexpr = SPARSITY_M - SPARSITY_N # index of N-th largest in ascending order - - # Extract the threshold value at KTH_IDX via masked sum - # Use 0.0 as fill (not -inf) so sum equals just the KTH element - cols = tl.arange(0, 8)[None, None, :] - threshold = tl.sum(tl.where(cols == KTH_IDX, sorted_vals, 0.0), axis=2) - - # Mask: keep elements >= threshold (may keep >N on ties — acceptable) - mask = reshaped >= tl.expand_dims(threshold, 2) - return tl.reshape(tl.where(mask, reshaped, MASK_VAL), (BLOCK_M, BLOCK_N)) - - -# --------------------------------------------------------------------------- -# Sink/window dense-region check -# --------------------------------------------------------------------------- -@triton.jit -def _is_dense_region( - kv_start, - tile_q, - seq_len_q, - seq_len_kv, - BLOCK_M: tl.constexpr, - NUM_SINK_TOKENS: tl.constexpr, - DENSE_WINDOW_SIZE: tl.constexpr, -): - """Check if a KV tile falls in a dense region (sink tokens or local window). - - Uses absolute token positions so the result is BLOCK_N-independent, - ensuring forward and backward (which may use different BLOCK_N) agree. - - Returns: - True if the tile should be kept dense (skip N:M sparsification). - """ - is_sink = kv_start < NUM_SINK_TOKENS - causal_offset = seq_len_kv - seq_len_q - q_abs_pos = tile_q * BLOCK_M + causal_offset - token_distance = q_abs_pos - kv_start - is_local = (token_distance >= 0) and (token_distance < DENSE_WINDOW_SIZE) - return is_sink or is_local - - # --------------------------------------------------------------------------- # Masking helper # --------------------------------------------------------------------------- @@ -327,55 +221,22 @@ def _attn_fwd( scores, BLOCK_M, BLOCK_N, SPARSITY_N, SPARSITY_M ) + # Optional skip-softmax decision — the decision logic (and optional + # atomic counter updates) lives in sparsity/attention; this kernel + # just consults it under its constexpr guard. + skip_tile = False if APPLY_SKIP_SOFTMAX: - # --- Skip-softmax (BLASST, https://arxiv.org/pdf/2512.12087) --- - # - # Algorithm: During FlashAttention's block-wise computation, we - # maintain a running maximum m_i^(j) across blocks. If a block's - # local maximum ~m_i^(j) is significantly smaller than the running - # maximum m_i^(j): - # - # ~m_i^(j) - m_i^(j) < ln(lambda) - # - # then exp(~m_i^(j) - m_i^(j)) < lambda ≈ 0, meaning the block's - # contribution to the final output is negligible. We skip the - # softmax computation, V load, and BMM2 computation entirely. - # - # The threshold is pre-scaled by qk_scale in the Python wrapper so - # it can be compared directly against scaled scores (matching the - # BLASST reference semantics on unscaled scores). - tile_row_max = tl.max(scores, 1) # [BLOCK_M] — ~m_i^(j) (scaled) - # Per-row: True if row's tile max is negligible vs running max - can_skip = tile_row_max < (row_max + SKIP_THRESHOLD_LOG2) - # Per-tile: skip entire tile only if ALL rows are negligible - skip_tile = tl.min(can_skip.to(tl.int32)) == 1 - - # Optional runtime sparsity measurement via atomic counters - if MEASURE_SPARSITY: - tl.atomic_add(Sparsity_total, 1) # count every tile - if skip_tile: - tl.atomic_add(Sparsity_skipped, 1) # count skipped tiles - - if not skip_tile: - m_new = tl.maximum(row_max, tile_row_max) - p = tl.math.exp2(scores - m_new[:, None]) - l_new = tl.sum(p, 1) - correction = tl.math.exp2(row_max - m_new) - row_sum = row_sum * correction + l_new - acc = acc * correction[:, None] - - v_offs = (kv_offset + kv_start + kv_pos[:, None]) * stride_vbs + dim_pos[None, :] - v = tl.load( - v_base + v_offs, - mask=((kv_start + kv_pos[:, None]) < seq_len_kv) & d_mask[None, :], - other=0.0, - ) - acc = tl.dot(p.to(v.dtype), v, acc) - row_max = m_new - # else: tile skipped: no softmax computation, V load, and BMM2 computation - else: - # --- Standard path: no skip check --- - # Online softmax update + skip_tile = _skip_softmax_decision( + scores, + row_max, + SKIP_THRESHOLD_LOG2, + Sparsity_total, + Sparsity_skipped, + MEASURE_SPARSITY, + ) + + if not skip_tile: + # --- Online softmax update --- m_new = tl.maximum(row_max, tl.max(scores, 1)) p = tl.math.exp2(scores - m_new[:, None]) l_new = tl.sum(p, 1) @@ -392,6 +253,7 @@ def _attn_fwd( ) acc = tl.dot(p.to(v.dtype), v, acc) row_max = m_new + # else: tile skipped — no softmax, no V load, no BMM2 for this tile # --- Final normalization: output = acc / row_sum --- # Clamp denominator to avoid 0/0 NaN when skip-softmax skips all KV tiles. @@ -1092,6 +954,7 @@ def attention( Returns: Output tensor [total_q_tokens, num_q_heads, head_dim]. """ + _load_sparsity_helpers() sm_scale = 1.0 / (q.shape[2] ** 0.5) if softmax_scale is None else softmax_scale return _Attention.apply( q, @@ -1115,266 +978,4 @@ def attention( ) -# --------------------------------------------------------------------------- -# Calibration kernel: collect multi-threshold skip-softmax sparsity stats -# --------------------------------------------------------------------------- -@triton.jit -def _attn_fwd_calibrate( - Q, - K, - V, - qk_scale, - b_start_loc, - b_seq_len, - b_start_loc_k, - b_seq_len_k, - Out, - stride_qbs, - stride_qh, - stride_kbs, - stride_kh, - stride_vbs, - stride_vh, - stride_obs, - stride_oh, - Threshold_trials, # [NUM_THRESHOLDS] float32 — pre-scaled to log2 space - Per_program_totals, # [num_programs * NUM_THRESHOLDS] int32 — per-program tile counts - Per_program_skipped, # [num_programs * NUM_THRESHOLDS] int32 — per-program skip counts - kv_group_num: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_D: tl.constexpr, - BLOCK_N: tl.constexpr, - IS_CAUSAL: tl.constexpr, - HEAD_DIM: tl.constexpr, - NUM_THRESHOLDS: tl.constexpr, - PADDED_THRESHOLDS: tl.constexpr, # next_power_of_2(NUM_THRESHOLDS) for tl.arange -): - """Forward kernel with multi-threshold sparsity measurement. - - Computes full attention (no skipping) while counting how many KV tiles - would be skipped at each threshold. Each program writes its local counts - to ``Per_program_totals`` and ``Per_program_skipped``; the Python wrapper - sums across programs afterward. This avoids global atomic contention. - """ - batch_idx = tl.program_id(0) - head_idx = tl.program_id(1) - tile_q = tl.program_id(2) - kv_head_idx = head_idx // kv_group_num - - seq_len_q = tl.load(b_seq_len + batch_idx) - seq_len_kv = tl.load(b_seq_len_k + batch_idx) - q_offset = tl.load(b_start_loc + batch_idx) - kv_offset = tl.load(b_start_loc_k + batch_idx) - - if tile_q * BLOCK_M >= seq_len_q: - return - - q_pos = tile_q * BLOCK_M + tl.arange(0, BLOCK_M) - kv_pos = tl.arange(0, BLOCK_N) - dim_pos = tl.arange(0, BLOCK_D) - d_mask = dim_pos < HEAD_DIM - - q_ptrs = (q_offset + q_pos[:, None]) * stride_qbs + head_idx * stride_qh + dim_pos[None, :] - q = tl.load(Q + q_ptrs, mask=(q_pos[:, None] < seq_len_q) & d_mask[None, :], other=0.0) - - k_base = K + kv_head_idx * stride_kh - v_base = V + kv_head_idx * stride_vh - - row_max = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") - row_sum = tl.zeros([BLOCK_M], dtype=tl.float32) - acc = tl.zeros([BLOCK_M, BLOCK_D], dtype=tl.float32) - - # Pre-load all thresholds once (vectorized, stays in registers). - # tl.arange requires power-of-2 size, so use PADDED_THRESHOLDS with masking. - thresh_offs = tl.arange(0, PADDED_THRESHOLDS) - thresh_mask = thresh_offs < NUM_THRESHOLDS - thresholds = tl.load(Threshold_trials + thresh_offs, mask=thresh_mask, other=float("inf")) - - # Per-program local counters: avoid global atomic contention in inner loop. - # Each program accumulates locally, then writes once to Per_program buffers. - local_skipped = tl.zeros([PADDED_THRESHOLDS], dtype=tl.int32) - num_tiles = 0 - - kv_bound = seq_len_kv if not IS_CAUSAL else tl.minimum((tile_q + 1) * BLOCK_M, seq_len_kv) - - for kv_start in range(0, kv_bound, BLOCK_N): - kv_start = tl.multiple_of(kv_start, BLOCK_N) - - k_offs = (kv_offset + kv_start + kv_pos[None, :]) * stride_kbs + dim_pos[:, None] - k = tl.load( - k_base + k_offs, - mask=((kv_start + kv_pos[None, :]) < seq_len_kv) & d_mask[:, None], - other=0.0, - ) - - scores = tl.dot(q, k) * qk_scale - scores = _apply_mask(scores, q_pos, kv_pos, seq_len_q, seq_len_kv, kv_start, IS_CAUSAL) - - tile_row_max = tl.max(scores, 1) - - # --- Vectorized multi-threshold sparsity measurement --- - # A tile is skipped iff ALL Q rows satisfy: tile_row_max < row_max + thresh. - # Equivalently: max(tile_row_max - row_max) < thresh (worst-case row - # must still be below threshold for the tile to be skippable). - max_gap = tl.max(tile_row_max - row_max) # scalar - skip_mask = (max_gap < thresholds).to(tl.int32) # [PADDED_THRESHOLDS] - local_skipped += skip_mask - num_tiles += 1 - - # --- Always compute full attention (no skipping) --- - m_new = tl.maximum(row_max, tile_row_max) - p = tl.math.exp2(scores - m_new[:, None]) - l_new = tl.sum(p, 1) - correction = tl.math.exp2(row_max - m_new) - row_sum = row_sum * correction + l_new - acc = acc * correction[:, None] - - v_offs = (kv_offset + kv_start + kv_pos[:, None]) * stride_vbs + dim_pos[None, :] - v = tl.load( - v_base + v_offs, - mask=((kv_start + kv_pos[:, None]) < seq_len_kv) & d_mask[None, :], - other=0.0, - ) - acc = tl.dot(p.to(v.dtype), v, acc) - row_max = m_new - - # --- Write per-program counters (no atomics, just stores) --- - # Compute unique flat program index for this (batch, head, q_tile) - num_q_tiles = tl.cdiv(tl.load(b_seq_len + 0), BLOCK_M) # conservative upper bound - num_heads = tl.num_programs(1) - prog_idx = batch_idx * num_heads * num_q_tiles + head_idx * num_q_tiles + tile_q - base = prog_idx * NUM_THRESHOLDS - tl.store( - Per_program_totals + base + thresh_offs, - tl.full([PADDED_THRESHOLDS], num_tiles, dtype=tl.int32), - mask=thresh_mask, - ) - tl.store( - Per_program_skipped + base + thresh_offs, - local_skipped, - mask=thresh_mask, - ) - - acc = acc / tl.maximum(row_sum[:, None], 1e-6) - o_ptrs = (q_offset + q_pos[:, None]) * stride_obs + head_idx * stride_oh + dim_pos[None, :] - tl.store(Out + o_ptrs, acc, mask=(q_pos[:, None] < seq_len_q) & d_mask[None, :]) - - -def attention_calibrate( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - b_start_loc: torch.Tensor, - b_seq_len: torch.Tensor, - max_input_len: int, - is_causal: bool = True, - softmax_scale: float | None = None, - b_start_loc_k: torch.Tensor | None = None, - b_seq_len_k: torch.Tensor | None = None, - max_input_len_k: int | None = None, - *, - threshold_trials: list[float] | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - """Flash attention with multi-threshold skip-softmax sparsity measurement. - - Computes full attention (identical output to dense attention) while - measuring how many KV tiles would be skipped at each threshold in - ``threshold_trials``. No autograd — forward only. - - Args: - q, k, v, b_start_loc, b_seq_len, max_input_len, is_causal, - softmax_scale, b_start_loc_k, b_seq_len_k, max_input_len_k: - Same as :func:`attention`. - threshold_trials: List of threshold values to measure sparsity for. - Each value is converted to log2-scaled space for the kernel. - - Returns: - Tuple of (output, sparsity_counters): - - output: ``[total_q_tokens, num_q_heads, head_dim]`` - - sparsity_counters: ``[num_thresholds, 2]`` int64 tensor where - ``[:, 0]`` = total tile evaluations, ``[:, 1]`` = skipped tiles. - Sparsity per threshold = ``counters[:, 1] / counters[:, 0]``. - """ - if threshold_trials is None or len(threshold_trials) == 0: - raise ValueError("threshold_trials must be a non-empty list") - - HEAD_DIM = q.shape[2] - num_q_heads = q.shape[1] - num_kv_heads = k.shape[1] - kv_group_num = num_q_heads // num_kv_heads - batch = b_seq_len.shape[0] - sm_scale = 1.0 / (HEAD_DIM**0.5) if softmax_scale is None else softmax_scale - qk_scale = sm_scale * LOG2E - BLOCK_D = triton.next_power_of_2(HEAD_DIM) - BLOCK_M = 128 - BLOCK_N = 64 - - if b_seq_len_k is None: - b_seq_len_k = b_seq_len - b_start_loc_k = b_start_loc - - num_thresholds = len(threshold_trials) - - # Convert thresholds to log2-scaled space: log2(lambda) * sm_scale - threshold_tensor = torch.tensor( - [math.log2(t) * sm_scale for t in threshold_trials], - dtype=torch.float32, - device=q.device, - ) - - o = torch.empty_like(q) - - num_q_tiles = triton.cdiv(max_input_len, BLOCK_M) - grid = (batch, num_q_heads, num_q_tiles) - num_programs = batch * num_q_heads * num_q_tiles - - # Per-program output buffers (no atomics needed — each program writes its own row) - per_program_totals = torch.zeros( - num_programs * num_thresholds, dtype=torch.int32, device=q.device - ) - per_program_skipped = torch.zeros( - num_programs * num_thresholds, dtype=torch.int32, device=q.device - ) - - _attn_fwd_calibrate[grid]( - q, - k, - v, - qk_scale, - b_start_loc, - b_seq_len, - b_start_loc_k, - b_seq_len_k, - o, - q.stride(0), - q.stride(1), - k.stride(0), - k.stride(1), - v.stride(0), - v.stride(1), - o.stride(0), - o.stride(1), - threshold_tensor, - per_program_totals, - per_program_skipped, - kv_group_num=kv_group_num, - BLOCK_M=BLOCK_M, - BLOCK_D=BLOCK_D, - BLOCK_N=BLOCK_N, - IS_CAUSAL=is_causal, - HEAD_DIM=HEAD_DIM, - NUM_THRESHOLDS=num_thresholds, - PADDED_THRESHOLDS=triton.next_power_of_2(num_thresholds), - num_warps=4, - num_stages=1, - ) - - # Reduce across programs: sum per-program counts → [num_thresholds] - totals = per_program_totals.view(num_programs, num_thresholds).sum(dim=0).to(torch.int64) - skipped = per_program_skipped.view(num_programs, num_thresholds).sum(dim=0).to(torch.int64) - sparsity_counters = torch.stack([totals, skipped], dim=1) # [num_thresholds, 2] - - return o, sparsity_counters - - -__all__ = ["attention", "attention_calibrate"] +__all__ = ["LOG2E", "_apply_mask", "attention"] diff --git a/modelopt/torch/kernels/quantization/__init__.py b/modelopt/torch/kernels/quantization/__init__.py new file mode 100644 index 00000000000..1ae6845c90f --- /dev/null +++ b/modelopt/torch/kernels/quantization/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Quantization kernels: conv (implicit GEMM) and gemm (tensor_quant + Triton FP4/FP8).""" diff --git a/modelopt/torch/kernels/quantization/attention/__init__.py b/modelopt/torch/kernels/quantization/attention/__init__.py new file mode 100644 index 00000000000..ee64a4dd673 --- /dev/null +++ b/modelopt/torch/kernels/quantization/attention/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Quantization-specific attention kernel pieces (placeholder for combined sparse+quant path).""" diff --git a/modelopt/torch/quantization/src/conv/README.md b/modelopt/torch/kernels/quantization/conv/README.md similarity index 95% rename from modelopt/torch/quantization/src/conv/README.md rename to modelopt/torch/kernels/quantization/conv/README.md index 6b14fd5953b..ae61235514b 100644 --- a/modelopt/torch/quantization/src/conv/README.md +++ b/modelopt/torch/kernels/quantization/conv/README.md @@ -32,7 +32,7 @@ When NVFP4 quantization is configured on a `Conv3d` layer via ModelOpt PTQ, the ```python import torch -from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda +from modelopt.torch.kernels.quantization.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda from modelopt.torch.quantization.tensor_quant import dynamic_block_quantize_op x = torch.randn(1, 128, 21, 60, 106, device="cuda") @@ -75,7 +75,7 @@ out_q = conv3d_implicit_gemm_cuda( ### `conv3d_implicit_gemm_cuda` -`from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda` +`from modelopt.torch.kernels.quantization.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda` | Parameter | Description | |-----------|-------------| @@ -91,7 +91,7 @@ out_q = conv3d_implicit_gemm_cuda( ### `fp4_fake_quant` -`from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import fp4_fake_quant` +`from modelopt.torch.kernels.quantization.conv.implicit_gemm_cuda import fp4_fake_quant` Standalone FP4 (E2M1) blockwise fake quantization with FP8 E4M3 scale quantization. Uses the same CUDA device functions as the fused path inside the GEMM kernel. diff --git a/modelopt/torch/kernels/quantization/conv/__init__.py b/modelopt/torch/kernels/quantization/conv/__init__.py new file mode 100644 index 00000000000..ac5091fa9d2 --- /dev/null +++ b/modelopt/torch/kernels/quantization/conv/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Implicit-GEMM CUDA kernel for quantized 3D convolution.""" diff --git a/modelopt/torch/quantization/src/conv/bench_implicit_gemm.py b/modelopt/torch/kernels/quantization/conv/bench_implicit_gemm.py similarity index 98% rename from modelopt/torch/quantization/src/conv/bench_implicit_gemm.py rename to modelopt/torch/kernels/quantization/conv/bench_implicit_gemm.py index 807ce178387..66e3f968517 100644 --- a/modelopt/torch/quantization/src/conv/bench_implicit_gemm.py +++ b/modelopt/torch/kernels/quantization/conv/bench_implicit_gemm.py @@ -94,7 +94,9 @@ def bench_fn(fn, warmup: int, iters: int) -> float: def run_benchmark(shapes_name: str, warmup: int, iters: int, fp4_block_size: int): """Run latency benchmark for the given shapes.""" - from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda + from modelopt.torch.kernels.quantization.conv.implicit_gemm_cuda import ( + conv3d_implicit_gemm_cuda, + ) shapes = get_shapes(shapes_name) diff --git a/modelopt/torch/quantization/src/conv/implicit_gemm_binding.cpp b/modelopt/torch/kernels/quantization/conv/implicit_gemm_binding.cpp similarity index 100% rename from modelopt/torch/quantization/src/conv/implicit_gemm_binding.cpp rename to modelopt/torch/kernels/quantization/conv/implicit_gemm_binding.cpp diff --git a/modelopt/torch/quantization/src/conv/implicit_gemm_cuda.py b/modelopt/torch/kernels/quantization/conv/implicit_gemm_cuda.py similarity index 100% rename from modelopt/torch/quantization/src/conv/implicit_gemm_cuda.py rename to modelopt/torch/kernels/quantization/conv/implicit_gemm_cuda.py diff --git a/modelopt/torch/quantization/src/conv/implicit_gemm_kernel.cu b/modelopt/torch/kernels/quantization/conv/implicit_gemm_kernel.cu similarity index 100% rename from modelopt/torch/quantization/src/conv/implicit_gemm_kernel.cu rename to modelopt/torch/kernels/quantization/conv/implicit_gemm_kernel.cu diff --git a/modelopt/torch/quantization/triton/__init__.py b/modelopt/torch/kernels/quantization/gemm/__init__.py similarity index 98% rename from modelopt/torch/quantization/triton/__init__.py rename to modelopt/torch/kernels/quantization/gemm/__init__.py index def70e5914f..39b07b4faa9 100644 --- a/modelopt/torch/quantization/triton/__init__.py +++ b/modelopt/torch/kernels/quantization/gemm/__init__.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python # SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/modelopt/torch/quantization/triton/fp4_kernel.py b/modelopt/torch/kernels/quantization/gemm/fp4_kernel.py similarity index 100% rename from modelopt/torch/quantization/triton/fp4_kernel.py rename to modelopt/torch/kernels/quantization/gemm/fp4_kernel.py diff --git a/modelopt/torch/quantization/triton/fp4_kernel_hopper.py b/modelopt/torch/kernels/quantization/gemm/fp4_kernel_hopper.py similarity index 100% rename from modelopt/torch/quantization/triton/fp4_kernel_hopper.py rename to modelopt/torch/kernels/quantization/gemm/fp4_kernel_hopper.py diff --git a/modelopt/torch/quantization/triton/fp8_kernel.py b/modelopt/torch/kernels/quantization/gemm/fp8_kernel.py similarity index 100% rename from modelopt/torch/quantization/triton/fp8_kernel.py rename to modelopt/torch/kernels/quantization/gemm/fp8_kernel.py diff --git a/modelopt/torch/quantization/triton/gptq_fused_kernel.py b/modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py similarity index 100% rename from modelopt/torch/quantization/triton/gptq_fused_kernel.py rename to modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py diff --git a/modelopt/torch/quantization/triton/nvfp4_quant.py b/modelopt/torch/kernels/quantization/gemm/nvfp4_quant.py similarity index 100% rename from modelopt/torch/quantization/triton/nvfp4_quant.py rename to modelopt/torch/kernels/quantization/gemm/nvfp4_quant.py diff --git a/modelopt/torch/quantization/src/tensor_quant.cpp b/modelopt/torch/kernels/quantization/gemm/tensor_quant.cpp similarity index 100% rename from modelopt/torch/quantization/src/tensor_quant.cpp rename to modelopt/torch/kernels/quantization/gemm/tensor_quant.cpp diff --git a/modelopt/torch/quantization/src/tensor_quant.h b/modelopt/torch/kernels/quantization/gemm/tensor_quant.h similarity index 100% rename from modelopt/torch/quantization/src/tensor_quant.h rename to modelopt/torch/kernels/quantization/gemm/tensor_quant.h diff --git a/modelopt/torch/quantization/src/tensor_quant_gpu.cu b/modelopt/torch/kernels/quantization/gemm/tensor_quant_gpu.cu similarity index 100% rename from modelopt/torch/quantization/src/tensor_quant_gpu.cu rename to modelopt/torch/kernels/quantization/gemm/tensor_quant_gpu.cu diff --git a/modelopt/torch/quantization/src/tensor_quant_gpu_fp8.cu b/modelopt/torch/kernels/quantization/gemm/tensor_quant_gpu_fp8.cu similarity index 100% rename from modelopt/torch/quantization/src/tensor_quant_gpu_fp8.cu rename to modelopt/torch/kernels/quantization/gemm/tensor_quant_gpu_fp8.cu diff --git a/modelopt/torch/quantization/src/tensor_quant_mx.cu b/modelopt/torch/kernels/quantization/gemm/tensor_quant_mx.cu similarity index 100% rename from modelopt/torch/quantization/src/tensor_quant_mx.cu rename to modelopt/torch/kernels/quantization/gemm/tensor_quant_mx.cu diff --git a/modelopt/torch/quantization/src/tensor_quant_mx.h b/modelopt/torch/kernels/quantization/gemm/tensor_quant_mx.h similarity index 100% rename from modelopt/torch/quantization/src/tensor_quant_mx.h rename to modelopt/torch/kernels/quantization/gemm/tensor_quant_mx.h diff --git a/modelopt/torch/kernels/sparsity/__init__.py b/modelopt/torch/kernels/sparsity/__init__.py new file mode 100644 index 00000000000..ca2bfdb1282 --- /dev/null +++ b/modelopt/torch/kernels/sparsity/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Sparsity kernels: attention (Triton skip-softmax backends) and gemm (placeholder).""" diff --git a/modelopt/torch/sparsity/attention_sparsity/kernels/__init__.py b/modelopt/torch/kernels/sparsity/attention/__init__.py similarity index 95% rename from modelopt/torch/sparsity/attention_sparsity/kernels/__init__.py rename to modelopt/torch/kernels/sparsity/attention/__init__.py index 0cc4a202f57..b45f4f27ae5 100644 --- a/modelopt/torch/sparsity/attention_sparsity/kernels/__init__.py +++ b/modelopt/torch/kernels/sparsity/attention/__init__.py @@ -18,7 +18,11 @@ import contextlib import threading -from modelopt.torch.kernels import IS_AVAILABLE, attention, register_triton_attention +from modelopt.torch.kernels.common.attention import ( + IS_AVAILABLE, + attention, + register_triton_attention, +) # --------------------------------------------------------------------------- # Optional backend registrations (depend on diffusers / ltx_core) diff --git a/modelopt/torch/kernels/sparsity/attention/calibrate.py b/modelopt/torch/kernels/sparsity/attention/calibrate.py new file mode 100644 index 00000000000..37c5fccd6bf --- /dev/null +++ b/modelopt/torch/kernels/sparsity/attention/calibrate.py @@ -0,0 +1,323 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + + +"""Skip-softmax multi-threshold calibration kernel and Python API. + +Runs a full attention forward (identical to dense attention) while measuring +how many KV tiles would be skipped at each candidate threshold. Used by the +sparse-attention calibration workflow in +``modelopt.torch.sparsity.attention_sparsity`` to fit a skip threshold. +""" + +import math + +import torch +import triton +import triton.language as tl + +from modelopt.torch.kernels.common.attention.triton_fa import LOG2E, _apply_mask + + +# --------------------------------------------------------------------------- +# Calibration kernel: collect multi-threshold skip-softmax sparsity stats +# --------------------------------------------------------------------------- +@triton.jit +def _attn_fwd_calibrate( + Q, + K, + V, + qk_scale, + b_start_loc, + b_seq_len, + b_start_loc_k, + b_seq_len_k, + Out, + stride_qbs, + stride_qh, + stride_kbs, + stride_kh, + stride_vbs, + stride_vh, + stride_obs, + stride_oh, + Threshold_trials, # [NUM_THRESHOLDS] float32 — pre-scaled to log2 space + Per_program_totals, # [num_programs * NUM_THRESHOLDS] int32 — per-program tile counts + Per_program_skipped, # [num_programs * NUM_THRESHOLDS] int32 — per-program skip counts + kv_group_num: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_N: tl.constexpr, + IS_CAUSAL: tl.constexpr, + HEAD_DIM: tl.constexpr, + NUM_THRESHOLDS: tl.constexpr, + PADDED_THRESHOLDS: tl.constexpr, # next_power_of_2(NUM_THRESHOLDS) for tl.arange +): + """Forward kernel with multi-threshold sparsity measurement. + + Computes full attention (no skipping) while counting how many KV tiles + would be skipped at each threshold. Each program writes its local counts + to ``Per_program_totals`` and ``Per_program_skipped``; the Python wrapper + sums across programs afterward. This avoids global atomic contention. + """ + batch_idx = tl.program_id(0) + head_idx = tl.program_id(1) + tile_q = tl.program_id(2) + kv_head_idx = head_idx // kv_group_num + + seq_len_q = tl.load(b_seq_len + batch_idx) + seq_len_kv = tl.load(b_seq_len_k + batch_idx) + q_offset = tl.load(b_start_loc + batch_idx) + kv_offset = tl.load(b_start_loc_k + batch_idx) + + if tile_q * BLOCK_M >= seq_len_q: + return + + q_pos = tile_q * BLOCK_M + tl.arange(0, BLOCK_M) + kv_pos = tl.arange(0, BLOCK_N) + dim_pos = tl.arange(0, BLOCK_D) + d_mask = dim_pos < HEAD_DIM + + q_ptrs = (q_offset + q_pos[:, None]) * stride_qbs + head_idx * stride_qh + dim_pos[None, :] + q = tl.load(Q + q_ptrs, mask=(q_pos[:, None] < seq_len_q) & d_mask[None, :], other=0.0) + + k_base = K + kv_head_idx * stride_kh + v_base = V + kv_head_idx * stride_vh + + row_max = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + row_sum = tl.zeros([BLOCK_M], dtype=tl.float32) + acc = tl.zeros([BLOCK_M, BLOCK_D], dtype=tl.float32) + + # Pre-load all thresholds once (vectorized, stays in registers). + # tl.arange requires power-of-2 size, so use PADDED_THRESHOLDS with masking. + thresh_offs = tl.arange(0, PADDED_THRESHOLDS) + thresh_mask = thresh_offs < NUM_THRESHOLDS + thresholds = tl.load(Threshold_trials + thresh_offs, mask=thresh_mask, other=float("inf")) + + # Per-program local counters: avoid global atomic contention in inner loop. + # Each program accumulates locally, then writes once to Per_program buffers. + local_skipped = tl.zeros([PADDED_THRESHOLDS], dtype=tl.int32) + num_tiles = 0 + + kv_bound = seq_len_kv if not IS_CAUSAL else tl.minimum((tile_q + 1) * BLOCK_M, seq_len_kv) + + for kv_start in range(0, kv_bound, BLOCK_N): + kv_start = tl.multiple_of(kv_start, BLOCK_N) + + k_offs = (kv_offset + kv_start + kv_pos[None, :]) * stride_kbs + dim_pos[:, None] + k = tl.load( + k_base + k_offs, + mask=((kv_start + kv_pos[None, :]) < seq_len_kv) & d_mask[:, None], + other=0.0, + ) + + scores = tl.dot(q, k) * qk_scale + scores = _apply_mask(scores, q_pos, kv_pos, seq_len_q, seq_len_kv, kv_start, IS_CAUSAL) + + tile_row_max = tl.max(scores, 1) + + # --- Vectorized multi-threshold sparsity measurement --- + # A tile is skipped iff ALL Q rows satisfy: tile_row_max < row_max + thresh. + # Equivalently: max(tile_row_max - row_max) < thresh (worst-case row + # must still be below threshold for the tile to be skippable). + max_gap = tl.max(tile_row_max - row_max) # scalar + skip_mask = (max_gap < thresholds).to(tl.int32) # [PADDED_THRESHOLDS] + local_skipped += skip_mask + num_tiles += 1 + + # --- Always compute full attention (no skipping) --- + m_new = tl.maximum(row_max, tile_row_max) + p = tl.math.exp2(scores - m_new[:, None]) + l_new = tl.sum(p, 1) + correction = tl.math.exp2(row_max - m_new) + row_sum = row_sum * correction + l_new + acc = acc * correction[:, None] + + v_offs = (kv_offset + kv_start + kv_pos[:, None]) * stride_vbs + dim_pos[None, :] + v = tl.load( + v_base + v_offs, + mask=((kv_start + kv_pos[:, None]) < seq_len_kv) & d_mask[None, :], + other=0.0, + ) + acc = tl.dot(p.to(v.dtype), v, acc) + row_max = m_new + + # --- Write per-program counters (no atomics, just stores) --- + # Compute unique flat program index for this (batch, head, q_tile). + # Use tl.num_programs(2) (grid z dim = cdiv(max_input_len, BLOCK_M)) so the + # stride matches the wrapper's buffer layout for any batch order. Loading + # b_seq_len[0] would collide with later batches when batch 0 is shorter. + num_q_tiles = tl.num_programs(2) + num_heads = tl.num_programs(1) + prog_idx = batch_idx * num_heads * num_q_tiles + head_idx * num_q_tiles + tile_q + base = prog_idx * NUM_THRESHOLDS + tl.store( + Per_program_totals + base + thresh_offs, + tl.full([PADDED_THRESHOLDS], num_tiles, dtype=tl.int32), + mask=thresh_mask, + ) + tl.store( + Per_program_skipped + base + thresh_offs, + local_skipped, + mask=thresh_mask, + ) + + acc = acc / tl.maximum(row_sum[:, None], 1e-6) + o_ptrs = (q_offset + q_pos[:, None]) * stride_obs + head_idx * stride_oh + dim_pos[None, :] + tl.store(Out + o_ptrs, acc, mask=(q_pos[:, None] < seq_len_q) & d_mask[None, :]) + + +def attention_calibrate( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + b_start_loc: torch.Tensor, + b_seq_len: torch.Tensor, + max_input_len: int, + is_causal: bool = True, + softmax_scale: float | None = None, + b_start_loc_k: torch.Tensor | None = None, + b_seq_len_k: torch.Tensor | None = None, + max_input_len_k: int | None = None, + *, + threshold_trials: list[float] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Flash attention with multi-threshold skip-softmax sparsity measurement. + + Computes full attention (identical output to dense attention) while + measuring how many KV tiles would be skipped at each threshold in + ``threshold_trials``. No autograd — forward only. + + Args: + q, k, v, b_start_loc, b_seq_len, max_input_len, is_causal, + softmax_scale, b_start_loc_k, b_seq_len_k, max_input_len_k: + Same as :func:`modelopt.torch.kernels.common.attention.attention`. + threshold_trials: List of threshold values to measure sparsity for. + Each value is converted to log2-scaled space for the kernel. + + Returns: + Tuple of (output, sparsity_counters): + - output: ``[total_q_tokens, num_q_heads, head_dim]`` + - sparsity_counters: ``[num_thresholds, 2]`` int64 tensor where + ``[:, 0]`` = total tile evaluations, ``[:, 1]`` = skipped tiles. + Sparsity per threshold = ``counters[:, 1] / counters[:, 0]``. + """ + if threshold_trials is None or len(threshold_trials) == 0: + raise ValueError("threshold_trials must be a non-empty list") + + # Calibration has only been validated with uniform-length batches (current + # diffusion + RULER paths). Varlen inputs would exercise code paths in the + # kernel that have not been tested — fail loudly rather than silently + # produce wrong sparsity counts. + if b_seq_len.numel() > 1 and not torch.all(b_seq_len == b_seq_len[0]).item(): + raise NotImplementedError( + "attention_calibrate currently supports only uniform-length batches. " + f"Got b_seq_len={b_seq_len.tolist()}. Varlen calibration is untested — " + "validate the kernel against a reference before removing this guard." + ) + if int(b_seq_len[0].item()) != max_input_len: + raise ValueError( + "attention_calibrate expects max_input_len to equal b_seq_len[0] " + f"(uniform batching). Got max_input_len={max_input_len}, " + f"b_seq_len[0]={int(b_seq_len[0].item())}." + ) + if ( + b_seq_len_k is not None + and b_seq_len_k.data_ptr() != b_seq_len.data_ptr() + and b_seq_len_k.numel() > 1 + and not torch.all(b_seq_len_k == b_seq_len_k[0]).item() + ): + raise NotImplementedError( + "attention_calibrate currently supports only uniform-length batches. " + f"Got b_seq_len_k={b_seq_len_k.tolist()}. Varlen calibration is untested." + ) + + HEAD_DIM = q.shape[2] + num_q_heads = q.shape[1] + num_kv_heads = k.shape[1] + kv_group_num = num_q_heads // num_kv_heads + batch = b_seq_len.shape[0] + sm_scale = 1.0 / (HEAD_DIM**0.5) if softmax_scale is None else softmax_scale + qk_scale = sm_scale * LOG2E + BLOCK_D = triton.next_power_of_2(HEAD_DIM) + BLOCK_M = 128 + BLOCK_N = 64 + + if b_seq_len_k is None: + b_seq_len_k = b_seq_len + b_start_loc_k = b_start_loc + + num_thresholds = len(threshold_trials) + + # Convert thresholds to log2-scaled space: log2(lambda) * sm_scale + threshold_tensor = torch.tensor( + [math.log2(t) * sm_scale for t in threshold_trials], + dtype=torch.float32, + device=q.device, + ) + + o = torch.empty_like(q) + + num_q_tiles = triton.cdiv(max_input_len, BLOCK_M) + grid = (batch, num_q_heads, num_q_tiles) + num_programs = batch * num_q_heads * num_q_tiles + + # Per-program output buffers (no atomics needed — each program writes its own row) + per_program_totals = torch.zeros( + num_programs * num_thresholds, dtype=torch.int32, device=q.device + ) + per_program_skipped = torch.zeros( + num_programs * num_thresholds, dtype=torch.int32, device=q.device + ) + + _attn_fwd_calibrate[grid]( + q, + k, + v, + qk_scale, + b_start_loc, + b_seq_len, + b_start_loc_k, + b_seq_len_k, + o, + q.stride(0), + q.stride(1), + k.stride(0), + k.stride(1), + v.stride(0), + v.stride(1), + o.stride(0), + o.stride(1), + threshold_tensor, + per_program_totals, + per_program_skipped, + kv_group_num=kv_group_num, + BLOCK_M=BLOCK_M, + BLOCK_D=BLOCK_D, + BLOCK_N=BLOCK_N, + IS_CAUSAL=is_causal, + HEAD_DIM=HEAD_DIM, + NUM_THRESHOLDS=num_thresholds, + PADDED_THRESHOLDS=triton.next_power_of_2(num_thresholds), + num_warps=4, + num_stages=1, + ) + + # Reduce across programs: sum per-program counts → [num_thresholds] + totals = per_program_totals.view(num_programs, num_thresholds).sum(dim=0).to(torch.int64) + skipped = per_program_skipped.view(num_programs, num_thresholds).sum(dim=0).to(torch.int64) + sparsity_counters = torch.stack([totals, skipped], dim=1) # [num_thresholds, 2] + + return o, sparsity_counters diff --git a/modelopt/torch/sparsity/attention_sparsity/kernels/diffusers_triton_attention.py b/modelopt/torch/kernels/sparsity/attention/diffusers_triton_attention.py similarity index 95% rename from modelopt/torch/sparsity/attention_sparsity/kernels/diffusers_triton_attention.py rename to modelopt/torch/kernels/sparsity/attention/diffusers_triton_attention.py index 2923447cf02..434c4824f8e 100644 --- a/modelopt/torch/sparsity/attention_sparsity/kernels/diffusers_triton_attention.py +++ b/modelopt/torch/kernels/sparsity/attention/diffusers_triton_attention.py @@ -36,7 +36,10 @@ attention_backend, ) -from modelopt.torch.kernels import attention, attention_calibrate +# ``attention`` and ``attention_calibrate`` are resolved lazily inside the +# call-site functions below. Capturing them at module top-level would fetch +# ``None`` from the partially-loaded ``common.attention`` package during the +# sparsity↔common circular import chain. _BACKEND_NAME = "modelopt_triton" _BACKEND_REGISTERED = False @@ -166,6 +169,8 @@ def _diffusers_triton_attention( calib_mode = getattr(_thread_local, "calibration_mode", False) if calib_mode: trials = getattr(_thread_local, "threshold_trials", None) + from modelopt.torch.kernels.common.attention import attention_calibrate + if trials and attention_calibrate is not None: o, counters = attention_calibrate(q, k, v, **kw, threshold_trials=trials) @@ -196,6 +201,8 @@ def _diffusers_triton_attention( if threshold is not None and threshold > 0.0: kw["skip_softmax_threshold"] = threshold + from modelopt.torch.kernels.common.attention import attention + assert attention is not None, "Triton attention kernel not available (requires CUDA + triton)" do_measure = getattr(_thread_local, "measure_sparsity", False) if do_measure: diff --git a/modelopt/torch/sparsity/attention_sparsity/kernels/ltx_triton_attention.py b/modelopt/torch/kernels/sparsity/attention/ltx_triton_attention.py similarity index 94% rename from modelopt/torch/sparsity/attention_sparsity/kernels/ltx_triton_attention.py rename to modelopt/torch/kernels/sparsity/attention/ltx_triton_attention.py index fd53e7f9f49..90601dc2cae 100644 --- a/modelopt/torch/sparsity/attention_sparsity/kernels/ltx_triton_attention.py +++ b/modelopt/torch/kernels/sparsity/attention/ltx_triton_attention.py @@ -25,7 +25,10 @@ import torch -from modelopt.torch.kernels import attention, attention_calibrate +# ``attention`` and ``attention_calibrate`` are resolved lazily inside the +# call-site functions below. Capturing them at module top-level would fetch +# ``None`` from the partially-loaded ``common.attention`` package during the +# sparsity↔common circular import chain. from modelopt.torch.utils.logging import warn_rank_0 # Thread-local storage for skip-softmax configuration @@ -126,6 +129,8 @@ def _ltx_triton_attention( calib_mode = getattr(_thread_local, "calibration_mode", False) if calib_mode: trials = getattr(_thread_local, "threshold_trials", None) + from modelopt.torch.kernels.common.attention import attention_calibrate + if trials and attention_calibrate is not None: o, counters = attention_calibrate(q_flat, k_flat, v_flat, **kw, threshold_trials=trials) @@ -150,6 +155,8 @@ def _ltx_triton_attention( elif threshold is not None and threshold > 0.0: kw["skip_softmax_threshold"] = threshold + from modelopt.torch.kernels.common.attention import attention + assert attention is not None, "Triton attention kernel not available (requires CUDA + triton)" o = attention(q_flat, k_flat, v_flat, **kw) return o.view(b, seq_q, heads * dim_head) diff --git a/modelopt/torch/kernels/sparsity/attention/skip_softmax_helpers.py b/modelopt/torch/kernels/sparsity/attention/skip_softmax_helpers.py new file mode 100644 index 00000000000..f066f9c4b7d --- /dev/null +++ b/modelopt/torch/kernels/sparsity/attention/skip_softmax_helpers.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + + +"""Skip-softmax / N:M sparse attention helpers. + +These ``@triton.jit`` helpers are called conditionally from the baseline +flash-attention forward kernel in ``common/attention/triton_fa.py`` when the +user requests N:M sparsity or sink/window-aware dense regions. +""" + +import triton +import triton.language as tl + + +# --------------------------------------------------------------------------- +# N:M sparse softmax helpers +# --------------------------------------------------------------------------- +@triton.jit +def _sparse_nm_masks_m4(x0, x1, x2, x3, N: tl.constexpr): + """Top-N of 4 selection via pure boolean logic (6 comparisons, no int casts). + + Uses ``>=`` so that ties are broken by index (lower index wins). + Guarantees exactly N masks are True for any input including all-equal. + + Boolean formulas for "at least K of 3 wins": + K=3 (N=1): AND of all — must beat all 3 others + K=2 (N=2): majority — must beat at least 2 (sorting network) + K=1 (N=3): OR of all — must beat at least 1 + """ + c01 = x0 >= x1 + c02 = x0 >= x2 + c03 = x0 >= x3 + c12 = x1 >= x2 + c13 = x1 >= x3 + c23 = x2 >= x3 + + nc01 = ~c01 + nc02 = ~c02 + nc03 = ~c03 + nc12 = ~c12 + nc13 = ~c13 + nc23 = ~c23 + + if N == 1: + # Keep max only: must beat all 3 + m0 = c01 & c02 & c03 + m1 = nc01 & c12 & c13 + m2 = nc02 & nc12 & c23 + m3 = nc03 & nc13 & nc23 + elif N == 2: + # Majority vote: must beat at least 2 of 3 + m0 = (c01 & c02) | (c01 & c03) | (c02 & c03) + m1 = (nc01 & c12) | (nc01 & c13) | (c12 & c13) + m2 = (nc02 & nc12) | (nc02 & c23) | (nc12 & c23) + m3 = (nc03 & nc13) | (nc03 & nc23) | (nc13 & nc23) + elif N == 3: + # Keep all but min: must beat at least 1 + m0 = c01 | c02 | c03 + m1 = nc01 | c12 | c13 + m2 = nc02 | nc12 | c23 + m3 = nc03 | nc13 | nc23 + else: + tl.static_assert(False, "N must be 1, 2, or 3 for M=4") + + return m0, m1, m2, m3 + + +@triton.jit +def _apply_sparse_nm_to_qk_tile( + qk, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + SPARSITY_N: tl.constexpr, + SPARSITY_M: tl.constexpr, +): + """Apply N:M sparse softmax to a QK score tile. + + For every ``SPARSITY_M`` consecutive elements along the N (key) dimension, + keeps the top ``SPARSITY_N`` values and sets the rest to ``-inf``. + ``BLOCK_N`` must be divisible by ``SPARSITY_M``. + + For M=4, exactly N values are retained (ties broken by position). + For M=8, a threshold-based approach (``tl.sort``) may retain more + than N values when ties straddle the threshold boundary. + """ + tl.static_assert(SPARSITY_M == 4 or SPARSITY_M == 8, "SPARSITY_M must be 4 or 8") # noqa: PLR1714 + MASK_VAL: tl.constexpr = float("-inf") + + if SPARSITY_M == 4: + tl.static_assert(BLOCK_N % 4 == 0, "BLOCK_N must be divisible by 4") + reshaped = tl.reshape(qk, (BLOCK_M, BLOCK_N // 4, 4)) + cols = tl.arange(0, 4)[None, None, :] + x0 = tl.sum(tl.where(cols == 0, reshaped, 0.0), axis=2) + x1 = tl.sum(tl.where(cols == 1, reshaped, 0.0), axis=2) + x2 = tl.sum(tl.where(cols == 2, reshaped, 0.0), axis=2) + x3 = tl.sum(tl.where(cols == 3, reshaped, 0.0), axis=2) + + m0, m1, m2, m3 = _sparse_nm_masks_m4(x0, x1, x2, x3, SPARSITY_N) + + out = tl.full((BLOCK_M, BLOCK_N // 4, 4), 0.0, dtype=qk.dtype) + out = tl.where(cols == 0, tl.expand_dims(tl.where(m0, x0, MASK_VAL), 2), out) + out = tl.where(cols == 1, tl.expand_dims(tl.where(m1, x1, MASK_VAL), 2), out) + out = tl.where(cols == 2, tl.expand_dims(tl.where(m2, x2, MASK_VAL), 2), out) + out = tl.where(cols == 3, tl.expand_dims(tl.where(m3, x3, MASK_VAL), 2), out) + return tl.reshape(out, (BLOCK_M, BLOCK_N)) + + else: # SPARSITY_M == 8 + tl.static_assert(BLOCK_N % 8 == 0, "BLOCK_N must be divisible by 8") + reshaped = tl.reshape(qk, (BLOCK_M, BLOCK_N // 8, 8)) + + # Sort each group of 8 ascending; N-th largest is at index (8 - N) + sorted_vals = tl.sort(reshaped, dim=2) + KTH_IDX: tl.constexpr = SPARSITY_M - SPARSITY_N # index of N-th largest in ascending order + + # Extract the threshold value at KTH_IDX via masked sum + # Use 0.0 as fill (not -inf) so sum equals just the KTH element + cols = tl.arange(0, 8)[None, None, :] + threshold = tl.sum(tl.where(cols == KTH_IDX, sorted_vals, 0.0), axis=2) + + # Mask: keep elements >= threshold (may keep >N on ties — acceptable) + mask = reshaped >= tl.expand_dims(threshold, 2) + return tl.reshape(tl.where(mask, reshaped, MASK_VAL), (BLOCK_M, BLOCK_N)) + + +# --------------------------------------------------------------------------- +# BLASST skip-softmax per-tile decision +# --------------------------------------------------------------------------- +@triton.jit +def _skip_softmax_decision( + scores, + row_max, + SKIP_THRESHOLD_LOG2: tl.constexpr, + Sparsity_total, + Sparsity_skipped, + MEASURE_SPARSITY: tl.constexpr, +): + """BLASST skip-softmax per-tile decision (https://arxiv.org/pdf/2512.12087). + + During FlashAttention's block-wise computation we maintain a running + maximum ``m_i^(j)`` across blocks. If a block's local maximum + ``~m_i^(j)`` is significantly smaller than the running maximum + (``~m_i^(j) - m_i^(j) < ln(lambda)``), then ``exp(~m_i^(j) - m_i^(j)) + < lambda ~= 0`` and the block's contribution to the output is negligible. + The caller may then skip the softmax computation, V load, and BMM2. + + The threshold is pre-scaled to log2 space by the Python wrapper so it can + be compared directly against the already-scaled scores. + + Returns: + True when *all* Q rows in the tile satisfy the skip criterion. + + When ``MEASURE_SPARSITY`` is set, also records total/skipped tile counts + via atomic adds on ``Sparsity_total`` / ``Sparsity_skipped``. + """ + tile_row_max = tl.max(scores, 1) # [BLOCK_M] — ~m_i^(j) (scaled) + # Per-row: True if row's tile max is negligible vs running max + can_skip = tile_row_max < (row_max + SKIP_THRESHOLD_LOG2) + # Per-tile: skip entire tile only if ALL rows are negligible + skip_tile = tl.min(can_skip.to(tl.int32)) == 1 + + if MEASURE_SPARSITY: + tl.atomic_add(Sparsity_total, 1) # count every tile + if skip_tile: + tl.atomic_add(Sparsity_skipped, 1) # count skipped tiles + + return skip_tile + + +# --------------------------------------------------------------------------- +# Sink/window dense-region check +# --------------------------------------------------------------------------- +@triton.jit +def _is_dense_region( + kv_start, + tile_q, + seq_len_q, + seq_len_kv, + BLOCK_M: tl.constexpr, + NUM_SINK_TOKENS: tl.constexpr, + DENSE_WINDOW_SIZE: tl.constexpr, +): + """Check if a KV tile falls in a dense region (sink tokens or local window). + + Uses absolute token positions so the result is BLOCK_N-independent, + ensuring forward and backward (which may use different BLOCK_N) agree. + + Returns: + True if the tile should be kept dense (skip N:M sparsification). + """ + is_sink = kv_start < NUM_SINK_TOKENS + causal_offset = seq_len_kv - seq_len_q + q_abs_pos = tile_q * BLOCK_M + causal_offset + token_distance = q_abs_pos - kv_start + is_local = (token_distance >= 0) and (token_distance < DENSE_WINDOW_SIZE) + return is_sink or is_local diff --git a/modelopt/torch/kernels/sparsity/gemm/__init__.py b/modelopt/torch/kernels/sparsity/gemm/__init__.py new file mode 100644 index 00000000000..5a366019db7 --- /dev/null +++ b/modelopt/torch/kernels/sparsity/gemm/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Sparsity GEMM kernels (placeholder for future implementations).""" diff --git a/modelopt/torch/quantization/extensions.py b/modelopt/torch/quantization/extensions.py index 003703567a3..a65396d64ff 100644 --- a/modelopt/torch/quantization/extensions.py +++ b/modelopt/torch/quantization/extensions.py @@ -22,6 +22,7 @@ __all__ = ["get_cuda_ext", "get_cuda_ext_fp8", "get_cuda_ext_mx", "precompile"] path = Path(__file__).parent +kernels_gemm = path.parent / "kernels" / "quantization" / "gemm" def get_cuda_ext(raise_if_failed: bool = False): @@ -29,7 +30,7 @@ def get_cuda_ext(raise_if_failed: bool = False): if not hasattr(get_cuda_ext, "extension"): get_cuda_ext.extension = load_cpp_extension( # type:ignore[attr-defined] name="modelopt_cuda_ext", - sources=[path / "src/tensor_quant.cpp", path / "src/tensor_quant_gpu.cu"], + sources=[kernels_gemm / "tensor_quant.cpp", kernels_gemm / "tensor_quant_gpu.cu"], cuda_version_specifiers=">=11", raise_if_failed=raise_if_failed, ) @@ -41,7 +42,7 @@ def get_cuda_ext_fp8(raise_if_failed: bool = False): if not hasattr(get_cuda_ext_fp8, "extension"): get_cuda_ext_fp8.extension = load_cpp_extension( # type:ignore[attr-defined] name="modelopt_cuda_ext_fp8", - sources=[path / "src/tensor_quant_gpu_fp8.cu"], + sources=[kernels_gemm / "tensor_quant_gpu_fp8.cu"], cuda_version_specifiers=">=11.8", fail_msg=( "CUDA extension for FP8 quantization could not be built and loaded, FP8 simulated" @@ -58,7 +59,7 @@ def get_cuda_ext_mx(raise_if_failed: bool = False): get_cuda_ext_mx.extension = load_cpp_extension( # type:ignore[attr-defined] name="modelopt_cuda_ext_mx", sources=[ - path / "src/tensor_quant_mx.cu", + kernels_gemm / "tensor_quant_mx.cu", ], cuda_version_specifiers=">=11.8", fail_msg=( diff --git a/modelopt/torch/quantization/nn/modules/quant_conv.py b/modelopt/torch/quantization/nn/modules/quant_conv.py index ed165556249..375fdc96e0c 100644 --- a/modelopt/torch/quantization/nn/modules/quant_conv.py +++ b/modelopt/torch/quantization/nn/modules/quant_conv.py @@ -19,7 +19,7 @@ import torch.nn as nn -from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda +from modelopt.torch.kernels.quantization.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda from ... import tensor_quant from .quant_module import QuantLinearConvBase, QuantModuleRegistry, _LegacyQuantLinearConvBaseMixin diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 59bcd215bbc..92eaf12ece5 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -30,6 +30,7 @@ from torch.nn.functional import linear from transformers.models.t5.modeling_t5 import T5Attention +from modelopt.torch.kernels.quantization.gemm import IS_AVAILABLE as IS_TRITON_AVAILABLE from modelopt.torch.opt.dynamic import DynamicModule from modelopt.torch.utils.distributed import ParallelState @@ -37,7 +38,6 @@ from ..conversion import register from ..nn import QuantInputBase, QuantModule, QuantModuleRegistry, TensorQuantizer from ..nn.modules.quant_linear import _QuantLinear -from ..triton import IS_AVAILABLE as IS_TRITON_AVAILABLE from ..utils import replace_function, sync_moe_expert_amax from ..utils.layerwise_calib import LayerActivationCollector from .attention import register_attention_for_kv_quant @@ -58,7 +58,7 @@ kitchen = None if IS_TRITON_AVAILABLE: - from ..triton import weight_dequant + from modelopt.torch.kernels.quantization.gemm import weight_dequant else: weight_dequant = None diff --git a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py index 6ff31424c77..fe30e283c2d 100644 --- a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py +++ b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py @@ -346,7 +346,7 @@ def _unpack_tensor(input: torch.Tensor): ) from e if fast: - from ..triton.fp4_kernel import fp4_dequantize + from modelopt.torch.kernels.quantization.gemm.fp4_kernel import fp4_dequantize return fp4_dequantize( self._quantized_data, diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index 16b9d32997e..15d782c4a79 100644 --- a/modelopt/torch/quantization/tensor_quant.py +++ b/modelopt/torch/quantization/tensor_quant.py @@ -21,7 +21,7 @@ from torch.autograd import Function from torch.onnx import symbolic_helper -import modelopt.torch.quantization.triton as triton_kernel +import modelopt.torch.kernels.quantization.gemm as triton_kernel from .config import QuantizerAttributeConfig from .extensions import get_cuda_ext, get_cuda_ext_fp8, get_cuda_ext_mx diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index ac2ec7a2553..aadfe40d24a 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -292,7 +292,7 @@ def gptq_blockwise_update_fused_scalar( block_size: Number of columns to process per GPTQ block. quant_block_size: Number of elements sharing one quantization scale factor. """ - from modelopt.torch.quantization.triton.gptq_fused_kernel import gptq_fused_block_scalar + from modelopt.torch.kernels.quantization.gemm.gptq_fused_kernel import gptq_fused_block_scalar num_cols = weight.shape[1] for bs in range(0, num_cols, block_size): diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrate.py b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrate.py index f63feac69ed..51df5bb4d4a 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrate.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrate.py @@ -153,8 +153,12 @@ def create_decode_calibration_forward_loop( ) -> Callable: """Create forward loop for decode phase calibration. - Uses flash attention for fast prefill, then switches to eager attention - for decode token generation with softmax hook measurement. + Uses SDPA for fast prefill, then switches to eager attention for decode + token generation with softmax hook measurement. (Previously used + ``flash_attention_2`` for prefill, but transformers>=5.0's FA2 path + unconditionally calls ``s_aux.to(query.dtype)`` on the attention-sinks + tensor and crashes for models without sinks. SDPA is just as fast for + prefill, has no softmax hook, and is version-stable.) Args: calibration_data: List of samples with 'input' and 'length' fields @@ -180,8 +184,8 @@ def forward_loop(model: nn.Module) -> None: with torch.no_grad(): try: - # Step 1: Fast prefill with flash attention (no measurement) - model.config._attn_implementation = "flash_attention_2" + # Step 1: Fast prefill with SDPA (no measurement) + model.config._attn_implementation = "sdpa" outputs = model(input_ids, use_cache=True) past_key_values = outputs.past_key_values next_token = outputs.logits[:, -1:, :].argmax(dim=-1) diff --git a/modelopt/torch/sparsity/attention_sparsity/conversion.py b/modelopt/torch/sparsity/attention_sparsity/conversion.py index cc928198509..f0c33520c49 100644 --- a/modelopt/torch/sparsity/attention_sparsity/conversion.py +++ b/modelopt/torch/sparsity/attention_sparsity/conversion.py @@ -79,7 +79,7 @@ def _set_attn_implementation(model: nn.Module, config: SparseAttentionConfig) -> ) if "triton" in backends: - from .kernels import register_triton_attention + from modelopt.torch.kernels.sparsity.attention import register_triton_attention if register_triton_attention is None: raise ImportError( @@ -128,7 +128,9 @@ def _register_diffusers_backends_if_needed(model: nn.Module) -> None: from diffusers.models.modeling_utils import ModelMixin if isinstance(model, ModelMixin): - from .kernels import register_diffusers_triton_attention + from modelopt.torch.kernels.sparsity.attention import ( + register_diffusers_triton_attention, + ) if register_diffusers_triton_attention is not None: register_diffusers_triton_attention() @@ -137,7 +139,7 @@ def _register_diffusers_backends_if_needed(model: nn.Module) -> None: # Patch ltx_core Attention modules if present (independent of diffusers) try: - from .kernels import register_ltx_triton_attention + from modelopt.torch.kernels.sparsity.attention import register_ltx_triton_attention except (ImportError, RuntimeError): return diff --git a/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py b/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py index 117e337809f..c1d6465ba66 100644 --- a/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py +++ b/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py @@ -384,7 +384,7 @@ def sparse_softmax(input, dim=-1, *args, **kwargs): input = self.apply_sparsity(input, sparse_mask) return original_softmax(input, dim, *args, **kwargs) - from ..kernels import set_skip_softmax_context + from modelopt.torch.kernels.sparsity.attention import set_skip_softmax_context stack = ExitStack() set_skip_softmax_context(True) diff --git a/modelopt/torch/sparsity/attention_sparsity/methods/triton_skip_softmax.py b/modelopt/torch/sparsity/attention_sparsity/methods/triton_skip_softmax.py index 1e2f3905e7a..ff74d13fae9 100644 --- a/modelopt/torch/sparsity/attention_sparsity/methods/triton_skip_softmax.py +++ b/modelopt/torch/sparsity/attention_sparsity/methods/triton_skip_softmax.py @@ -168,7 +168,9 @@ def _get_scale_factor(self) -> float | None: def _get_diffusers_backend_context(): """Activate the modelopt_triton diffusers backend if registered.""" try: - from ..kernels.diffusers_triton_attention import get_triton_attention_backend + from modelopt.torch.kernels.sparsity.attention.diffusers_triton_attention import ( + get_triton_attention_backend, + ) with get_triton_attention_backend(): yield @@ -178,13 +180,17 @@ def _get_diffusers_backend_context(): def _set_triton_backends(self, **kwargs): """Set config on both diffusers and LTX Triton backends.""" try: - from ..kernels.diffusers_triton_attention import set_triton_skip_softmax_config + from modelopt.torch.kernels.sparsity.attention.diffusers_triton_attention import ( + set_triton_skip_softmax_config, + ) set_triton_skip_softmax_config(**kwargs) except ImportError: pass try: - from ..kernels.ltx_triton_attention import set_ltx_triton_context + from modelopt.torch.kernels.sparsity.attention.ltx_triton_attention import ( + set_ltx_triton_context, + ) set_ltx_triton_context(active=True, **kwargs) except ImportError: @@ -193,13 +199,17 @@ def _set_triton_backends(self, **kwargs): def _clear_triton_backends(self): """Clear config on both Triton backends.""" try: - from ..kernels.diffusers_triton_attention import clear_triton_skip_softmax_config + from modelopt.torch.kernels.sparsity.attention.diffusers_triton_attention import ( + clear_triton_skip_softmax_config, + ) clear_triton_skip_softmax_config() except ImportError: pass try: - from ..kernels.ltx_triton_attention import clear_ltx_triton_context + from modelopt.torch.kernels.sparsity.attention.ltx_triton_attention import ( + clear_ltx_triton_context, + ) clear_ltx_triton_context() except ImportError: @@ -211,7 +221,7 @@ def _collect_calibration_stats(self, module): seq_k = None try: - from ..kernels.diffusers_triton_attention import ( + from modelopt.torch.kernels.sparsity.attention.diffusers_triton_attention import ( get_calibration_counters, get_calibration_seq_k, ) @@ -223,7 +233,7 @@ def _collect_calibration_stats(self, module): if counters is None: try: - from ..kernels.ltx_triton_attention import ( + from modelopt.torch.kernels.sparsity.attention.ltx_triton_attention import ( get_calibration_counters, get_calibration_seq_k, ) @@ -288,7 +298,9 @@ def get_sparsity_counters(self) -> tuple[int, int]: def _collect_sparsity_counters(self) -> None: """Read runtime sparsity counters from the backend and accumulate.""" try: - from ..kernels.diffusers_triton_attention import get_sparsity_counters + from modelopt.torch.kernels.sparsity.attention.diffusers_triton_attention import ( + get_sparsity_counters, + ) total, skipped = get_sparsity_counters() self._sparsity_total += total diff --git a/pyproject.toml b/pyproject.toml index fdd60b5193a..bace52dff9c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -227,8 +227,8 @@ extend-ignore = [ "SIM", "UP", ] # TODO: Disabled for now, will enable later, once all puzzletron code is migrated -"modelopt/torch/quantization/triton/*" = ["N803", "N806", "E731"] # triton style -"modelopt/torch/sparsity/attention_sparsity/kernels/*" = [ +"modelopt/torch/kernels/quantization/gemm/*" = ["N803", "N806", "E731"] # triton style +"modelopt/torch/kernels/sparsity/attention/*" = [ "N803", "N806", ] # triton kernel style diff --git a/tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa.py similarity index 98% rename from tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa.py rename to tests/gpu/torch/kernels/common/attention/test_triton_fa.py index a5174496cf6..7fc3a554c7a 100644 --- a/tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa.py @@ -26,10 +26,10 @@ pytest.mark.filterwarnings("ignore::DeprecationWarning"), ] -from modelopt.torch.kernels import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE +from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE if TRITON_KERNEL_AVAILABLE: - from modelopt.torch.kernels import attention, register_triton_attention + from modelopt.torch.kernels.common.attention import attention, register_triton_attention if register_triton_attention is not None: register_triton_attention() diff --git a/tests/gpu/torch/sparsity/attention_sparsity/conftest.py b/tests/gpu/torch/kernels/conftest.py similarity index 100% rename from tests/gpu/torch/sparsity/attention_sparsity/conftest.py rename to tests/gpu/torch/kernels/conftest.py diff --git a/tests/gpu/torch/quantization/kernels/test_implicit_gemm.py b/tests/gpu/torch/kernels/quantization/conv/test_implicit_gemm.py similarity index 99% rename from tests/gpu/torch/quantization/kernels/test_implicit_gemm.py rename to tests/gpu/torch/kernels/quantization/conv/test_implicit_gemm.py index 56ceaacc01f..96cc24c2b98 100644 --- a/tests/gpu/torch/quantization/kernels/test_implicit_gemm.py +++ b/tests/gpu/torch/kernels/quantization/conv/test_implicit_gemm.py @@ -28,7 +28,9 @@ @pytest.fixture(scope="module") def cuda_conv3d(): """Import and return the CUDA implicit GEMM conv3d function.""" - from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda + from modelopt.torch.kernels.quantization.conv.implicit_gemm_cuda import ( + conv3d_implicit_gemm_cuda, + ) return conv3d_implicit_gemm_cuda @@ -36,7 +38,7 @@ def cuda_conv3d(): def _triton_fp4_available(): """Check if the Triton FP4 fake quant kernel is available (requires compute >= 8.9).""" try: - import modelopt.torch.quantization.triton as triton_kernel + import modelopt.torch.kernels.quantization.gemm as triton_kernel return hasattr(triton_kernel, "fp4_fake_quant_block") except ImportError: @@ -305,7 +307,7 @@ def test_deterministic(self, cuda_conv3d): @pytest.fixture(scope="module") def cuda_fp4(): """Import and return the CUDA FP4 fake quant function.""" - from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import fp4_fake_quant + from modelopt.torch.kernels.quantization.conv.implicit_gemm_cuda import fp4_fake_quant return fp4_fake_quant @@ -780,7 +782,7 @@ class TestFP4FakeQuantVsTriton: @pytest.mark.parametrize("num_blocks", [4, 16, 64]) def test_vs_triton(self, cuda_fp4, block_size, num_blocks): """CUDA kernel should match the Triton fp4_fake_quant_block.""" - from modelopt.torch.quantization.triton import fp4_fake_quant_block + from modelopt.torch.kernels.quantization.gemm import fp4_fake_quant_block torch.manual_seed(42) x = torch.randn(num_blocks, block_size, device="cuda", dtype=torch.float32) * 10 @@ -865,7 +867,7 @@ class TestFP4FakeQuantVsModelopt: @pytest.mark.parametrize("seed", [42, 123, 999]) def test_vs_triton_fp4_fake_quant_block(self, cuda_fp4, block_size, seed): """Compare against modelopt Triton fp4_fake_quant_block.""" - from modelopt.torch.quantization.triton import fp4_fake_quant_block + from modelopt.torch.kernels.quantization.gemm import fp4_fake_quant_block torch.manual_seed(seed) num_blocks = 16 @@ -957,7 +959,7 @@ def test_vs_triton_realistic_shape(self, cuda_fp4): x = torch.randn(num_blocks, block_size, device="cuda", dtype=torch.float32) * 5 global_amax = x.abs().max() - from modelopt.torch.quantization.triton import fp4_fake_quant_block + from modelopt.torch.kernels.quantization.gemm import fp4_fake_quant_block ours = cuda_fp4(x, global_amax.unsqueeze(0), block_size) theirs = fp4_fake_quant_block( @@ -983,7 +985,7 @@ def test_vs_triton_input_dtypes(self, cuda_fp4, dtype): Our kernel casts to float32 internally, so the result should match Triton's output when both receive the same dtype input. """ - from modelopt.torch.quantization.triton import fp4_fake_quant_block + from modelopt.torch.kernels.quantization.gemm import fp4_fake_quant_block torch.manual_seed(42) block_size = 16 diff --git a/tests/gpu/torch/sparsity/attention_sparsity/test_diffusers_triton_attention.py b/tests/gpu/torch/kernels/sparsity/attention/test_diffusers_triton_attention.py similarity index 96% rename from tests/gpu/torch/sparsity/attention_sparsity/test_diffusers_triton_attention.py rename to tests/gpu/torch/kernels/sparsity/attention/test_diffusers_triton_attention.py index f479b8883f9..54f66a279e2 100644 --- a/tests/gpu/torch/sparsity/attention_sparsity/test_diffusers_triton_attention.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_diffusers_triton_attention.py @@ -26,11 +26,9 @@ diffusers = pytest.importorskip("diffusers") -from modelopt.torch.kernels import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE -from modelopt.torch.sparsity.attention_sparsity.kernels import ( - diffusers_triton_attention as diffusers_mod, -) -from modelopt.torch.sparsity.attention_sparsity.kernels import ltx_triton_attention as ltx_mod +from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE +from modelopt.torch.kernels.sparsity.attention import diffusers_triton_attention as diffusers_mod +from modelopt.torch.kernels.sparsity.attention import ltx_triton_attention as ltx_mod @pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") diff --git a/tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_calibrate.py b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py similarity index 98% rename from tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_calibrate.py rename to tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py index 37c4da9969c..eaa1f5e3258 100644 --- a/tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_calibrate.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py @@ -29,10 +29,10 @@ pytest.mark.filterwarnings("ignore::DeprecationWarning"), ] -from modelopt.torch.kernels import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE +from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE if TRITON_KERNEL_AVAILABLE: - from modelopt.torch.kernels import attention, attention_calibrate + from modelopt.torch.kernels.common.attention import attention, attention_calibrate @pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") diff --git a/tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_skip_softmax.py b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py similarity index 98% rename from tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_skip_softmax.py rename to tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py index 21b2a12ca71..56f0a9e9d86 100644 --- a/tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_skip_softmax.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py @@ -25,10 +25,10 @@ pytest.mark.filterwarnings("ignore::DeprecationWarning"), ] -from modelopt.torch.kernels import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE +from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE if TRITON_KERNEL_AVAILABLE: - from modelopt.torch.kernels import attention, register_triton_attention + from modelopt.torch.kernels.common.attention import attention, register_triton_attention if register_triton_attention is not None: register_triton_attention() diff --git a/tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_sparse_nm.py b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_sparse_nm.py similarity index 98% rename from tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_sparse_nm.py rename to tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_sparse_nm.py index 4eec5799a52..ff215a6ff8a 100644 --- a/tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_sparse_nm.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_sparse_nm.py @@ -27,14 +27,16 @@ pytest.mark.filterwarnings("ignore::DeprecationWarning"), ] -from modelopt.torch.kernels import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE +from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE if TRITON_KERNEL_AVAILABLE: import triton import triton.language as tl - from modelopt.torch.kernels import attention - from modelopt.torch.kernels.triton_fa import _apply_sparse_nm_to_qk_tile + from modelopt.torch.kernels.common.attention import attention + from modelopt.torch.kernels.sparsity.attention.skip_softmax_helpers import ( + _apply_sparse_nm_to_qk_tile, + ) @triton.jit def _test_apply_sparse_nm( diff --git a/tests/gpu/torch/quantization/conftest.py b/tests/gpu/torch/quantization/conftest.py index 9e34e5ef680..c6f9d23b6cc 100644 --- a/tests/gpu/torch/quantization/conftest.py +++ b/tests/gpu/torch/quantization/conftest.py @@ -16,7 +16,7 @@ import pytest -from modelopt.torch.quantization import triton as triton_kernel +from modelopt.torch.kernels.quantization import gemm as triton_kernel @pytest.fixture(autouse=True) diff --git a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py index 1a28d229f45..d2503669ac9 100644 --- a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py +++ b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py @@ -20,7 +20,7 @@ from _test_utils.torch.quantization.quant_utils import quant from _test_utils.torch.quantization.tensor_quant_common import FakeTensorQuantTester -import modelopt.torch.quantization.triton as triton_kernel +import modelopt.torch.kernels.quantization.gemm as triton_kernel import modelopt.torch.quantization.utils as quant_utils from modelopt.torch.quantization import tensor_quant from modelopt.torch.quantization.extensions import get_cuda_ext, get_cuda_ext_mx diff --git a/tests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.py b/tests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.py index 72b20df9329..0c267ee2123 100644 --- a/tests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.py +++ b/tests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.py @@ -33,7 +33,7 @@ diffusers = pytest.importorskip("diffusers") -from modelopt.torch.kernels import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE +from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE if TRITON_KERNEL_AVAILABLE: import modelopt.torch.sparsity.attention_sparsity as mtsa diff --git a/tests/unit/torch/kernels/test_triton_fa.py b/tests/unit/torch/kernels/common/attention/test_triton_fa.py similarity index 88% rename from tests/unit/torch/kernels/test_triton_fa.py rename to tests/unit/torch/kernels/common/attention/test_triton_fa.py index ac054e10e6b..6969ae0a0e3 100644 --- a/tests/unit/torch/kernels/test_triton_fa.py +++ b/tests/unit/torch/kernels/common/attention/test_triton_fa.py @@ -33,7 +33,8 @@ def test_triton_fa_importable_on_cpu(): except ImportError: pytest.skip("triton is not installed") - from modelopt.torch.kernels import triton_fa + from modelopt.torch.kernels.common.attention import triton_fa + from modelopt.torch.kernels.sparsity.attention import calibrate assert "attention" in triton_fa.__all__ - assert "attention_calibrate" in triton_fa.__all__ + assert callable(calibrate.attention_calibrate) diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_kernel_backends.py b/tests/unit/torch/kernels/sparsity/attention/test_kernel_backends.py similarity index 85% rename from tests/unit/torch/sparsity/attention_sparsity/test_kernel_backends.py rename to tests/unit/torch/kernels/sparsity/attention/test_kernel_backends.py index 775723e66c4..f997c98be64 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_kernel_backends.py +++ b/tests/unit/torch/kernels/sparsity/attention/test_kernel_backends.py @@ -34,12 +34,12 @@ class TestSkipSoftmaxContext: def test_default_is_false(self): - from modelopt.torch.sparsity.attention_sparsity.kernels import get_skip_softmax_context + from modelopt.torch.kernels.sparsity.attention import get_skip_softmax_context assert get_skip_softmax_context() is False def test_set_and_get(self): - from modelopt.torch.sparsity.attention_sparsity.kernels import ( + from modelopt.torch.kernels.sparsity.attention import ( get_skip_softmax_context, set_skip_softmax_context, ) @@ -60,9 +60,7 @@ class TestDiffusersTritonBackend: @pytest.fixture(autouse=True) def _reset(self): - from modelopt.torch.sparsity.attention_sparsity.kernels import ( - diffusers_triton_attention as mod, - ) + from modelopt.torch.kernels.sparsity.attention import diffusers_triton_attention as mod mod._BACKEND_REGISTERED = False mod.clear_triton_skip_softmax_config() @@ -70,7 +68,7 @@ def _reset(self): mod.clear_triton_skip_softmax_config() def test_set_clear_config(self): - from modelopt.torch.sparsity.attention_sparsity.kernels.diffusers_triton_attention import ( + from modelopt.torch.kernels.sparsity.attention.diffusers_triton_attention import ( clear_triton_skip_softmax_config, set_triton_skip_softmax_config, ) @@ -79,7 +77,7 @@ def test_set_clear_config(self): clear_triton_skip_softmax_config() def test_register_idempotent(self): - from modelopt.torch.sparsity.attention_sparsity.kernels.diffusers_triton_attention import ( + from modelopt.torch.kernels.sparsity.attention.diffusers_triton_attention import ( register_diffusers_triton_attention, ) @@ -87,7 +85,7 @@ def test_register_idempotent(self): register_diffusers_triton_attention() # Should be a no-op def test_get_backend_before_register_raises(self): - from modelopt.torch.sparsity.attention_sparsity.kernels.diffusers_triton_attention import ( + from modelopt.torch.kernels.sparsity.attention.diffusers_triton_attention import ( get_triton_attention_backend, ) @@ -113,12 +111,10 @@ def test_with_diffusers_model(self): """A ModelMixin subclass triggers diffusers backend registration.""" from diffusers.models.modeling_utils import ModelMixin + from modelopt.torch.kernels.sparsity.attention import diffusers_triton_attention as mod from modelopt.torch.sparsity.attention_sparsity.conversion import ( _register_diffusers_backends_if_needed, ) - from modelopt.torch.sparsity.attention_sparsity.kernels import ( - diffusers_triton_attention as mod, - ) mod._BACKEND_REGISTERED = False diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_ltx_triton_attention.py b/tests/unit/torch/kernels/sparsity/attention/test_ltx_triton_attention.py similarity index 96% rename from tests/unit/torch/sparsity/attention_sparsity/test_ltx_triton_attention.py rename to tests/unit/torch/kernels/sparsity/attention/test_ltx_triton_attention.py index 4751fbae353..6ed2b3f20b8 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_ltx_triton_attention.py +++ b/tests/unit/torch/kernels/sparsity/attention/test_ltx_triton_attention.py @@ -32,7 +32,7 @@ @pytest.fixture def ltx_mod(): """Import ltx_triton_attention and ensure thread-local state is reset.""" - from modelopt.torch.sparsity.attention_sparsity.kernels import ltx_triton_attention as mod + from modelopt.torch.kernels.sparsity.attention import ltx_triton_attention as mod mod.clear_ltx_triton_context() try: @@ -119,7 +119,7 @@ def __init__(self): parent = Parent() ltx_mod.register_ltx_triton_attention(parent) - from modelopt.torch.sparsity.attention_sparsity.kernels.ltx_triton_attention import ( + from modelopt.torch.kernels.sparsity.attention.ltx_triton_attention import ( _TritonLTXAttentionWrapper, ) diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py index 93389a46105..8e68c28e19c 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py @@ -270,7 +270,7 @@ def test_triton_backend_sets_attn_impl(self): "*": {"method": "triton_skip_softmax", "backend": "triton"}, } with patch( - "modelopt.torch.sparsity.attention_sparsity.kernels.register_triton_attention", + "modelopt.torch.kernels.sparsity.attention.register_triton_attention", MagicMock(return_value=True), ): _set_attn_implementation(model, config) @@ -287,7 +287,7 @@ def test_triton_backend_register_failure_raises(self): config.sparse_cfg = {"*": {"method": "triton_skip_softmax", "backend": "triton"}} with ( patch( - "modelopt.torch.sparsity.attention_sparsity.kernels.register_triton_attention", + "modelopt.torch.kernels.sparsity.attention.register_triton_attention", MagicMock(return_value=False), ), pytest.raises(RuntimeError, match="Failed to register"), @@ -305,7 +305,7 @@ def test_triton_backend_no_triton_raises(self): config.sparse_cfg = {"*": {"method": "triton_skip_softmax", "backend": "triton"}} with ( patch( - "modelopt.torch.sparsity.attention_sparsity.kernels.register_triton_attention", + "modelopt.torch.kernels.sparsity.attention.register_triton_attention", None, ), pytest.raises(ImportError, match="Triton backend requires"),