nx_vit_encoder_layer.nx
buildroot/runtime/nx_vit_encoder_layer.nx
about
nx_vit_encoder_layer.nx -- the faithful f32 PRE-NORM ViT encoder block (the ViTPose repeated unit), composing the
gated blocks: h = x + attn_dense(MHSA(LN_before(x))); out = h + fc2(GELU(fc1(LN_after(h)))). MHSA has explicit
qkv_bias and HF [out,in] weights (via nx_f32_matmul_t = x@W^T), 12 heads x head_dim 64, scale 1/sqrt(head_dim),
non-causal. Weights are passed as a 16-pointer array to stay under the arg cap. license_tier: ORIGINAL
dependencies 7 imports · 3 importers
imports: nx_syscalls.nxnx_f32.nxnx_f32_matmul_t.nxnx_f32_softmax.nxnx_f32_layernorm.nxnx_f32_gelu.nxnx_thread_pool.nx
imported by: nx_bert_ce_lib.nxnx_vit_encoder_layer_gate.nxnx_vitpose_forward.nx
structs
| 25 | struct NxVitAttnCtx |
consts
| 36 | const NX_VIT_ATTN_CTX_BYTES: i64 = 72 |
functions
| 14 | func vit_bias_add(x: *i64, bias: *i64, n_rows: i64, dim: i64) -> i64 |
| 39 | func vit_attn_head(Q: *i64, K: *i64, V: *i64, ctx: *i64, T: i64, D: i64, hd: i64, scale: i64, h: i64) -> i64 |
| 68 | func _vit_attn_task(ctx_i: i64) -> i64 |
| 74 | func vit_mhsa(x: *i64, Wq: *i64, bq: *i64, Wk: *i64, bk: *i64, Wv: *i64, bv: *i64, Wo: *i64, bo: *i64, |
| 95 | func vit_mhsa_pool(pool: *NxThreadPool, x: *i64, Wq: *i64, bq: *i64, Wk: *i64, bk: *i64, Wv: *i64, bv: *i64, Wo: *i64, bo: *i64, |
| 131 | func vit_encoder_layer(x: *i64, wts: *i64, T: i64, D: i64, nh: i64, hd: i64, mlp: i64, eps: i64, scale: i64, out: *i64) -> i64 |