From 3af8063bf143b6acec21c3f1de0b7f2654d878c9 Mon Sep 17 00:00:00 2001 From: Nikhil Mehta <51523958+starskyreverie@users.noreply.github.com> Date: Sun, 22 Sep 2024 15:57:21 -0700 Subject: [PATCH 1/2] remove markdown formatting in ScalarE activation_accum code block just for extra clarity when rendered since it's in a code block :) can be confused w/ exponentiation, multiplies, etc otherwise --- general/nki/trainium_inferentia2_arch.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/general/nki/trainium_inferentia2_arch.rst b/general/nki/trainium_inferentia2_arch.rst index fd36ba46..56052cf9 100644 --- a/general/nki/trainium_inferentia2_arch.rst +++ b/general/nki/trainium_inferentia2_arch.rst @@ -439,13 +439,13 @@ in a pipeline fashion. Mathematically, ScalarE implements: # Output: 2D out_tile for lane_id in range(in_tile.shape[0]): for k in range(in_tile.shape[1]) - out_tile[lane_id][k] = func(in_tile[lane_id][k] * **scale****[lane_id]** + out_tile[lane_id][k] = func(in_tile[lane_id][k] * scale[lane_id] + bias[lane_id]) # Case 2: scale is a compile constant in the instruction for lane_id in range(in_tile.shape[0]): for k in range(in_tile.shape[1]) - out_tile[lane_id][k] = func(in_tile[lane_id][k] * **scale** + out_tile[lane_id][k] = func(in_tile[lane_id][k] * scale + bias[lane_id]) This functionality can be invoked using the :doc:`nki.isa.activation ` From 0b2735a365544633dc304ddef5e9e5394bf6024f Mon Sep 17 00:00:00 2001 From: Nikhil Mehta <51523958+starskyreverie@users.noreply.github.com> Date: Sun, 22 Sep 2024 15:58:27 -0700 Subject: [PATCH 2/2] correctly format ScalarE activation accum code --- general/nki/trainium_inferentia2_arch.rst | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/general/nki/trainium_inferentia2_arch.rst b/general/nki/trainium_inferentia2_arch.rst index 56052cf9..47887100 100644 --- a/general/nki/trainium_inferentia2_arch.rst +++ b/general/nki/trainium_inferentia2_arch.rst @@ -438,14 +438,14 @@ in a pipeline fashion. Mathematically, ScalarE implements: # Input: 2D in_tile, 1D scale, 1D bias # Output: 2D out_tile for lane_id in range(in_tile.shape[0]): - for k in range(in_tile.shape[1]) - out_tile[lane_id][k] = func(in_tile[lane_id][k] * scale[lane_id] + for k in range(in_tile.shape[1]): + out_tile[lane_id][k] = func(in_tile[lane_id][k] * scale[lane_id] + bias[lane_id]) # Case 2: scale is a compile constant in the instruction for lane_id in range(in_tile.shape[0]): - for k in range(in_tile.shape[1]) - out_tile[lane_id][k] = func(in_tile[lane_id][k] * scale + for k in range(in_tile.shape[1]): + out_tile[lane_id][k] = func(in_tile[lane_id][k] * scale + bias[lane_id]) This functionality can be invoked using the :doc:`nki.isa.activation `