Build failure on custom/newer architectures (sm_121): NameError: name 'num' is not defined

[mamorett]

Description
Hello, thanks for this great library!
I encountered a build error while trying to install SageAttention on a Spark DGX system.
The installation fails with a NameError because the variable num in setup.py is defined inside an if/elif chain that doesn't cover my specific GPU architecture, but it is accessed unconditionally afterwards.

Environment

• Hardware: Spark DGX
• Reported Compute Capability: 12.1 (Detected as sm_121)
• Python: 3.12
• CUDA: 12.1

The Error

File "setup.py", line 115, in <module>
    NVCC_FLAGS += ["-gencode", f"arch=compute_{num},code=sm_{num}"]
NameError: name 'num' is not defined. Did you mean: 'sum'?

Root Cause
In setup.py, the variable num is only assigned if capability matches specific prefixes (8.0, 8.6, 8.9, 9.0, 12.0).
My system reports a capability starting with 12.1. As a result, the elif chain is exhausted, num is never initialized, and the script crashes at line 115 when trying to append NVCC_FLAGS.

Proposed Solution
I was able to compile successfully by explicitly handling the 12.1 case (treating it similarly to 12.0/Next-Gen).

Here is the patch I applied to setup.py around line 110:

    elif capability.startswith("12.0"):
        HAS_SM120 = True
        num = "120" # need to use sm120a to use mxfp8/mxfp4/nvfp4 instructions.
    
    # --- PATCH START ---
    elif capability.startswith("12.1"): 
        # Handling for sm_121 (Spark DGX / Custom Arch)
        HAS_SM120 = True
        num = "121" 
    # --- PATCH END ---

    # SAFEGUARD:
    if 'num' not in locals():
        print(f"Warning: Unsupported capability {capability}, skipping arch flags.")
    else:
        NVCC_FLAGS += ["-gencode", f"arch=compute_{num},code=sm_{num}"]
        if capability.endswith("+PTX"):
            NVCC_FLAGS += ["-gencode", f"arch=compute_{num},code=compute_{num}"]

This ensures that:

1. New architectures like 12.1 are handled.
2. If an unknown architecture appears in the future, the script won't crash with a NameError.

Thanks!