Skip to content

[RWKV7] Drop some kernels to enhance speed #346

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zhiyuan1i merged 10 commits into from
Apr 11, 2025
Merged

[RWKV7] Drop some kernels to enhance speed #346

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
8e3cdef
[RWKV] remove unnecessary kernels
zhiyuan1i Apr 10, 2025
d18c3df
reduce register
zhiyuan1i Apr 10, 2025
4bc1896
support fp16/32
zhiyuan1i Apr 11, 2025
35cdcd4
drop precision setting
zhiyuan1i Apr 11, 2025
8f7c7fa
[Attn] Remove `head_first` & rename `offsets` to `cu_seqlens` (#345)
yzhangcs Apr 10, 2025
683bc45
Merge branch 'main' into rwkv-speedup
zhiyuan1i Apr 11, 2025
2b3318e
fix
zhiyuan1i Apr 11, 2025
97b0979
update
zhiyuan1i Apr 11, 2025
32210ec
use bf16 in rwkv7 tests
zhiyuan1i Apr 11, 2025
11ac5bd
remove lines
zhiyuan1i Apr 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion fla/layers/rwkv7.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -206,7 +206,6 @@ def forward(
output_final_state=use_cache,
cu_seqlens=cu_seqlens,
head_first=False,
input_precision=r.dtype,
)

if past_key_values is not None:
Expand Down
6 changes: 5 additions & 1 deletion fla/ops/generalized_delta_rule/dplr/chunk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -330,7 +330,11 @@ def chunk_dplr_delta_rule(
"when head_first=False was specified. "
"Please verify your input tensor format matches the expected shape [B, T, H, ...]."
)

if q.dtype == torch.float32:
warnings.warn(
"""ChunkDeltaRuleFunction does not support float32. Please use bfloat16.
If you want to use float32, please solve the issue by yourself."""
)
if cu_seqlens is not None:
if q.shape[0] != 1:
raise ValueError(
Expand Down
198 changes: 62 additions & 136 deletions fla/ops/generalized_delta_rule/dplr/chunk_A_bwd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
for num_warps in [2, 4, 8, 16, 32]
for num_stages in [2, 3, 4]
],
key=['BK', 'NC', 'BT', 'K'],
key=['BK', 'BT', 'K'],
use_cuda_graph=use_cuda_graph,
)
@triton.jit(do_not_specialize=['T'])
Expand Down Expand Up @@ -55,20 +55,19 @@ def chunk_dplr_bwd_kernel_intra(
BT: tl.constexpr,
BC: tl.constexpr,
BK: tl.constexpr,
NC: tl.constexpr,
IS_VARLEN: tl.constexpr,
GATHER_SUPPORTED: tl.constexpr
):
i_k, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
i_b, i_h = i_bh // H, i_bh % H
i_t, i_i = i_c // NC, i_c % NC
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
T = eos - bos
if i_t * BT + i_i * BC >= T:
bos, eos = (i_b * T).to(tl.int32), (i_b * T + T).to(tl.int32)

if i_t * BT >= T:
return

