Upload folder using huggingface_hub

.gitattributes

@@ -37,3 +37,4 @@ Qwen-Image-Edit-2509/processor/tokenizer.json filter=lfs diff=lfs merge=lfs -tex
 Qwen-Image-i2i_merged/processor/tokenizer.json filter=lfs diff=lfs merge=lfs -text
 flash-attn3-kernel/torch28-cxx11-cu128-x86_64-linux/vllm_flash_attn3/_vllm_flash_attn3_28fbd26_dirty.abi3.so filter=lfs diff=lfs merge=lfs -text
 vllm_flash_attn3/_flash_attn3_1d39a44.abi3.so filter=lfs diff=lfs merge=lfs -text
+flash-attn3-kernel/torch28-cxx11-cu128-x86_64-linux/vllm_flash_attn3/_flash_attn3_1d39a44.abi3.so filter=lfs diff=lfs merge=lfs -text

flash-attn3-kernel/torch28-cxx11-cu128-x86_64-linux/vllm_flash_attn3/_flash_attn3_1d39a44.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:17fbfaef77c6560c59ca5751ace86ca16a3af0471da3d265c4b101e3b1b2ec30
+size 779759224

flash-attn3-kernel/torch28-cxx11-cu128-x86_64-linux/vllm_flash_attn3/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _vllm_flash_attn3_28fbd26_dirty
-ops = torch.ops._vllm_flash_attn3_28fbd26_dirty
+from . import _flash_attn3_1d39a44
+ops = torch.ops._flash_attn3_1d39a44
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_vllm_flash_attn3_28fbd26_dirty::{op_name}"
+    return f"_flash_attn3_1d39a44::{op_name}"

flash-attn3-kernel/torch28-cxx11-cu128-x86_64-linux/vllm_flash_attn3/flash_attn3/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+from pathlib import Path
+from types import ModuleType
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

flash-attn3-kernel/torch28-cxx11-cu128-x86_64-linux/vllm_flash_attn3/flash_attn_interface.py

@@ -5,12 +5,7 @@ from typing import Optional, Union
 import torch
 import torch.nn as nn
 
-# isort: off
-# We need to import the CUDA kernels after importing torch
-from ._ops import ops
-
-# isort: on
-
+from ._ops import ops as flash_attn_3_cuda
 
 def maybe_contiguous(x):
     return x.contiguous() if x is not None and x.stride(-1) != 1 else x
@@ -43,13 +38,13 @@ def _flash_attn_forward(
         softmax_scale,
         causal,
         window_size=(-1, -1),
+        attention_chunk=0,
         softcap=0.0,
         rotary_interleaved=True,
         scheduler_metadata=None,
         num_splits=1,
         pack_gqa=None,
-        sm_margin=0,
-        s_aux=None):
+        sm_margin=0):
     q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)]
     v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v
     cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [
@@ -61,7 +56,7 @@ def _flash_attn_forward(
     ]
     rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)]
     seqlens_rotary = maybe_contiguous(seqlens_rotary)
-    out, softmax_lse, *rest = ops.fwd(
+    out, softmax_lse, *rest = flash_attn_3_cuda.fwd(
         q,
         k,
         v,
@@ -89,13 +84,13 @@ def _flash_attn_forward(
         causal,
         window_size[0],
         window_size[1],
+        attention_chunk,
         softcap,
         rotary_interleaved,
         scheduler_metadata,
         num_splits,
         pack_gqa,
         sm_margin,
-        s_aux
     )
     return out, softmax_lse, *rest
 
@@ -125,7 +120,7 @@ def _flash_attn_backward(
 ):
     # dq, dk, dv are allocated by us so they should already be contiguous
     dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
-    dq, dk, dv, softmax_d, *rest = ops.bwd(
+    dq, dk, dv, softmax_d, *rest = flash_attn_3_cuda.bwd(
         dout,
         q,
         k,
@@ -161,9 +156,11 @@ class FlashAttnQKVPackedFunc(torch.autograd.Function):
         causal,
         q_descale=None, k_descale=None, v_descale=None,
         window_size=(-1, -1),
+        attention_chunk=0,
         softcap=0.0,
         deterministic=False,
         num_heads_q=None,
+        sm_margin=0,
     ):
         if softmax_scale is None:
             softmax_scale = qkv.shape[-1] ** (-0.5)
@@ -192,22 +189,27 @@ class FlashAttnQKVPackedFunc(torch.autograd.Function):
             softmax_scale,
             causal=causal,
             window_size=window_size,
+            attention_chunk=attention_chunk,
             softcap=softcap,
+            sm_margin=sm_margin,
         )
         # ctx.save_for_backward(q, k, v, out_padded, softmax_lse)
         ctx.save_for_backward(q, k, v, out, softmax_lse)
         ctx.softmax_scale = softmax_scale
         ctx.causal = causal
         ctx.window_size = window_size
+        ctx.attention_chunk = attention_chunk
         ctx.softcap = softcap
         ctx.deterministic = deterministic
         ctx.ndim = qkv.dim()
+        ctx.sm_margin = sm_margin
         # return out, softmax_lse
         return out
 
     @staticmethod
     def backward(ctx, dout, *args):
         q, k, v, out, softmax_lse = ctx.saved_tensors
+        assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk"
         if ctx.ndim == 5:
             qkv_shape = q.shape[:-2] + (3, *q.shape[-2:])
             dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device)
@@ -235,10 +237,11 @@ class FlashAttnQKVPackedFunc(torch.autograd.Function):
             ctx.causal,
             ctx.window_size,
             ctx.softcap,
-            ctx.deterministic, 
+            ctx.deterministic,
+            ctx.sm_margin,
         )
         dqkv = dqkv[..., : dout.shape[-1]]  # We could have padded the head dimension
