Motivation

Without this change, we will see many NaN from certain shape matmul, for example 32, 128 x 128, 1024

Fix Persistent Matmul Kernel: Remove Incorrect tile_id_c Logic

Problem

The persistent matmul kernel had a critical bug where computation and storage used misaligned tile IDs, causing incorrect results to be written to wrong memory locations, which caused gibberish result for our Grok 10B and v6 Mini model.

Root Cause

# Buggy: compute tile N, store to tile (N + NUM_SMS)
for tile_id in tl.range(start_pid, num_tiles, NUM_SMS):
    pid_m, pid_n = _compute_pid(tile_id, ...)      # Compute tile N
    # ... accumulate ...
    
    tile_id_c += NUM_SMS
    pid_m, pid_n = _compute_pid(tile_id_c, ...)    # ❌ Store to different tile!
    tl.store(c_ptrs, c, mask=c_mask)

Solution

Use the same pid_m, pid_n for both computation and storage:

# Fixed: compute and store to same tile
for tile_id in tl.range(start_pid, num_tiles, NUM_SMS):
    pid_m, pid_n = _compute_pid(tile_id, ...)      # Compute tile N
    # ... accumulate ...
    # Reuse pid_m, pid_n for storage
    tl.store(c_ptrs, c, mask=c_mask)               # ✅ Store to same tile

Additional Changes

• Added explicit bfloat16 and float32 dtype handling

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.