# offset calculation
Expand Down Expand Up @@ -96,8 +95,8 @@ def chunk_dplr_bwd_kernel_intra(
stride_qk = H*K
stride_A = H*BT

p_ge = tl.make_block_ptr(ge, (T, K), (stride_qk, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
p_gi = tl.make_block_ptr(gi, (T, K), (stride_qk, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
p_ge = tl.make_block_ptr(ge, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0))
p_gi = tl.make_block_ptr(gi, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0))
# [BC, BK]
b_ge = tl.load(p_ge, boundary_check=(0, 1))
b_gi = tl.load(p_gi, boundary_check=(0, 1))
Expand All @@ -106,101 +105,29 @@ def chunk_dplr_bwd_kernel_intra(
b_dk = tl.zeros([BC, BK], dtype=tl.float32)
b_db = tl.zeros([BC, BK], dtype=tl.float32)
# intra chunk gradient calculation
p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (stride_A, 1), (i_t*BT + i_i*BC, i_i*BC), (BC, BC), (1, 0))
p_dAab = tl.make_block_ptr(dAab, (T, BT), (stride_A, 1), (i_t*BT + i_i*BC, i_i*BC), (BC, BC), (1, 0))
p_dAqb = tl.make_block_ptr(dAqb, (T, BT), (stride_A, 1), (i_t*BT + i_i*BC, i_i*BC), (BC, BC), (1, 0))
p_dAak = tl.make_block_ptr(dAak, (T, BT), (stride_A, 1), (i_t*BT + i_i*BC, i_i*BC), (BC, BC), (1, 0))
p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (stride_A, 1), (i_t*BT, 0), (BC, BC), (1, 0))
p_dAab = tl.make_block_ptr(dAab, (T, BT), (stride_A, 1), (i_t*BT, 0), (BC, BC), (1, 0))
p_dAqb = tl.make_block_ptr(dAqb, (T, BT), (stride_A, 1), (i_t*BT, 0), (BC, BC), (1, 0))
p_dAak = tl.make_block_ptr(dAak, (T, BT), (stride_A, 1), (i_t*BT, 0), (BC, BC), (1, 0))
o_i = tl.arange(0, BC)
p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t*BT + i_i*BC, i_k*BK), (BC, BK), (1, 0))
p_b = tl.make_block_ptr(b, (T, K), (stride_qk, 1), (i_t*BT + i_i*BC, i_k*BK), (BC, BK), (1, 0))
p_a = tl.make_block_ptr(a, (T, K), (stride_qk, 1), (i_t*BT + i_i*BC, i_k*BK), (BC, BK), (1, 0))
p_q = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_t*BT + i_i*BC, i_k*BK), (BC, BK), (1, 0))
b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32)
b_b = tl.load(p_b, boundary_check=(0, 1)).to(tl.float32)
b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32)
b_a = tl.load(p_a, boundary_check=(0, 1)).to(tl.float32)
b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32)
b_dAab = tl.load(p_dAab, boundary_check=(0, 1)).to(tl.float32)
b_dAqb = tl.load(p_dAqb, boundary_check=(0, 1)).to(tl.float32)
b_dAak = tl.load(p_dAak, boundary_check=(0, 1)).to(tl.float32)
p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t*BT, i_k*BK), (BC, BK), (1, 0))
p_b = tl.make_block_ptr(b, (T, K), (stride_qk, 1), (i_t*BT, i_k*BK), (BC, BK), (1, 0))
p_a = tl.make_block_ptr(a, (T, K), (stride_qk, 1), (i_t*BT, i_k*BK), (BC, BK), (1, 0))
p_q = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_t*BT, i_k*BK), (BC, BK), (1, 0))
b_k = tl.load(p_k, boundary_check=(0, 1))
b_b = tl.load(p_b, boundary_check=(0, 1))
b_q = tl.load(p_q, boundary_check=(0, 1))
b_a = tl.load(p_a, boundary_check=(0, 1))
b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1))
b_dAab = tl.load(p_dAab, boundary_check=(0, 1))
b_dAqb = tl.load(p_dAqb, boundary_check=(0, 1))
b_dAak = tl.load(p_dAak, boundary_check=(0, 1))

