Clarify on why TRITON_MAX_TILE_NUMEL limits (2**17)

jafraustro · August 12, 2026, 12:05am

What is the reason for the max numel limit 2 **17 number? pytorch/test/test_sparse_csr.py at 44fae102774216bbba8df91f140bba54a30e5f62 · pytorch/pytorch · GitHub

github.com/pytorch/pytorch

torch/sparse/_triton_ops.py

44fae1027


      
              crow_indices[..., :-1]: (0, None, -1),
              values: (0, None, None),
          }
          
          def kernel(grid, *sliced_tensors):
              _bsr_softmax_kernel[grid](
                  *ptr_stride_extractor(*sliced_tensors),
                  row_block,
                  col_block,
                  max_row_nnz,
                  # Triton's max numel is bounded by 2 ** 17.
                  min(2**17, max_row_nnz),
              )
          
          launch_kernel(kernel, tensor_dims_map, full_grid, grid_blocks)
          
          values = (
              values.reshape(-1, row_block, nnz, col_block)
              .transpose(-3, -2)
              .reshape(*input.values().shape)
          )

            res_tri = bsr_softmax(bsr)
            res_coo = torch.sparse.softmax(coo, -1)
            self.assertEqual(res_tri, res_coo.to(input.dtype))

        # Test long rows which exceed Triton's max numel limit set to 2 ** 17
        input = tensor(b + (1, 150000))
        bsr = input.to_sparse_bsr(1)
        self.assertEqual(input.softmax(-1), bsr_softmax(bsr))

Should this number be device specific rather than hardcoded?
Is it tied to a specific hardware constraint (e.g., shared memory size, max block size)?