Kevin LinandClaude Opus 5.5
5e61accc37
Split the FastWan-QAD Wan2.1-1.3B examples into per-recipe scripts (fp8, nvfp4_qat, nvfp4_sa2) with a shared _qad_common.py and a standalone multi-run benchmark harness, replacing FastWan_QAD_TAEHV.py. The PR's fastvideo-kernel build changes (sm_121a arch probe and build.sh mapping) were dropped during maintainer review (Swarm order #664) in favour of #1598, which owns the kernel arch policy; the RTX 5090 (sm_120a) recipes need nothing beyond main's kernels. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
32 lines
1.3 KiB
Python
32 lines
1.3 KiB
Python
"""Shared GPU-arch helpers for the QAD 5090 example scripts.
|
|
|
|
The default (non ``--bf16``) path of the NVFP4 scripts runs DiT linears through
|
|
flashinfer's cutlass FP4 gemm, which has no fallback and targets consumer
|
|
Blackwell only (sm_120a RTX 5090-class, or sm_121a GB10/DGX Spark, which also
|
|
needs CUDA >= 12.9). These helpers pick the right FLASHINFER_CUDA_ARCH_LIST and
|
|
fail fast with a readable error instead of an opaque flashinfer ValueError.
|
|
"""
|
|
|
|
import torch
|
|
|
|
# Compute capabilities the flashinfer FP4 gemm path is gated on.
|
|
FP4_CAPABILITIES = ((12, 0), (12, 1))
|
|
|
|
|
|
def flashinfer_arch_list() -> str:
|
|
"""FLASHINFER_CUDA_ARCH_LIST derived from the local GPU, defaulting to 12.0a."""
|
|
if torch.cuda.is_available():
|
|
major, minor = torch.cuda.get_device_capability()
|
|
if (major, minor) in FP4_CAPABILITIES:
|
|
return f"{major}.{minor}a"
|
|
return "12.0a"
|
|
|
|
|
|
def require_fp4_capable_gpu() -> None:
|
|
"""Exit with a one-line error if the GPU cannot run the NVFP4 gemm path."""
|
|
cap = torch.cuda.get_device_capability() if torch.cuda.is_available() else None
|
|
if cap not in FP4_CAPABILITIES:
|
|
raise SystemExit(
|
|
f"The default NVFP4 path requires an sm_120a GPU (RTX 5090-class Blackwell, "
|
|
f"compute capability 12.0/12.1); found {cap}. Re-run with --bf16 on other GPUs.")
|