Add sm_90a arch gencode to docker (#2338)

Summary:
We need to override the upstream config and add `sm_90a` gencode to the docker.

Pull Request resolved: #2338

Test Plan:
https://github.com/pytorch/benchmark/actions/runs/9704484783

https://github.com/pytorch/benchmark/actions/runs/9704484783/job/26784787989#step:5:2057

```
 #11 900.6 + /usr/local/cuda-12.4/bin/nvprune -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90 -gencode arch=compute_90a,code=sm_90a /usr/local/cuda-12.4/lib64/libcublas_static.a -o /usr/local/cuda-12.4/lib64/libcublas_static.a
#11 901.5 + /usr/local/cuda-12.4/bin/nvprune -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90 -gencode arch=compute_90a,code=sm_90a /usr/local/cuda-12.4/lib64/libcublasLt_static.a -o /usr/local/cuda-12.4/lib64/libcublasLt_static.a
```

Nightly docker build:
https://github.com/pytorch/benchmark/actions/runs/9704719476

Reviewed By: aaronenyeshi

Differential Revision: D59136248

Pulled By: xuzhao9

fbshipit-source-id: 09a71a509b2ac448d66f950e92c33effc173d48f

docker/gcp-a100-runner-dind.dockerfile

@@ -1,9 +1,12 @@
 # default base image: ghcr.io/actions/actions-runner:latest
 # base image: Ubuntu 22.04 jammy
+# Prune CUDA to only keep gencode >= A100
 ARG BASE_IMAGE=ghcr.io/actions/actions-runner:latest
 FROM ${BASE_IMAGE}
 
 ENV LANG=C.UTF-8 LC_ALL=C.UTF-8
+ARG OVERRIDE_GENCODE="-gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90 -gencode arch=compute_90a,code=sm_90a"
+ARG OVERRIDE_GENCODE_CUDNN="-gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90 -gencode arch=compute_90a,code=sm_90a"
 
 RUN sudo apt-get -y update && sudo apt -y update
 # fontconfig: required by model doctr_det_predictor
@@ -30,7 +33,7 @@ RUN sudo mkdir -p /workspace; sudo chown runner:runner /workspace
 # Use the CUDA installation scripts from pytorch/builder
 # Install CUDA 12.4 only to reduce docker size
 RUN cd /workspace; git clone https://github.com/pytorch/builder.git
-RUN sudo bash -c 'source /workspace/builder/common/install_cuda.sh; install_124; prune_124'
+RUN sudo bash -c "set -x; source /workspace/builder/common/install_cuda.sh; install_124; export OVERRIDE_GENCODE=\"${OVERRIDE_GENCODE}\" OVERRIDE_GENCODE_CUDNN=\"${OVERRIDE_GENCODE_CUDNN}\"; prune_124"
 
 # Install miniconda
 RUN wget -q https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /workspace/Miniconda3-latest-Linux-x86_64.sh