Fix and support scalar and vector scale/bias symmetrically in Activation

[RuifengHua]

Summary

blas.Activation (and the compact blas.activation) computes dst = op(src * scale + bias). scale and bias are both per-partition terms that nisa.activation broadcasts (P, 1) -> (P, F) (each partition's single value is applied across that partition's whole row). The instruction supports, for either term, a scalar or a (P, 1) vector — and on NeuronCore-v3 a (P, 1) vector may span multiple partition tiles. The primitive, however, handled scale/bias with an under-developed, asymmetric path that under-delivered against the hardware in three ways:

1. A raw (P, 1) tensor / TensorView scale crashed.
Scale of type nl.ndarray / TensorView was passed to nisa.activation unchanged. But a tensor allocated via alloc_logical((P, 1), ...) has the container layout (pdim, n_p_tiles, F) = (P, 1, 1) — 3D, carrying the internal n_p_tiles dimension — whereas nisa.activation expects a 2D (P, 1) view. Handed the 3D container, the broadcast produces the wrong element count and raises ValueError: shape mismatch. Only a TileStream scale worked, because TileStream.get_tile().get_view() strips the n_p_tiles dimension down to the 2D (P, 1) view. So the natural thing a caller has (a raw alloc_logical tensor) was exactly the thing that failed — despite the type hint advertising nl.ndarray / TensorView as accepted.

2. A multi-tile (P > pmax) vector scale was rejected.
The scale was resolved once, before the tile loop, and __init__ asserted the scale TileStream was a single tile. So a (P, 1) scale with P > pmax (which must tile along the partition loop like dst/src do) hit the assert and was refused — even though bias, an identically-shaped (P, 1) term, was fetched inside the loop and tiled correctly. The scale path simply hadn't been generalized to the multi-tile case the way bias had.

3. bias could not be a scalar.
scale accepted a scalar float, but bias was typed and handled as a tensor/TileStream only. nisa.activation accepts a scalar bias on NeuronCore-v3, so this was a gratuitous asymmetry — a caller wanting op(src * scale + c) for a constant c could not express it directly.

This PR makes scale and bias first-class and symmetric — each may be None, a scalar float, or a (P, 1) vector (raw nl.ndarray / TensorView / TileStream), single- or multi-tile — by handling both through one path: normalize a raw tensor / TensorView into a TileStream (resolving the correct 2D (P, 1) per-tile view), resolve each term per tile inside the loop (so a (P, 1) vector advances in lockstep with dst/src, enabling P > pmax), accept a scalar for either term, and validate a vector term is a (P, 1) vector whose partition size and tile count match dst (a clear boundary error instead of a cryptic downstream one).

A (P, 1) per-partition scale/bias is a common ML operation — per-row / per-channel dequantization and RMSNorm-style scaling — and P routinely exceeds pmax (long sequences, large hidden dims), which is exactly the multi-tile case (2) unblocks.

One hardware asymmetry remains and is left to nisa.activation to enforce (documented, not guarded): a vector scale must be float32, whereas a vector bias may be any supported dtype.

Repro (before fix)

A (P, 1) tensor scale — the natural thing a caller has from alloc_logical — crashed:

import nki, nki.language as nl
import numpy as np
from nkilib_src.nkilib.experimental.primitives import blas, dma, tile_stream

@nki.jit
def k(x, s):  # s is a (P, 1) per-partition scale (e.g. a per-row dequant scale)
    p, f = x.shape
    y = nl.ndarray((p, f), dtype=x.dtype, buffer=nl.shared_hbm)
    x_sb = tile_stream.alloc_logical((p, f), p, x.dtype, "x")
    s_sb = tile_stream.alloc_logical((p, 1), p, nl.float32, "s")
    dma.load(x_sb, x)
    dma.load(s_sb, s)
    blas.activation(x_sb, op=nl.copy, scale=s_sb)  # (P, 1) tensor scale
    dma.store(y, x_sb)
    return y

x = np.ones((4, 8), dtype=np.float32)
s = np.arange(1, 5, dtype=np.float32).reshape(4, 1)  # per-row scales 1, 2, 3, 4
out = nki.simulate(k)(x, s)
# before fix: the raw (P, 1) scale reaches nisa.activation as its 3D alloc_logical
#             container (P, 1, 1) instead of a 2D (P, 1) view, so the broadcast
#             produces the wrong element count -> ValueError: shape mismatch
# after fix:  each row is scaled by its own value -> out[i, :] == s[i]
#             (rows come out all-1, all-2, all-3, all-4)

Test plan

Rewrote test/integration/nkilib/experimental/primitives/blas/test_activation.py, verified in simulation mode (--test-mode simulation):

• Compact API fast smoke tier — plain, scalar/vector/mixed scale+bias, a needs-positive op (rsqrt), and in-place
• Compact API pairwise sweep — op x scale-kind x bias-kind x inplace x P x F x dtype (P <= pmax), covering every scale/bias type pairing including mixes (e.g. scalar scale + vector bias)
• Class API multi-tile (P > pmax, incl. partial last tiles) — per-partition (P, 1) vector scale and/or bias tiled in lockstep with the data; verified each row (including the partial last tile) gets its own scale/bias value
• Shape guard — a scale that is not a (P, 1) vector matching the data partition dim is rejected at the primitive boundary
• Confirmed the multi-tile vector-scale and shape-guard tests fail against the pre-fix code and pass after the fix

Note: validated in simulation only. Tolerances: 1e-4 (fp32) / 1e-2 (bf16).

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.