Reimplement tiled matmul with a 1D launch grid and a grouped program-ID remap so adjacent programs share A and B tiles and hit warm L2 lines. For each output entry:
C[i,j]=k=0∑K−1A[i,k]⋅B[k,j]The arithmetic is identical to the plain matmul. The single change: switch from a 2D grid to a 1D grid of the same total size, then derive (pid_m,pid_n) from pid with a GROUP_SIZE_M constexpr so the schedule walks the output in M-major groups. This is exactly the pattern Triton's official matmul tutorial uses.
Why it matters: a strictly row-major or column-major linear schedule streams one matrix and re-reads the other from DRAM on every tile. Grouped scheduling reuses each loaded A row-block across the group and each B column-block across the same group, multiplying L2 hit rate by roughly GROUP_SIZE_M.
Input: A = [[1, 2], [3, 4]], B = [[5, 6], [7, 8]]
Output: [[19, 22], [43, 50]]
Input: A = [[1, 0, 1], [0, 1, 0]], B = [[1, 2], [3, 4], [5, 6]]
Output: [[6, 8], [3, 4]]
Compute num_pid_in_group=GROUP_SIZE_M⋅num_pid_n, then group_id=pid/num_pid_in_group, and first_pid_m=group_id⋅GROUP_SIZE_M. Clamp the group height for the partial last group.
Inside the group, decode pid_m and pid_n so adjacent values of pid keep the same rows of A warm across consecutive programs.
Once the two program IDs are decoded, the rest of the kernel is identical to a standard 2D-grid matmul: same K-loop, same tl.dot, same masked store.
torch.float32 on cuda.atol=1e-2 + rtol=1e-2 * abs(expected) for K≥128 to absorb fp32 accumulation drift.Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.
Reimplement tiled matmul with a 1D launch grid and a grouped program-ID remap so adjacent programs share A and B tiles and hit warm L2 lines. For each output entry:
C[i,j]=k=0∑K−1A[i,k]⋅B[k,j]The arithmetic is identical to the plain matmul. The single change: switch from a 2D grid to a 1D grid of the same total size, then derive (pid_m,pid_n) from pid with a GROUP_SIZE_M constexpr so the schedule walks the output in M-major groups. This is exactly the pattern Triton's official matmul tutorial uses.
Why it matters: a strictly row-major or column-major linear schedule streams one matrix and re-reads the other from DRAM on every tile. Grouped scheduling reuses each loaded A row-block across the group and each B column-block across the same group, multiplying L2 hit rate by roughly GROUP_SIZE_M.
Input: A = [[1, 2], [3, 4]], B = [[5, 6], [7, 8]]
Output: [[19, 22], [43, 50]]
Input: A = [[1, 0, 1], [0, 1, 0]], B = [[1, 2], [3, 4], [5, 6]]
Output: [[6, 8], [3, 4]]
Compute num_pid_in_group=GROUP_SIZE_M⋅num_pid_n, then group_id=pid/num_pid_in_group, and first_pid_m=group_id⋅GROUP_SIZE_M. Clamp the group height for the partial last group.
Inside the group, decode pid_m and pid_n so adjacent values of pid keep the same rows of A warm across consecutive programs.
Once the two program IDs are decoded, the rest of the kernel is identical to a standard 2D-grid matmul: same K-loop, same tl.dot, same masked store.
torch.float32 on cuda.atol=1e-2 + rtol=1e-2 * abs(expected) for K≥128 to absorb fp32 accumulation drift.Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.