-        return dqkv, None, None, None, None, None, None, None, None, None, None
+        return dqkv, None, None, None, None, None, None, None, None, None, None, None
 
 
 class FlashAttnFunc(torch.autograd.Function):
@@ -254,12 +257,12 @@ class FlashAttnFunc(torch.autograd.Function):
         qv=None,
         q_descale=None, k_descale=None, v_descale=None,
         window_size=(-1, -1),
+        attention_chunk=0,
         softcap=0.0,
         num_splits=1,
         pack_gqa=None,
         deterministic=False,
         sm_margin=0,
-        s_aux=None,
     ):
         if softmax_scale is None:
             softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5)
@@ -280,17 +283,18 @@ class FlashAttnFunc(torch.autograd.Function):
             softmax_scale,
             causal=causal,
             window_size=window_size,
+            attention_chunk=attention_chunk,
             softcap=softcap,
             num_splits=num_splits,
             pack_gqa=pack_gqa,
             sm_margin=sm_margin,
-            s_aux=s_aux,
         )
         # ctx.save_for_backward(q, k, v, out_padded, softmax_lse)
         ctx.save_for_backward(q, k, v, out, softmax_lse)
         ctx.softmax_scale = softmax_scale
         ctx.causal = causal
         ctx.window_size = window_size
+        ctx.attention_chunk = attention_chunk
         ctx.softcap = softcap
         ctx.deterministic = deterministic
         ctx.sm_margin = sm_margin
@@ -299,6 +303,7 @@ class FlashAttnFunc(torch.autograd.Function):
     @staticmethod
     def backward(ctx, dout, *args):
         q, k, v, out, softmax_lse = ctx.saved_tensors
+        assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk"
         dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
         _flash_attn_backward(
             dout,
@@ -320,9 +325,9 @@ class FlashAttnFunc(torch.autograd.Function):
             ctx.deterministic,
             ctx.sm_margin,
         )
-        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
-        dk = dk[..., : dout.shape[-1]]
-        dv = dv[..., : dout.shape[-1]]
+        dq = dq[..., : q.shape[-1]]  # We could have padded the head dimension
+        dk = dk[..., : k.shape[-1]]
+        dv = dv[..., : v.shape[-1]]
         return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None
 
 
@@ -345,12 +350,12 @@ class FlashAttnVarlenFunc(torch.autograd.Function):
         qv=None,
         q_descale=None, k_descale=None, v_descale=None,
         window_size=(-1, -1),
+        attention_chunk=0,
         softcap=0.0,
         num_splits=1,
         pack_gqa=None,
         deterministic=False,
         sm_margin=0,
-        s_aux=None,
     ):
         if softmax_scale is None:
             softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5)
@@ -375,11 +380,11 @@ class FlashAttnVarlenFunc(torch.autograd.Function):
             softmax_scale,
             causal=causal,
             window_size=window_size,
+            attention_chunk=attention_chunk,
             softcap=softcap,
             num_splits=num_splits,
             pack_gqa=pack_gqa,
             sm_margin=sm_margin,
-            s_aux=s_aux,
         )
         # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k)
         ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k)
@@ -388,6 +393,7 @@ class FlashAttnVarlenFunc(torch.autograd.Function):
         ctx.softmax_scale = softmax_scale
         ctx.causal = causal
         ctx.window_size = window_size
+        ctx.attention_chunk = attention_chunk
         ctx.softcap = softcap
         ctx.deterministic = deterministic
         ctx.sm_margin = sm_margin
@@ -396,6 +402,7 @@ class FlashAttnVarlenFunc(torch.autograd.Function):
     @staticmethod
     def backward(ctx, dout, *args):
         q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors
+        assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk"
         dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
         _flash_attn_backward(
             dout,
@@ -420,9 +427,9 @@ class FlashAttnVarlenFunc(torch.autograd.Function):
             ctx.deterministic,
             ctx.sm_margin,
         )
-        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
-        dk = dk[..., : dout.shape[-1]]
-        dv = dv[..., : dout.shape[-1]]
+        dq = dq[..., : q.shape[-1]]  # We could have padded the head dimension
+        dk = dk[..., : k.shape[-1]]
+        dv = dv[..., : v.shape[-1]]
         return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None
 
 
