File size: 21,664 Bytes
c855122 90f3d2f c855122 ce0e45f c855122 ce0e45f c855122 ce0e45f c855122 ce0e45f c855122 90f3d2f c855122 90f3d2f c855122 ce0e45f c855122 ce0e45f c855122 90f3d2f c855122 ce0e45f c855122 ce0e45f c855122 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 | """Attention for LNC=1: ONE program over ALL head groups.
The 2-core kernel in ``attention_block.py`` owns one 8-head group per program. At LNC=1 that
means one launch per group per layer (2 for ViT-L, 3 for ViT-H+), each re-transposing x, each
re-loading the rope tables, each writing an [N, H] bf16 out-projection PARTIAL to HBM that the
LN-fold kernel then sums.
This kernel, per image (``nimg`` images stacked on the row axis, one launch):
* transposes x once and loads the rope tables once (shared by every group: tables are
quad-tiled per head, identical for all heads);
* for each group: projects Q/K/V + RoPE (same producer as attention_block), runs the vendor
attention per head, and writes O feature-major into a PERSISTENT per-group SBUF slab;
* after all groups: ONE out-projection contracting over ALL features (G*4 chunks accumulate
in one fp32 PSUM tile), adds the FULL bias once, writes ONE [N, H] result.
No partials, no host sum; the LN-fold kernel receives a single branch.
``resident=1``: Q/K/V/O weights, biases and rope tables are staged ONCE per launch (bottom of
SBUF) and reused by every image. Host decides by an SBUF estimate (does not fit at N>=1024 for
ViT-L); ``resident=0`` restages each group's weights per image.
N<=2048 only (the shipping single-section vendor path). Numerics: identical producer and vendor
attention; the out-projection now accumulates all groups in fp32 PSUM instead of summing bf16
per-group partials, so results differ from attention_block by bf16 rounding of the partials.
"""
from __future__ import annotations
import torch
from nki.framework.kernel import wrap_nki_to_torch_hop
from torch_neuronx import nki_op
import nki
import nki.isa as nisa
import nki.language as nl
import os as _osr
_ROPE_BF16 = _osr.environ.get("VIT_ENC_ROPE_BF16", "0") == "1"
_ROPE_DT = nl.bfloat16 if _ROPE_BF16 else nl.float32
from nki.isa import engine
from nkilib.core.utils.modular_allocator import ModularAllocator
from .attention_inner import _attention_cte
from .attention_block import pack_qk_bias as _pack_qk_bias
_P_MAX = 128
_PSUM_FMAX = 512
_PSUM_BANK = 2048
_K_TILE = 512
_V_TILE = 128
def _kernel_assert(cond: bool, msg: str):
assert cond, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {msg}"
def sbuf_estimate(H: int, Fp: int, N: int, D: int, resident: bool) -> int:
"""Conservative bytes/lane at the kernel's peak (host-side decision helper, not exact)."""
G = Fp // _PSUM_FMAX
n_chunks = H // _P_MAX
n_tt = (N + _PSUM_FMAX - 1) // _PSUM_FMAX
n_ktiles = (N + _K_TILE - 1) // _K_TILE
lh = _PSUM_FMAX // D
n_och = (lh * D) // _P_MAX
w_grp = 3 * n_chunks * _PSUM_FMAX * 2 + _PSUM_FMAX * 4 + 64 # wqk + wv + bv_bc + biases
w_out = G * n_och * H * 2 + H * 4 # wo + bo_bc
tab = 4 * n_tt * _PSUM_FMAX * 4
o_feat = G * n_och * N * 2
x_t = n_chunks * n_tt * _PSUM_FMAX * 2
kv = lh * (n_ktiles * _K_TILE * 2 + (N // _V_TILE) * D * 2)
work = 2 * _PSUM_FMAX * 2 * 2 + 2 * _PSUM_FMAX * 4 + _PSUM_FMAX * 4
inner = 48 * 1024
resid = (G * w_grp + w_out + tab) if resident else 0
base = resid + o_feat + x_t + (0 if resident else tab)
p1 = kv + (0 if resident else w_grp) + work
p2 = kv + inner
p3 = (0 if resident else w_out) + 2048
return base + max(p1, p2, p3)
@nki.jit
def _attention_block_l1_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
wot_hbm, bo_hbm, scale, head_dim, n_valid, nimg=1, resident=0):
NT, H = h_hbm.shape
N = NT // nimg
Fp = wqk_hbm.shape[2]
D = head_dim
half = D // 2
quad = _P_MAX // half
local_heads = _PSUM_FMAX // D
local_quads = local_heads // quad
G = Fp // _PSUM_FMAX # head groups, ALL handled by this one program
n_chunks = H // _P_MAX
n_tt = (N + _PSUM_FMAX - 1) // _PSUM_FMAX
n_ktiles = (N + _K_TILE - 1) // _K_TILE
n_vtiles = N // _V_TILE
n_ochunks = (local_heads * D) // _P_MAX # feature chunks per group
v_per_tt = _PSUM_FMAX // _V_TILE
_kernel_assert(wqk_hbm.shape == (2, H, Fp), "bad wqk")
_kernel_assert(bqk_hbm.shape == (2, _P_MAX, Fp // _P_MAX), "bqk must be pack_qk_bias(...)")
_kernel_assert(tab_hbm.shape == (4, _P_MAX, N), "bad tab")
_kernel_assert(wvt_hbm.shape == (H, Fp) and bv_hbm.shape == (1, Fp), "bad V weights")
_kernel_assert(wot_hbm.shape == (Fp, H) and bo_hbm.shape == (1, H), "bad O weights")
_kernel_assert(H % _P_MAX == 0, "H must be a multiple of 128")
_kernel_assert(quad * half == _P_MAX, "D must give quad*half=128")
_kernel_assert(Fp % _PSUM_FMAX == 0, "padded features must be a multiple of 512")
_kernel_assert(N % _P_MAX == 0, "N must be a multiple of 128")
_kernel_assert(N <= _PSUM_FMAX * 4, "N>2048 not supported")
_kernel_assert(NT == nimg * N, "rows must be nimg x N")
result = nl.ndarray((1, NT, H), dtype=nl.bfloat16, buffer=nl.shared_hbm)
# ---------------- resident (per-launch invariant) staging ----------------
res_top = 0
r_w = []
r_wv = []
r_b = []
r_bv = []
if resident:
ra = ModularAllocator(initial_address=0)
for g in range(G):
col0 = g * _PSUM_FMAX
w_sb_ = ra.alloc_sbuf_tensor(shape=(_P_MAX, local_quads * 2 * _P_MAX), dtype=nl.bfloat16,
block_dim=[2, n_chunks], num_free_tiles=[2, n_chunks])
for qk in range(2):
for c in range(n_chunks):
nisa.dma_copy(dst=w_sb_[qk][c],
src=wqk_hbm[qk, c * _P_MAX:(c + 1) * _P_MAX,
col0:col0 + local_quads * 2 * _P_MAX])
wv_sb_ = ra.alloc_sbuf_tensor(shape=(_P_MAX, local_heads * D), dtype=nl.bfloat16,
block_dim=[n_chunks], num_free_tiles=[n_chunks])
for c in range(n_chunks):
nisa.dma_copy(dst=wv_sb_[c],
src=wvt_hbm[c * _P_MAX:(c + 1) * _P_MAX, col0:col0 + local_heads * D])
# Q/K biases: host-packed fp32 [2, 128, Fp/128] (pack_qk_bias); this program's groups land in
# ONE DMA per q/k on the Scalar HW-DGE queue (was one [128,1] DMA + one cast per group).
b_sb_ = ra.alloc_sbuf_tensor(shape=(_P_MAX, local_quads * 2), dtype=nl.float32,
block_dim=[2], num_free_tiles=[2], align_to=4)
for qk in range(2):
nisa.dma_copy(dst=b_sb_[qk], src=bqk_hbm[qk, 0:_P_MAX, g * (_PSUM_FMAX // _P_MAX):g * (_PSUM_FMAX // _P_MAX) + local_quads * 2],
dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar)
bv_bc_ = ra.alloc_sbuf_tensor(shape=(_V_TILE, local_heads * D), dtype=nl.float32, align_to=4)
nisa.dma_copy(dst=bv_bc_,
src=bv_hbm.ap(pattern=[[0, _V_TILE], [1, local_heads * D]], offset=col0))
r_w.append(w_sb_)
r_wv.append(wv_sb_)
r_b.append(b_sb_)
r_bv.append(bv_bc_)
r_tab = ra.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=_ROPE_DT,
block_dim=[4, n_tt], num_free_tiles=[4, n_tt], align_to=4)
for s in range(n_tt):
for t_i in range(4):
_tw = min(_PSUM_FMAX, N - s * _PSUM_FMAX)
nisa.dma_copy(dst=r_tab[t_i][s][:, 0:_tw],
src=tab_hbm[t_i, 0:_P_MAX, s * _PSUM_FMAX:s * _PSUM_FMAX + _tw])
r_wo = ra.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.bfloat16,
block_dim=[G * n_ochunks], num_free_tiles=[G * n_ochunks])
for oc in range(G * n_ochunks):
nisa.dma_copy(dst=r_wo[oc], src=wot_hbm[oc * _P_MAX:(oc + 1) * _P_MAX, 0:H])
r_bo = ra.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.float32, align_to=4)
nisa.dma_copy(dst=r_bo, src=bo_hbm.ap(pattern=[[0, _P_MAX], [1, H]], offset=0))
res_top = ra.get_current_address()
n_grp = N // _P_MAX
for img in range(nimg):
rofs = img * N * H
allocator = ModularAllocator(initial_address=res_top)
# O for ALL groups stays in SBUF until the single out-projection.
o_feat = []
for g in range(G):
of_ = allocator.alloc_sbuf_tensor(shape=(_P_MAX, N), dtype=nl.bfloat16,
block_dim=[n_ochunks], num_free_tiles=[n_ochunks],
align_to=32)
o_feat.append(of_)
# x^T: shared by every group (transposed ONCE per image).
x_t = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=nl.bfloat16,
block_dim=[n_chunks, n_tt], num_free_tiles=[n_chunks, n_tt],
align_to=32)
for s in range(n_tt):
for c in range(n_chunks):
nisa.dma_transpose(
dst=x_t[c][s].ap([[_PSUM_FMAX, _P_MAX], [1, 1], [1, 1],
[1, min(_PSUM_FMAX, N - s * _PSUM_FMAX)]], offset=0),
src=h_hbm.ap([[H, min(_PSUM_FMAX, N - s * _PSUM_FMAX)], [1, 1], [1, 1], [1, _P_MAX]],
offset=rofs + (s * _PSUM_FMAX) * H + c * _P_MAX),
)
# rope tables: shared by every group.
if resident:
tab = r_tab
else:
tab = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=_ROPE_DT,
block_dim=[4, n_tt], num_free_tiles=[4, n_tt], align_to=4)
for s in range(n_tt):
for t_i in range(4):
_tw = min(_PSUM_FMAX, N - s * _PSUM_FMAX)
nisa.dma_copy(dst=tab[t_i][s][:, 0:_tw],
src=tab_hbm[t_i, 0:_P_MAX, s * _PSUM_FMAX:s * _PSUM_FMAX + _tw])
img_base = allocator.get_current_address()
for g in range(G):
allocator.set_current_address(img_base)
col0 = g * _PSUM_FMAX
k_tiles = []
v_tiles = []
for _lh in range(local_heads):
kt_ = allocator.alloc_sbuf_tensor(shape=(D, _K_TILE), dtype=nl.bfloat16,
block_dim=[n_ktiles], num_free_tiles=[n_ktiles],
align_to=32)
k_tiles.append(kt_)
vt_ = allocator.alloc_sbuf_tensor(shape=(_V_TILE, D), dtype=nl.bfloat16,
block_dim=[n_vtiles], num_free_tiles=[n_vtiles])
v_tiles.append(vt_)
phase2_base = allocator.get_current_address()
if resident:
w_sb = r_w[g]
wv_sb = r_wv[g]
b_sb = r_b[g]
bv_bc = r_bv[g]
else:
w_sb = allocator.alloc_sbuf_tensor(shape=(_P_MAX, local_quads * 2 * _P_MAX),
dtype=nl.bfloat16,
block_dim=[2, n_chunks], num_free_tiles=[2, n_chunks])
for qk in range(2):
for c in range(n_chunks):
nisa.dma_copy(dst=w_sb[qk][c],
src=wqk_hbm[qk, c * _P_MAX:(c + 1) * _P_MAX,
col0:col0 + local_quads * 2 * _P_MAX])
wv_sb = allocator.alloc_sbuf_tensor(shape=(_P_MAX, local_heads * D), dtype=nl.bfloat16,
block_dim=[n_chunks], num_free_tiles=[n_chunks])
for c in range(n_chunks):
nisa.dma_copy(dst=wv_sb[c],
src=wvt_hbm[c * _P_MAX:(c + 1) * _P_MAX, col0:col0 + local_heads * D])
# Q/K biases: host-packed fp32 [2, 128, Fp/128] (pack_qk_bias); this program's groups land in
# ONE DMA per q/k on the Scalar HW-DGE queue (was one [128,1] DMA + one cast per group).
b_sb = allocator.alloc_sbuf_tensor(shape=(_P_MAX, local_quads * 2), dtype=nl.float32,
block_dim=[2], num_free_tiles=[2], align_to=4)
for qk in range(2):
nisa.dma_copy(dst=b_sb[qk], src=bqk_hbm[qk, 0:_P_MAX, g * (_PSUM_FMAX // _P_MAX):g * (_PSUM_FMAX // _P_MAX) + local_quads * 2],
dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar)
bv_bc = allocator.alloc_sbuf_tensor(shape=(_V_TILE, local_heads * D), dtype=nl.float32,
align_to=4)
nisa.dma_copy(dst=bv_bc,
src=bv_hbm.ap(pattern=[[0, _V_TILE], [1, local_heads * D]], offset=col0))
proj = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=nl.bfloat16,
block_dim=[2], num_free_tiles=[2])
acc = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=_ROPE_DT, align_to=4)
tmp = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=_ROPE_DT, align_to=4)
roped = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=nl.bfloat16,
block_dim=[2], num_free_tiles=[2])
vout = allocator.alloc_sbuf_tensor(shape=(_V_TILE, local_heads * D), dtype=nl.float32,
align_to=4)
q_priv = []
for _lh in range(local_heads):
q_priv.append(nl.ndarray((1, D, N), dtype=nl.bfloat16, buffer=nl.private_hbm))
# ---- producer (identical to attention_block's N<=2048 path) ----
for tt in range(n_tt):
t0 = tt * _PSUM_FMAX
tw = min(_PSUM_FMAX, N - t0)
for vsub in range(tw // _V_TILE):
ps_v = nl.ndarray((_V_TILE, local_heads * D), dtype=nl.float32,
buffer=nl.psum, address=(0, 7 * _PSUM_BANK))
for c in range(n_chunks):
nisa.nc_matmul(dst=ps_v,
stationary=x_t[c][tt][:, vsub * _V_TILE:(vsub + 1) * _V_TILE],
moving=wv_sb[c], accumulate=(c > 0))
nisa.tensor_tensor(dst=vout, data1=ps_v, data2=bv_bc, op=nl.add)
for lh in range(local_heads):
nisa.tensor_copy(dst=v_tiles[lh][tt * v_per_tt + vsub],
src=vout[:, lh * D:(lh + 1) * D], engine=engine.scalar)
for qk in range(2):
for j in range(local_quads):
for par in range(2):
gq = j * 2 + par
cols = gq * _P_MAX
ps = nl.ndarray((_P_MAX, _PSUM_FMAX), dtype=nl.float32, buffer=nl.psum,
address=(0, (5 + par) * _PSUM_BANK))
for c in range(n_chunks):
nisa.nc_matmul(dst=ps[:, 0:tw],
stationary=w_sb[qk][c][:, cols:cols + _P_MAX],
moving=x_t[c][tt][:, 0:tw], accumulate=(c > 0))
nisa.tensor_scalar(dst=proj[par][:, 0:tw], data=ps[:, 0:tw],
op0=nl.add, operand0=b_sb[qk][:, gq:gq + 1])
nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[0][:, 0:tw],
data2=tab[0][tt][:, 0:tw], op=nl.multiply)
nisa.tensor_tensor(dst=tmp[:, 0:tw], data1=proj[1][:, 0:tw],
data2=tab[2][tt][:, 0:tw], op=nl.multiply)
nisa.tensor_tensor(dst=roped[0][:, 0:tw], data1=acc[:, 0:tw],
data2=tmp[:, 0:tw], op=nl.add)
nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[1][:, 0:tw],
data2=tab[1][tt][:, 0:tw], op=nl.multiply)
nisa.tensor_tensor(dst=tmp[:, 0:tw], data1=proj[0][:, 0:tw],
data2=tab[3][tt][:, 0:tw], op=nl.multiply)
nisa.tensor_tensor(dst=roped[1][:, 0:tw], data1=acc[:, 0:tw],
data2=tmp[:, 0:tw], op=nl.add)
for b in range(quad):
lh = j * quad + b
r0 = b * half
if qk == 1:
for par2 in range(2):
nisa.tensor_copy(
dst=k_tiles[lh][tt][par2 * half:(par2 + 1) * half, 0:tw],
src=roped[par2][r0:r0 + half, 0:tw])
else:
for par2 in range(2):
nisa.dma_copy(
dst=q_priv[lh][0, par2 * half:(par2 + 1) * half, t0:t0 + tw],
src=roped[par2][r0:r0 + half, 0:tw])
# ---- attention for this group's heads -> o_feat[g] (SBUF, feature-major) ----
allocator.set_current_address(phase2_base)
o_sb = []
for _lh in range(local_heads):
oc = (_lh * D) // _P_MAX
rr = (_lh * D) % _P_MAX
grp = []
for _g in range(n_grp):
grp.append(o_feat[g][oc][rr:rr + D, _g * _P_MAX:(_g + 1) * _P_MAX])
o_sb.append(grp)
attn_base = allocator.get_current_address()
for lh in range(local_heads):
_attention_cte(
q_priv[lh], None, None,
scale=scale, causal_mask=False,
tp_q=False, tp_k=False, tp_out=True,
cache_softmax=False,
k_cache_sbuf=k_tiles[lh],
v_cache_sbuf=v_tiles[lh],
out_o_hbm=o_sb[lh],
init_sbuf_addr=attn_base,
seqlen_k_valid=n_valid,
)
allocator.set_current_address(attn_base)
# ---- ONE out-projection over ALL G*n_ochunks feature chunks, full bias once ----
allocator.set_current_address(img_base)
if resident:
wo_sb = r_wo
bo_bc = r_bo
else:
wo_sb = allocator.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.bfloat16,
block_dim=[G * n_ochunks], num_free_tiles=[G * n_ochunks])
for oc in range(G * n_ochunks):
nisa.dma_copy(dst=wo_sb[oc], src=wot_hbm[oc * _P_MAX:(oc + 1) * _P_MAX, 0:H])
bo_bc = allocator.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.float32, align_to=4)
nisa.dma_copy(dst=bo_bc, src=bo_hbm.ap(pattern=[[0, _P_MAX], [1, H]], offset=0))
out_sb = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=nl.bfloat16, align_to=32)
n_ob = (H + _PSUM_FMAX - 1) // _PSUM_FMAX
for tb in range(n_grp):
for ob in range(n_ob):
ow = min(_PSUM_FMAX, H - ob * _PSUM_FMAX)
ps_out = nl.ndarray((_P_MAX, _PSUM_FMAX), dtype=nl.float32, buffer=nl.psum,
address=(0, 0))
for g in range(G):
for oc in range(n_ochunks):
nisa.nc_matmul(
dst=ps_out[:, 0:ow],
stationary=o_feat[g][oc][:, tb * _P_MAX:(tb + 1) * _P_MAX],
moving=wo_sb[g * n_ochunks + oc][:, ob * _PSUM_FMAX:ob * _PSUM_FMAX + ow],
accumulate=((g > 0) or (oc > 0)),
)
nisa.tensor_tensor(dst=out_sb[:, 0:ow], data1=ps_out[:, 0:ow],
data2=bo_bc[:, ob * _PSUM_FMAX:ob * _PSUM_FMAX + ow], op=nl.add)
nisa.dma_copy(
dst=result.ap(pattern=[[H, _P_MAX], [1, ow]],
offset=rofs + tb * _P_MAX * H + ob * _PSUM_FMAX),
src=out_sb[:, 0:ow])
return result
_ATTN_L1 = wrap_nki_to_torch_hop(_attention_block_l1_kernel[1])
@nki_op("vit_enc::attention_block_l1", mutates_args={})
def attention_block_l1(
h: torch.Tensor,
wqk: torch.Tensor,
bqk: torch.Tensor,
tab: torch.Tensor,
wvt: torch.Tensor,
bv: torch.Tensor,
wot: torch.Tensor,
bo: torch.Tensor,
scale: float,
head_dim: int,
n_valid: int,
nimg: int = 1,
resident: int = 0,
) -> torch.Tensor:
"""All head groups in one program. Returns ONE [1, nimg*N, H] bf16 result (full bias)."""
for t in (h, wqk, bqk, tab, wvt, bv, wot, bo):
if t is not None and t.requires_grad:
raise RuntimeError("attention_block_l1 is inference-only; use torch.no_grad().")
if h.shape[0] // nimg > 2048:
raise NotImplementedError("attention_block_l1 supports N<=2048 only")
if bqk.dim() == 3 and bqk.shape[1] != 128:
bqk = _pack_qk_bias(bqk)
return _ATTN_L1(h, wqk, bqk, tab, wvt, bv, wot, bo, scale, head_dim, n_valid, nimg, resident)
|