diff --git a/problems/pmpp/conv2d_py/reference.py b/problems/pmpp/conv2d_py/reference.py
index 0eb48e6f8..1876e9e2f 100644
--- a/problems/pmpp/conv2d_py/reference.py
+++ b/problems/pmpp/conv2d_py/reference.py
@@ -1,14 +1,13 @@
 from utils import verbose_allclose
 import torch
 import torch.nn.functional as F
-from task import input_t, output_t, KernelSpec
+from task import input_t, output_t
 
-def ref_kernel(data: input_t, spec: KernelSpec) -> output_t:
+def ref_kernel(data: input_t) -> output_t:
     """
     Reference implementation of 2D convolution using PyTorch.
     Args:
         data: Tuple of (input tensor, kernel tensor)
-        spec: Convolution specifications (stride, padding)
     Returns:
         Output tensor after convolution
     """
@@ -16,11 +15,14 @@ def ref_kernel(data: input_t, spec: KernelSpec) -> output_t:
     return F.conv2d(
         input_tensor, 
         kernel,
-        stride=spec.stride,
-        padding=spec.padding
+
+        # No padding and no striding
+        # TODO: Can revisit this in future problems
+        stride=1,
+        padding=0
     )
 
-def generate_input(size: int, kernel_size: int, channels: int, batch: int, seed: int) -> input_t:
+def generate_input(size: int, kernelsize: int, channels: int, batch: int, seed: int) -> input_t:
     """
     Generates random input and kernel tensors.
     Returns:
@@ -40,7 +42,7 @@ def generate_input(size: int, kernel_size: int, channels: int, batch: int, seed:
     # Generate kernel tensor: [out_channels, in_channels, kernel_height, kernel_width]
     # Here we use same number of output channels as input channels for simplicity
     kernel = torch.randn(
-        channels, channels, kernel_size, kernel_size,
+        channels, channels, kernelsize, kernelsize,
         device='cuda',
         dtype=torch.float32,
         generator=gen
@@ -50,10 +52,9 @@ def generate_input(size: int, kernel_size: int, channels: int, batch: int, seed:
 
 def check_implementation(
     data: input_t,
-    spec: KernelSpec,
     output: output_t,
 ) -> str:
-    expected = ref_kernel(data, spec)
+    expected = ref_kernel(data)
     reasons = verbose_allclose(output, expected, rtol=1e-3, atol=1e-3)
     
     if len(reasons) > 0:
diff --git a/problems/pmpp/conv2d_py/submission.py b/problems/pmpp/conv2d_py/submission.py
new file mode 100644
index 000000000..a1b7d16d4
--- /dev/null
+++ b/problems/pmpp/conv2d_py/submission.py
@@ -0,0 +1,21 @@
+from task import input_t, output_t
+import torch
+import torch.nn.functional as F
+
+
+def custom_kernel(data: input_t) -> output_t:
+    """
+    Implementation of 2D convolution using PyTorch with no padding and no striding.
+    Args:
+        data: Tuple of (input tensor, kernel tensor)
+        spec: Convolution specifications
+    Returns:
+        Output tensor after convolution
+    """
+    input_tensor, kernel = data
+    return F.conv2d(
+        input_tensor, 
+        kernel,
+        stride=1,
+        padding=0
+    )
\ No newline at end of file
diff --git a/problems/pmpp/conv2d_py/task.py b/problems/pmpp/conv2d_py/task.py
index 6cce0e6eb..397332abd 100644
--- a/problems/pmpp/conv2d_py/task.py
+++ b/problems/pmpp/conv2d_py/task.py
@@ -1,18 +1,13 @@
 from typing import TypedDict, TypeVar, Tuple
 import torch
-from dataclasses import dataclass
 
 input_t = TypeVar("input_t", bound=Tuple[torch.Tensor, torch.Tensor])
 output_t = TypeVar("output_t", bound=torch.Tensor)
 
-@dataclass
-class KernelSpec:
-    stride: int
-    padding: int
 
 class TestSpec(TypedDict):
     size: int
-    kernel_size: int
+    kernelsize: int
     channels: int
     batch: int
-    seed: int 
\ No newline at end of file
+    seed: int   
\ No newline at end of file
diff --git a/problems/pmpp/conv2d_py/task.yml b/problems/pmpp/conv2d_py/task.yml
index 5b0955ede..0f8d075ae 100644
--- a/problems/pmpp/conv2d_py/task.yml
+++ b/problems/pmpp/conv2d_py/task.yml
@@ -11,21 +11,30 @@ lang: "py"
 
 description: |
   Implement a 2D convolution kernel that matches the reference implementation.
-  The kernel should perform 2D convolution with the given specifications (stride and padding).
+  The kernel should perform 2D convolution with the given specifications
+  We will benchmark different sizes, kernel sizes, channels and batch sizes but they will all 
+  be even numbers with the exception of batch size which can sometimes be 1
+  We assume no padding and striding and instead vary the size of the input and kernel,
+  number of channels, and batch size.
+
+  Input: Tuple of (input_tensor, kernel)
+    - input_tensor: 4D tensor of shape (batch, channels, height, width) with arbitrary values
+    - kernel: 4D tensor of shape (channels, channels, kernelsize, kernelsize) with arbitrary values
+  Output: 4D tensor of shape (batch, channels, height-kernelsize+1, width-kernelsize+1) with convolved values
 
 config:
   main: "eval.py" 
 
 tests:
-  - {"size": 32, "kernel_size": 3, "channels": 16, "batch": 1, "seed": 4242}
-  - {"size": 32, "kernel_size": 5, "channels": 16, "batch": 2, "seed": 5236}
-  - {"size": 64, "kernel_size": 3, "channels": 32, "batch": 1, "seed": 1001}
-  - {"size": 64, "kernel_size": 5, "channels": 32, "batch": 2, "seed": 5531}
-  - {"size": 128, "kernel_size": 3, "channels": 64, "batch": 1, "seed": 9173}
+  - {"size": 32, "kernelsize": 4, "channels": 16, "batch": 1, "seed": 4242}
+  - {"size": 32, "kernelsize": 4, "channels": 16, "batch": 2, "seed": 5236}
+  - {"size": 64, "kernelsize": 4, "channels": 32, "batch": 1, "seed": 1001}
+  - {"size": 64, "kernelsize": 8, "channels": 32, "batch": 2, "seed": 5531}
+  - {"size": 128, "kernelsize": 8, "channels": 64, "batch": 1, "seed": 9173}
 
 benchmarks:
-  - {"size": 128, "kernel_size": 3, "channels": 64, "batch": 4, "seed": 54352}
-  - {"size": 128, "kernel_size": 5, "channels": 64, "batch": 4, "seed": 93246}
-  - {"size": 256, "kernel_size": 3, "channels": 128, "batch": 2, "seed": 6256}
-  - {"size": 256, "kernel_size": 5, "channels": 128, "batch": 2, "seed": 8841}
-  - {"size": 512, "kernel_size": 3, "channels": 256, "batch": 1, "seed": 6252} 
\ No newline at end of file
+  - {"size": 128, "kernelsize": 8, "channels": 64, "batch": 4, "seed": 54352}
+  - {"size": 128, "kernelsize": 16, "channels": 64, "batch": 4, "seed": 93246}
+  - {"size": 256, "kernelsize": 16, "channels": 128, "batch": 2, "seed": 6256}
+  - {"size": 256, "kernelsize": 16, "channels": 128, "batch": 2, "seed": 8841}
+  - {"size": 512, "kernelsize": 32, "channels": 256, "batch": 1, "seed": 6252} 
\ No newline at end of file
diff --git a/problems/pmpp/grayscale_py/submission.py b/problems/pmpp/grayscale_py/submission.py
new file mode 100644
index 000000000..de0c14945
--- /dev/null
+++ b/problems/pmpp/grayscale_py/submission.py
@@ -0,0 +1,8 @@
+from task import input_t, output_t
+import torch
+
+def custom_kernel(data: input_t) -> output_t:
+    weights = torch.tensor([0.2989, 0.5870, 0.1140], 
+                         device=data.device, 
+                         dtype=data.dtype)
+    return torch.sum(data * weights, dim=-1)
diff --git a/problems/pmpp/grayscale_py/task.yml b/problems/pmpp/grayscale_py/task.yml
index 9ba74756d..62c44e0c1 100644
--- a/problems/pmpp/grayscale_py/task.yml
+++ b/problems/pmpp/grayscale_py/task.yml
@@ -9,7 +9,7 @@ lang: "py"
 
 description: |
   Implement an RGB to grayscale conversion kernel that matches the reference implementation.
-  The kernel should convert RGB images to grayscale using the standard coefficients:
+  The kernel should convert square RGB images with even sizes to grayscale using the standard coefficients:
   Y = 0.2989 R + 0.5870 G + 0.1140 B
   
   Input: RGB tensor of shape (H, W, 3) with values in [0, 1]
@@ -19,15 +19,16 @@ config:
   main: "eval.py" 
 
 tests:
-  - {"size": 127, "seed": 4242}
-  - {"size": 128, "seed": 5236}
-  - {"size": 129, "seed": 1001}
+
+  - {"size": 128, "seed": 1001}
   - {"size": 256, "seed": 5531}
   - {"size": 512, "seed": 9173}
 
 benchmarks:
-  - {"size": 1024, "seed": 54352}
-  - {"size": 2048, "seed": 93246}
-  - {"size": 4096, "seed": 6256}
-  - {"size": 8192, "seed": 8841}
-  - {"size": 16384, "seed": 6252} 
\ No newline at end of file
+  
+  - {"size": 512, "seed": 54352}
+  - {"size": 1024, "seed": 93246}
+  - {"size": 2048, "seed": 6256}
+  - {"size": 4096, "seed": 8841}
+  - {"size": 8192, "seed": 6252}
+  - {"size": 16384, "seed": 54352}
\ No newline at end of file
diff --git a/problems/pmpp/histogram_py/reference.py b/problems/pmpp/histogram_py/reference.py
index 8fb766b88..54d699387 100644
--- a/problems/pmpp/histogram_py/reference.py
+++ b/problems/pmpp/histogram_py/reference.py
@@ -1,47 +1,62 @@
 from utils import verbose_allclose
 import torch
-from task import input_t, output_t, HistogramSpec
+from task import input_t, output_t
 
-def ref_kernel(data: input_t, spec: HistogramSpec) -> output_t:
+def ref_kernel(data: input_t) -> output_t:
     """
     Reference implementation of histogram using PyTorch.
     Args:
