Skip to content

Commit ad8f9ab

Browse files
thomas-hiddenpeakthomas-hiddenpeak
committed
feat: add SM87 and SM110 compute capability support
- Add (87, 87) and (110, 110) match arms in compute_cap_matching - Add ampere-87 and blackwell-110 entries to Docker image matrix - Add SM87 and SM110 binary builds to Dockerfile-cuda-all - Route SM87 and SM110 devices to dedicated binaries in cuda-all-entrypoint.sh - Add comprehensive test cases for both compute capabilities
1 parent f016879 commit ad8f9ab

4 files changed

Lines changed: 79 additions & 1 deletion

File tree

‎.github/workflows/matrix.json‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,15 @@
2727
"grpc": true,
2828
"dockerfile": "Dockerfile-cuda"
2929
},
30+
{
31+
"name": "ampere-87",
32+
"imageNamePrefix": "87-",
33+
"runOn": "always",
34+
"sccache": true,
35+
"cudaComputeCap": 87,
36+
"grpc": true,
37+
"dockerfile": "Dockerfile-cuda"
38+
},
3039
{
3140
"name": "ada-lovelace",
3241
"imageNamePrefix": "89-",
@@ -54,6 +63,15 @@
5463
"grpc": true,
5564
"dockerfile": "Dockerfile-cuda"
5665
},
66+
{
67+
"name": "blackwell-110",
68+
"imageNamePrefix": "110-",
69+
"runOn": "always",
70+
"sccache": true,
71+
"cudaComputeCap": 110,
72+
"grpc": true,
73+
"dockerfile": "Dockerfile-cuda"
74+
},
5775
{
5876
"name": "blackwell-120",
5977
"imageNamePrefix": "120-",

‎Dockerfile-cuda-all‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,10 @@ RUN --mount=type=secret,id=actions_results_url,env=ACTIONS_RESULTS_URL \
6161
--mount=type=secret,id=actions_runtime_token,env=ACTIONS_RUNTIME_TOKEN \
6262
CUDA_COMPUTE_CAP=80 cargo chef cook --release --features candle-cuda --recipe-path recipe.json && sccache -s;
6363

64+
RUN --mount=type=secret,id=actions_results_url,env=ACTIONS_RESULTS_URL \
65+
--mount=type=secret,id=actions_runtime_token,env=ACTIONS_RUNTIME_TOKEN \
66+
CUDA_COMPUTE_CAP=87 cargo chef cook --release --features candle-cuda --recipe-path recipe.json && sccache -s;
67+
6468
RUN --mount=type=secret,id=actions_results_url,env=ACTIONS_RESULTS_URL \
6569
--mount=type=secret,id=actions_runtime_token,env=ACTIONS_RUNTIME_TOKEN \
6670
CUDA_COMPUTE_CAP=90 cargo chef cook --release --features candle-cuda --recipe-path recipe.json && sccache -s;
@@ -69,6 +73,10 @@ RUN --mount=type=secret,id=actions_results_url,env=ACTIONS_RESULTS_URL \
6973
--mount=type=secret,id=actions_runtime_token,env=ACTIONS_RUNTIME_TOKEN \
7074
CUDA_COMPUTE_CAP=100 cargo chef cook --release --features candle-cuda --recipe-path recipe.json && sccache -s;
7175

76+
RUN --mount=type=secret,id=actions_results_url,env=ACTIONS_RESULTS_URL \
77+
--mount=type=secret,id=actions_runtime_token,env=ACTIONS_RUNTIME_TOKEN \
78+
CUDA_COMPUTE_CAP=110 cargo chef cook --release --features candle-cuda --recipe-path recipe.json && sccache -s;
79+
7280
RUN --mount=type=secret,id=actions_results_url,env=ACTIONS_RESULTS_URL \
7381
--mount=type=secret,id=actions_runtime_token,env=ACTIONS_RUNTIME_TOKEN \
7482
CUDA_COMPUTE_CAP=120 cargo chef cook --release --features candle-cuda --recipe-path recipe.json && sccache -s;
@@ -91,6 +99,12 @@ RUN --mount=type=secret,id=actions_results_url,env=ACTIONS_RESULTS_URL \
9199

92100
RUN mv /usr/src/target/release/text-embeddings-router /usr/src/target/release/text-embeddings-router-80
93101

102+
RUN --mount=type=secret,id=actions_results_url,env=ACTIONS_RESULTS_URL \
103+
--mount=type=secret,id=actions_runtime_token,env=ACTIONS_RUNTIME_TOKEN \
104+
CUDA_COMPUTE_CAP=87 cargo build --release --bin text-embeddings-router -F candle-cuda && sccache -s;
105+
106+
RUN mv /usr/src/target/release/text-embeddings-router /usr/src/target/release/text-embeddings-router-87
107+
94108
RUN --mount=type=secret,id=actions_results_url,env=ACTIONS_RESULTS_URL \
95109
--mount=type=secret,id=actions_runtime_token,env=ACTIONS_RUNTIME_TOKEN \
96110
CUDA_COMPUTE_CAP=90 cargo build --release --bin text-embeddings-router -F candle-cuda && sccache -s;
@@ -103,6 +117,12 @@ RUN --mount=type=secret,id=actions_results_url,env=ACTIONS_RESULTS_URL \
103117

104118
RUN mv /usr/src/target/release/text-embeddings-router /usr/src/target/release/text-embeddings-router-100
105119

120+
RUN --mount=type=secret,id=actions_results_url,env=ACTIONS_RESULTS_URL \
121+
--mount=type=secret,id=actions_runtime_token,env=ACTIONS_RUNTIME_TOKEN \
122+
CUDA_COMPUTE_CAP=110 cargo build --release --bin text-embeddings-router -F candle-cuda && sccache -s;
123+
124+
RUN mv /usr/src/target/release/text-embeddings-router /usr/src/target/release/text-embeddings-router-110
125+
106126
RUN --mount=type=secret,id=actions_results_url,env=ACTIONS_RESULTS_URL \
107127
--mount=type=secret,id=actions_runtime_token,env=ACTIONS_RUNTIME_TOKEN \
108128
CUDA_COMPUTE_CAP=120 cargo build --release --bin text-embeddings-router -F candle-cuda && sccache -s;
@@ -127,8 +147,10 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-ins
127147

128148
COPY --from=builder /usr/src/target/release/text-embeddings-router-75 /usr/local/bin/text-embeddings-router-75
129149
COPY --from=builder /usr/src/target/release/text-embeddings-router-80 /usr/local/bin/text-embeddings-router-80
150+
COPY --from=builder /usr/src/target/release/text-embeddings-router-87 /usr/local/bin/text-embeddings-router-87
130151
COPY --from=builder /usr/src/target/release/text-embeddings-router-90 /usr/local/bin/text-embeddings-router-90
131152
COPY --from=builder /usr/src/target/release/text-embeddings-router-100 /usr/local/bin/text-embeddings-router-100
153+
COPY --from=builder /usr/src/target/release/text-embeddings-router-110 /usr/local/bin/text-embeddings-router-110
132154
COPY --from=builder /usr/src/target/release/text-embeddings-router-120 /usr/local/bin/text-embeddings-router-120
133155

134156
COPY --chmod=775 cuda-all-entrypoint.sh entrypoint.sh

‎backends/candle/src/compute_cap.rs‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,11 @@ fn compute_cap_matching(runtime_compute_cap: usize, compile_compute_cap: usize)
2828
(75, 75) => true,
2929
(80..=89, 80) => true,
3030
(86..=89, 80..=86) => true,
31+
(87, 87) => true,
3132
(89, 89) => true,
3233
(90, 90) => true,
3334
(100, 100) => true,
35+
(110, 110) => true,
3436
(120, 120) => true,
3537
(_, _) => false,
3638
}
@@ -54,47 +56,79 @@ mod tests {
5456
assert!(compute_cap_matching(75, 75));
5557
assert!(compute_cap_matching(80, 80));
5658
assert!(compute_cap_matching(86, 86));
59+
assert!(compute_cap_matching(87, 87));
5760
assert!(compute_cap_matching(89, 89));
5861
assert!(compute_cap_matching(90, 90));
62+
assert!(compute_cap_matching(110, 110));
5963
assert!(compute_cap_matching(120, 120));
6064

6165
assert!(compute_cap_matching(86, 80));
66+
assert!(compute_cap_matching(87, 80));
6267
assert!(compute_cap_matching(89, 80));
68+
assert!(compute_cap_matching(87, 86));
6369
assert!(compute_cap_matching(89, 86));
6470

6571
assert!(!compute_cap_matching(75, 80));
6672
assert!(!compute_cap_matching(75, 86));
73+
assert!(!compute_cap_matching(75, 87));
6774
assert!(!compute_cap_matching(75, 89));
6875
assert!(!compute_cap_matching(75, 90));
76+
assert!(!compute_cap_matching(75, 110));
6977

7078
assert!(!compute_cap_matching(80, 75));
7179
assert!(!compute_cap_matching(80, 86));
80+
assert!(!compute_cap_matching(80, 87));
7281
assert!(!compute_cap_matching(80, 89));
7382
assert!(!compute_cap_matching(80, 90));
83+
assert!(!compute_cap_matching(80, 110));
7484

7585
assert!(!compute_cap_matching(86, 75));
86+
assert!(!compute_cap_matching(86, 87));
7687
assert!(!compute_cap_matching(86, 89));
7788
assert!(!compute_cap_matching(86, 90));
89+
assert!(!compute_cap_matching(86, 110));
90+
91+
assert!(!compute_cap_matching(87, 75));
92+
assert!(!compute_cap_matching(87, 89));
93+
assert!(!compute_cap_matching(87, 90));
94+
assert!(!compute_cap_matching(87, 110));
7895

7996
assert!(!compute_cap_matching(89, 75));
97+
assert!(!compute_cap_matching(89, 87));
8098
assert!(!compute_cap_matching(89, 90));
99+
assert!(!compute_cap_matching(89, 110));
81100

82101
assert!(!compute_cap_matching(90, 75));
83102
assert!(!compute_cap_matching(90, 80));
84103
assert!(!compute_cap_matching(90, 86));
104+
assert!(!compute_cap_matching(90, 87));
85105
assert!(!compute_cap_matching(90, 89));
106+
assert!(!compute_cap_matching(90, 110));
86107

87108
assert!(!compute_cap_matching(100, 75));
88109
assert!(!compute_cap_matching(100, 80));
89110
assert!(!compute_cap_matching(100, 86));
111+
assert!(!compute_cap_matching(100, 87));
90112
assert!(!compute_cap_matching(100, 89));
91113
assert!(!compute_cap_matching(100, 90));
114+
assert!(!compute_cap_matching(100, 110));
92115

93116
assert!(!compute_cap_matching(120, 75));
94117
assert!(!compute_cap_matching(120, 80));
95118
assert!(!compute_cap_matching(120, 86));
119+
assert!(!compute_cap_matching(120, 87));
96120
assert!(!compute_cap_matching(120, 89));
97121
assert!(!compute_cap_matching(120, 90));
98122
assert!(!compute_cap_matching(120, 100));
123+
assert!(!compute_cap_matching(120, 110));
124+
125+
assert!(!compute_cap_matching(110, 75));
126+
assert!(!compute_cap_matching(110, 80));
127+
assert!(!compute_cap_matching(110, 86));
128+
assert!(!compute_cap_matching(110, 87));
129+
assert!(!compute_cap_matching(110, 89));
130+
assert!(!compute_cap_matching(110, 90));
131+
assert!(!compute_cap_matching(110, 100));
132+
assert!(!compute_cap_matching(110, 120));
99133
}
100134
}

‎cuda-all-entrypoint.sh‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,12 +30,16 @@ compute_cap=$(nvidia-smi --query-gpu=compute_cap --format=csv | sed -n '2p' | se
3030

3131
if [ ${compute_cap} -eq 75 ]; then
3232
exec text-embeddings-router-75 "$@"
33-
elif [ ${compute_cap} -ge 80 -a ${compute_cap} -lt 90 ]; then
33+
elif [ ${compute_cap} -ge 80 -a ${compute_cap} -lt 87 ]; then
3434
exec text-embeddings-router-80 "$@"
35+
elif [ ${compute_cap} -eq 87 ]; then
36+
exec text-embeddings-router-87 "$@"
3537
elif [ ${compute_cap} -eq 90 ]; then
3638
exec text-embeddings-router-90 "$@"
3739
elif [ ${compute_cap} -eq 100 ]; then
3840
exec text-embeddings-router-100 "$@"
41+
elif [ ${compute_cap} -eq 110 ]; then
42+
exec text-embeddings-router-110 "$@"
3943
elif [ ${compute_cap} -eq 120 ]; then
4044
exec text-embeddings-router-120 "$@"
4145
else

0 commit comments

Comments
 (0)