fix(op_builder): avoid duplicate/wrong -gencode flags (#7974)
## Summary
- Fix duplicate/wrong `-gencode=` flags in both JIT and non-JIT
compilation paths (`op_builder/builder.py`)
- Fix `TORCH_CUDA_ARCH_LIST` env-var restore logic in
`OpBuilder.jit_load()`
DeepSpeed's `compute_capability_args()` generates its own `-gencode`
flags, but PyTorch (`load()` in JIT mode, `BuildExtension` in non-JIT
mode) *also* reads `TORCH_CUDA_ARCH_LIST` and generates `-gencode`
flags. This causes two problems:
1. **JIT mode**: `jit_load()` set `TORCH_CUDA_ARCH_LIST=""`, which
PyTorch treats as *unset* and falls back to auto-detection — resulting
in every flag appearing **twice**.
2. **Non-JIT mode**: subclasses that override `filter_ccs()` (e.g.
`FPQuantizerBuilder`, `EvoformerAttnBuilder`) remove certain archs, but
`BuildExtension` re-reads the **unfiltered** `TORCH_CUDA_ARCH_LIST` and
adds them back — **undermining the filter**.
The fix synchronises `TORCH_CUDA_ARCH_LIST` with the filtered arch list
in `compute_capability_args()`, for both JIT and non-JIT paths.
Fixes #7972
## Before / After
<details>
<summary>Before (buggy behavior)</summary>
**JIT mode** — `TORCH_CUDA_ARCH_LIST` cleared to `""`, PyTorch
auto-detects and adds flags, DeepSpeed also adds the same flags:
```
nvcc ... -gencode=arch=compute_80,code=compute_80 -gencode=arch=compute_80,code=sm_80
... -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_80,code=compute_80
```
Plus a spurious warning:
```
UserWarning: TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation.
```
**Non-JIT mode** — `FPQuantizerBuilder.filter_ccs()` removes `< 8.0`,
but `BuildExtension` re-adds them from the unfiltered env var:
```
# FPQuantizer compiled for sm_70 even though filter_ccs() removed it
nvcc ... -gencode=arch=compute_80,code=sm_80 # from DeepSpeed (correct)
... -gencode=arch=compute_70,code=sm_70 # from BuildExtension (wrong!)
```
</details>
<details>
<summary>After (fixed behavior)</summary>
**JIT mode** — `TORCH_CUDA_ARCH_LIST` is set to the detected
architectures, PyTorch generates flags once, no duplicates:
```
nvcc ... -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_80,code=compute_80
```
No spurious warning. Env var is properly restored/removed after build.
**Non-JIT mode** — `TORCH_CUDA_ARCH_LIST` is updated to the filtered
list. Each extension keeps its own `-gencode` flags, and
`BuildExtension` reads the filtered env var:
```
# FPQuantizer: only sm_80+ as intended
nvcc ... -gencode=arch=compute_80,code=sm_80 # from DeepSpeed
... -gencode=arch=compute_80,code=sm_80 # from BuildExtension (harmless dup)
```
> **Note:** in multi-builder `setup.py` builds, the last builder's
filtered arch list wins for `TORCH_CUDA_ARCH_LIST`. This may cause
harmless duplicates for some extensions, but will never reintroduce
archs that any builder's `filter_ccs()` removed — a strict improvement
over the current behavior where the unfiltered original is always used.
</details>
## Changes
- `op_builder/builder.py`
- `CUDAOpBuilder.compute_capability_args()`:
- Always sync `TORCH_CUDA_ARCH_LIST` with the filtered arch list
- JIT mode: return `[]` (PyTorch generates flags via `load()`)
- Non-JIT mode: return `-gencode` args as before (per-builder flags in
`extra_compile_args`)
- `OpBuilder.jit_load()`: simplified stash/restore — properly `del` the
env var if it was not originally set
---------
Signed-off-by: Cursx <674760201@qq.com>
