Update 02-fused-softmax.py

anmyachev · web-flow · commit ba1254e727e3 · 2025-06-03T19:34:05.000+02:00
diff --git a/python/tutorials/02-fused-softmax.py b/python/tutorials/02-fused-softmax.py
@@ -126,7 +126,7 @@ def softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride, n
 max_num_resident_warps = NUM_SM * warps_per_sm
 kernels = {}
 # Possible SLM allocation sizes in kB
-tg_slm_sizes = [2**i for i in [0, 1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128]]  # TODO: Get from properties
+tg_slm_sizes = [i * 2**10 for i in [0, 1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128]]  # TODO: Get from properties
 
 
 def softmax(x):
