Fix a bug in argmax function

trident/kernel/argmax.py

@@ -18,6 +18,7 @@
 
 class Argmax:
     @staticmethod
+    @triton.heuristics({"require_x_boundary_check": lambda args: args["x_size"] % args["x_block_size"]})
     @triton.jit
     def forward(
         output_ptr: tl.tensor,
@@ -28,8 +29,10 @@ def forward(
         x_stride: tl.int32,
         dtype: tl.constexpr,
         x_block_size: tl.constexpr,
+        require_x_boundary_check: tl.constexpr,
     ):
         y_offset = tl.program_id(0)
+
         output_block_ptr = tl.make_block_ptr(
             output_ptr,
             shape=(y_size,),
@@ -46,6 +49,13 @@ def forward(
             block_shape=(1, x_block_size),
             order=(1, 0),
         )
-        input = tl.load(input_block_ptr, boundary_check=(1, 0), padding_option="zero")
+
+        if require_x_boundary_check:
+            input = tl.load(input_block_ptr, boundary_check=(1,))
+            condition = tl.arange(0, x_block_size) < x_size
+            input = tl.where(condition, input, float("-inf"))
+        else:
+            input = tl.load(input_block_ptr)
+
         output = tl.argmax(input, 1)
         tl.store(output_block_ptr, output.to(dtype))