ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

build: parametrize CUDA version for wheels and docker

Browse Source 
- Add cu128/cu130 build matrix to release workflow
- Parametrize Dockerfile and docker-compose with CUDA_TAG build arg
- Allow csrc/ and setup.py in docker context via .dockerignore
- Add nvcc/torch CUDA version mismatch preflight warning in setup.py
- Add cuda_toolkit_version() helper in csrc/build.py
- Use at::IntArrayRef explicitly to fix ATen overload ambiguity
- Guard kernels with CUDART_VERSION >= 11020 check
- Remove invalid [tool.pip] section from pyproject.toml
This commit is contained in:
ViperEkura
2026-07-31 14:10:55 +08:00
parent 738cb8f128
commit 5756054d38
9 changed files with 88 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/attn_entry_utils.cuh
+2 -2
View File
@@ -23,8 +23,8 @@ using bf16 = __nv_bfloat16;
template<typename P>
inline void alloc_split_partials(P& p) {
auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
auto o_part = torch::empty({p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt);
auto ml_part = torch::empty({p.batch, p.q_head, MAX_SPLITS, 2}, fopt);
auto o_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt);
auto ml_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, 2}, fopt);
p.o_part = (float*)o_part.data_ptr();
p.ml_part = (float*)ml_part.data_ptr();
}