-        data: Input tensor to compute histogram on
-        spec: Histogram specifications (num_bins, min_val, max_val)
+        data: tensor of shape (size,)
     Returns:
         Tensor containing bin counts
     """
-    # Clip values to range
-    clipped = torch.clamp(data, spec.min_val, spec.max_val)
+    # Fixed range [0, 100]
+    min_val, max_val = 0, 100
+    
+    # Number of bins is input size / 16
+    num_bins = data.shape[0] // 16
+
+    clipped = torch.clamp(data, min_val, max_val)
     
     # Scale to bin indices
-    bin_width = (spec.max_val - spec.min_val) / spec.num_bins
-    indices = ((clipped - spec.min_val) / bin_width).long()
-    indices = torch.clamp(indices, 0, spec.num_bins - 1)
+    bin_width = (max_val - min_val) / num_bins
+    indices = ((clipped - min_val) / bin_width).long()
+    indices = torch.clamp(indices, 0, num_bins - 1)
     
     # Count values in each bin
-    return torch.bincount(indices, minlength=spec.num_bins).to(torch.float32)
+    return torch.bincount(indices, minlength=num_bins).to(torch.float32)
 
 def generate_input(size: int, seed: int) -> input_t:
     """
-    Generates random input tensor with values roughly in [0, 1].
+    Generates random input tensor for histogram.
+    The number of bins is automatically set to size/16.
+    
+    Args:
+        size: Size of the input tensor (must be multiple of 16)
+        seed: Random seed
     Returns:
-        Tensor to compute histogram on
+        The input tensor with values in [0, 100]
     """
     gen = torch.Generator(device='cuda')
     gen.manual_seed(seed)
-    # Generate values with normal distribution for interesting histograms
-    return torch.randn(size, device='cuda', dtype=torch.float32, generator=gen).contiguous()
+    
+    # Generate integer values between 0 and 100
+    data = torch.randint(0, 101, (size,), device='cuda', dtype=torch.int32, generator=gen)
+    
+    # Convert to float since the histogram implementation expects float input
+    return data.float().contiguous()
 
 def check_implementation(
     data: input_t,
-    spec: HistogramSpec,
     output: output_t,
 ) -> str:
-    expected = ref_kernel(data, spec)
+    """
+    Compare custom implementation's output to the reference output.
+    """
+    expected = ref_kernel(data)
     reasons = verbose_allclose(output, expected)
     
     if len(reasons) > 0:
         return "mismatch found! custom implementation doesn't match reference: " + reasons[0]
     
-    return '' 
\ No newline at end of file
+    return ''
diff --git a/problems/pmpp/histogram_py/submission.py b/problems/pmpp/histogram_py/submission.py
new file mode 100644
index 000000000..338585f7c
--- /dev/null
+++ b/problems/pmpp/histogram_py/submission.py
@@ -0,0 +1,27 @@
+import torch
+from task import input_t, output_t
+
+def custom_kernel(data: input_t) -> output_t:
+    """
+    Reference implementation of histogram using PyTorch.
+    Args:
+        data: tensor of shape (size,)
+    Returns:
+        Tensor containing bin counts
+    """
+    # Fixed range [0, 100]
+    min_val, max_val = 0, 100
+    
+    # Number of bins is input size / 16
+    num_bins = data.shape[0] // 16
+    
+    # Clip values to range
+    clipped = torch.clamp(data, min_val, max_val)
+    
+    # Scale to bin indices
+    bin_width = (max_val - min_val) / num_bins
+    indices = ((clipped - min_val) / bin_width).long()
+    indices = torch.clamp(indices, 0, num_bins - 1)
+    
+    # Count values in each bin
+    return torch.bincount(indices, minlength=num_bins).to(torch.float32)
\ No newline at end of file
diff --git a/problems/pmpp/histogram_py/task.py b/problems/pmpp/histogram_py/task.py
index e9d7fadfc..61fd33787 100644
--- a/problems/pmpp/histogram_py/task.py
+++ b/problems/pmpp/histogram_py/task.py
@@ -1,16 +1,10 @@
 from typing import TypedDict, TypeVar
 import torch
-from dataclasses import dataclass
 
 input_t = TypeVar("input_t", bound=torch.Tensor)
 output_t = TypeVar("output_t", bound=torch.Tensor)
 
-@dataclass
-class HistogramSpec:
-    num_bins: int
-    min_val: float
-    max_val: float
-
 class TestSpec(TypedDict):
     size: int
-    seed: int 
\ No newline at end of file
+    seed: int 
+
diff --git a/problems/pmpp/histogram_py/task.yml b/problems/pmpp/histogram_py/task.yml
index 9d30c0c34..667ad8ea3 100644
--- a/problems/pmpp/histogram_py/task.yml
+++ b/problems/pmpp/histogram_py/task.yml
@@ -10,22 +10,24 @@ files:
 lang: "py"
 
 description: |
-  Implement a histogram kernel using CUDA inline function that matches the reference implementation.
-  The kernel should count the number of elements falling into each bin across the specified range.
+  Implement a histogram kernel that counts the number of elements falling into each bin across the specified range.
+  The minimum and maximum values of the range are fixed to 0 and 100 respectively.
+  All sizes are multiples of 16 and the number of bins is set to the size of the input tensor divided by 16.
+
+  Input:
+    - data: a tensor of shape (size,)
 
 config:
-  main: "eval.py" 
+  main: "eval.py"
 
 tests:
-  - {"size": 1023, "seed": 4242}
-  - {"size": 1024, "seed": 5236}
-  - {"size": 1025, "seed": 1001}
-  - {"size": 2048, "seed": 5531}
-  - {"size": 4096, "seed": 9173}
+  - {"size": 5120,   "seed": 9991}
+  - {"size": 7840,   "seed": 2105}
+  - {"size": 30080,  "seed": 9999}
+  - {"size": 100000, "seed": 1212}
 
 benchmarks:
-  - {"size": 8192, "seed": 54352}
-  - {"size": 16384, "seed": 93246}
-  - {"size": 32768, "seed": 6256}
-  - {"size": 65536, "seed": 8841}
-  - {"size": 131072, "seed": 6252} 
\ No newline at end of file
+  - {"size": 1310720, "seed": 6252}
+  - {"size": 2621440, "seed": 8841}
+  - {"size": 5242880, "seed": 6252}
+  - {"size": 10485760, "seed": 8841}
diff --git a/problems/pmpp/matmul_py/task.yml b/problems/pmpp/matmul_py/task.yml
index 6f14aa5bc..e8cd416e3 100644
--- a/problems/pmpp/matmul_py/task.yml
+++ b/problems/pmpp/matmul_py/task.yml
@@ -12,6 +12,7 @@ lang: "py"
 description: |
   Implement a custom matmul function that matches the reference implementation.
   The function should handle a tuple of input tensors and apply matmul
+  The shapes of all outer and inner dimensions of tensors are multiples of 16
 
 config:
   main: "eval.py"
diff --git a/problems/pmpp/prefixsum_py/reference.py b/problems/pmpp/prefixsum_py/reference.py
index bce90273a..3850a9aaa 100644
--- a/problems/pmpp/prefixsum_py/reference.py
+++ b/problems/pmpp/prefixsum_py/reference.py
@@ -22,14 +22,18 @@ def generate_input(size: int, seed: int) -> input_t:
     gen.manual_seed(seed)
     return torch.randn(size, device='cuda', dtype=torch.float32, generator=gen).contiguous()
 
-def check_implementation(
-    data: input_t,
-    output: output_t,
-) -> str:
+# This algorithm is very sensitive to the tolerance and the error is magnified by the input size
+# The tolerance is scaled by the square root of the input size
+def check_implementation(data: input_t, output: output_t) -> str:
     expected = ref_kernel(data)
-    reasons = verbose_allclose(output, expected, rtol=1e-5, atol=1e-5)
     
+    # Then get the size for scaling the tolerance
+    n = data.numel()
+    
+    scale_factor = n ** 0.5  # Square root of input size
+    rtol = 1e-5 * scale_factor
+    atol = 1e-5 * scale_factor
+    reasons = verbose_allclose(output, expected, rtol=rtol, atol=atol)
     if len(reasons) > 0:
         return "mismatch found! custom implementation doesn't match reference: " + reasons[0]
-    
-    return '' 
\ No newline at end of file
+    return ''
\ No newline at end of file
diff --git a/problems/pmpp/prefixsum_py/submission.py b/problems/pmpp/prefixsum_py/submission.py
new file mode 100644
index 000000000..6ccdf4ad6
--- /dev/null
+++ b/problems/pmpp/prefixsum_py/submission.py
@@ -0,0 +1,12 @@
+import torch
+from task import input_t, output_t
+
+def custom_kernel(data: input_t) -> output_t:
+    """
+    Reference implementation of inclusive prefix sum using PyTorch.
+    Args:
+        data: Input tensor to compute prefix sum on
+    Returns:
+        Tensor containing the inclusive prefix sum
+    """
+    return torch.cumsum(data, dim=0)
\ No newline at end of file
diff --git a/problems/pmpp/prefixsum_py/task.yml b/problems/pmpp/prefixsum_py/task.yml
index ff13e391e..8762ee349 100644
--- a/problems/pmpp/prefixsum_py/task.yml
+++ b/problems/pmpp/prefixsum_py/task.yml
@@ -10,8 +10,14 @@ files:
 lang: "py"
 
 description: |
-  Implement an inclusive prefix sum (scan) kernel using CUDA inline function that matches the reference implementation.
+  Implement an inclusive prefix sum (scan) kernel that matches the reference implementation.
   The kernel should compute the cumulative sum of all elements up to each position.
+  Because of numerical instability, the tolerance is scaled by the square root of the input size.
+
+  Input:
+  - `data`: A 1D tensor of size `n`
+  Output:
+  - `output`: A 1D tensor of size `n`
 
 config:
   main: "eval.py" 
@@ -24,8 +30,21 @@ tests:
   - {"size": 4096, "seed": 9173}
 
 benchmarks:
-  - {"size": 8192, "seed": 54352}
-  - {"size": 16384, "seed": 93246}
-  - {"size": 32768, "seed": 6256}
-  - {"size": 65536, "seed": 8841}
-  - {"size": 131072, "seed": 6252} 
\ No newline at end of file
+  # - {"size": 8192, "seed": 54352}
+  # - {"size": 16384, "seed": 93246}
+  # - {"size": 32768, "seed": 6256}
+  # - {"size": 65536, "seed": 8841}
+  # - {"size": 131072, "seed": 6252} 
+  - {"size": 262144, "seed": 12345}
+  - {"size": 524288, "seed": 67890}
+  - {"size": 1048576, "seed": 13579}
+  - {"size": 2097152, "seed": 24680}
+  - {"size": 4194304, "seed": 35791}
+  - {"size": 8388608, "seed": 46802}
+  - {"size": 16777216, "seed": 57913}
+  - {"size": 33554432, "seed": 68024}
+  - {"size": 67108864, "seed": 79135}
+  - {"size": 134217728, "seed": 80246} # fits on T4
+  - {"size": 268435456, "seed": 91357}
+  # - {"size": 536870912, "seed": 102468}
+  # - {"size": 1073741824, "seed": 113579}
\ No newline at end of file
diff --git a/problems/pmpp/sort_py/reference.py b/problems/pmpp/sort_py/reference.py
index c8d05f777..a1dbf2c35 100644
--- a/problems/pmpp/sort_py/reference.py
+++ b/problems/pmpp/sort_py/reference.py
@@ -12,15 +12,34 @@ def ref_kernel(data: input_t) -> output_t:
     """
     return torch.sort(data)[0]
 
