From 683c257c56db44d29ba3ee2e18775292cd409321 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 11 Mar 2026 15:08:13 -0700 Subject: [PATCH 1/2] Add a Triton flash attention kernel with HuggingFace integration Add a Triton Flash Attention kernel that supports variable-length batching, GQA, causal/non-causal masking, and autograd-compatible forward/backward. Register it as attn_implementation="modelopt_triton" for HuggingFace models. Signed-off-by: Kai Xu --- .../llm_sparsity/attention_sparsity/README.md | 5 +- .../llm_sparsity/attention_sparsity/hf_sa.py | 37 +- modelopt/torch/kernels/__init__.py | 47 ++ modelopt/torch/kernels/hf_triton_attention.py | 143 ++++ modelopt/torch/kernels/triton_fa.py | 716 ++++++++++++++++++ .../sparsity/attention_sparsity/config.py | 14 +- .../sparsity/attention_sparsity/conversion.py | 58 ++ .../attention_sparsity/kernels/__init__.py | 24 + .../methods/flash_skip_softmax.py | 16 + .../attention_sparsity/methods/registry.py | 12 + .../attention_sparsity/plugins/huggingface.py | 38 +- .../attention_sparsity/sparse_attention.py | 51 +- pyproject.toml | 1 + .../attention_sparsity/test_triton_fa.py | 439 +++++++++++ 14 files changed, 1505 insertions(+), 96 deletions(-) create mode 100644 modelopt/torch/kernels/__init__.py create mode 100644 modelopt/torch/kernels/hf_triton_attention.py create mode 100644 modelopt/torch/kernels/triton_fa.py create mode 100644 modelopt/torch/sparsity/attention_sparsity/kernels/__init__.py create mode 100644 tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa.py diff --git a/examples/llm_sparsity/attention_sparsity/README.md b/examples/llm_sparsity/attention_sparsity/README.md index e9d50ae10c9..a55abaa6ac7 100644 --- a/examples/llm_sparsity/attention_sparsity/README.md +++ b/examples/llm_sparsity/attention_sparsity/README.md @@ -1,6 +1,9 @@ # Attention Sparsity for HuggingFace Models -In this tutorial, we demonstrate how to use NVIDIA Model Optimizer to apply attention sparsity to HuggingFace models. Attention sparsity reduces computational cost by skipping near-zero attention scores during the softmax computation. +In this tutorial, we demonstrate how to use NVIDIA Model Optimizer to apply attention sparsity to HuggingFace models. Attention sparsity reduces computational cost by skipping near-zero attention scores during the softmax computation. Two attention backends are supported: + +- **pytorch** (default): Patches `F.softmax` to apply skip-softmax sparsity (requires `attn_implementation="eager"`) +- **triton**: Uses a fused Triton Flash Attention kernel with in-kernel sparsity (uses `attn_implementation="modelopt_triton"`) ## Getting Started diff --git a/examples/llm_sparsity/attention_sparsity/hf_sa.py b/examples/llm_sparsity/attention_sparsity/hf_sa.py index 0b97298f5fe..8115d4aaff1 100644 --- a/examples/llm_sparsity/attention_sparsity/hf_sa.py +++ b/examples/llm_sparsity/attention_sparsity/hf_sa.py @@ -144,9 +144,8 @@ def main(args): print(f"Loading model: {args.pyt_ckpt_path}") - # Load model and tokenizer - # Note: attn_implementation="eager" is required for calibration to work properly - # (flash_attention_2 or sdpa would bypass the softmax patching needed for stats collection) + # No need to specify attn_implementation here — mtsa.sparsify() sets it + # automatically ("eager" for pytorch backend, "modelopt_triton" for triton). model = AutoModelForCausalLM.from_pretrained( args.pyt_ckpt_path, attn_implementation="eager", @@ -164,21 +163,21 @@ def main(args): output_before, test_prompt, input_ids = generate_sample_output(model, tokenizer, args) # Apply sparse attention with optional calibration - print(f"\nApplying sparse attention: {args.sparse_attn}") - sparse_config = SPARSE_ATTN_CFG_CHOICES[args.sparse_attn] - - # Override calibration options if provided via CLI + print(f"\nApplying sparse attention: {args.sparse_attn} (backend={args.backend})") + sparse_config = copy.deepcopy(SPARSE_ATTN_CFG_CHOICES[args.sparse_attn]) + + # Apply CLI overrides to sparse_cfg + sparse_cfg = sparse_config.get("sparse_cfg", {}) + for layer_cfg in sparse_cfg.values(): + if isinstance(layer_cfg, dict) and "method" in layer_cfg: + layer_cfg["backend"] = args.backend if args.target_sparse_ratio is not None: - sparse_config = copy.deepcopy(sparse_config) - sparse_cfg = sparse_config.get("sparse_cfg", {}) - if isinstance(sparse_cfg, dict) and "calibration" in sparse_cfg: - calibration_cfg = sparse_cfg["calibration"] - if isinstance(calibration_cfg, dict): - calibration_cfg["target_sparse_ratio"] = { - "prefill": args.target_sparse_ratio, - "decode": args.target_sparse_ratio, - } - print(f"Overriding target_sparse_ratio to {args.target_sparse_ratio}") + calib = sparse_cfg.setdefault("calibration", {}) + assert isinstance(calib, dict) + calib["target_sparse_ratio"] = { + "prefill": args.target_sparse_ratio, + "decode": args.target_sparse_ratio, + } model = mtsa.sparsify(model, config=sparse_config) print("Sparse attention applied successfully!") @@ -242,8 +241,8 @@ def main(args): "--backend", type=str, default="pytorch", - choices=["pytorch"], - help="Backend for sparse attention (default: pytorch). More backends coming soon.", + choices=["pytorch", "triton"], + help="Backend for sparse attention (default: pytorch). 'triton' uses the fused Triton kernel.", ) # Sequence length arguments diff --git a/modelopt/torch/kernels/__init__.py b/modelopt/torch/kernels/__init__.py new file mode 100644 index 00000000000..66dbad7e792 --- /dev/null +++ b/modelopt/torch/kernels/__init__.py @@ -0,0 +1,47 @@ +# 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 +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 + with import_plugin("transformers"): + from .hf_triton_attention import register_triton_attention as _register_triton_attention + + register_triton_attention = _register_triton_attention + +__all__ = [ + "IS_AVAILABLE", + "attention", + "register_triton_attention", +] diff --git a/modelopt/torch/kernels/hf_triton_attention.py b/modelopt/torch/kernels/hf_triton_attention.py new file mode 100644 index 00000000000..5f10df2502e --- /dev/null +++ b/modelopt/torch/kernels/hf_triton_attention.py @@ -0,0 +1,143 @@ +# 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. + +"""HuggingFace attention backend using the Triton flash attention kernel. + +Registers as attn_implementation="modelopt_triton" so HF models dispatch to the +Triton kernel natively. Handles format conversion between HF's [batch, heads, seq, dim] +and the kernel's flat packed [total_tokens, heads, dim] varlen format. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from modelopt.torch.kernels.triton_fa import attention + + +def _seq_lens_from_mask( + attention_mask: torch.Tensor | None, + fallback: int, + device: torch.device, +) -> tuple[torch.Tensor | None, bool]: + """Derive per-sequence lengths from attention mask. + + Returns (b_seq_len, has_padding). If the mask is not a usable 2D format, + returns (None, False). + """ + if attention_mask is not None and attention_mask.dim() == 2: + mask = attention_mask.bool() if attention_mask.dtype != torch.bool else attention_mask + b_seq_len = mask.sum(dim=1).to(torch.int32).to(device) + has_padding = bool((b_seq_len != fallback).any()) + return b_seq_len, has_padding + return None, False + + +def triton_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs, +) -> tuple[torch.Tensor, None]: + """Attention forward compatible with HF AttentionInterface. + + Converts HF tensors to varlen format, calls the Triton kernel, converts back. + Handles both prefill (seq_len > 1) and decode (seq_len == 1). + + Args: + module: The attention module (LlamaAttention etc.). + query: [batch, num_heads, seq_len, head_dim]. + key: [batch, num_kv_heads, seq_k, head_dim]. + value: [batch, num_kv_heads, seq_k, head_dim]. + attention_mask: Optional; kernel handles causal masking internally. + 2D [batch, seq_len] masks are used to derive per-sequence lengths. + Other formats (e.g. 4D causal masks) are ignored. + scaling: Softmax scale (e.g. 1/sqrt(head_dim)). + dropout: Ignored (kernel has no dropout); use 0 for eval. + **kwargs: Reserved for future extensions. + + Returns: + (attn_output, None) with attn_output [batch, seq_len, num_heads, head_dim]. + """ + batch, num_heads, seq_len, head_dim = query.shape + seq_k = key.shape[2] + num_kv_heads = key.shape[1] + device = query.device + is_decode = seq_len <= 1 + + # Reshape from HF [batch, heads, seq, dim] -> flat [batch*seq, heads, dim] + q = query.permute(0, 2, 1, 3).reshape(batch * seq_len, num_heads, head_dim).contiguous() + k = key.permute(0, 2, 1, 3).reshape(batch * seq_k, num_kv_heads, head_dim).contiguous() + v = value.permute(0, 2, 1, 3).reshape(batch * seq_k, num_kv_heads, head_dim).contiguous() + + # Build varlen metadata + b_seq_len_q, has_padding = _seq_lens_from_mask(attention_mask, seq_len, device) + if b_seq_len_q is None: + b_seq_len_q = torch.full((batch,), seq_len, device=device, dtype=torch.int32) + + kw = { + "b_start_loc": torch.arange(batch, device=device, dtype=torch.int32) * seq_len, + "b_seq_len": b_seq_len_q, + "max_input_len": seq_len, + "is_causal": not is_decode, + "softmax_scale": scaling, + } + # Decode: Q has 1 token, K/V have seq_k tokens (KV cache, no padding) + if is_decode: + kw["b_start_loc_k"] = torch.arange(batch, device=device, dtype=torch.int32) * seq_k + kw["b_seq_len_k"] = torch.full((batch,), seq_k, device=device, dtype=torch.int32) + kw["max_input_len_k"] = seq_k + + o = attention(q, k, v, **kw) + + attn_output = o.view(batch, seq_len, num_heads, head_dim) + + # Zero out padding positions (kernel produces NaN for all-padding rows due to 0/0). + # Assumes right-padding (valid tokens at positions 0..n-1), which is the HF + # convention during prefill. Left-padded inputs are not supported. + if has_padding: + pad_mask = torch.arange(seq_len, device=device)[None, :] >= b_seq_len_q[:, None] + attn_output = attn_output.masked_fill(pad_mask[:, :, None, None], 0.0) + + return (attn_output, None) + + +def register_triton_attention() -> bool: + """Register the Triton backend with HF AttentionInterface. + + Called by _set_attn_implementation() during sparsification. Must run before + the model's first forward pass so HF dispatches to the Triton kernel. + + Returns: + True if registration succeeded, False if transformers API not available. + """ + try: + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + except (ImportError, AttributeError): + return False + + ALL_ATTENTION_FUNCTIONS.register("modelopt_triton", triton_attention_forward) + return True + + +__all__ = [ + "register_triton_attention", + "triton_attention_forward", +] diff --git a/modelopt/torch/kernels/triton_fa.py b/modelopt/torch/kernels/triton_fa.py new file mode 100644 index 00000000000..b9184788bf1 --- /dev/null +++ b/modelopt/torch/kernels/triton_fa.py @@ -0,0 +1,716 @@ +# 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. + +# ruff: noqa: N803, N806 — Triton kernels use uppercase for constexpr and tensor args by convention + +"""Triton flash attention kernel with variable-length sequences and GQA. + +Based on the Flash Attention v2 algorithm (https://arxiv.org/abs/2307.08691). + +Input format: flat packed [total_tokens, num_heads, head_dim] with per-sequence +metadata (b_start_loc, b_seq_len). Supports causal masking and autograd. +""" + +import torch +import triton +import triton.language as tl + +LOG2E: float = 1.44269504088896 + +# --------------------------------------------------------------------------- +# Autotune configs for forward kernel +# --------------------------------------------------------------------------- +_FWD_CONFIGS = [ + triton.Config({"BLOCK_M": bm, "BLOCK_N": bn}, num_stages=s, num_warps=w) + for bm in [64, 128] + for bn in [32, 64, 128] + for s in [1, 2, 3] + for w in [4, 8] +] + +# Use a single config in testing for reproducibility +if "PYTEST_VERSION" in __import__("os").environ: + _FWD_CONFIGS = [triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_stages=1, num_warps=4)] + + +# --------------------------------------------------------------------------- +# Masking helper +# --------------------------------------------------------------------------- +@triton.jit +def _apply_mask( + scores, + q_pos, + kv_pos, + seq_len_q, + seq_len_kv, + kv_start, + IS_CAUSAL: tl.constexpr, +): + """Apply causal mask and padding mask to a score tile.""" + if IS_CAUSAL: + scores += tl.where( + (kv_start + kv_pos[None, :] < seq_len_kv) + & (q_pos[:, None] >= (kv_start + kv_pos[None, :])), + 0, + float("-inf"), + ) + else: + scores += tl.where((kv_start + kv_pos[None, :]) < seq_len_kv, 0, float("-inf")) + return scores + + +# --------------------------------------------------------------------------- +# Forward kernel +# --------------------------------------------------------------------------- +@triton.autotune(configs=_FWD_CONFIGS, key=["N_CTX", "HEAD_DIM"]) +@triton.jit +def _attn_fwd( + Q, # [total_q, num_q_heads, head_dim] query tensor + K, # [total_kv, num_kv_heads, head_dim] key tensor + V, # [total_kv, num_kv_heads, head_dim] value tensor + qk_scale, # softmax_scale * log2(e) + b_start_loc, # [batch] start offset of each Q sequence + b_seq_len, # [batch] length of each Q sequence + b_start_loc_k, # [batch] start offset of each KV sequence + b_seq_len_k, # [batch] length of each KV sequence + Out, # [total_q, num_q_heads, head_dim] output tensor + Lse, # [total_q, num_q_heads] log-sum-exp + stride_qbs, + stride_qh, # Q strides: per-token, per-head + stride_kbs, + stride_kh, # K strides: per-token, per-head + stride_vbs, + stride_vh, # V strides: per-token, per-head + stride_obs, + stride_oh, # Output strides: per-token, per-head + stride_lse_tok, + stride_lse_head, # LSE strides: per-token, per-head + N_CTX, # Max Q sequence length (autotune cache key only) + kv_group_num: tl.constexpr, # GQA ratio: num_q_heads // num_kv_heads + BLOCK_M: tl.constexpr, # Q tile size (autotuned) + BLOCK_D: tl.constexpr, # Head dim tile size (next_power_of_2(HEAD_DIM)) + BLOCK_N: tl.constexpr, # KV tile size (autotuned) + IS_CAUSAL: tl.constexpr, # Whether to apply causal mask + HEAD_DIM: tl.constexpr, # Actual head dimension (for d_mask) + STORE_LSE: tl.constexpr, # Whether to save LSE for backward pass +): + # --- Grid: (batch, num_q_heads, num_q_tiles) --- + # Example: batch=2, num_q_heads=32, seq_len=256, BLOCK_M=128 + # grid = (2, 32, 2), 128 thread blocks launched in parallel + # block (1, 5, 0) handles: batch 1, Q head 5, tokens 0-127 + batch_idx = tl.program_id(0) # 0..batch-1 + head_idx = tl.program_id(1) # 0..num_q_heads-1 + tile_q = tl.program_id(2) # 0..ceil(seq_len/BLOCK_M)-1 + kv_head_idx = head_idx // kv_group_num # GQA: map Q head to shared KV head + + # --- Load Q and KV varlen metadata --- + 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 # This Q tile is past the sequence end + + # --- Tile position indices --- + q_pos = tile_q * BLOCK_M + tl.arange(0, BLOCK_M) # Absolute Q token positions + kv_pos = tl.arange(0, BLOCK_N) # Relative KV positions within a tile + dim_pos = tl.arange(0, BLOCK_D) # Head dimension positions + d_mask = dim_pos < HEAD_DIM # Mask for non-power-of-2 head dims + + # --- Load Q tile [BLOCK_M, BLOCK_D]: stays in SRAM for the entire KV loop --- + 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) + + # Base pointers for K and V at this KV head (per-tile offset added in loop) + k_base = K + kv_head_idx * stride_kh + v_base = V + kv_head_idx * stride_vh + + # --- Online softmax state (per Q row) --- + row_max = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") # Running max for stability + row_sum = tl.zeros([BLOCK_M], dtype=tl.float32) # Running sum of exp(scores) + acc = tl.zeros([BLOCK_M, BLOCK_D], dtype=tl.float32) # Running weighted sum of V + + # Causal bound: Q position i only attends to KV positions 0..i + kv_bound = seq_len_kv if not IS_CAUSAL else tl.minimum((tile_q + 1) * BLOCK_M, seq_len_kv) + + # --- Main loop: iterate over KV tiles --- + for kv_start in range(0, kv_bound, BLOCK_N): + kv_start = tl.multiple_of(kv_start, BLOCK_N) # Compiler hint for alignment + + # Load K^T [BLOCK_D, BLOCK_N] (transposed layout for Q @ K^T matmul) + 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 = Q @ K^T * scale [BLOCK_M, BLOCK_N] + 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) + + # --- Online softmax update --- + # 1. Update running max + m_new = tl.maximum(row_max, tl.max(scores, 1)) + # 2. Compute unnormalized attention weights + p = tl.math.exp2(scores - m_new[:, None]) + l_new = tl.sum(p, 1) + # 3. Correction factor: rescale previous tiles when max changes + correction = tl.math.exp2(row_max - m_new) + row_sum = row_sum * correction + l_new + acc = acc * correction[:, None] + + # Load V [BLOCK_N, BLOCK_D] and accumulate: acc += attn_weights @ V + 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 + + # --- Final normalization: output = acc / row_sum --- + acc = acc / row_sum[:, None] + + # Save LSE for backward pass (log2-space: lse = max + log2(sum)) + if STORE_LSE: + lse = row_max + tl.math.log2(row_sum) + lse = tl.where(row_sum == 0.0, float("-inf"), lse) + lse_ptrs = (q_offset + q_pos) * stride_lse_tok + head_idx * stride_lse_head + tl.store(Lse + lse_ptrs, lse, mask=q_pos < seq_len_q) + + # --- Store output [BLOCK_M, BLOCK_D] --- + 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, :]) + + +# --------------------------------------------------------------------------- +# Backward kernels +# --------------------------------------------------------------------------- +@triton.jit +def _attn_bwd_preprocess( + Out, + dO, + Delta, + stride_obs, + stride_oh, + stride_dobs, + stride_doh, + stride_delta_tok, + stride_delta_head, + total_tokens, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_M: tl.constexpr, +): + """Phase 1 of backward: compute delta_i = rowsum(O_i * dO_i). + + Delta is used in the dS computation: dS = P * (dP - delta). + This avoids recomputing O in the dQ/dK/dV kernels. + """ + head = tl.program_id(0) + offs_tok = tl.program_id(1) * BLOCK_M + tl.arange(0, BLOCK_M) + dim_pos = tl.arange(0, BLOCK_D) + mask_tok = offs_tok < total_tokens + mask_d = dim_pos < HEAD_DIM + + # Load O and dO tiles [BLOCK_M, BLOCK_D] + o = tl.load( + Out + offs_tok[:, None] * stride_obs + head * stride_oh + dim_pos[None, :], + mask=mask_tok[:, None] & mask_d[None, :], + other=0.0, + ) + do = tl.load( + dO + offs_tok[:, None] * stride_dobs + head * stride_doh + dim_pos[None, :], + mask=mask_tok[:, None] & mask_d[None, :], + other=0.0, + ) + + # delta_i = sum_d(O[i,d] * dO[i,d]) per token position + delta = tl.sum(o * do, axis=1) + tl.store(Delta + offs_tok * stride_delta_tok + head * stride_delta_head, delta, mask=mask_tok) + + +@triton.jit +def _attn_bwd_dq( + Q, + K, + V, + dO, + dQ, + Lse, + Delta, + b_start_loc, + b_seq_len, + b_start_loc_k, + b_seq_len_k, + qk_scale, + sm_scale, + stride_qbs, + stride_qh, + stride_kbs, + stride_kh, + stride_vbs, + stride_vh, + stride_dobs, + stride_doh, + stride_dqbs, + stride_dqh, + stride_lse_tok, + stride_lse_head, + 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, +): + """Phase 3 of backward: compute dQ for one Q tile, looping over KV tiles. + + For each KV tile, recomputes attention scores S = Q @ K^T, then: + P = softmax(S) (via exp2 and saved LSE) + dP = dO @ V^T + dS = P * (dP - delta) + dQ += dS @ K + """ + # --- Grid: each thread block handles one (batch, q_head, q_tile) --- + 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 + + # --- Load per-sequence varlen metadata --- + 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_mask = q_pos < seq_len_q + + # --- Load Q, dO tiles: stay in SRAM for the entire KV loop --- + q_ptrs = (q_offset + q_pos[:, None]) * stride_qbs + head_idx * stride_qh + dim_pos[None, :] + q = tl.load(Q + q_ptrs, mask=q_mask[:, None] & d_mask[None, :], other=0.0) + do_ptrs = (q_offset + q_pos[:, None]) * stride_dobs + head_idx * stride_doh + dim_pos[None, :] + do = tl.load(dO + do_ptrs, mask=q_mask[:, None] & d_mask[None, :], other=0.0) + + # Load saved LSE and delta from forward pass (same [total_tokens, heads] layout) + row_ptrs = (q_offset + q_pos) * stride_lse_tok + head_idx * stride_lse_head + lse = tl.load(Lse + row_ptrs, mask=q_mask, other=0.0) + row_delta = tl.load(Delta + row_ptrs, mask=q_mask, other=0.0) + + dq = tl.zeros([BLOCK_M, BLOCK_D], dtype=tl.float32) + kv_bound = seq_len_kv if not IS_CAUSAL else tl.minimum((tile_q + 1) * BLOCK_M, seq_len_kv) + + # --- Loop over KV tiles: recompute S, then compute dQ contribution --- + for kv_start in range(0, kv_bound, BLOCK_N): + kv_mask = (kv_start + kv_pos) < seq_len_kv + + # Load K^T and V for this KV tile + k_ptrs = ( + (kv_offset + kv_start + kv_pos[None, :]) * stride_kbs + + kv_head_idx * stride_kh + + dim_pos[:, None] + ) + kT = tl.load(K + k_ptrs, mask=kv_mask[None, :] & d_mask[:, None], other=0.0) + v_ptrs = ( + (kv_offset + kv_start + kv_pos[:, None]) * stride_vbs + + kv_head_idx * stride_vh + + dim_pos[None, :] + ) + v = tl.load(V + v_ptrs, mask=kv_mask[:, None] & d_mask[None, :], other=0.0) + + # Recompute attention: S = Q @ K^T, P = exp2(S - LSE) + scores = tl.dot(q, kT) * qk_scale + scores = _apply_mask(scores, q_pos, kv_pos, seq_len_q, seq_len_kv, kv_start, IS_CAUSAL) + p = tl.math.exp2(scores - lse[:, None]) + + # dP = dO @ V^T, dS = P * (dP - delta), dQ += dS @ K + dp = tl.dot(do, tl.trans(v)) + ds = p * (dp - row_delta[:, None]) + dq += tl.dot(ds.to(kT.dtype), tl.trans(kT)) + + # --- Store dQ (scaled by sm_scale since scores were pre-scaled by qk_scale) --- + dq *= sm_scale + dq_ptrs = (q_offset + q_pos[:, None]) * stride_dqbs + head_idx * stride_dqh + dim_pos[None, :] + tl.store(dQ + dq_ptrs, dq.to(q.dtype), mask=q_mask[:, None] & d_mask[None, :]) + + +@triton.jit +def _attn_bwd_dkdv( + Q, + K, + V, + dO, + dK, + dV, + Lse, + Delta, + b_start_loc, + b_seq_len, + b_start_loc_k, + b_seq_len_k, + qk_scale, + sm_scale, + stride_qbs, + stride_qh, + stride_kbs, + stride_kh, + stride_vbs, + stride_vh, + stride_dobs, + stride_doh, + stride_dkbs, + stride_dkh, + stride_dvbs, + stride_dvh, + stride_lse_tok, + stride_lse_head, + 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, +): + """Phase 2 of backward: compute dK, dV for one KV tile. + + Loops over all Q tiles (and GQA heads sharing this KV head), accumulating: + dV += P^T @ dO + dK += dS^T @ Q where dS = P * (dO @ V^T - delta) + """ + # --- Grid: each thread block handles one (batch, kv_head, kv_tile) --- + batch_idx = tl.program_id(0) + kv_head_idx = tl.program_id(1) + tile_kv = tl.program_id(2) + + # --- Load per-sequence varlen metadata --- + 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) + + kv_start = tile_kv * BLOCK_N + if kv_start >= seq_len_kv: + return + + kv_pos = tl.arange(0, BLOCK_N) # Relative positions within this KV tile + dim_pos = tl.arange(0, BLOCK_D) + d_mask = dim_pos < HEAD_DIM + kv_abs = kv_start + kv_pos # Absolute positions for memory access + kv_mask = kv_abs < seq_len_kv + + # --- Load K, V tiles: stay in SRAM throughout the Q loop --- + kv_k_ptrs = ( + (kv_offset + kv_abs[:, None]) * stride_kbs + kv_head_idx * stride_kh + dim_pos[None, :] + ) + kv_v_ptrs = ( + (kv_offset + kv_abs[:, None]) * stride_vbs + kv_head_idx * stride_vh + dim_pos[None, :] + ) + k_tile = tl.load(K + kv_k_ptrs, mask=kv_mask[:, None] & d_mask[None, :], other=0.0) + v_tile = tl.load(V + kv_v_ptrs, mask=kv_mask[:, None] & d_mask[None, :], other=0.0) + kT = tl.trans(k_tile) + + # --- Accumulate dK, dV across all Q tiles --- + dk = tl.zeros([BLOCK_N, BLOCK_D], dtype=tl.float32) + dv = tl.zeros([BLOCK_N, BLOCK_D], dtype=tl.float32) + + n_q_tiles = (seq_len_q + BLOCK_M - 1) // BLOCK_M + # Causal: Q position i attends to KV 0..i, so this KV tile (at kv_start) + # only receives gradients from Q tiles where q_pos >= kv_start. Skip earlier ones. + first_q_tile = kv_start // BLOCK_M if IS_CAUSAL else 0 + q_pos_base = tl.arange(0, BLOCK_M) + + for qi in range(first_q_tile, n_q_tiles): + q_pos = qi * BLOCK_M + q_pos_base + q_mask = q_pos < seq_len_q + + # GQA: accumulate contributions from all Q heads sharing this KV head + for g in range(kv_group_num): + head_idx = kv_head_idx * kv_group_num + g + + # Load Q, dO, LSE, delta for this Q tile and head + q_ptrs = ( + (q_offset + q_pos[:, None]) * stride_qbs + head_idx * stride_qh + dim_pos[None, :] + ) + q_tile = tl.load(Q + q_ptrs, mask=q_mask[:, None] & d_mask[None, :], other=0.0) + do_ptrs = ( + (q_offset + q_pos[:, None]) * stride_dobs + head_idx * stride_doh + dim_pos[None, :] + ) + do_tile = tl.load(dO + do_ptrs, mask=q_mask[:, None] & d_mask[None, :], other=0.0) + lse_ptrs = (q_offset + q_pos) * stride_lse_tok + head_idx * stride_lse_head + lse = tl.load(Lse + lse_ptrs, mask=q_mask, other=0.0) + row_delta = tl.load(Delta + lse_ptrs, mask=q_mask, other=0.0) + + # Recompute attention: S = Q @ K^T, P = exp2(S - LSE) + scores = tl.dot(q_tile, kT) * qk_scale + scores = _apply_mask(scores, q_pos, kv_pos, seq_len_q, seq_len_kv, kv_start, IS_CAUSAL) + p = tl.math.exp2(scores - lse[:, None]) + + # dV += P^T @ dO + dv += tl.dot(tl.trans(p.to(do_tile.dtype)), do_tile) + # dS = P * (dO @ V^T - delta), dK += dS^T @ Q + dp = tl.dot(do_tile, tl.trans(v_tile)) + ds = p * (dp - row_delta[:, None]) + dk += tl.dot(tl.trans(ds.to(q_tile.dtype)), q_tile) + + # --- Store dK, dV (dK scaled by sm_scale) --- + dk *= sm_scale + tl.store(dK + kv_k_ptrs, dk.to(k_tile.dtype), mask=kv_mask[:, None] & d_mask[None, :]) + tl.store(dV + kv_v_ptrs, dv.to(v_tile.dtype), mask=kv_mask[:, None] & d_mask[None, :]) + + +# --------------------------------------------------------------------------- +# Autograd wrapper + public API +# --------------------------------------------------------------------------- +class _Attention(torch.autograd.Function): + @staticmethod + def forward( + ctx, + q, + k, + v, + b_start_loc, + b_seq_len, + max_input_len, + is_causal, + sm_scale, + b_start_loc_k, + b_seq_len_k, + max_input_len_k, + ): + 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] + + # Prefill: Q/K/V are the same packed tensor, reuse Q offsets for K/V. + # Decode: K/V is a separate KV cache tensor, caller must pass explicit metadata. + if b_seq_len_k is None: + b_seq_len_k = b_seq_len + b_start_loc_k = b_start_loc + max_input_len_k = max_input_len + + # Pre-multiply scale by log2(e) so the kernel can use exp2() + # exp(score * sm_scale) = exp2(score * sm_scale * log2(e)) + qk_scale = sm_scale * LOG2E + # Triton tiles must be powers of 2; pad head dim + BLOCK_D = triton.next_power_of_2(HEAD_DIM) + + o = torch.empty_like(q) + lse = torch.empty(q.shape[0], num_q_heads, device=q.device, dtype=torch.float32) + + # Grid: (batch, q_heads, q_tiles). Uses a function because BLOCK_M is autotuned. + def grid(META): + return (batch, num_q_heads, triton.cdiv(max_input_len, META["BLOCK_M"])) + + _attn_fwd[grid]( + q, + k, + v, + qk_scale, + b_start_loc, + b_seq_len, + b_start_loc_k, + b_seq_len_k, + o, + lse, + q.stride(0), + q.stride(1), + k.stride(0), + k.stride(1), + v.stride(0), + v.stride(1), + o.stride(0), + o.stride(1), + lse.stride(0), + lse.stride(1), + N_CTX=max_input_len, + kv_group_num=kv_group_num, + BLOCK_D=BLOCK_D, + IS_CAUSAL=is_causal, + HEAD_DIM=HEAD_DIM, + STORE_LSE=True, + # BLOCK_M, BLOCK_N, num_warps, num_stages chosen by autotune + ) + + ctx.save_for_backward(q, k, v, o, lse, b_start_loc, b_seq_len, b_start_loc_k, b_seq_len_k) + ctx.max_input_len = max_input_len + ctx.max_input_len_k = max_input_len_k + ctx.sm_scale = sm_scale + ctx.qk_scale = qk_scale + ctx.is_causal = is_causal + ctx.HEAD_DIM = HEAD_DIM + ctx.kv_group_num = kv_group_num + ctx.num_q_heads = num_q_heads + ctx.num_kv_heads = num_kv_heads + ctx.batch = batch + return o + + @staticmethod + def backward(ctx, grad_output): + q, k, v, o, lse, b_start_loc, b_seq_len, b_start_loc_k, b_seq_len_k = ctx.saved_tensors + HEAD_DIM = ctx.HEAD_DIM + BLOCK = 64 # smaller block for backward to reduce shared memory pressure + BLOCK_D = triton.next_power_of_2(HEAD_DIM) + do = grad_output.contiguous() + num_warps = 4 + + # Phase 1: delta = rowsum(O * dO) + delta = torch.empty_like(lse) + _attn_bwd_preprocess[(ctx.num_q_heads, triton.cdiv(q.shape[0], BLOCK))]( + o, + do, + delta, + o.stride(0), + o.stride(1), + do.stride(0), + do.stride(1), + delta.stride(0), + delta.stride(1), + q.shape[0], + HEAD_DIM=HEAD_DIM, + BLOCK_D=BLOCK_D, + BLOCK_M=BLOCK, + ) + + dq = torch.zeros_like(q) + dk = torch.zeros_like(k) + dv = torch.zeros_like(v) + + bwd_args = ( + q, + k, + v, + do, + lse, + delta, + b_start_loc, + b_seq_len, + b_start_loc_k, + b_seq_len_k, + ctx.qk_scale, + ctx.sm_scale, + q.stride(0), + q.stride(1), + k.stride(0), + k.stride(1), + v.stride(0), + v.stride(1), + do.stride(0), + do.stride(1), + ) + + # Phase 2: dK, dV + _attn_bwd_dkdv[(ctx.batch, ctx.num_kv_heads, triton.cdiv(ctx.max_input_len_k, BLOCK))]( + *bwd_args[:4], + dk, + dv, + *bwd_args[4:], + dk.stride(0), + dk.stride(1), + dv.stride(0), + dv.stride(1), + lse.stride(0), + lse.stride(1), + kv_group_num=ctx.kv_group_num, + BLOCK_M=BLOCK, + BLOCK_D=BLOCK_D, + BLOCK_N=BLOCK, + IS_CAUSAL=ctx.is_causal, + HEAD_DIM=HEAD_DIM, + num_warps=num_warps, + num_stages=1, + ) + + # Phase 3: dQ + _attn_bwd_dq[(ctx.batch, ctx.num_q_heads, triton.cdiv(ctx.max_input_len, BLOCK))]( + *bwd_args[:4], + dq, + *bwd_args[4:], + dq.stride(0), + dq.stride(1), + lse.stride(0), + lse.stride(1), + kv_group_num=ctx.kv_group_num, + BLOCK_M=BLOCK, + BLOCK_D=BLOCK_D, + BLOCK_N=BLOCK, + IS_CAUSAL=ctx.is_causal, + HEAD_DIM=HEAD_DIM, + num_warps=num_warps, + num_stages=1, + ) + + return dq, dk, dv, None, None, None, None, None, None, None, None + + +def attention( + 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, +) -> torch.Tensor: + """Variable-length flash attention with GQA and autograd support. + + Args: + q: [total_q_tokens, num_q_heads, head_dim] + k: [total_kv_tokens, num_kv_heads, head_dim] + v: [total_kv_tokens, num_kv_heads, head_dim] + b_start_loc: [batch] start offset of each Q sequence in the flat tensor. + b_seq_len: [batch] length of each Q sequence. + max_input_len: Maximum Q sequence length (for grid sizing). + is_causal: Whether to apply causal masking. + softmax_scale: Scale factor (default: 1/sqrt(head_dim)). + b_start_loc_k: [batch] start offset for K/V (None = same as Q). + b_seq_len_k: [batch] length for K/V (None = same as Q). + max_input_len_k: Maximum K/V sequence length (None = same as Q). + + Returns: + Output tensor [total_q_tokens, num_q_heads, head_dim]. + """ + sm_scale = 1.0 / (q.shape[2] ** 0.5) if softmax_scale is None else softmax_scale + return _Attention.apply( + q, + k, + v, + b_start_loc, + b_seq_len, + max_input_len, + is_causal, + sm_scale, + b_start_loc_k, + b_seq_len_k, + max_input_len_k, + ) + + +__all__ = ["attention"] diff --git a/modelopt/torch/sparsity/attention_sparsity/config.py b/modelopt/torch/sparsity/attention_sparsity/config.py index 180ff89b0a5..4baf5bbe632 100644 --- a/modelopt/torch/sparsity/attention_sparsity/config.py +++ b/modelopt/torch/sparsity/attention_sparsity/config.py @@ -74,8 +74,8 @@ class SparseAttentionAttributeConfig(ModeloptBaseConfig): title="Backend implementation.", description=( "Backend to use for sparse attention computation. " - "Only 'pytorch' is supported, which uses softmax patching with F.softmax. " - "Requires model to be loaded with attn_implementation='eager'." + "'pytorch' uses softmax patching with F.softmax (requires attn_implementation='eager'). " + "'triton' uses the fused Triton kernel (requires attn_implementation='modelopt_triton')." ), ) @@ -91,6 +91,7 @@ class SparseAttentionAttributeConfig(ModeloptBaseConfig): description=( "Whether the model uses causal (autoregressive) attention. " "If True, sparsity statistics are calculated over the lower triangle only. " + "Set to False for cross-attention models. " "Defaults to True for decoder-only models like GPT, LLaMA, etc." ), ) @@ -106,11 +107,12 @@ def validate_method(cls, v): @field_validator("backend") @classmethod def validate_backend(cls, v): - """Validate backend is pytorch.""" - if v != "pytorch": + """Validate backend is pytorch or triton.""" + if v not in ("pytorch", "triton"): raise ValueError( - f"Invalid backend: {v}. Only 'pytorch' backend is supported. " - f"Model must be loaded with attn_implementation='eager'." + f"Invalid backend: {v}. Supported backends: 'pytorch' (requires " + f"attn_implementation='eager'), 'triton' (requires " + f"attn_implementation='modelopt_triton')." ) return v diff --git a/modelopt/torch/sparsity/attention_sparsity/conversion.py b/modelopt/torch/sparsity/attention_sparsity/conversion.py index 2155a13d0d8..cdc2aed948e 100644 --- a/modelopt/torch/sparsity/attention_sparsity/conversion.py +++ b/modelopt/torch/sparsity/attention_sparsity/conversion.py @@ -32,6 +32,61 @@ from .utils import get_named_sparse_attention_modules, get_sparse_attention_modules +def _set_attn_implementation(model: nn.Module, config: SparseAttentionConfig) -> None: + """Set the correct attn_implementation based on the sparse attention backend. + + - ``backend="triton"``: registers the Triton kernel with HF and sets + ``attn_implementation="modelopt_triton"``. + - ``backend="pytorch"`` (default): sets ``attn_implementation="eager"`` so that + softmax-patching methods (e.g. skip-softmax) work correctly. FlashAttention + and SDPA bypass ``F.softmax``, so eager is required. + + This is called automatically during ``mtsa.sparsify()`` so users never need + to manually set ``attn_implementation``. + """ + sparse_cfg = config.sparse_cfg if hasattr(config, "sparse_cfg") else {} + + # Collect backends only from layer configs (identified by having a "method" key). + # Other dict entries (e.g. "calibration") are not layer configs. + backends = { + v.get("backend", "pytorch") + for v in sparse_cfg.values() + if isinstance(v, dict) and "method" in v + } + + if "triton" in backends and "pytorch" in backends: + raise ValueError( + "Mixed backends ('triton' and 'pytorch') in the same model are not " + "supported. All sparse attention layers must use the same backend." + ) + + model_config = getattr(model, "config", None) + + if "triton" in backends: + from .kernels import register_triton_attention + + if register_triton_attention is None: + raise ImportError( + "Triton backend requires 'triton' and 'transformers' packages. " + "Install with: pip install triton transformers" + ) + if not register_triton_attention(): + raise RuntimeError( + "Failed to register Triton attention with HuggingFace. " + "Check that your transformers version supports ALL_ATTENTION_FUNCTIONS." + ) + + # Set attn_implementation so HF dispatches to the Triton kernel. + # HF's ALL_ATTENTION_FUNCTIONS is checked at forward time, not construction + # time, so this works even after the model is already loaded. + if model_config is not None: + model_config._attn_implementation = "modelopt_triton" + elif model_config is not None: + # For pytorch backend, force eager for softmax patching. + # TODO: Add the triton backend support for skip-softmax. + model_config._attn_implementation = "eager" + + def is_attn_sparsified(model: nn.Module) -> bool: """Check if a model has sparse attention applied. @@ -61,6 +116,9 @@ def convert_to_sparse_attention_model( # Initialize the true module if necessary model = model.init_modellike() if isinstance(model, ModelLikeModule) else model + # Set the correct attn_implementation for the chosen backend + _set_attn_implementation(model, config) + # Apply custom model plugins register_custom_model_plugins_on_the_fly(model) diff --git a/modelopt/torch/sparsity/attention_sparsity/kernels/__init__.py b/modelopt/torch/sparsity/attention_sparsity/kernels/__init__.py new file mode 100644 index 00000000000..dee1bc472a2 --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/kernels/__init__.py @@ -0,0 +1,24 @@ +# 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. + +"""Re-exports from modelopt.torch.kernels for backward compatibility.""" + +from modelopt.torch.kernels import IS_AVAILABLE, attention, register_triton_attention + +__all__ = [ + "IS_AVAILABLE", + "attention", + "register_triton_attention", +] 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 9bfd6a9549d..2501b58f659 100644 --- a/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py +++ b/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py @@ -24,6 +24,9 @@ import numpy as np import torch +import torch.nn.functional as F + +from modelopt.torch.quantization.utils import replace_function from . import SparseAttentionMethod, register_sparse_method @@ -365,6 +368,19 @@ def get_threshold_info(self) -> dict[str, Any]: "value": self.thresholds_config, } + def get_sparse_context(self, module: torch.nn.Module): + """Return a context manager that patches F.softmax with sparse masking.""" + original_softmax = F.softmax + + def sparse_softmax(input, dim=-1, *args, **kwargs): + sparse_mask, stats = self.calculate_sparsity(input) + module._last_stats = stats + if not self._calibration_mode: + input = self.apply_sparsity(input, sparse_mask) + return original_softmax(input, dim, *args, **kwargs) + + return replace_function(torch.nn.functional, "softmax", sparse_softmax) + @property def name(self) -> str: """Method identifier.""" diff --git a/modelopt/torch/sparsity/attention_sparsity/methods/registry.py b/modelopt/torch/sparsity/attention_sparsity/methods/registry.py index 6329e4446f9..3f3e78db6d4 100644 --- a/modelopt/torch/sparsity/attention_sparsity/methods/registry.py +++ b/modelopt/torch/sparsity/attention_sparsity/methods/registry.py @@ -70,6 +70,18 @@ def apply_sparsity( Masked attention scores with sparse elements set to -inf """ + def get_sparse_context(self, module: torch.nn.Module): + """Return a context manager that activates this method's sparsity during forward. + + Each method subclass implements its own activation mechanism: + - Softmax-patching methods replace F.softmax during the forward pass. + - Kernel-fused methods set flags on ``module`` that the kernel reads. + + Args: + module: The SparseAttentionModule wrapping the attention layer. + """ + raise NotImplementedError(f"{type(self).__name__} must implement get_sparse_context()") + def get_threshold_info(self) -> dict[str, Any]: """Get threshold information for display/debugging. diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/huggingface.py b/modelopt/torch/sparsity/attention_sparsity/plugins/huggingface.py index 828d126e864..6cb47fd2f18 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/huggingface.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/huggingface.py @@ -15,7 +15,7 @@ """Dynamic sparse attention registration for HuggingFace models.""" -import warnings +import logging import torch.nn as nn import transformers @@ -25,6 +25,8 @@ from ..sparse_attention import SparseAttentionModule, SparseAttentionRegistry from . import CUSTOM_MODEL_PLUGINS +logger = logging.getLogger(__name__) + class _GenericSparseAttention(SparseAttentionModule): """Generic sparse attention that works with any HF attention module. @@ -93,10 +95,12 @@ def register_sparse_attention_on_the_fly(model: nn.Module) -> bool: SparseAttentionRegistry.register({module_type: type_name})(_GenericSparseAttention) attention_types.add(module_type) registered_count += 1 - print(f"Registered {type_name} for sparse attention optimization") + logger.info("Registered %s for sparse attention optimization", type_name) if registered_count > 0: - print(f"Dynamically registered {registered_count} attention module types for sparsity") + logger.info( + "Dynamically registered %d attention module types for sparsity", registered_count + ) return registered_count > 0 @@ -123,31 +127,5 @@ def _is_supported_model(model: nn.Module) -> bool: return isinstance(model, nn.Module) -def validate_eager_attention(model: nn.Module) -> None: - """Validate and enforce eager attention for HuggingFace models. - - Sparse attention requires attn_implementation='eager' because it - patches torch.nn.functional.softmax, which is only called in eager mode. - - Args: - model: Model to validate - """ - if not isinstance(model, transformers.PreTrainedModel): - return - - attn_impl = getattr(model.config, "_attn_implementation", None) - if attn_impl and attn_impl != "eager": - warnings.warn( - f"Sparse attention requires attn_implementation='eager', but model uses '{attn_impl}'. " - "Forcing eager attention implementation." - ) - model.config._attn_implementation = "eager" - - # Register plugins -CUSTOM_MODEL_PLUGINS.extend( - [ - validate_eager_attention, - register_sparse_attention_on_the_fly, - ] -) +CUSTOM_MODEL_PLUGINS.append(register_sparse_attention_on_the_fly) diff --git a/modelopt/torch/sparsity/attention_sparsity/sparse_attention.py b/modelopt/torch/sparsity/attention_sparsity/sparse_attention.py index 281e11e7d95..4333d12430b 100644 --- a/modelopt/torch/sparsity/attention_sparsity/sparse_attention.py +++ b/modelopt/torch/sparsity/attention_sparsity/sparse_attention.py @@ -17,11 +17,7 @@ from typing import Any -import torch -import torch.nn.functional as F - from modelopt.torch.opt.dynamic import DynamicModule, _DMRegistryCls -from modelopt.torch.quantization.utils import replace_function from .config import SparseAttentionAttributeConfig from .methods import get_sparse_method @@ -32,21 +28,16 @@ class SparseAttentionModule(DynamicModule): """Generic sparse attention module wrapper for applying sparsity to attention layers. This module wraps existing attention implementations to add sparse attention - capabilities by patching torch.nn.functional.softmax. + capabilities. The activation mechanism is delegated to the configured method + via ``method.get_sparse_context(module)``, so each method defines how it + integrates with the forward pass (e.g. softmax patching, kernel flags). Forward Flow: ------------- 1. Check if sparse attention is enabled (pass-through if disabled) - 2. Create softmax patch context with sparse_softmax function - 3. Apply sparse attention by patching F.softmax: - - Patches torch.nn.functional.softmax with sparse_softmax - - sparse_softmax applies method's sparsity logic before softmax - 4. Forward through original attention with sparsity applied - - Requirements: - ------------- - - Model must be loaded with attn_implementation="eager" for proper softmax interception - - Only PyTorch backend is supported (patches F.softmax) + 2. Obtain method-specific context via ``_sparse_method_instance.get_sparse_context(self)`` + 3. Run the original forward inside the context + 4. Collect statistics if stats manager is enabled Attributes: ----------- @@ -190,32 +181,12 @@ def forward(self, *args, **kwargs): return result def _get_sparse_context(self): - """Get the softmax patch context for applying sparse attention.""" - return self._create_softmax_patch_context() - - def _create_softmax_patch_context(self): - """Create context manager for patching softmax function.""" - return replace_function(torch.nn.functional, "softmax", self._create_sparse_softmax()) - - def _create_sparse_softmax(self): - """Create sparse softmax function for current method.""" - original_softmax = F.softmax + """Get the context manager for applying sparse attention. - def sparse_softmax(input, dim=-1, *args, **kwargs): - # Calculate sparsity mask and collect statistics - sparse_mask, stats = self._sparse_method_instance.calculate_sparsity(input) - - # Store stats for collection - self._last_stats = stats - - # Only apply sparsity mask after calibration (not during calibration) - # During calibration, we measure sparsity without modifying the output - if not self._sparse_method_instance._calibration_mode: - input = self._sparse_method_instance.apply_sparsity(input, sparse_mask) - - return original_softmax(input, dim, *args, **kwargs) - - return sparse_softmax + Delegates to the method instance so each method defines its own + activation mechanism (softmax patching, kernel flags, etc.). + """ + return self._sparse_method_instance.get_sparse_context(self) # Create registry for sparse attention modules diff --git a/pyproject.toml b/pyproject.toml index e287975b42d..06e9479ef81 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -204,6 +204,7 @@ extend-ignore = [ "E501", ] # Ignore missing docstrings or line length for Jupyter notebooks "modelopt/torch/quantization/triton/*" = ["N803", "N806", "E731"] # triton style +"modelopt/torch/sparsity/attention_sparsity/kernels/*" = ["N803", "N806"] # triton kernel style "examples/deepseek/ds_kernel.py" = ["N803", "N806", "E731"] # triton style [tool.ruff.lint.pycodestyle] diff --git a/tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa.py b/tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa.py new file mode 100644 index 00000000000..c86a4131edd --- /dev/null +++ b/tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa.py @@ -0,0 +1,439 @@ +# 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. + +"""GPU tests for Triton flash attention kernel.""" + +import pytest +import torch +import torch.nn.functional as F + +pytestmark = [ + pytest.mark.filterwarnings("ignore::UserWarning"), + pytest.mark.filterwarnings("ignore::RuntimeWarning"), + pytest.mark.filterwarnings("ignore::DeprecationWarning"), +] + +from modelopt.torch.kernels import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE + +if TRITON_KERNEL_AVAILABLE: + from modelopt.torch.kernels import attention, register_triton_attention + + if register_triton_attention is not None: + register_triton_attention() + + +def _sdpa_reference(q, k, v, b_start_loc, b_seq_len): + """SDPA causal reference. Supports GQA. Returns [total_tokens, num_heads, dim].""" + batch = b_seq_len.shape[0] + num_q, num_kv = q.shape[1], k.shape[1] + parts = [] + for b in range(batch): + s, n = int(b_start_loc[b].item()), int(b_seq_len[b].item()) + qb = q[s : s + n].unsqueeze(0).permute(0, 2, 1, 3) + kb = k[s : s + n].unsqueeze(0).permute(0, 2, 1, 3) + vb = v[s : s + n].unsqueeze(0).permute(0, 2, 1, 3) + if num_q != num_kv: + r = num_q // num_kv + kb = kb.repeat_interleave(r, dim=1) + vb = vb.repeat_interleave(r, dim=1) + ob = F.scaled_dot_product_attention(qb, kb, vb, is_causal=True) + parts.append(ob.permute(0, 2, 1, 3).squeeze(0)) + return torch.cat(parts, dim=0) + + +@pytest.fixture(scope="module") +def tiny_llama_dir(tmp_path_factory): + """Tiny Llama: 2 layers, 64 hidden, 4 q-heads, 2 kv-heads, head_dim=16.""" + from _test_utils.torch.transformers_models import create_tiny_llama_dir + + return create_tiny_llama_dir( + tmp_path_factory.mktemp("tiny_llama"), + with_tokenizer=True, + num_hidden_layers=2, + hidden_size=64, + num_attention_heads=4, + num_key_value_heads=2, + intermediate_size=64, + max_position_embeddings=64, + ) + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") +class TestTritonFaVsSdpa: + """Triton flash attention matches PyTorch SDPA for prefill and decode.""" + + @pytest.mark.parametrize( + ("dtype", "num_heads", "num_kv_heads", "head_dim"), + [ + (torch.float32, 2, 2, 32), + (torch.float16, 4, 2, 64), + (torch.bfloat16, 4, 2, 128), + ], + ids=["fp32_mha", "fp16_gqa", "bf16_gqa_hdim128"], + ) + def test_prefill_matches_sdpa(self, dtype, num_heads, num_kv_heads, head_dim): + """Prefill matches SDPA.""" + seq_lens = [8, 12] + total = sum(seq_lens) + scale = 1.0 / (head_dim**0.5) + + torch.manual_seed(123) + q = torch.randn(total, num_heads, head_dim, device="cuda", dtype=dtype) + k = torch.randn(total, num_kv_heads, head_dim, device="cuda", dtype=dtype) + v = torch.randn(total, num_kv_heads, head_dim, device="cuda", dtype=dtype) + locs = torch.tensor([0, seq_lens[0]], device="cuda", dtype=torch.int32) + lens = torch.tensor(seq_lens, device="cuda", dtype=torch.int32) + + o = attention( + q, + k, + v, + b_start_loc=locs, + b_seq_len=lens, + max_input_len=max(seq_lens), + is_causal=True, + softmax_scale=scale, + ) + torch.testing.assert_close(o, _sdpa_reference(q, k, v, locs, lens), rtol=1e-3, atol=1e-3) + + def test_decode_matches_sdpa(self): + """Decode matches SDPA.""" + batch = 2 + seq_lens_k = [5, 9] # KV lengths (context + current token) + num_heads, num_kv_heads, head_dim = 4, 2, 32 + scale = 1.0 / (head_dim**0.5) + + torch.manual_seed(103) + # Q: one token per batch element -> flat [batch, num_heads, head_dim] + q_flat = torch.randn(batch, num_heads, head_dim, device="cuda", dtype=torch.float32) + + # K/V: variable-length, packed into flat tensors + total_kv = sum(seq_lens_k) + k_flat = torch.randn(total_kv, num_kv_heads, head_dim, device="cuda", dtype=torch.float32) + v_flat = torch.randn(total_kv, num_kv_heads, head_dim, device="cuda", dtype=torch.float32) + + cumsum = [0] + for sl in seq_lens_k: + cumsum.append(cumsum[-1] + sl) + b_start_loc_q = torch.arange(batch, device="cuda", dtype=torch.int32) + b_seq_len_q = torch.ones(batch, device="cuda", dtype=torch.int32) + b_start_loc_k = torch.tensor(cumsum[:-1], device="cuda", dtype=torch.int32) + b_seq_len_k = torch.tensor(seq_lens_k, device="cuda", dtype=torch.int32) + + out = attention( + q_flat, + k_flat, + v_flat, + b_start_loc=b_start_loc_q, + b_seq_len=b_seq_len_q, + max_input_len=1, + is_causal=False, + softmax_scale=scale, + b_start_loc_k=b_start_loc_k, + b_seq_len_k=b_seq_len_k, + max_input_len_k=max(seq_lens_k), + ) + + for i in range(batch): + sl = seq_lens_k[i] + s = cumsum[i] + qb = q_flat[i : i + 1].unsqueeze(2) # [1, heads, 1, dim] + kb = k_flat[s : s + sl].unsqueeze(0).permute(0, 2, 1, 3) + vb = v_flat[s : s + sl].unsqueeze(0).permute(0, 2, 1, 3) + kb = kb.repeat_interleave(num_heads // num_kv_heads, dim=1) + vb = vb.repeat_interleave(num_heads // num_kv_heads, dim=1) + ref = F.scaled_dot_product_attention(qb, kb, vb, is_causal=False).squeeze(2) + torch.testing.assert_close(out[i : i + 1], ref, rtol=1e-3, atol=1e-3) + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") +class TestSparseAttentionIntegration: + """HF model + mtsa.sparsify integration.""" + + def test_triton_matches_eager(self, tiny_llama_dir): + """Triton attention produces same logits and generated tokens as eager.""" + pytest.importorskip("transformers") + from transformers import AutoModelForCausalLM, AutoTokenizer + + tok = AutoTokenizer.from_pretrained(tiny_llama_dir) + if tok.pad_token_id is None: + tok.pad_token_id = tok.eos_token_id + ids = tok("The capital of France is", return_tensors="pt").input_ids.to("cuda") + + # Eager baseline + model_eager = AutoModelForCausalLM.from_pretrained( + tiny_llama_dir, + attn_implementation="eager", + torch_dtype=torch.bfloat16, + device_map="cuda", + ) + model_eager.eval() + with torch.no_grad(): + logits_eager = model_eager(input_ids=ids).logits + out_eager = model_eager.generate( + ids, + max_new_tokens=5, + do_sample=False, + pad_token_id=tok.pad_token_id, + ) + del model_eager + + # Triton + model_triton = AutoModelForCausalLM.from_pretrained( + tiny_llama_dir, + attn_implementation="modelopt_triton", + torch_dtype=torch.bfloat16, + device_map="cuda", + ) + model_triton.eval() + with torch.no_grad(): + logits_triton = model_triton(input_ids=ids).logits + out_triton = model_triton.generate( + ids, + max_new_tokens=5, + do_sample=False, + pad_token_id=tok.pad_token_id, + ) + + # Logits should be close (bf16 tolerance) + torch.testing.assert_close(logits_triton, logits_eager, rtol=2e-2, atol=2e-2) + # Generated tokens must be identical (greedy decoding is deterministic) + assert torch.equal(out_triton, out_eager), ( + f"Generated tokens differ:\n eager: {out_eager}\n triton: {out_triton}" + ) + + def test_triton_padded_batch(self, tiny_llama_dir): + """Padded batch (2D attention mask) produces valid logits for each sequence.""" + pytest.importorskip("transformers") + from transformers import AutoModelForCausalLM, AutoTokenizer + + model = AutoModelForCausalLM.from_pretrained( + tiny_llama_dir, + attn_implementation="modelopt_triton", + torch_dtype=torch.bfloat16, + device_map="cuda", + ) + model.eval() + tok = AutoTokenizer.from_pretrained(tiny_llama_dir) + if tok.pad_token_id is None: + tok.pad_token_id = tok.eos_token_id + tok.padding_side = "right" + + inputs = tok( + ["Hello world", "The capital of France is Paris and"], + return_tensors="pt", + padding=True, + ).to("cuda") + with torch.no_grad(): + logits = model(**inputs).logits + assert not torch.isnan(logits).any() and not torch.isinf(logits).any() + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") +class TestBackward: + """Backward pass gradient correctness tests.""" + + def _sdpa_backward_ref(self, q, k, v, scale, is_causal=True): + """Run SDPA forward+backward, return output and gradients.""" + q_ref = q.clone().unsqueeze(0).permute(0, 2, 1, 3).requires_grad_(True) + k_ref = k.clone().unsqueeze(0).permute(0, 2, 1, 3).requires_grad_(True) + v_ref = v.clone().unsqueeze(0).permute(0, 2, 1, 3).requires_grad_(True) + num_q, num_kv = q_ref.shape[1], k_ref.shape[1] + if num_q != num_kv: + r = num_q // num_kv + k_exp = k_ref.repeat_interleave(r, dim=1) + v_exp = v_ref.repeat_interleave(r, dim=1) + else: + k_exp, v_exp = k_ref, v_ref + o_ref = F.scaled_dot_product_attention( + q_ref, k_exp, v_exp, is_causal=is_causal, scale=scale + ) + o_ref.sum().backward() + dq = q_ref.grad.permute(0, 2, 1, 3).squeeze(0) + dk = k_ref.grad.permute(0, 2, 1, 3).squeeze(0) + dv = v_ref.grad.permute(0, 2, 1, 3).squeeze(0) + return o_ref.permute(0, 2, 1, 3).squeeze(0).detach(), dq.detach(), dk.detach(), dv.detach() + + def test_backward_causal_matches_sdpa(self): + """dQ, dK, dV match SDPA backward for causal self-attention.""" + from modelopt.torch.kernels import attention + + seq_len = 16 + num_heads, num_kv_heads, head_dim = 2, 2, 32 + scale = 1.0 / (head_dim**0.5) + + torch.manual_seed(42) + q = torch.randn( + seq_len, num_heads, head_dim, device="cuda", dtype=torch.float32, requires_grad=True + ) + k = torch.randn( + seq_len, num_kv_heads, head_dim, device="cuda", dtype=torch.float32, requires_grad=True + ) + v = torch.randn( + seq_len, num_kv_heads, head_dim, device="cuda", dtype=torch.float32, requires_grad=True + ) + + o = attention( + q, + k, + v, + b_start_loc=torch.tensor([0], device="cuda", dtype=torch.int32), + b_seq_len=torch.tensor([seq_len], device="cuda", dtype=torch.int32), + max_input_len=seq_len, + is_causal=True, + softmax_scale=scale, + ) + o.sum().backward() + + _, dq_ref, dk_ref, dv_ref = self._sdpa_backward_ref( + q.detach(), k.detach(), v.detach(), scale, is_causal=True + ) + + torch.testing.assert_close(q.grad, dq_ref, rtol=5e-3, atol=5e-3) + torch.testing.assert_close(k.grad, dk_ref, rtol=5e-3, atol=5e-3) + torch.testing.assert_close(v.grad, dv_ref, rtol=5e-3, atol=5e-3) + + def test_backward_gqa(self): + """Backward with GQA (4 q-heads, 2 kv-heads), multi-tile (seq_len=256).""" + from modelopt.torch.kernels import attention + + seq_len = 256 + num_heads, num_kv_heads, head_dim = 4, 2, 32 + scale = 1.0 / (head_dim**0.5) + + torch.manual_seed(43) + q = torch.randn( + seq_len, num_heads, head_dim, device="cuda", dtype=torch.float32, requires_grad=True + ) + k = torch.randn( + seq_len, num_kv_heads, head_dim, device="cuda", dtype=torch.float32, requires_grad=True + ) + v = torch.randn( + seq_len, num_kv_heads, head_dim, device="cuda", dtype=torch.float32, requires_grad=True + ) + + o = attention( + q, + k, + v, + b_start_loc=torch.tensor([0], device="cuda", dtype=torch.int32), + b_seq_len=torch.tensor([seq_len], device="cuda", dtype=torch.int32), + max_input_len=seq_len, + is_causal=True, + softmax_scale=scale, + ) + o.sum().backward() + + _, dq_ref, dk_ref, dv_ref = self._sdpa_backward_ref( + q.detach(), k.detach(), v.detach(), scale, is_causal=True + ) + + torch.testing.assert_close(q.grad, dq_ref, rtol=5e-3, atol=5e-3) + torch.testing.assert_close(k.grad, dk_ref, rtol=5e-3, atol=5e-3) + torch.testing.assert_close(v.grad, dv_ref, rtol=5e-3, atol=5e-3) + + def test_backward_multi_batch_variable_length(self): + """Multi-batch variable-length causal backward matches per-sample SDPA.""" + from modelopt.torch.kernels import attention + + seq_lens = [8, 12] + total = sum(seq_lens) + num_heads, num_kv_heads, head_dim = 2, 2, 32 + scale = 1.0 / (head_dim**0.5) + + torch.manual_seed(45) + q = torch.randn( + total, num_heads, head_dim, device="cuda", dtype=torch.float32, requires_grad=True + ) + k = torch.randn( + total, num_kv_heads, head_dim, device="cuda", dtype=torch.float32, requires_grad=True + ) + v = torch.randn( + total, num_kv_heads, head_dim, device="cuda", dtype=torch.float32, requires_grad=True + ) + locs = torch.tensor([0, seq_lens[0]], device="cuda", dtype=torch.int32) + lens = torch.tensor(seq_lens, device="cuda", dtype=torch.int32) + + o = attention( + q, + k, + v, + b_start_loc=locs, + b_seq_len=lens, + max_input_len=max(seq_lens), + is_causal=True, + softmax_scale=scale, + ) + o.sum().backward() + + # Per-sample SDPA reference + dq_ref = torch.zeros_like(q) + dk_ref = torch.zeros_like(k) + dv_ref = torch.zeros_like(v) + for b in range(len(seq_lens)): + s, n = int(locs[b].item()), seq_lens[b] + _, dq_b, dk_b, dv_b = self._sdpa_backward_ref( + q.detach()[s : s + n], + k.detach()[s : s + n], + v.detach()[s : s + n], + scale, + is_causal=True, + ) + dq_ref[s : s + n] = dq_b + dk_ref[s : s + n] = dk_b + dv_ref[s : s + n] = dv_b + + torch.testing.assert_close(q.grad, dq_ref, rtol=5e-3, atol=5e-3) + torch.testing.assert_close(k.grad, dk_ref, rtol=5e-3, atol=5e-3) + torch.testing.assert_close(v.grad, dv_ref, rtol=5e-3, atol=5e-3) + + def test_backward_longer_sequences(self): + """Backward with seq_len=512, GQA, exercises multi-tile loops.""" + from modelopt.torch.kernels import attention + + seq_len = 512 + num_heads, num_kv_heads, head_dim = 4, 2, 64 + scale = 1.0 / (head_dim**0.5) + + torch.manual_seed(49) + q = torch.randn( + seq_len, num_heads, head_dim, device="cuda", dtype=torch.float32, requires_grad=True + ) + k = torch.randn( + seq_len, num_kv_heads, head_dim, device="cuda", dtype=torch.float32, requires_grad=True + ) + v = torch.randn( + seq_len, num_kv_heads, head_dim, device="cuda", dtype=torch.float32, requires_grad=True + ) + + o = attention( + q, + k, + v, + b_start_loc=torch.tensor([0], device="cuda", dtype=torch.int32), + b_seq_len=torch.tensor([seq_len], device="cuda", dtype=torch.int32), + max_input_len=seq_len, + is_causal=True, + softmax_scale=scale, + ) + o.sum().backward() + + _, dq_ref, dk_ref, dv_ref = self._sdpa_backward_ref( + q.detach(), k.detach(), v.detach(), scale, is_causal=True + ) + + torch.testing.assert_close(q.grad, dq_ref, rtol=5e-3, atol=5e-3) + torch.testing.assert_close(k.grad, dk_ref, rtol=5e-3, atol=5e-3) + torch.testing.assert_close(v.grad, dv_ref, rtol=5e-3, atol=5e-3) From 6198a22fa018d76014c4d568bce5e6f3a6fadf81 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Tue, 17 Mar 2026 14:13:39 -0700 Subject: [PATCH 2/2] Address review feedbacks Signed-off-by: Kai Xu --- .../sparsity/attention_sparsity/plugins/huggingface.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/huggingface.py b/modelopt/torch/sparsity/attention_sparsity/plugins/huggingface.py index 6cb47fd2f18..599832943dc 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/huggingface.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/huggingface.py @@ -15,8 +15,6 @@ """Dynamic sparse attention registration for HuggingFace models.""" -import logging - import torch.nn as nn import transformers @@ -25,8 +23,6 @@ from ..sparse_attention import SparseAttentionModule, SparseAttentionRegistry from . import CUSTOM_MODEL_PLUGINS -logger = logging.getLogger(__name__) - class _GenericSparseAttention(SparseAttentionModule): """Generic sparse attention that works with any HF attention module. @@ -95,12 +91,10 @@ def register_sparse_attention_on_the_fly(model: nn.Module) -> bool: SparseAttentionRegistry.register({module_type: type_name})(_GenericSparseAttention) attention_types.add(module_type) registered_count += 1 - logger.info("Registered %s for sparse attention optimization", type_name) + print(f"Registered {type_name} for sparse attention optimization") if registered_count > 0: - logger.info( - "Dynamically registered %d attention module types for sparsity", registered_count - ) + print(f"Dynamically registered {registered_count} attention module types for sparsity") return registered_count > 0