[biofoolgreen]

Problem and behavior

Cosmos3-Nano-Sim-Auto is rejected at model construction on SM120 because replay admits only SM90 and SM10.x. Admitting (12, 0) alone exposes a CUDA illegal-memory-access failure inside FA4's gapped-varlen forward on the published seven-camera joint RGB/LiDAR workload with the pinned flash-attn-4==4.0.0b24.

This PR offers a framework-side alternative that retains the existing dependency pins and still executes FA4's SM120 kernels. On SM120, segments with the same query/key lengths are stacked into dense batches, run through FA4 dense forward, and scattered back into their original output/LSE rows. The existing visibility partitions, GQA folding and FP32 LSE merge are retained. Other architectures keep their existing execution path. SM120 backward explicitly raises because training has not been qualified.

Related work and tradeoffs

• fix: admit SM120 and use FA4 with invalid-tile fix #295 admits SM120 and upgrades FA4 to the release containing Dao-AILab/flash-attention#2763, which guards invalid SM80/SM120 varlen tiles. That dependency-upgrade approach may be preferable once hardware-qualified.
• fix(mot): gapped replay on GPUs FlashAttention-4 has no kernels for (sm120) #302 provides a gathered PyTorch FA2 inference path. This alternative keeps FA4 and uses equal-length dense batches instead.

This is a competing compatibility approach, not a claim that the existing PRs miss another independent FA4 bug. It adds host synchronization and temporary tensor copies per launch. It is intended for evaluation while the project still pins b24; it does not claim improved performance, general SM12x support, or validated training/backward support.

Validation

Manual hardware validation completed on 2026-10-26 against the exact PR commit f2de74f022aa21de160e372802ae0dc75375f24b on 4 x NVIDIA RTX 6000D (SM120). The dedicated checkout was clean and the imported framework module path was checked before testing.

• 31 pytest cases passed, 0 failures/errors/skips; 273 unrelated cases deselected. This includes 13 new CPU/SM120 tests and 18 existing causal replay tests.
• 24 additional numerical comparisons passed: the six BF16 MHA/GQA cases were run independently on physical GPUs 0, 1, 2 and 3. Maximum absolute output error was 0.00406915 and maximum LSE error was 9.53674e-7 on every tested GPU; gap rows remained zero and the backward guard rejected unsupported execution.
• Ruff lint, format checks and the source/lockfile identity checks passed.
• The unmodified published seven-camera joint RGB/LiDAR reference scene completed on the final PR commit, exit code 0, with CUDA_LAUNCH_BLOCKING unset. All output checks passed: seven 1280x720 videos, 297 frames per view at 30 FPS, and 99 LiDAR sweeps. Range/intensity tensors had shape [99,128,1800], finite values, and frame IDs 0,3,...,294.

Measurement	Seconds
Full inference process, including startup/teardown	590.192
Framework batch pipeline	573.603
Encode + generation + decode	559.896
Generation + decode	546.911
Input preparation	12.984
Output saving	12.614

Observed device memory peaks from nvidia-smi, sampled every 1 second over the inference process:

Physical GPU	Sampled peak used memory (MiB)	GiB
0	50995	49.80
1	49721	48.56
2	51197	50.00
3	48229	47.10

Environment: Python 3.13, torch 2.10.0+cu130, torchcodec 0.10.0+cu130, flash-attn-4 4.0.0b24, nvidia-cutlass-dsl 4.6.0.dev0, driver 580.95.05. Model revision: c3fa1378e93710242501a6288d7c9e7c9084a580; shared sample revision: b8dcc2f70421e31db86d1b81d194c791d65eaad2. Sampling, model weights and the published request were unchanged; guardrails were disabled for this public driving reference scene. The early/middle/late multiview preview was inspected as a basic visual check.

This is one functional full-scene run, with existing runtime caches, rather than a repeated steady-state performance benchmark. Device memory includes CUDA/context/runtime allocations; 1-second sampling may miss short peaks. These manual results do not represent upstream GitHub Actions CI. SM90/SM10.x hardware regressions, SM121, and training/backward qualification remain untested. The earlier 732-second run used synchronous CUDA debugging and is not used to claim a speedup.

The retained evidence includes JUnit XML, per-case timing CSV, per-GPU numeric results, a 1-second GPU sampling CSV, stage timing JSON, output verification, source/lockfile SHA256s and SHA256-verified previews.

Regression command:

COSMOS_INTERNAL=0 COSMOS_TRAINING=0 python -m pytest --levels 0 -q \
  cosmos_framework/model/generator/mot/maskless_attention_test.py \
  cosmos_framework/model/generator/omni_mot_causal_model_test.py \
  -k 'replay or gapped or FlashAttention'