Flash mla interface

src/transformers/integrations/__init__.py

@@ -66,6 +66,7 @@
         "replace_with_higgs_linear",
     ],
     "hqq": ["prepare_for_hqq_linear"],
+    "flash_mla": ["flash_mla_attention_forward"],
     "hub_kernels": [
         "LayerRepository",
         "lazy_load_kernel",
@@ -214,6 +215,7 @@
     )
     from .higgs import HiggsLinear, dequantize_higgs, quantize_with_higgs, replace_with_higgs_linear
     from .hqq import prepare_for_hqq_linear
+    from .flash_mla import flash_mla_attention_forward
     from .hub_kernels import (
         LayerRepository,
         lazy_load_kernel,

src/transformers/integrations/finegrained_fp8.py

@@ -412,55 +412,55 @@ def w8a8_block_fp8_matmul(
     Otherwise falls back to Triton.
     """
 
-    if _supports_cutlass(block_size, output_dtype):
-        kernel = _get_quantization_kernel()
-        if kernel is not None:
-            try:
-                # CUTLASS expects:
-                # - A: [M, K] row-major, float8_e4m3fn
-                # - B: [K, N] column-major, float8_e4m3fn
-                # - As: [M, K//128] M-major (activation scales)
-                # - Bs: [K//128, N//128] K-major (weight scales)
-
-                # Reshape A to 2D if needed
-                original_shape = A.shape
-                M = A.numel() // A.shape[-1]
-                K = A.shape[-1]
-                N = B.shape[0]
-
-                # CUTLASS requires dimensions divisible by 16
-                if K % 16 != 0 or N % 16 != 0:
-                    raise ValueError(f"CUTLASS requires K ({K}) and N ({N}) divisible by 16")
-
-                A_2d = A.view(M, K).contiguous()
-                # B needs to be column-major for CUTLASS: [K, N] with stride(0)==1
-                # Our B is [N, K] row-major. Make it contiguous first, then transpose.
-                # B.contiguous() gives [N, K] with stride=(K,1)
-                # B.contiguous().t() gives [K, N] with stride=(1,K) which is column-major
-                # Do NOT call .contiguous() after .t() as it would make it row-major!
-                B_col_major = B.contiguous().t()
-
-                # Scales need proper layout for CUTLASS blockwise:
-                # As should be [M, K//128] with M-major layout (stride(0)==1)
-                # Bs should be [K//128, N//128] with K-major layout (stride(0)==1)
-
-                # As: reshape to [M, K//128], then make M-major via t().contiguous().t()
-                As_2d = As.view(M, -1).contiguous()
-                As_2d = As_2d.t().contiguous().t()  # [M, K//128] with stride(0)==1
-
-                # Bs: our input is [N//128, K//128], need [K//128, N//128] with stride(0)==1
-                # Transpose to get [K//128, N//128], then make K-major via t().contiguous().t()
-                Bs_km = Bs.contiguous().t()  # [K//128, N//128]
-                Bs_km = Bs_km.t().contiguous().t()  # Make K-major (stride(0)==1)
-
-                # Call CUTLASS kernel - it returns the output tensor
-                # Signature: cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias=None) -> Tensor
-                C = kernel.cutlass_scaled_mm(A_2d, B_col_major, As_2d, Bs_km, output_dtype, None)
-                # Reshape output back
-                C_shape = original_shape[:-1] + (N,)
-                return C.view(C_shape)
-            except Exception as e:
-                logger.warning_once(f"CUTLASS kernel failed: {e}. Falling back to Triton.")
+    # if _supports_cutlass(block_size, output_dtype):
+    #     kernel = _get_quantization_kernel()
+    #     if kernel is not None:
+    #         try:
+    #             # CUTLASS expects:
+    #             # - A: [M, K] row-major, float8_e4m3fn
+    #             # - B: [K, N] column-major, float8_e4m3fn
+    #             # - As: [M, K//128] M-major (activation scales)
+    #             # - Bs: [K//128, N//128] K-major (weight scales)
+
+    #             # Reshape A to 2D if needed
+    #             original_shape = A.shape
+    #             M = A.numel() // A.shape[-1]
+    #             K = A.shape[-1]
+    #             N = B.shape[0]
+
+    #             # CUTLASS requires dimensions divisible by 16
+    #             if K % 16 != 0 or N % 16 != 0:
+    #                 raise ValueError(f"CUTLASS requires K ({K}) and N ({N}) divisible by 16")
+
+    #             A_2d = A.view(M, K).contiguous()
+    #             # B needs to be column-major for CUTLASS: [K, N] with stride(0)==1
+    #             # Our B is [N, K] row-major. Make it contiguous first, then transpose.
+    #             # B.contiguous() gives [N, K] with stride=(K,1)
+    #             # B.contiguous().t() gives [K, N] with stride=(1,K) which is column-major
+    #             # Do NOT call .contiguous() after .t() as it would make it row-major!
+    #             B_col_major = B.contiguous().t()
+
+    #             # Scales need proper layout for CUTLASS blockwise:
+    #             # As should be [M, K//128] with M-major layout (stride(0)==1)
+    #             # Bs should be [K//128, N//128] with K-major layout (stride(0)==1)
+
+    #             # As: reshape to [M, K//128], then make M-major via t().contiguous().t()
+    #             As_2d = As.view(M, -1).contiguous()
+    #             As_2d = As_2d.t().contiguous().t()  # [M, K//128] with stride(0)==1
+
+    #             # Bs: our input is [N//128, K//128], need [K//128, N//128] with stride(0)==1
+    #             # Transpose to get [K//128, N//128], then make K-major via t().contiguous().t()
+    #             Bs_km = Bs.contiguous().t()  # [K//128, N//128]
+    #             Bs_km = Bs_km.t().contiguous().t()  # Make K-major (stride(0)==1)
+
+    #             # Call CUTLASS kernel - it returns the output tensor
+    #             # Signature: cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias=None) -> Tensor
+    #             C = kernel.cutlass_scaled_mm(A_2d, B_col_major, As_2d, Bs_km, output_dtype, None)
+    #             # Reshape output back
+    #             C_shape = original_shape[:-1] + (N,)
+    #             return C.view(C_shape)
+    #         except Exception as e:
+    #             logger.warning_once(f"CUTLASS kernel failed: {e}. Falling back to Triton.")
 
     # Fall back to Triton
     return w8a8_block_fp8_matmul_triton(A, B, As, Bs, block_size, output_dtype)

src/transformers/integrations/flash_mla.py

@@ -0,0 +1,237 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Flash-MLA attention integration for sparse attention with Dynamic Sparse Attention (DSA).
+
+This module provides a wrapper around the flash-mla kernel from kernels-community/flash-mla,
+with automatic fallback to flash_attention_2 when input tokens < 2048.
+"""
+
+import torch
+
+from ..utils import logging
+from .flash_attention import flash_attention_forward, get_target_dtype
+
+
+logger = logging.get_logger(__name__)
+
+# Minimum sequence length to use flash-mla sparse attention
+# Below this threshold, we fall back to flash_attention_2
+FLASH_MLA_MIN_SEQ_LEN = 2048
+
+
+def flash_mla_attention_forward(
+    module: torch.nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    dropout: float = 0.0,
+    scaling: float | None = None,
+    sliding_window: int | None = None,
+    softcap: float | None = None,
+    is_causal: bool | None = None,
+    **kwargs,
+) -> tuple[torch.Tensor, None]:
+    """
+    Flash-MLA attention forward pass with automatic fallback to flash_attention_2.
+
+    This wrapper handles:
+    - Fallback to flash_attention_2 when sequence length < 2048
+    - Sparse attention via topk_indices when sequence length >= 2048
+    - Tensor layout conversion (transformers BHSD to flash-mla BSHD)
+    - head_dim_v padding if needed (flash-mla requires specific dimensions)
+
+    Args:
+        module (`torch.nn.Module`):
+            The attention module containing config and layer information.
+        query (`torch.Tensor`):
+            Query tensor of shape `[B, H, S, D]` (BHSD format).
+        key (`torch.Tensor`):
+            Key tensor of shape `[B, H, T, D]` (BHSD format).
+        value (`torch.Tensor`):
+            Value tensor of shape `[B, H, T, D_v]` (BHSD format).
+        attention_mask (`torch.Tensor | None`):
+            Combined attention mask (causal + DSA sparse mask). Used for flash_attention_2 fallback.
+        dropout (`float`, optional):
+            Dropout probability. Defaults to 0.0.
+        scaling (`float | None`, optional):
+            Scaling factor for attention scores. Defaults to None.
+        sliding_window (`int | None`, optional):
+            Sliding window size. Defaults to None.
+        softcap (`float | None`, optional):
+            Soft cap for attention logits. Defaults to None.
+        is_causal (`bool | None`, optional):
+            Whether attention is causal. Defaults to None.
+        **kwargs:
+            Additional keyword arguments, including:
+            - topk_indices (`torch.Tensor | None`): Indices for sparse attention from DSA indexer.
+
+    Returns:
+        `tuple[torch.Tensor, None]`: Attention output tensor and None (no attention weights).
+    """
+    # Extract topk_indices from kwargs (used for sparse attention)
+    topk_indices = kwargs.pop("topk_indices", None)
+
+    # Get total sequence length from key tensor
+    # key shape is [B, H, T, D] in BHSD format
+    seq_len = key.shape[2]
+
+    # Fallback to flash_attention_2 when sequence length is below threshold
+    # This is because flash-mla sparse attention is optimized for longer sequences
+    if seq_len < FLASH_MLA_MIN_SEQ_LEN:
+        logger.debug(
+            f"Sequence length {seq_len} < {FLASH_MLA_MIN_SEQ_LEN}, falling back to flash_attention_2"
+        )
+        return flash_attention_forward(
+            module=module,
+            query=query,
+            key=key,
+            value=value,
+            attention_mask=attention_mask,
+            dropout=dropout,
+            scaling=scaling,
+            sliding_window=sliding_window,
+            softcap=softcap,
+            is_causal=is_causal,
+            **kwargs,
+        )
+
+    # Use flash-mla sparse attention with topk_indices
+    return _flash_mla_sparse_forward(
+        module=module,
+        query=query,
+        key=key,
+        value=value,
+        topk_indices=topk_indices,
+        dropout=dropout,
+        scaling=scaling,
+        **kwargs,
+    )
+
+
+def _flash_mla_sparse_forward(
+    module: torch.nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    topk_indices: torch.Tensor | None,
+    dropout: float = 0.0,
+    scaling: float | None = None,
+    **kwargs,
+) -> tuple[torch.Tensor, None]:
+    """
+    Internal function for flash-mla sparse attention computation.
+
+    This function handles the actual flash-mla kernel call with sparse attention
+    via topk_indices from the DSA indexer.
+
+    Args:
+        module (`torch.nn.Module`):
+            The attention module containing config and layer information.
+        query (`torch.Tensor`):
+            Query tensor of shape `[B, H, S, D]` (BHSD format).
+        key (`torch.Tensor`):
+            Key tensor of shape `[B, H, T, D]` (BHSD format).
+        value (`torch.Tensor`):
+            Value tensor of shape `[B, H, T, D_v]` (BHSD format).
+        topk_indices (`torch.Tensor | None`):
+            Indices for sparse attention from DSA indexer, shape `[B, S, topk]`.
+        dropout (`float`, optional):
+            Dropout probability. Defaults to 0.0.
+        scaling (`float | None`, optional):
+            Scaling factor for attention scores. Defaults to None.
+        **kwargs:
+            Additional keyword arguments.
+
+    Returns:
+        `tuple[torch.Tensor, None]`: Attention output tensor and None (no attention weights).
+    """
+    if kwargs.get("output_attentions", False):
+        logger.warning_once(
+            "Flash-MLA does not support `output_attentions=True`. "
+            "Please set your attention to `eager` if you want this feature."
+        )
+
+    # Get batch size and sequence lengths
+    batch_size, num_heads, q_len, head_dim = query.shape
+    _, _, kv_len, _ = key.shape
+
+    # Convert from BHSD (transformers) to BSHD (flash-mla) format
+    # query: [B, H, S, D] -> [B, S, H, D]
+    # key: [B, H, T, D] -> [B, T, H, D]
+    # value: [B, H, T, D_v] -> [B, T, H, D_v]
+    query = query.transpose(1, 2).contiguous()
+    key = key.transpose(1, 2).contiguous()
+    value = value.transpose(1, 2).contiguous()
+
+    # Handle dtype conversion for flash attention compatibility
+    target_dtype = get_target_dtype(query, module)
+    if target_dtype is not None:
+        query = query.to(target_dtype)
+        key = key.to(target_dtype)
+        value = value.to(target_dtype)
+
+    # Get the flash-mla kernel function
+    # This is loaded via hub_kernels infrastructure
+    try:
+        from ..integrations.hub_kernels import get_kernel
+
+        flash_mla_kernel = get_kernel("kernels-community/flash-mla")
+        flash_mla_sparse_fwd = flash_mla_kernel.flash_mla_sparse_fwd
+    except (ImportError, AttributeError) as e:
+        raise RuntimeError(
+            f"Failed to load flash-mla kernel. Make sure kernels-community/flash-mla is available. Error: {e}"
+        )
+
+    # Prepare scaling factor
+    if scaling is None:
+        scaling = head_dim**-0.5
+
+    # Get value head dimension (may differ from query/key head dimension in MLA)
+    v_head_dim = value.shape[-1]
+
+    # Flash-MLA may require specific head_dim_v (e.g., 512)
+    # Pad if necessary
+    flash_mla_v_head_dim = 512
+    needs_v_padding = v_head_dim < flash_mla_v_head_dim
+    if needs_v_padding:
+        value = torch.nn.functional.pad(value, (0, flash_mla_v_head_dim - v_head_dim))
+
+    # Call flash-mla kernel with sparse attention
+    # The kernel expects:
+    # - q: [B, S, H, D]
+    # - k_cache: [B, T, H, D] (or compressed format)
+    # - v_cache: [B, T, H, D_v]
+    # - topk_indices: [B, S, topk] for sparse attention
+    attn_output = flash_mla_sparse_fwd(
+        q=query,
+        kv=torch.cat([key, value], dim=-1),
+        indices=topk_indices,
+        sm_scale=scaling,
+        topk_length = module.top_k_length if hasattr(module, "top_k_length") else None,
+    )
+
+    # Remove padding if we added it
+    if needs_v_padding:
+        attn_output = attn_output[..., :v_head_dim]
+
+    # Convert back from BSHD to BHSD format
+    # attn_output: [B, S, H, D_v] -> [B, H, S, D_v]
+    attn_output = attn_output.transpose(1, 2)
+
+    return attn_output, None
+
+
+__all__ = ["flash_mla_attention_forward"]

src/transformers/integrations/hub_kernels.py

@@ -277,6 +277,7 @@ def register_kernel_mapping_transformers(*args, **kwargs):
     "causal-conv1d": {"repo_id": "kernels-community/causal-conv1d", "version": 1},
     "mamba-ssm": {"repo_id": "kernels-community/mamba-ssm", "version": 1},
     "falcon_mamba-ssm": {"repo_id": "kernels-community/mamba-ssm", "version": 1},
+    "flash-mla": {"repo_id": "kernels-community/flash-mla"},
 }
 
 _KERNEL_MODULE_MAPPING: dict[str, ModuleType | None] = {}
@@ -338,6 +339,12 @@ def load_and_register_attn_kernel(
         if attention_wrapper is None:
             attention_wrapper = flash_attention_forward
         kernel_function = attention_wrapper
+    if hasattr(kernel, "flash_mla_sparse_fwd"):
+        from .flash_mla import flash_mla_attention_forward
+
+        if attention_wrapper is None:
+            attention_wrapper = flash_mla_attention_forward
+        kernel_function = attention_wrapper
     elif kernel_name is not None:
         kernel_function = getattr(kernel, kernel_name)