-def generate_input(size: int, seed: int) -> input_t:
+def generate_input(size: int, seed: int) -> torch.Tensor:
     """
-    Generates random input tensor.
+    Generates random input tensor where elements are drawn from different distributions.
+    
+    Args:
+        size: Total size of the final 1D tensor
+        seed: Base seed for random generation
+    
     Returns:
-        Tensor to be sorted
+        1D tensor of size `size` containing flattened values from different distributions
     """
+    # Calculate dimensions for a roughly square 2D matrix
+    rows = int(size ** 0.5)  # Square root for roughly square shape
+    cols = (size + rows - 1) // rows  # Ceiling division to ensure total size >= requested size
+    
     gen = torch.Generator(device='cuda')
-    gen.manual_seed(seed)
-    return torch.randn(size, device='cuda', dtype=torch.float32, generator=gen).contiguous()
+    result = torch.empty((rows, cols), device='cuda', dtype=torch.float32)
+    
+    # Different seed for each row!
+    for i in range(rows):
+        row_seed = seed + i
+        gen.manual_seed(row_seed)
+        
+        # Generate values for this row with mean=row_seed
+        result[i, :] = torch.randn(cols, device='cuda', dtype=torch.float32, generator=gen) + row_seed
+    
+    # Flatten and trim to exact size requested
+    return result.flatten()[:size].contiguous()
 
 def check_implementation(
     data: input_t,
diff --git a/problems/pmpp/sort_py/task.yml b/problems/pmpp/sort_py/task.yml
index 9fcb8fe1f..659191b24 100644
--- a/problems/pmpp/sort_py/task.yml
+++ b/problems/pmpp/sort_py/task.yml
@@ -11,7 +11,11 @@ lang: "py"
 
 description: |
   Implement a sort kernel that matches the reference implementation.
-  The kernel should sort the input array in ascending order using the merge sort algorithm.
+  The kernel should sort the input array in ascending order using a sort algorithm of your choice.
+  
+  Input arrays are generated as random floating-point numbers, where each row of a roughly square matrix
+  is drawn from a normal distribution with a different mean value per row based on the seed and then flattened into a 1D array.
+  
 
 config:
   main: "eval.py" 
@@ -24,8 +28,8 @@ tests:
   - {"size": 4096, "seed": 9173}
 
 benchmarks:
-  - {"size": 8192, "seed": 54352}
-  - {"size": 16384, "seed": 93246}
-  - {"size": 32768, "seed": 6256}
-  - {"size": 65536, "seed": 8841}
-  - {"size": 131072, "seed": 6252} 
\ No newline at end of file
+  - {"size": 100000, "seed": 54352}
+  - {"size": 500000, "seed": 93246}
+  - {"size": 1000000, "seed": 6256}
+  - {"size": 10000000, "seed": 8841}
+  - {"size": 100000000, "seed": 6252} 
\ No newline at end of file