Update rotary_nki_kernels.py

apatwal72 · web-flow · commit 352e6f7174b2 · 2025-03-19T14:32:50.000-04:00
specify indicies for other dims when storing tiles
diff --git a/src/nki_samples/tutorials/rotary/rotary_nki_kernels.py b/src/nki_samples/tutorials/rotary/rotary_nki_kernels.py
@@ -312,12 +312,12 @@ def nki_apply_rotary_embedding(q, k, cos, sin):
 
         nl.store(
             output_q_hbm_tile[seq_batch_id * nl.tile_size.pmax + i_p, i_f],
-            output_tile[0],
+            output_tile[0, :, :],
             mask=(seq_batch_id * nl.tile_size.pmax + i_p < seq_len),
         )
         nl.store(
             output_k_hbm_tile[seq_batch_id * nl.tile_size.pmax + i_p, i_f],
-            output_tile[1],
+            output_tile[1, :, :],
             mask=(seq_batch_id * nl.tile_size.pmax + i_p < seq_len),
         )
 

Original file line number	Diff line number	Diff line change
`@@ -312,12 +312,12 @@ def nki_apply_rotary_embedding(q, k, cos, sin):`
`312`	`312`
`313`	`313`	`nl.store(`
`314`	`314`	`output_q_hbm_tile[seq_batch_id * nl.tile_size.pmax + i_p, i_f],`
`315`		`- output_tile[0],`
	`315`	`+ output_tile[0, :, :],`
`316`	`316`	`mask=(seq_batch_id * nl.tile_size.pmax + i_p < seq_len),`
`317`	`317`	`)`
`318`	`318`	`nl.store(`
`319`	`319`	`output_k_hbm_tile[seq_batch_id * nl.tile_size.pmax + i_p, i_f],`
`320`		`- output_tile[1],`
	`320`	`+ output_tile[1, :, :],`
`321`	`321`	`mask=(seq_batch_id * nl.tile_size.pmax + i_p < seq_len),`
`322`	`322`	`)`
`323`	`323`