@@ -432,9 +439,11 @@ def flash_attn_qkvpacked_func(
     causal=False,
     q_descale=None, k_descale=None, v_descale=None,
     window_size=(-1, -1),
+    attention_chunk=0,
     softcap=0.0,
     deterministic=False,
     num_heads_q=None,
+    sm_margin=0,
 ):
     """dropout_p should be set to 0.0 during evaluation
     If Q, K, V are already stacked into 1 tensor, this function will be faster than
@@ -476,9 +485,11 @@ def flash_attn_qkvpacked_func(
         causal,
         q_descale, k_descale, v_descale,
         window_size,
+        attention_chunk,
         softcap,
         deterministic,
         num_heads_q,
+        sm_margin,
     )
 
 
@@ -491,12 +502,12 @@ def flash_attn_func(
     qv=None,
     q_descale=None, k_descale=None, v_descale=None,
     window_size=(-1, -1),
+    attention_chunk=0,
     softcap=0.0,
     num_splits=1,
     pack_gqa=None,
     deterministic=False,
     sm_margin=0,
-    s_aux=None,
 ):
     """dropout_p should be set to 0.0 during evaluation
     Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
@@ -552,12 +563,12 @@ def flash_attn_func(
         qv,
         q_descale, k_descale, v_descale,
         window_size,
+        attention_chunk,
         softcap,
         num_splits,
         pack_gqa,
         deterministic,
         sm_margin,
-        s_aux,
     )
 
 
@@ -576,12 +587,12 @@ def flash_attn_varlen_func(
     qv=None,
     q_descale=None, k_descale=None, v_descale=None,
     window_size=(-1, -1),
+    attention_chunk=0,
     softcap=0.0,
     num_splits=1,
     pack_gqa=None,
     deterministic=False,
     sm_margin=0,
-    s_aux=None,
 ):
     return FlashAttnVarlenFunc.apply(
         q,
@@ -598,17 +609,17 @@ def flash_attn_varlen_func(
         qv,
         q_descale, k_descale, v_descale,
         window_size,
+        attention_chunk,
         softcap,
         num_splits,
         pack_gqa,
         deterministic,
         sm_margin,
-        s_aux,
     )
 
 
 def flash_attn_combine(out_partial, lse_partial, out=None, out_dtype=None):
-    return ops.fwd_combine(out_partial, lse_partial, out, out_dtype)
+    return flash_attn_3_cuda.fwd_combine(out_partial, lse_partial, out, out_dtype)
 
 
 def flash_attn_with_kvcache(
@@ -634,6 +645,7 @@ def flash_attn_with_kvcache(
     softmax_scale=None,
     causal=False,
     window_size=(-1, -1),  # -1 means infinite context window
+    attention_chunk=0,
     softcap=0.0, # 0.0 means deactivated
     rotary_interleaved=True,
     scheduler_metadata=None,
@@ -641,7 +653,6 @@ def flash_attn_with_kvcache(
     pack_gqa=None,   # Can be tuned for speed
     sm_margin=0,     # Can be tuned if some SMs are used for communication
     return_softmax_lse=False,
-    s_aux=None,
 ):
     """
     If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from
@@ -762,13 +773,13 @@ def flash_attn_with_kvcache(
         softmax_scale,
         causal=causal,
         window_size=window_size,
+        attention_chunk=attention_chunk,
         softcap=softcap,
         rotary_interleaved=rotary_interleaved,
         scheduler_metadata=scheduler_metadata,
         num_splits=num_splits,
         pack_gqa=pack_gqa,
         sm_margin=sm_margin,
-        s_aux=s_aux,
     )
     # return (out, softmax_lse) if return_softmax_lse else out
     return (out, softmax_lse, *rest) if return_softmax_lse else out
@@ -786,6 +797,7 @@ def get_scheduler_metadata(
     max_seqlen_k_new=0,
     causal=False,
     window_size=(-1, -1),  # -1 means infinite context window
+    attention_chunk=0,
     has_softcap=False,
     num_splits=0,    # Can be tuned for speed
     pack_gqa=None,   # Can be tuned for speed
@@ -794,7 +806,7 @@ def get_scheduler_metadata(
     cache_seqlens = maybe_contiguous(cache_seqlens)
     if headdim_v is None:
         headdim_v = headdim
-    scheduler_metadata = ops.get_scheduler_metadata(
+    scheduler_metadata = flash_attn_3_cuda.get_scheduler_metadata(
         batch_size, max_seqlen_q, max_seqlen_k, num_heads_q, num_heads_kv, headdim, headdim_v,
         qkv_dtype,
         cache_seqlens,
@@ -807,6 +819,7 @@ def get_scheduler_metadata(
         max_seqlen_k_new,
         causal,
         window_size[0], window_size[1],
+        attention_chunk,
         has_softcap,
         num_splits,
         pack_gqa,