v1.1: K4 (mlp/mlp_res_ln) + K1/K3 (res_layernorm) DMA/engine-schedule rework

cabaa8f 6 days ago

15.1 kB

	"""K4 -- fused MLP (fc1 + gelu + fc2), x-prestage ("XP") in prologue mode.

	WHAT IT COMPUTES
	----------------
	``out = gelu(x @ w1.T + b1) @ w2.T + b2``, with the rows sharded across LNC cores and
	the weights pre-transposed to bf16. Returns ``[R, K]`` bf16.

	WHY IT IS SHAPED THIS WAY
	-------------------------
	This is the effective shipping kernel: the ``b1pre`` MLP (fc1 bias staged in the prologue
	so every gelu's bias is resident from t=0) with the x-prestage applied. The schedule was
	rebuilt from the device profile (external contribution by jlongen-aws, PR #2 on
	``jburtoft/vjepa2-neuron-kernels``). Four design choices carry the performance; the
	instruction order and numerics are unchanged, so the output is bit-identical to the
	previous version:

	* x BEFORE weights: row-block 0's ``x`` load + transpose is issued first, then W1, then the
	later row blocks' ``x``, then W2. The first matmul no longer waits on ~16 MB of weights it
	does not need yet (the first fc1 matmul used to retire at ~57 us).
	* b1 in ONE DMA on the Scalar engine's HW-DGE queue: ``b1`` is host-packed to ``[128, H/128]``
	(:func:`pack_b1`) and lands in a single ~16 KB DMA instead of ``n_h`` ``[128, 1]`` fp32 DMAs
	that trickled in and stalled each gelu (and, through its PSUM bank, the fc1 matmuls behind
	it). Routing it on ``engine=scalar`` HW-DGE keeps gelu's wait off the W1/W2 traffic.
	* row-block 0's ``x`` is transposed on the TENSOR engine (``nc_transpose``) in the prologue,
	while the Tensor engine is otherwise dark in the pre-K4 seam. Later row blocks keep
	``dma_transpose`` (they stream under rb0's matmuls).
	* x prestaged in the prologue: EVERY row block's ``x`` staging is issued before the ``rb``
	loop, in row-block order, so a later block's serialised DMA chain runs inside the pre-K4
	seam instead of in the idle window at the rb boundary. Appending later blocks after rb=0's
	stages means rb=0's readiness cannot regress -- the first matmul cannot be delayed.

	The x staging is emitted through a module-level emitter (``_stage_xt``) so the ``x`` load,
	the ``nc_transpose`` and the ``dma_transpose`` each keep exactly one source site.

	PSUM: fc1 uses 1 bank, fc2 uses up to 4 banks; the ``y1_psum`` / ``out_psums`` tiles are
	declared inside their loops so the ring never accumulates onto a stale PSUM value. The b2
	bias broadcast is kept (removing it regressed the sibling kernel by 23-30%). Weight DMAs stay
	one-per-tile: merging them into a single access-pattern DMA lost ~8 us (one completion signal
	per 8 MB weight).

	This kernel is LNC-parametric (``LNC`` from ``NEURON_LOGICAL_NC_CONFIG``) and keeps the
	``min()`` remainder handling, so it is NOT restricted to exact-multiple shapes. INFERENCE
	ONLY: no autograd rule.
	"""

	from __future__ import annotations

	import math
	import os as _os

	import torch
	from torch_neuronx import nki_op

	import nki
	import nki.isa as nisa
	import nki.language as nl

	#: PE array partition limit, and the moving-operand free-axis limit.
	_P_MAX = 128
	_F_MOV_MAX = 512

	_BF16 = nl.bfloat16

	#: logical-core degree the kernels are built for; must match NEURON_LOGICAL_NC_CONFIG.
	LNC = int(_os.environ.get("NEURON_LOGICAL_NC_CONFIG", "2"))

	_INV_SQRT2 = 1.0 / math.sqrt(2.0)
	_INV_SQRT_2PI = 1.0 / math.sqrt(2.0 * math.pi)


	def _div_ceil(n: int, d: int) -> int:
	"""Ceiling division (never write (n + d - 1) // d inline)."""
	return (n + d - 1) // d


	def _kernel_assert(condition: bool, error_text: str):
	"""Assert with an NKI-formatted error message."""
	assert condition, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {error_text}"


	def _stage_xt(x_hbm, r_abs, pr, n_k, K, rb_size, tensor_transpose):
	"""Emit one row block's x staging and return its ``n_k`` transposed bf16 tiles.

	``n_rs`` loads of ``[<=128, K]`` bf16, each followed by ``n_k`` transposes into the row
	block's ``xt_tiles``. The transpose runs on the Tensor engine (``nc_transpose``, via a PSUM
	tile) when ``tensor_transpose`` is set -- used for row block 0 while the Tensor engine is
	dark in the pre-K4 seam -- and on the DMA engine (``dma_transpose``) otherwise. Kept as a
	module-level emitter so the ``x`` load and each transpose keep a single source site.

	Args:
	x_hbm: [R, K] @ HBM bf16, the MLP input.
	r_abs: absolute starting row of this row block.
	pr: number of rows in this row block (<= rb_size).
	n_k: number of K tiles (fc1 contraction dimension tiles).
	K: hidden-in dimension.
	rb_size: row-block width (free axis of the transposed tiles).
	tensor_transpose: transpose on the Tensor engine (True) or the DMA engine (False).

	Returns:
	list of ``n_k`` transposed ``[128, rb_size]`` bf16 SBUF tiles.
	"""
	n_rs = _div_ceil(pr, _P_MAX)
	xt_tiles = []
	for _ in range(n_k):
	xt_tiles.append(nl.ndarray((_P_MAX, rb_size), dtype=_BF16, buffer=nl.sbuf))
	for rs in range(n_rs):
	rr0 = rs * _P_MAX
	prs = min(_P_MAX, pr - rr0)
	x_b16 = nl.ndarray((_P_MAX, K), dtype=_BF16, buffer=nl.sbuf)
	nisa.dma_copy(
	dst=x_b16[0:prs, 0:K],
	src=x_hbm[r_abs + rr0:r_abs + rr0 + prs, 0:K],
	)
	for kt in range(n_k):
	k0 = kt * _P_MAX
	pk = min(_P_MAX, K - k0)
	if tensor_transpose:
	# Transpose on the (idle) Tensor engine via PSUM, then evict to the xt tile.
	tp = nl.ndarray((_P_MAX, _P_MAX), dtype=_BF16, buffer=nl.psum)
	nisa.nc_transpose(dst=tp[0:pk, 0:prs], data=x_b16[0:prs, k0:k0 + pk],
	engine=nisa.engine.tensor)
	nisa.tensor_copy(dst=xt_tiles[kt][0:pk, rr0:rr0 + prs], src=tp[0:pk, 0:prs])
	else:
	nisa.dma_transpose(dst=xt_tiles[kt][0:pk, rr0:rr0 + prs],
	src=x_b16[0:prs, k0:k0 + pk])
	return xt_tiles


	@nki.jit
	def _mlp_k4_xp_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm):
	"""gelu(x @ w1.T + b1) @ w2.T + b2, rows sharded, b1 AND both x blocks resident early.

	Args:
	x_hbm (nl.ndarray): [R, K] @ HBM, bfloat16 (kernel 1 emits bf16).
	w1t_hbm (nl.ndarray): [K, H] @ HBM, bf16, == fc1.weight.T.
	b1_hbm (nl.ndarray): [128, H/128] @ HBM, float32, == pack_b1(fc1 bias).
	w2t_hbm (nl.ndarray): [H, K] @ HBM, bf16, == fc2.weight.T.
	b2_hbm (nl.ndarray): [K] @ HBM, float32.

	Returns:
	nl.ndarray: [R, K] @ HBM, bfloat16.
	"""
	R, K = x_hbm.shape
	K_w1t, H = w1t_hbm.shape
	H_w2t, K_w2t = w2t_hbm.shape
	_kernel_assert(R >= 1 and K >= 1 and H >= 1, f"R({R}) K({K}) H({H}) must be >= 1")
	_kernel_assert(K_w1t == K, f"w1t in-dim {K_w1t} != K {K}; expected w1t = w1.T of shape [K, H]")
	_kernel_assert(K_w2t == K, f"w2t out-dim {K_w2t} != K {K}; expected w2t = w2.T of shape [H, K]")
	_kernel_assert(H_w2t == H, f"w2t in-dim {H_w2t} != H {H}")
	_kernel_assert(H % _P_MAX == 0, f"H={H} must be a multiple of {_P_MAX} for the packed b1 layout")
	_kernel_assert(tuple(b1_hbm.shape) == (_P_MAX, H // _P_MAX),
	f"b1 must be pack_b1(fc1 bias): [{_P_MAX}, {H // _P_MAX}], got {tuple(b1_hbm.shape)}")
	_kernel_assert(
	R % LNC == 0,
	f"the {LNC}-core kernel needs R divisible by {LNC}, got R={R}. "
	f"program_id is a symbolic value, so a partial final row block cannot be "
	f"clamped; use the one-core pre-transposed kernel for uneven R",
	)

	rows_per = R // LNC
	pid = nl.program_id(0)
	r_base = pid * rows_per

	out_hbm = nl.ndarray((R, K), dtype=_BF16, buffer=nl.shared_hbm)

	n_k = _div_ceil(K, _P_MAX) # K tiles (fc1 contraction)
	n_h = _div_ceil(H, _P_MAX) # H tiles (fc1 out-rows / fc2 contraction)

	# b1 FIRST, in ONE DMA on the Scalar engine's HW-DGE queue: the packed [128, H/128] bias
	# lands in a single ~16 KB transfer so every gelu's bias is resident long before the first
	# matmul retires, and gelu's wait no longer counts the W1/W2 traffic. b1_sb[:, ht] is the
	# bias column for H tile ht (b1_sb[p, ht] == fc1_bias[ht*128 + p]).
	b1_sb = nl.ndarray((_P_MAX, n_h), dtype=nl.float32, buffer=nl.sbuf)
	nisa.dma_copy(dst=b1_sb, src=b1_hbm, dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar)

	# Main fused loop geometry (needed to size the x prestage).
	rb_size = _F_MOV_MAX
	n_rb = _div_ceil(rows_per, rb_size)
	n_ko = _div_ceil(K, _F_MOV_MAX)

	# x BEFORE weights: stage row block 0's x first (Tensor-engine transpose, filling the dark
	# pre-K4 seam), then W1, then later row blocks' x (DMA transpose, streaming under rb0's
	# matmuls), then W2. Appending later blocks after rb=0's stages leaves rb=0's readiness fixed.
	xt_sets = [None] * n_rb
	xt_sets[0] = _stage_xt(x_hbm, r_base, min(rb_size, rows_per), n_k, K, rb_size,
	tensor_transpose=True)

	# W1 staging (pre-transposed bf16, one DMA per K tile).
	w1t_tiles = []
	for kt in range(n_k):
	k0 = kt * _P_MAX
	pk = min(_P_MAX, K - k0)
	t = nl.ndarray((_P_MAX, H), dtype=_BF16, buffer=nl.sbuf)
	nisa.dma_copy(dst=t[0:pk, :], src=w1t_hbm[k0:k0 + pk, :])
	w1t_tiles.append(t)

	# Later row blocks' x via DMA transpose, under rb0's matmuls.
	for rb_pre in range(1, n_rb):
	r_pre = rb_pre * rb_size
	xt_sets[rb_pre] = _stage_xt(x_hbm, r_base + r_pre, min(rb_size, rows_per - r_pre),
	n_k, K, rb_size, tensor_transpose=False)

	# W2 staging (one DMA per H tile, so fc2 can start on the first strips).
	w2t_tiles = []
	for ht in range(n_h):
	h0 = ht * _P_MAX
	ph = min(_P_MAX, H - h0)
	t = nl.ndarray((_P_MAX, K), dtype=_BF16, buffer=nl.sbuf)
	nisa.dma_copy(dst=t[0:ph, :], src=w2t_hbm[h0:h0 + ph, :])
	w2t_tiles.append(t)

	for rb in range(n_rb):
	r_local = rb * rb_size
	pr = min(rb_size, rows_per - r_local)
	n_rs = _div_ceil(pr, _P_MAX)
	r_abs = r_base + r_local

	# x_T bf16 tiles: consume the prologue-staged set for this row block.
	xt_tiles = xt_sets[rb]

	# ---- fc1: gt_tiles[ht] = gelu(w1 @ x.T + b1), bf16, SBUF-resident ----
	gt_tiles = []
	for ht in range(n_h):
	h0 = ht * _P_MAX
	ph = min(_P_MAX, H - h0)
	# Declared INSIDE the ht loop, so the ring never reuses a written PSUM tile.
	# nc_matmul's default accumulate=None auto-detects "previously written" and
	# would accumulate onto a stale value.
	y1_psum = nl.ndarray((_P_MAX, rb_size), dtype=nl.float32, buffer=nl.psum)
	for kt in nl.affine_range(n_k):
	k0 = kt * _P_MAX
	pk = min(_P_MAX, K - k0)
	nisa.nc_matmul(
	dst=y1_psum[0:ph, 0:pr],
	stationary=w1t_tiles[kt][0:pk, h0:h0 + ph],
	moving=xt_tiles[kt][0:pk, 0:pr],
	)
	# b1 comes from the prologue-resident packed tile, column ht.
	gt = nl.ndarray((_P_MAX, rb_size), dtype=_BF16, buffer=nl.sbuf)
	nisa.activation(dst=gt[0:ph, 0:pr], op=nl.gelu,
	data=y1_psum[0:ph, 0:pr], bias=b1_sb[0:ph, ht:ht + 1])
	gt_tiles.append(gt)

	# ---- fc2: out[r,k] = sum_h g[r,h] * w2[k,h] + b2[k] ----
	for ko in range(n_ko):
	ko0 = ko * _F_MOV_MAX
	nk = min(_F_MOV_MAX, K - ko0)
	out_psums = []
	for _ in range(n_rs):
	out_psums.append(
	nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.psum)
	)
	for ht in nl.affine_range(n_h):
	h0 = ht * _P_MAX
	ph = min(_P_MAX, H - h0)
	for rs in range(n_rs):
	rr0 = rs * _P_MAX
	prs = min(_P_MAX, pr - rr0)
	nisa.nc_matmul(
	dst=out_psums[rs][0:prs, 0:nk],
	stationary=gt_tiles[ht][0:ph, rr0:rr0 + prs],
	moving=w2t_tiles[ht][0:ph, ko0:ko0 + nk],
	)
	# b2 broadcast, hoisted out of the rs loop (removing the broadcast itself
	# cost 23% and 30% in two separate attempts).
	b2_bc = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.sbuf)
	nisa.dma_copy(dst=b2_bc[0:_P_MAX, 0:nk],
	src=b2_hbm.ap(pattern=[[0, _P_MAX], [1, nk]], offset=ko0))
	for rs in range(n_rs):
	rr0 = rs * _P_MAX
	prs = min(_P_MAX, pr - rr0)
	out_sb = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=_BF16, buffer=nl.sbuf)
	nisa.tensor_tensor(dst=out_sb[0:prs, 0:nk],
	data1=out_psums[rs][0:prs, 0:nk],
	data2=b2_bc[0:prs, 0:nk],
	op=nl.add)
	nisa.dma_copy(
	dst=out_hbm[r_abs + rr0:r_abs + rr0 + prs, ko0:ko0 + nk],
	src=out_sb[0:prs, 0:nk],
	)

	return out_hbm


	def pack_b1(b1: torch.Tensor) -> torch.Tensor:
	"""Pack fc1 bias ``[H]`` fp32 into ``[128, H/128]`` so ``b1p[p, ht] == b1[ht*128 + p]``.

	Lets the kernel load the whole bias in one DMA instead of ``H/128`` column DMAs. ``H`` must
	be a multiple of 128.
	"""
	H = b1.shape[0]
	assert H % _P_MAX == 0, f"pack_b1 needs H % {_P_MAX} == 0, got H={H}"
	return b1.reshape(-1, _P_MAX).t().contiguous()


	def _no_grad(*tensors) -> None:
	"""Refuse gradient-tracking inputs -- this op registers no autograd rule."""
	for t in tensors:
	if t is not None and t.requires_grad:
	raise RuntimeError(
	"nki_mlp_k4_xp registers no autograd rule (INFERENCE ONLY). "
	"Use under torch.no_grad()/eval()."
	)


	@nki_op("vit_enc::nki_mlp_k4_xp", mutates_args={})
	def nki_mlp_k4_xp(
	x: torch.Tensor,
	w1t: torch.Tensor,
	b1: torch.Tensor,
	w2t: torch.Tensor,
	b2: torch.Tensor,
	) -> torch.Tensor:
	"""Fused MLP on LNC=2, pre-transposed bf16 weights, x prestaged in the prologue.

	Args:
	x : [R, K] bf16 on 'neuron'. ``R`` must be divisible by :data:`LNC`.
	w1t : [K, H] bf16 == fc1.weight.T.
	b1 : fc1 bias, either ``[H]`` float32 (packed on the fly) or pre-packed
	``[128, H/128]`` from :func:`pack_b1`.
	w2t : [H, K] bf16 == fc2.weight.T. b2 : [K] float32.

	Returns:
	[R, K] bf16.
	"""
	_no_grad(x, w1t, b1, w2t, b2)
	if b1.dim() == 1: # accept the pre-pack [H] bias; pack it on the fly
	b1 = pack_b1(b1)
	return _mlp_k4_xp_kernel[LNC](x, w1t, b1, w2t, b2)


	#: Requested export alias. Same object as the incumbent-named nki_op, so the runtime
	#: op-name string ("vit_enc::nki_mlp_k4_xp") is unchanged whichever name is used.
	mlp = nki_mlp_k4_xp