Release 2.8.3.post2: backport weights_only=True torch.load fix (CVE-2026-31253)

flash_attn/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.8.3"
+__version__ = "2.8.3.post2"
 
 from flash_attn.flash_attn_interface import (
     flash_attn_func,

flash_attn/models/llama.py

@@ -385,7 +385,7 @@ def state_dicts_from_checkpoint(
 ) -> List[dict]:
     # Need to sort, otherwise we mess up the ordering and the weights are wrong
     return [
-        torch.load(path, map_location="cpu")
+        torch.load(path, map_location="cpu", weights_only=True)
         for path in sorted((Path(checkpoint_path) / model_name).glob("consolidated.*.pth"))
     ]
 

flash_attn/utils/pretrained.py

@@ -59,7 +59,7 @@ def state_dict_from_pretrained(model_name, device=None, dtype=None):
     if load_safe:
         loader = partial(safe_load_file, device=mapped_device)
     else:
-        loader = partial(torch.load, map_location=mapped_device)
+        loader = partial(torch.load, map_location=mapped_device, weights_only=True)
 
     if is_sharded:
         # resolved_archive_file becomes a list of files that point to the different

hopper/setup.py

@@ -343,6 +343,44 @@ def open_url(url):
     return urllib.request.urlopen(request, timeout=300)
 
 
+def safe_extractall(tar, path):
+    """Extract a tar stream, refusing members that would escape ``path``.
+
+    Uses the PEP 706 ``data`` filter when the interpreter supports it (added in
+    3.12, backported to 3.10.12 / 3.11.4). On older interpreters the filter
+    keyword does not exist, so we validate each member's resolved path stays
+    inside ``path`` and reject link members before extracting. The tar is opened
+    in stream mode, so members are validated and extracted in a single pass.
+    """
+    dest = os.path.realpath(path)
+    if hasattr(tarfile, "data_filter"):
+        tar.extractall(path=dest, filter="data")
+        return
+
+    def stays_inside(resolved):
+        return resolved == dest or resolved.startswith(dest + os.sep)
+
+    for member in tar:
+        target = os.path.realpath(os.path.join(dest, member.name))
+        if not stays_inside(target):
+            raise RuntimeError(f"Refusing tar member outside extract dir: {member.name!r}")
+        # The data filter permits links whose target resolves inside the
+        # destination (the cuda_nvcc archives ship such intra-package symlinks,
+        # e.g. libnvvm.so -> libnvvm.so.4). Mirror that instead of rejecting all
+        # links, or real builds regress on interpreters without the data filter.
+        if member.issym() or member.islnk():
+            # A symlink target is relative to the link's own directory; a hard
+            # link's is relative to the archive root.
+            link_base = os.path.dirname(target) if member.issym() else dest
+            link_target = os.path.realpath(os.path.join(link_base, member.linkname))
+            if not stays_inside(link_target):
+                raise RuntimeError(
+                    f"Refusing link member pointing outside extract dir: "
+                    f"{member.name!r} -> {member.linkname!r}"
+                )
+        tar.extract(member, path=dest)
+
+
 def download_and_copy(name, src_func, dst_path, version, url_func):
     if is_offline_build():
         return
@@ -359,9 +397,21 @@ def download_and_copy(name, src_func, dst_path, version, url_func):
     src_path = os.path.join(tmp_path, src_path)
     download = not os.path.exists(src_path)
     if download:
+        # Refuse to extract into a pre-planted symlink: an attacker with write
+        # access to the predictable cache dir could otherwise redirect the
+        # extraction (and the shutil.copy below) to an arbitrary location.
+        # islink() only inspects the leaf, so also require the fully resolved
+        # path to stay under the cache root, catching a symlinked parent.
+        cache_root = os.path.realpath(flashattn_cache_path)
+        resolved_tmp = os.path.realpath(tmp_path)
+        if os.path.islink(tmp_path) or not (
+            resolved_tmp == cache_root or resolved_tmp.startswith(cache_root + os.sep)
+        ):
+            raise RuntimeError(f"Refusing to extract into symlinked cache path: {tmp_path}")
+        os.makedirs(tmp_path, exist_ok=True)
         print(f'downloading and extracting {url} ...')
         file = tarfile.open(fileobj=open_url(url), mode="r|*")
-        file.extractall(path=tmp_path)
+        safe_extractall(file, tmp_path)
     os.makedirs(os.path.split(dst_path)[0], exist_ok=True)
     print(f'copy {src_path} to {dst_path} ...')
     if os.path.isdir(src_path):

setup.py

@@ -454,9 +454,14 @@ def get_wheel_url():
         # We're using the CUDA version used to build torch, not the one currently installed
         # _, cuda_version_raw = get_cuda_bare_metal_version(CUDA_HOME)
         torch_cuda_version = parse(torch.version.cuda)
-        # For CUDA 11, we only compile for CUDA 11.8, and for CUDA 12 we only compile for CUDA 12.3
+        # For CUDA 11 we compile for 11.8, for CUDA 12 for 12.3, and for CUDA 13 for 13.0
         # to save CI time. Minor versions should be compatible.
-        torch_cuda_version = parse("11.8") if torch_cuda_version.major == 11 else parse("12.3")
+        if torch_cuda_version.major == 11:
+            torch_cuda_version = parse("11.8")
+        elif torch_cuda_version.major == 12:
+            torch_cuda_version = parse("12.3")
+        else:
+            torch_cuda_version = parse("13.0")
         # cuda_version = f"{cuda_version_raw.major}{cuda_version_raw.minor}"
         cuda_version = f"{torch_cuda_version.major}"
 

training/src/eval.py

@@ -31,7 +31,7 @@ def load_checkpoint(path, device='cpu'):
         path /= 'last.ckpt'
     # dst = f'cuda:{torch.cuda.current_device()}'
     log.info(f'Loading checkpoint from {str(path)}')
-    state_dict = torch.load(path, map_location=device)
+    state_dict = torch.load(path, map_location=device, weights_only=True)
     # T2T-ViT checkpoint is nested in the key 'state_dict_ema'
     if state_dict.keys() == {'state_dict_ema'}:
         state_dict = state_dict['state_dict_ema']

training/src/utils/checkpoint.py

@@ -17,7 +17,7 @@ def load_checkpoint(path, device='cpu'):
         else:
             raise ValueError(f"Unable to find 'latest' file at {latest_path}")
         path /= f'{tag}/mp_rank_00_model_states.pt'
-    state_dict = torch.load(path, map_location=device)
+    state_dict = torch.load(path, map_location=device, weights_only=True)
     if is_deepspeed:
         state_dict = state_dict['module']