Skip to content

fix: admit SM120 and use FA4 with invalid-tile fix #295

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
junghok-nv wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

fix: admit SM120 and use FA4 with invalid-tile fix #295

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
413ddf3
Add guarded experimental SM120 gapped forward replay
junghok-nv Oct 21, 2026
8afd9d4
Reduce SM120 replay fix to the existing attention file
junghok-nv Oct 21, 2026
7a8cd46
Limit SM120 patch to architecture admission only
junghok-nv Oct 21, 2026
2fefe2e
fix: use FA4 release with invalid-varlen-tile guard
junghok-nv Oct 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion cosmos_framework/model/generator/mot/maskless_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -432,7 +432,7 @@ def gapped_replay_unavailable_reason(device: torch.device) -> str | None:
if torch.version.hip is not None:
return "it needs NVIDIA CUDA, not ROCm"
major, minor = torch.cuda.get_device_capability(device)
if (major, minor) != (9, 0) and major != 10:
if (major, minor) not in ((9, 0), (12, 0)) and major != 10:

@pengcuo pengcuo Oct 22, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FA4 gates on arch // 10, so sm121 (DGX Spark / GB10) takes exactly the same kernels as sm120; admitting (12, 0) but refusing (12, 1) is an arbitrary split. Suggest major == 12 here (and major == 11 for Thor only if someone runs the tests there).

# Hopper and Blackwell use different FA4 kernels; the replay GPU tests cover
# the segment isolation and untouched gap rows required by both paths.
return f"it needs FlashAttention-4 on Hopper (sm90) or Blackwell (sm100/sm103), got sm{major}{minor}"

@pengcuo pengcuo Oct 22, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The message still lists only sm90/sm100/sm103 after the condition above started admitting sm120.

Suggested change
return f"it needs FlashAttention-4 on Hopper (sm90) or Blackwell (sm100/sm103), got sm{major}{minor}"
return f"it needs FlashAttention-4 on Hopper (sm90) or Blackwell (sm100/sm103/sm120), got sm{major}{minor}"

Expand Down
6 changes: 3 additions & 3 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -196,7 +196,7 @@ vllm = [
# Match nvcr.io/nvidia/pytorch:25.11-py3: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-26-01.html
cu128 = [
# FA4 CuTeDSL kernels; the cu13 extra is only needed by CUDA 13 environments.
"flash-attn-4==4.0.0b24",
"flash-attn-4==4.0.0b30",
"flash-attn-3-nv==1.0.3+cu128.torch210; platform_machine == 'x86_64'",
"flash-attn==2.7.4.post1+cu128.torch210; platform_machine == 'x86_64'",
"natten==0.21.6.dev6+cu128.torch210; platform_machine == 'x86_64'",
Expand Down Expand Up @@ -231,7 +231,7 @@ cu128-train = [
# Match nvcr.io/nvidia/pytorch:25.11-py3: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-25-11.html
cu130 = [
# Match imaginaire4's FA4 installation on both x86_64 and aarch64.

@pengcuo pengcuo Oct 22, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

After this bump the comment is no longer true unless imaginaire4's container moves to b30 as well. Please coordinate the pin with the i4 side (or drop the comment), since the two trees are kept in sync by the release flow.

"flash-attn-4[cu13]==4.0.0b24",
"flash-attn-4[cu13]==4.0.0b30",
"flash-attn-3-nv==1.0.3+cu130.torch210; platform_machine == 'x86_64'",
"flash-attn==2.7.4.post1+cu130.torch210; platform_machine == 'x86_64'",
# GB300/B300 (sm_103a): use the custom build that includes sm_103a kernels (both arches).
Expand Down Expand Up @@ -267,7 +267,7 @@ cu130-train = [
]
# torch 2.13 on CUDA 13.0. Same layout as cu130 but pinned to torch 2.13.
cu130-torch213 = [
"flash-attn-4[cu13]==4.0.0b24",
"flash-attn-4[cu13]==4.0.0b30",
#"flash-attn-3-nv==1.0.3+cu130.torch210; platform_machine == 'x86_64'", # TODO: torch213 build (cu130.torch210 wheels are cp313-only + torch210 ABI)
#"flash-attn==2.7.4.post1+cu130.torch210; platform_machine == 'x86_64'", # TODO: torch213 build (cu130.torch210 wheels are cp313-only + torch210 ABI)
# torch2.13 build (natten 0.21.6, sm 80/86/89/90/100/103/110/120/121) — matches this group's torch/ABI.
Expand Down
Loading