# inter chunk gradient calculation
o_k = i_k * BK + tl.arange(0, BK)
m_k = o_k < K
if i_i > 0:
p_gn = gi + (i_t * BT + i_i * BC - 1) * stride_qk + o_k
p_gn = tl.max_contiguous(tl.multiple_of(p_gn, BK), BK)
# [BK,]
b_gn = tl.load(p_gn, mask=m_k, other=0)
# [BK,]
for i_j in range(0, i_i):
p_kj = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0))
p_bj = tl.make_block_ptr(b, (T, K), (stride_qk, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0))
p_gkj = tl.make_block_ptr(gi, (T, K), (stride_qk, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0))
p_dAqikj = tl.make_block_ptr(dAqk, (T, BT), (stride_A, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0))
p_dAaibj = tl.make_block_ptr(dAab, (T, BT), (stride_A, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0))
p_dAqibj = tl.make_block_ptr(dAqb, (T, BT), (stride_A, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0))
p_dAaikj = tl.make_block_ptr(dAak, (T, BT), (stride_A, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0))
# [BC, BK]
b_kj = tl.load(p_kj, boundary_check=(0, 1))
b_bj = tl.load(p_bj, boundary_check=(0, 1))
b_gkj = tl.load(p_gkj, boundary_check=(0, 1))
tmp = exp(b_gn[None, :] - b_gkj)
b_kjg = b_kj * tmp
b_bjg = b_bj * tmp
# [BC, BC]
b_dAqikj = tl.load(p_dAqikj, boundary_check=(0, 1))
b_dAaibj = tl.load(p_dAaibj, boundary_check=(0, 1))
b_dAqibj = tl.load(p_dAqibj, boundary_check=(0, 1))
b_dAaikj = tl.load(p_dAaikj, boundary_check=(0, 1))
# [BC, BK]
b_dq += tl.dot(b_dAqikj, b_kjg)
b_dq += tl.dot(b_dAqibj, b_bjg)
# [BC, BC]
b_da += tl.dot(b_dAaibj, b_bjg)
b_da += tl.dot(b_dAaikj, b_kjg)
b_dq *= exp(b_gi - b_gn[None, :])
b_da *= exp(b_ge - b_gn[None, :])

NC = min(NC, tl.cdiv(T - i_t * BT, BC))
if i_i < NC - 1:
p_gn = gi + (min(i_t * BT + i_i * BC + BC, T) - 1)*stride_qk + o_k
p_gn = tl.max_contiguous(tl.multiple_of(p_gn, BK), BK)
# [BK,]
b_gn = tl.load(p_gn, mask=m_k, other=0)
for i_j in range(i_i + 1, NC):
m_j = (i_t * BT + i_j * BC + tl.arange(0, BC)) < T
p_qj = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0))
p_aj = tl.make_block_ptr(a, (T, K), (stride_qk, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0))
p_gij = tl.make_block_ptr(gi, (T, K), (stride_qk, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0))
p_gej = tl.make_block_ptr(ge, (T, K), (stride_qk, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0))
p_dAqjki = tl.make_block_ptr(dAqk, (BT, T), (1, stride_A), (i_i*BC, i_t*BT + i_j*BC), (BC, BC), (0, 1))
p_dAajbi = tl.make_block_ptr(dAab, (BT, T), (1, stride_A), (i_i*BC, i_t*BT + i_j*BC), (BC, BC), (0, 1))
p_dAqjbi = tl.make_block_ptr(dAqb, (BT, T), (1, stride_A), (i_i*BC, i_t*BT + i_j*BC), (BC, BC), (0, 1))
p_dAajki = tl.make_block_ptr(dAak, (BT, T), (1, stride_A), (i_i*BC, i_t*BT + i_j*BC), (BC, BC), (0, 1))
b_qj = tl.load(p_qj, boundary_check=(0, 1))
b_aj = tl.load(p_aj, boundary_check=(0, 1))
b_gij = tl.load(p_gij, boundary_check=(0, 1))
b_gej = tl.load(p_gej, boundary_check=(0, 1))
b_gij = tl.where(m_j[:, None] & m_k, b_gij, float('-inf'))
b_gej = tl.where(m_j[:, None] & m_k, b_gej, float('-inf'))
b_qjg = b_qj * exp(b_gij - b_gn[None, :])
b_ajg = b_aj * exp(b_gej - b_gn[None, :])
# [BC, BC]
b_dAqjki = tl.load(p_dAqjki, boundary_check=(0, 1))
b_dAajbi = tl.load(p_dAajbi, boundary_check=(0, 1))
b_dAqjbi = tl.load(p_dAqjbi, boundary_check=(0, 1))
b_dAajki = tl.load(p_dAajki, boundary_check=(0, 1))
b_dk += tl.dot(b_dAqjki, b_qjg)
b_dk += tl.dot(b_dAajki, b_ajg)
b_db += tl.dot(b_dAqjbi, b_qjg)
b_db += tl.dot(b_dAajbi, b_ajg)
tmp = exp(b_gn[None, :] - b_gi)
b_dk *= tmp
b_db *= tmp

