This document is relevant for: Trn2, Trn3

nki.isa.range_select#

nki.isa.range_select(dst: ~nki.language.NkiTensor, on_true_tile: ~nki.language.NkiTensor, comp_op0, comp_op1, bound0: ~nki.language.NkiTensor, bound1: ~nki.language.NkiTensor, reduce_cmd=reduce_cmd.reset_reduce, reduce_res: ~typing.Optional[~nki.language.NkiTensor] = None, reduce_op=<function maximum>, range_start=0, on_false_value=Ellipsis, name=None)[source]#

Select elements from on_true_tile based on comparison with bounds using Vector Engine.

Note

Available only on NeuronCore-v3 and newer.

For each element in on_true_tile, compares its free dimension index + range_start against bound0 and bound1 using the specified comparison operators (comp_op0 and comp_op1). If both comparisons evaluate to True, copies the element to the output; otherwise uses on_false_value.

Additionally performs a reduction operation specified by reduce_op on the results, storing the reduction result in reduce_res.

Note on numerical stability:

In self-attention, we often have this instruction sequence: range_select (VectorE) -> reduce_res -> activation (ScalarE). When range_select outputs a full row of fill_value, caution is needed to avoid NaN in the activation instruction that subtracts the output of range_select by reduce_res (max value):

  • If dst.dtype and reduce_res.dtype are both FP32, we should not hit any NaN issue since FP32_MIN - FP32_MIN = 0. Exponentiation on 0 is stable (1.0 exactly).

  • If dst.dtype is FP16/BF16/FP8, the fill_value in the output tile will become -INF since HW performs a downcast from FP32_MIN to a smaller dtype. In this case, you must make sure reduce_res.dtype is FP32 to avoid NaN in activation. NaN can be avoided because activation always upcasts input tiles to FP32 to perform math operations: -INF - FP32_MIN = -INF. Exponentiation on -INF is stable (0.0 exactly).

Constraints:

The comparison operators must be one of:

  • nl.equal

  • nl.less

  • nl.less_equal

  • nl.greater

  • nl.greater_equal

Partition dim sizes must match across on_true_tile, bound0, and bound1:

  • bound0 and bound1 must have one element per partition

  • on_true_tile must be one of the FP dtypes, and bound0/bound1 must be FP32 types.

The comparison with bound0, bound1, and free dimension index is done in FP32. Make sure range_start + free dimension index is within 2^24 range.

Numpy equivalent:

indices = np.zeros_like(on_true_tile, dtype=np.float32)
indices[:] = range_start + np.arange(on_true_tile[0].size)

mask = comp_op0(indices, bound0) & comp_op1(indices, bound1)
select_out_tile = np.where(mask, on_true_tile, on_false_value)
reduce_tile = reduce_op(select_out_tile, axis=1, keepdims=True)
Parameters:

  • dst – output tile with selected elements

  • on_true_tile – input tile containing elements to select from

  • on_false_value – constant value to use when selection condition is False. Due to hardware constraints, this must be FP32_MIN (-3.4028235e+38). See the numerical stability note above for guidance on output dtype selection.

  • comp_op0 – first comparison operator

  • comp_op1 – second comparison operator

  • bound0 – tile with one element per partition for first comparison

  • bound1 – tile with one element per partition for second comparison

  • reduce_op – reduction operator to apply on across the selected output. Currently only nl.maximum is supported.

  • reduce_cmd – controls the state of the Vector Engine accumulator registers. Defaults to reduce_cmd.reset_reduce. See nisa.reduce_cmd for supported values.

  • reduce_res –

    optional tile to store reduction results.

    Pass None to keep the reduction result in the Vector Engine’s internal accumulator without writing it out. This is useful when chaining multiple calls that reduce into the same accumulator — only the final call needs to pass a tile to retrieve the accumulated result.

  • range_start – starting base offset for index array for the free dimension of on_true_tile. Defaults to 0, and must be a compile-time integer.

This document is relevant for: Trn2, Trn3