Skip to content

Commit c0b114f

Browse files
azrael417mauro-bis
azrael417
and
mauro-bis
committed
fixing some issues with up and downsampling attention.
Co-authored-by: Mauro Bisson <maurob@nvidia.com>
1 parent fc85890 commit c0b114f

2 files changed

Lines changed: 8 additions & 151 deletions

File tree

‎torch_harmonics/attention/csrc/attention_cuda_bwd.cu‎

Lines changed: 6 additions & 150 deletions
Original file line numberDiff line numberDiff line change
@@ -54,150 +54,6 @@
5454

5555
namespace attention_kernels {
5656

57-
#if 0
58-
class ScopeTimer
59-
{
60-
public:
61-
explicit ScopeTimer(const std::string &label = "") :
62-
label_(label), start_(std::chrono::high_resolution_clock::now())
63-
{
64-
}
65-
66-
~ScopeTimer()
67-
{
68-
auto end = std::chrono::high_resolution_clock::now();
69-
auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(end - start_);
70-
std::cout << label_ << "Elapsed time: " << elapsed.count() << " ms" << std::endl;
71-
}
72-
73-
private:
74-
std::string label_;
75-
std::chrono::high_resolution_clock::time_point start_;
76-
};
77-
78-
// easier to understand version of manual shfl_xor_sync, performance appears similar
79-
static __device__ float __warp_sum_cub(float val)
80-
{
81-
// use cub to reduce within a warp
82-
__shared__ typename cub::WarpReduce<float>::TempStorage temp_storage;
83-
84-
// 1. Compute sum (initially only in lane 0)
85-
float sum = cub::WarpReduce<float>(temp_storage).Sum(val);
86-
// 2. Broadcast sum to all threads
87-
sum = __shfl_sync(0xFFFFFFFF, sum, 0);
88-
return sum;
89-
}
90-
91-
// This kernel computes the backward pass for the S2 attention mechanism, using
92-
// shared memory as a cache and one warp per output point, warp-parallel over
93-
// channels, which should be layed out in the fastest dimension for coalesced
94-
// memory access.
95-
template <int BDIM_X>
96-
__global__ __launch_bounds__(BDIM_X) void s2_attention_bwd_dkvq_kernel(
97-
int num_channels, int nlon_in, int nlat_out, int nlon_out,
98-
const torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> kx,
99-
const torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> vx,
100-
const torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> qy,
101-
const torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> dy,
102-
torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> dydk,
103-
torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> dydv,
104-
torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> dydq,
105-
const torch::PackedTensorAccessor64<int64_t, 1, torch::RestrictPtrTraits> psi_col_idx,
106-
const torch::PackedTensorAccessor64<int64_t, 1, torch::RestrictPtrTraits> psi_row_offset,
107-
const torch::PackedTensorAccessor32<float, 1, torch::RestrictPtrTraits> quad_weights)
108-
{
109-
110-
extern __shared__ float sh[];
111-
float *sh_alpha_k = sh + threadIdx.y * num_channels * 5;
112-
float *sh_alpha_vw = sh_alpha_k + num_channels;
113-
float *sh_alpha_kvw = sh_alpha_vw + num_channels;
114-
float *sh_dy = sh_alpha_kvw + num_channels;
115-
float *sh_qy = sh_dy + num_channels;
116-
// (optionally, could use more shared memory for other intermediates)
117-
118-
const uint64_t batchId = blockIdx.y;
119-
const uint64_t wid = uint64_t(blockIdx.x) * blockDim.y + threadIdx.y;
120-
if (wid >= uint64_t(nlat_out) * nlon_in) return;
121-
const int tidx = threadIdx.x;
122-
const int ho = wid / nlon_out;
123-
const int wo = wid - (ho * nlon_out);
124-
125-
// Zero shared memory
126-
for (int chan = tidx; chan < num_channels; chan += WARP_SIZE) {
127-
sh_alpha_k[chan] = 0.0f;
128-
sh_alpha_vw[chan] = 0.0f;
129-
sh_alpha_kvw[chan] = 0.0f;
130-
sh_dy[chan] = dy[batchId][chan][ho][wo];
131-
sh_qy[chan] = qy[batchId][chan][ho][wo];
132-
}
133-
float alpha_sum = 0.0f;
134-
float qdotk_max = -FLT_MAX;
135-
float integral = 0.0f;
136-
__syncthreads();
137-
138-
const int64_t rbeg = psi_row_offset[ho];
139-
const int64_t rend = psi_row_offset[ho + 1];
140-
const int rlen = rend - rbeg;
141-
142-
// 1st pass: accumulate alpha_sum, integral, and shared stats, along with a progressively computed qdotk_max.
143-
for (int off = 0; off < rlen; off++) {
144-
const int64_t col = psi_col_idx[rbeg + off];
145-
const int hi = col / nlon_in;
146-
const int wi = col - (hi * nlon_in);
147-
const int wip = (wi + wo) - ((wi + wo) / nlon_in) * nlon_in;
148-
float qdotk = 0.0f, gdotv = 0.0f;
149-
for (int chan = tidx; chan < num_channels; chan += WARP_SIZE) {
150-
qdotk += sh_qy[chan] * kx[batchId][chan][hi][wip];
151-
gdotv += sh_dy[chan] * vx[batchId][chan][hi][wip];
152-
}
153-
qdotk = __warp_sum_cub(qdotk);
154-
gdotv = __warp_sum_cub(gdotv);
155-
float qdotk_max_tmp = max(qdotk_max, qdotk);
156-
float alpha_inz = expf(qdotk - qdotk_max_tmp) * quad_weights[hi];
157-
float max_correction = expf(qdotk_max - qdotk_max_tmp);
158-
alpha_sum = alpha_sum * max_correction + alpha_inz;
159-
integral = integral * max_correction + alpha_inz * gdotv;
160-
for (int chan = tidx; chan < num_channels; chan += WARP_SIZE) {
161-
float kxval = kx[batchId][chan][hi][wip];
162-
sh_alpha_k[chan] = sh_alpha_k[chan] * max_correction + alpha_inz * kxval;
163-
sh_alpha_vw[chan] = sh_alpha_vw[chan] * max_correction + alpha_inz * gdotv;
164-
sh_alpha_kvw[chan] = sh_alpha_kvw[chan] * max_correction + alpha_inz * kxval * gdotv;
165-
}
166-
qdotk_max = qdotk_max_tmp;
167-
}
168-
169-
integral /= alpha_sum;
170-
171-
// Write dydq
172-
for (int chan = tidx; chan < num_channels; chan += WARP_SIZE) {
173-
dydq[batchId][chan][ho][wo]
174-
= (sh_alpha_kvw[chan] * alpha_sum - sh_alpha_vw[chan] * sh_alpha_k[chan]) / (alpha_sum * alpha_sum);
175-
}
176-
177-
// Third pass: accumulate gradients for k and v
178-
for (int off = 0; off < rlen; off++) {
179-
const int64_t col = psi_col_idx[rbeg + off];
180-
const int hi = col / nlon_in;
181-
const int wi = col - (hi * nlon_in);
182-
const int wip = (wi + wo) - ((wi + wo) / nlon_in) * nlon_in;
183-
float qdotk = 0.0f, gdotv = 0.0f;
184-
for (int chan = tidx; chan < num_channels; chan += WARP_SIZE) {
185-
qdotk += qy[batchId][chan][ho][wo] * kx[batchId][chan][hi][wip];
186-
gdotv += sh_dy[chan] * vx[batchId][chan][hi][wip];
187-
}
188-
qdotk = __warp_sum_cub(qdotk);
189-
gdotv = __warp_sum_cub(gdotv);
190-
float alpha_inz = expf(qdotk - qdotk_max) * quad_weights[hi];
191-
for (int chan = tidx; chan < num_channels; chan += WARP_SIZE) {
192-
float qyval = qy[batchId][chan][ho][wo];
193-
float dyval = sh_dy[chan];
194-
atomicAdd(&dydk[batchId][chan][hi][wip], qyval * (alpha_inz / alpha_sum) * (gdotv - integral));
195-
atomicAdd(&dydv[batchId][chan][hi][wip], (alpha_inz / alpha_sum) * dyval);
196-
}
197-
}
198-
}
199-
#endif
200-
20157
// BEGIN backward kernels and functions
20258

20359
// called with (blockDim.x=32 and blockDim.y>1, BDIM=blockDim.x*blockDim.y)
@@ -238,7 +94,7 @@ void s2_attn_bwd_generic_vec_k(int nchans_in, // no. of FLOATV_T elements along
23894
const int batch = blockIdx.y;
23995

24096
const uint64_t wid = uint64_t(blockIdx.x) * blockDim.y + threadIdx.y;
241-
if (wid >= uint64_t(nlat_out)*nlon_in) {
97+
if (wid >= uint64_t(nlat_out)*nlon_out) {
24298
return;
24399
}
244100

@@ -460,7 +316,7 @@ void s2_attn_bwd_special_vec_k(int nchan_in, // no. of FLOATV_T elements along
460316
const int batch = blockIdx.y;
461317
const uint64_t ctaid = uint64_t(blockIdx.x) * blockDim.y + threadIdx.y;
462318

463-
if (ctaid >= uint64_t(nlat_out)*nlon_in) {
319+
if (ctaid >= uint64_t(nlat_out)*nlon_out) {
464320
return;
465321
}
466322

@@ -834,15 +690,15 @@ void launch_spc_attn_bwd(int nloc, // "BDIM_X*nloc" >= nchans_out
834690
size_t shsize = sizeof(FLOATV_T)*(nchans_in+nchans_out) * block.y; // 2 arrays per cta, block.y > 1 iif block.x==32
835691

836692
// nloc determines the size of local arrays used to store
837-
// temporary buffers loc_k__[], loc_vw_[] and loc_kvw[],
693+
// temporary buffers loc_k__[], loc_vw_[] and loc_kvw[],
838694
// of size nchans_in each;
839695
// if nchans_out is >= BDIM_X*(nloc-1) and <= BDIM_X*nloc
840696
// then we can use the same compile-time known loops used
841-
// for input channels, with the execpetion of testing
697+
// for input channels, with the execpetion of testing
842698
// whether to execute the last iteration based on "nchans_out"
843-
// ibstead of "nchans_in"; in this way as long as the
699+
// instead of "nchans_in"; in this way as long as the
844700
// difference between the number of input and output channels
845-
// is <= BDIM_X we can use the faster path
701+
// is <= BDIM_X we can use the faster path
846702
if (nchans_out >= BDIM_X*(CUR_LOC_SIZE-1) &&
847703
nchans_out <= BDIM_X* CUR_LOC_SIZE ) {
848704
s2_attn_bwd_special_vec_k<BDIM_X, BDIM_Y, 1, CUR_LOC_SIZE>

‎torch_harmonics/attention/csrc/attention_cuda_fwd.cu‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -552,7 +552,8 @@ torch::Tensor s2_attention_fwd_cuda(at::Tensor kx,
552552
if (!vx_is_channels_last) { vxP = permute_4D_to0231(vxP); }
553553
if (!qy_is_channels_last) { qyP = permute_4D_to0231(qyP); }
554554

555-
torch::Tensor yP = torch::empty_like(vxP);
555+
int64_t out_dims[] = {batch_size, nlat_out, nlon_out, nchans_out};
556+
torch::Tensor yP = torch::empty(out_dims, kxP.options());
556557

557558
s2_attn_fwd_dispatch(batch_size,
558559
nchans_in,

0 commit comments

Comments
 (0)