# intra chunk gradient calculation
for j in range(0, min(BC, T - i_t * BT - i_i * BC)):
for j in range(0, min(BC, T - i_t * BT)):
# trick to index the block
if GATHER_SUPPORTED:
row_idx = tl.full([1, BK], j, dtype=tl.int16)
Expand Down Expand Up @@ -241,43 +168,45 @@ def chunk_dplr_bwd_kernel_intra(
# [1, BK] b_qj, b_aj
b_qj = tl.sum(tl.where(mask_idx[:, None], b_q, 0), 0)[None, :]
b_aj = tl.sum(tl.where(mask_idx[:, None], b_a, 0), 0)[None, :]
# tl.static_print(b_kj)
m_e = o_i[:, None] > j
m_i = o_i[:, None] >= j
tmp1 = exp(b_gi - b_gij)
tmp2 = exp(b_ge - b_gij)
b_dq += tl.where(m_i, b_dAqk_j * b_kj * tmp1, 0.)
b_dq += tl.where(m_i, b_dAqb_j * b_bj * tmp1, 0.)
b_da += tl.where(m_e, b_dAab_j * b_bj * tmp2, 0.)
b_da += tl.where(m_e, b_dAak_j * b_kj * tmp2, 0.)

m_i = o_i[:, None] <= j
m_e = o_i[:, None] < j
tmp1 = exp(b_gij - b_gi)
tmp2 = exp(b_gej - b_gi)
b_dk += tl.where(m_i, b_dA_qk_j * b_qj * tmp1, 0.)
b_dk += tl.where(m_e, b_dA_ak_j * b_aj * tmp2, 0.)
b_db += tl.where(m_i, b_dA_qb_j * b_qj * tmp1, 0.)
b_db += tl.where(m_e, b_dA_ab_j * b_aj * tmp2, 0.)
tmp1 = exp((b_gi - b_gij).to(tl.float32))
tmp2 = exp((b_ge - b_gij).to(tl.float32))

m_e1 = (o_i[:, None] > j).to(tl.int1)
m_i1 = (o_i[:, None] >= j).to(tl.int1)
b_dq += m_i1 * b_dAqk_j * b_kj * tmp1
b_dq += m_i1 * b_dAqb_j * b_bj * tmp1
b_da += m_e1 * b_dAab_j * b_bj * tmp2
b_da += m_e1 * b_dAak_j * b_kj * tmp2

m_i2 = (o_i[:, None] <= j).to(tl.int1)
m_e2 = (o_i[:, None] < j).to(tl.int1)
tmp1 = exp((b_gij - b_gi).to(tl.float32))
tmp2 = exp((b_gej - b_gi).to(tl.float32))
b_dk += m_i2 * b_dA_qk_j * b_qj * tmp1
b_dk += m_e2 * b_dA_ak_j * b_aj * tmp2
b_db += m_i2 * b_dA_qb_j * b_qj * tmp1
b_db += m_e2 * b_dA_ab_j * b_aj * tmp2

# post processing
p_dq = tl.make_block_ptr(dq, (T, K), (stride_qk, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
p_dk = tl.make_block_ptr(dk, (T, K), (stride_qk, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
p_da = tl.make_block_ptr(da, (T, K), (stride_qk, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
p_db = tl.make_block_ptr(db, (T, K), (stride_qk, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
p_dgk = tl.make_block_ptr(dgk, (T, K), (stride_qk, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
p_dgk_offset = tl.make_block_ptr(dgk_offset, (T, K), (stride_qk, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
p_dqg = tl.make_block_ptr(dqg, (T, K), (stride_qk, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
p_dkg = tl.make_block_ptr(dkg, (T, K), (stride_qk, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
p_dag = tl.make_block_ptr(dag, (T, K), (stride_qk, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
p_dbg = tl.make_block_ptr(dbg, (T, K), (stride_qk, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
p_dq = tl.make_block_ptr(dq, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0))
p_dk = tl.make_block_ptr(dk, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0))
p_da = tl.make_block_ptr(da, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0))
p_db = tl.make_block_ptr(db, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0))
p_dgk = tl.make_block_ptr(dgk, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0))
p_dgk_offset = tl.make_block_ptr(dgk_offset, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0))
p_dqg = tl.make_block_ptr(dqg, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0))
p_dkg = tl.make_block_ptr(dkg, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0))
p_dag = tl.make_block_ptr(dag, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0))
p_dbg = tl.make_block_ptr(dbg, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0))
p_gn = gi + (min(i_t * BT + BT, T) - 1)*stride_qk + o_k
p_gn = tl.max_contiguous(tl.multiple_of(p_gn, BK), BK)
b_gn = tl.load(p_gn, mask=m_k, other=0)
b_da += tl.load(p_dag, boundary_check=(0, 1)) * exp(b_ge)
b_dq += tl.load(p_dqg, boundary_check=(0, 1)) * exp(b_gi) * scale
tmp = exp(b_gn[None, :] - b_gi)
b_dk += tl.load(p_dkg, boundary_check=(0, 1)) * tmp
b_db += tl.load(p_dbg, boundary_check=(0, 1)) * tmp
b_dk += tl.load(p_dkg, boundary_check=(0, 1)).to(tl.float32) * tmp
b_db += tl.load(p_dbg, boundary_check=(0, 1)).to(tl.float32) * tmp
tl.store(p_dq, (b_dq).to(p_dq.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_da, b_da.to(p_da.dtype.element_ty), boundary_check=(0, 1))
Expand Down Expand Up @@ -326,9 +255,9 @@ def chunk_dplr_bwd_dgk_kernel(
NT = tl.cdiv(T, BT)
else:
NT = tl.cdiv(T, BT)
i_tg = i_b * NT + i_t
bos, eos = i_b * T, i_b * T + T
T = eos - bos
i_tg = (i_b * NT + i_t).to(tl.int32)
bos, eos = (i_b * T).to(tl.int32), (i_b * T + T).to(tl.int32)

stride_qk = H * K
dgk += (bos * H + i_h) * K
dgk_offset += (bos * H + i_h) * K
Expand Down Expand Up @@ -368,16 +297,14 @@ def chunk_dplr_bwd_dqk_intra(
dgk_last: torch.Tensor,
scale: float = 1.0,
cu_seqlens: Optional[torch.LongTensor] = None,
chunk_size: int = 64,
chunk_size: int = 16,
):
B, T, H, K = q.shape
BT = min(chunk_size, max(16, triton.next_power_of_2(T)))
BC = min(16, BT)
BK = min(64, triton.next_power_of_2(K)) if check_shared_mem() else min(32, triton.next_power_of_2(K))

chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
NC = triton.cdiv(BT, BC)
NK = triton.cdiv(K, BK)

dq = torch.empty_like(q)
Expand All @@ -387,7 +314,7 @@ def chunk_dplr_bwd_dqk_intra(
dgk = torch.empty_like(gi, dtype=torch.float)
dgk_offset = torch.empty_like(gi, dtype=torch.float)

grid = (NK, NT * NC, B * H)
grid = (NK, NT, B * H)
chunk_dplr_bwd_kernel_intra[grid](
q=q,
k=k,
Expand Down Expand Up @@ -416,9 +343,8 @@ def chunk_dplr_bwd_dqk_intra(
H=H,
K=K,
BT=BT,
BC=BC,
BC=BT,
BK=BK,
NC=NC,
GATHER_SUPPORTED=is_gather_supported
)

Expand Down
Loading