--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -0,0 +1,130 @@
+"""Triton softmax with fp32 Kahan-compensated summation.
+
+Produces results matching the fp64 reference within atol=rtol=1e-5
+even for the 256K-vocab and extreme-logit shapes.
+"""
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+
+
+OP_TYPE = "softmax"
+SUPPORTED_PRECISIONS = ["fp32"]
+HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
+
+
+# ---------------------------------------------------------------------------
+# Triton kernel
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def _softmax_kernel_fwd(
+    X_ptr,
+    Y_ptr,
+    stride_xr,
+    stride_yr,
+    N: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """One-row-per-program softmax with fp32 Kahan-compensated summation.
+
+    Block-N must divide N (N is power of 2, BLOCK_N is power of 2 -> fine).
+    """
+    row_id = tl.program_id(0)
+    x_base = X_ptr + row_id * stride_xr
+    y_base = Y_ptr + row_id * stride_yr
+
+    # ---- Pass 1: find row-max ------------------------------------------
+    row_max = tl.full([1], -float('inf'), dtype=tl.float32)
+    for off0 in range(0, N, BLOCK_N):
+        idx = off0 + tl.arange(0, BLOCK_N)
+        x = tl.load(x_base + idx).to(tl.float32)
+        local_max = tl.max(x)
+        row_max = tl.maximum(row_max, local_max)
+
+    row_max = row_max.reshape(()).to(tl.float32)
+
+    # ---- Pass 2: Kahan-sum exp(x - max) --------------------------------
+    #  Kahan sum:  running_sum and c (compensation), both fp32.
+    running_sum = tl.zeros([1], dtype=tl.float32)
+    comp = tl.zeros([1], dtype=tl.float32)
+
+    for off0 in range(0, N, BLOCK_N):
+        idx = off0 + tl.arange(0, BLOCK_N)
+        x = tl.load(x_base + idx).to(tl.float32)
+        e = tl.exp(x - row_max)
+        # Kahan step:  y = e - c;  new_sum = running_sum + y;  c = (new_sum - running_sum) - y
+        y_kahan = e - comp
+        tmp_sum = running_sum + y_kahan
+        comp = (tmp_sum - running_sum) - y_kahan
+        running_sum = tmp_sum
+
+    total_sum = running_sum.reshape(()).to(tl.float32)
+
+    # ---- Pass 3: write out exp(x - max) / total_sum --------------------
+    for off0 in range(0, N, BLOCK_N):
+        idx = off0 + tl.arange(0, BLOCK_N)
+        x = tl.load(x_base + idx).to(tl.float32)
+        e = tl.exp(x - row_max)
+        y = e / total_sum
+        tl.store(y_base + idx, y.to(tl.float32))
+
+
+def triton_softmax(x: torch.Tensor) -> torch.Tensor:
+    """x: (batch, vocab) fp32 on CUDA -> y: (batch, vocab) fp32 on CUDA."""
+    assert x.is_cuda and x.dtype == torch.float32 and x.ndim == 2
+    batch, vocab = x.shape
+    y = torch.empty_like(x)
+
+    # Pick BLOCK_N that divides vocab.  For small vocab use vocab itself.
+    # For large vocab use a reasonable block that keeps registers in check.
+    if vocab <= 2048:
+        BLOCK_N = vocab
+    elif vocab <= 8192:
+        BLOCK_N = 2048
+    elif vocab <= 32768:
+        BLOCK_N = 4096
+    else:
+        BLOCK_N = 4096
+
+    # Ensure BLOCK_N divides N
+    while vocab % BLOCK_N != 0 and BLOCK_N > 64:
+        BLOCK_N //= 2
+
+    grid = (batch,)
+    _softmax_kernel_fwd[grid](
+        x, y,
+        x.stride(0), y.stride(0),
+        vocab,
+        BLOCK_N=BLOCK_N,
+    )
+    return y
+
+
+# ---------------------------------------------------------------------------
+# Model interface
+# ---------------------------------------------------------------------------
+
+class Model(nn.Module):
+    def __init__(self, batch: int, vocab: int):
+        super().__init__()
+        self.batch = batch
+        self.vocab = vocab
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return triton_softmax(x)
+
+
+# Default shape; overridden per-iteration by check.py / benchmark.py.
+BATCH = 8
+VOCAB = 32768
+
+
+def get_inputs():
+    x = torch.randn(BATCH, VOCAB, dtype=torch.float32) * 4.0
+    return [x]
+
+
+def get_init_inputs():
+    return [BATCH, VOCAB]

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -46,16 +46,15 @@
     row_max = row_max.reshape(()).to(tl.float32)
 
     # ---- Pass 2: Kahan-sum exp(x - max) --------------------------------
-    #  Kahan sum:  running_sum and c (compensation), both fp32.
-    running_sum = tl.zeros([1], dtype=tl.float32)
-    comp = tl.zeros([1], dtype=tl.float32)
+    running_sum = tl.full([1], 0.0, dtype=tl.float32)
+    comp = tl.full([1], 0.0, dtype=tl.float32)
 
     for off0 in range(0, N, BLOCK_N):
         idx = off0 + tl.arange(0, BLOCK_N)
         x = tl.load(x_base + idx).to(tl.float32)
         e = tl.exp(x - row_max)
-        # Kahan step:  y = e - c;  new_sum = running_sum + y;  c = (new_sum - running_sum) - y
-        y_kahan = e - comp
+        e_chunk = tl.sum(e, keep_dims=True).to(tl.float32)
+        y_kahan = e_chunk - comp
         tmp_sum = running_sum + y_kahan
         comp = (tmp_sum - running_sum) - y_kahan
         running_sum = tmp_sum

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution_draft.py
@@ -0,0 +1,51 @@
+"""Triton softmax with fp32 Kahan-compensated summation.
+
+Two-phase kernel that splits each row into chunks processed by separate
+blocks, giving us batch * (N / chunk_size) blocks for GPU saturation.
+
+Phase 1: Each block computes local-max and Kahan-sum-exp for its chunk.
+Phase 2: Final kernel reads per-row max/sum and normalizes.
+"""
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+
+OP_TYPE = "softmax"
+SUPPORTED_PRECISIONS = ["fp32"]
+HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
+
+
+# ---------------------------------------------------------------------------
+# Phase 1: per-chunk max + Kahan sum of exp(x - max)
+# ---------------------------------------------------------------------------
+@triton.jit
+def _softmax_phase1(
+    X_ptr,
+    partial_max_ptr,
+    partial_sum_ptr,
+    stride_xr,
+    N: tl.constexpr,
+    CHUNK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    batch = N  # will be overridden; total vocab size
+    num_chunks = pid  # total blocks = batch * chunks_per_row
+    chunks_per_row = grid_x(1)  # not available; compute another way
+
+    # Better indexing: pid -> (row_id, chunk_id)
+    # We set grid = (batch * chunks_per_row,) and compute:
+    chunk_id = pid % chunks_per_row  # can't reference grid
+    row_id = pid // chunks_per_row
+
+    # Each program has already been assigned to a (row, chunk) pair.
+    # The host passes chunks_per_row as a constexpr.
+    pass
+
+
+# This approach has the problem that pid % chunks_per_row needs
+# chunks_per_row as constexpr. Let me use a simpler approach:
+# grid = (batch,), each program processes all chunks inline.
+#
+# Instead, let me use the simpler 3-pass approach but with better
+# block sizes and persistent scheduling.

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -1,109 +1,134 @@
 """Triton softmax with fp32 Kahan-compensated summation.
 
-Produces results matching the fp64 reference within atol=rtol=1e-5
-even for the 256K-vocab and extreme-logit shapes.
+Multi-block-per-row approach for GPU saturation:
+  Kernel 1: Each chunk computes local-max and Kahan-sum of exp(x - row_max_local).
+            Then atomically max-reduces into a per-row shared max, and a second
+            pass (same kernel, after grid sync) computes Kahan-sum of exp(x-global_max).
+  Kernel 2: Normalize: read x, subtract global max, exp, divide by global sum.
+
+Actually, let's keep it simple and correct:
+  Kernel 1 (max_reduce): multiple blocks per row, each processes a slice,
+    writes per-slice max. Then atomic-max reduces into per-row max.
+  Kernel 2 (sum_write): each block reads row_max, does Kahan-sum of exp(x-row_max),
+    writes normed result.
 """
 import torch
 import torch.nn as nn
 import triton
 import triton.language as tl
 
-
 OP_TYPE = "softmax"
 SUPPORTED_PRECISIONS = ["fp32"]
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
 
-# ---------------------------------------------------------------------------
-# Triton kernel
-# ---------------------------------------------------------------------------
+@triton.jit
+def _max_kernel(
+    X_ptr,
+    row_max_ptr,
+    stride_xr,
+    N: tl.constexpr,
+    CHUNK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    chunks_per_row = tl.num_programs(1) if tl.num_programs(1) > 1 else 1
+    
+    row_id = pid // chunks_per_row
+    chunk_id = pid % chunks_per_row
+    
+    x_base = X_ptr + row_id * stride_xr
+    
+    local_max = tl.full([1], -3.4028235e38, dtype=tl.float32)
+    
+    start = chunk_id * CHUNK
+    end = min(start + CHUNK, N)
+    
+    for off in range(start, end, CHUNK):
+        idx = off + tl.arange(0, CHUNK)
+        mask = idx < N
+        x = tl.load(x_base + idx, mask=mask, other=-3.4028235e38).to(tl.float32)
+        m = tl.max(x, keep_dims=True)
+        local_max = tl.maximum(local_max, m)
+    
+    lm = local_max.reshape(()).to(tl.float32)
+    
+    tl.atomic_max(row_max_ptr + row_id, lm)
 
-@triton.jit
-def _softmax_kernel_fwd(
+
+@triton.jit 
+def _norm_kernel(
     X_ptr,
     Y_ptr,
+    row_max_ptr,
     stride_xr,
     stride_yr,
     N: tl.constexpr,
-    BLOCK_N: tl.constexpr,
+    CHUNK: tl.constexpr,
 ):
-    """One-row-per-program softmax with fp32 Kahan-compensated summation.
-
-    Block-N must divide N (N is power of 2, BLOCK_N is power of 2 -> fine).
-    """
-    row_id = tl.program_id(0)
+    pid = tl.program_id(0)
+    
+    row_id = pid
+    chunk_id = 0  # for now, one block per row
+    
     x_base = X_ptr + row_id * stride_xr
     y_base = Y_ptr + row_id * stride_yr
-
-    # ---- Pass 1: find row-max ------------------------------------------
-    row_max = tl.full([1], -float('inf'), dtype=tl.float32)
-    for off0 in range(0, N, BLOCK_N):
-        idx = off0 + tl.arange(0, BLOCK_N)
-        x = tl.load(x_base + idx).to(tl.float32)
-        local_max = tl.max(x)
-        row_max = tl.maximum(row_max, local_max)
-
-    row_max = row_max.reshape(()).to(tl.float32)
-
-    # ---- Pass 2: Kahan-sum exp(x - max) --------------------------------
+    
+    row_max = tl.load(row_max_ptr + row_id).to(tl.float32)
+    
+    # Kahan sum all chunks
     running_sum = tl.full([1], 0.0, dtype=tl.float32)
     comp = tl.full([1], 0.0, dtype=tl.float32)
-
-    for off0 in range(0, N, BLOCK_N):
-        idx = off0 + tl.arange(0, BLOCK_N)
-        x = tl.load(x_base + idx).to(tl.float32)
+    
+    for off in range(0, N, CHUNK):
+        idx = off + tl.arange(0, CHUNK)
+        mask = idx < N
+        x = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32)
         e = tl.exp(x - row_max)
-        e_chunk = tl.sum(e, keep_dims=True).to(tl.float32)
-        y_kahan = e_chunk - comp
-        tmp_sum = running_sum + y_kahan
-        comp = (tmp_sum - running_sum) - y_kahan
-        running_sum = tmp_sum
-
-    total_sum = running_sum.reshape(()).to(tl.float32)
-
-    # ---- Pass 3: write out exp(x - max) / total_sum --------------------
-    for off0 in range(0, N, BLOCK_N):
-        idx = off0 + tl.arange(0, BLOCK_N)
-        x = tl.load(x_base + idx).to(tl.float32)
+        e_masked = tl.where(mask, e, 0.0)
+        e_sum = tl.sum(e_masked, keep_dims=True).to(tl.float32)
+        y_k = e_sum - comp
+        new_sum = running_sum + y_k
+        comp = (new_sum - running_sum) - y_k
+        running_sum = new_sum
+    
+    total = running_sum.reshape(()).to(tl.float32)
+    
+    for off in range(0, N, CHUNK):
+        idx = off + tl.arange(0, CHUNK)
+        mask = idx < N
+        x = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32)
         e = tl.exp(x - row_max)
-        y = e / total_sum
-        tl.store(y_base + idx, y.to(tl.float32))
+        y = tl.where(mask, e / total, 0.0)
+        tl.store(y_base + idx, y, mask=mask)
 
 
 def triton_softmax(x: torch.Tensor) -> torch.Tensor:
-    """x: (batch, vocab) fp32 on CUDA -> y: (batch, vocab) fp32 on CUDA."""
     assert x.is_cuda and x.dtype == torch.float32 and x.ndim == 2
     batch, vocab = x.shape
     y = torch.empty_like(x)
-
-    # Pick BLOCK_N that divides vocab.  For small vocab use vocab itself.
-    # For large vocab use a reasonable block that keeps registers in check.
-    if vocab <= 2048:
-        BLOCK_N = vocab
-    elif vocab <= 8192:
-        BLOCK_N = 2048
+    
+    row_max = torch.full((batch,), -float('inf'), dtype=torch.float32, device=x.device)
+    
+    if vocab <= 4096:
+        CHUNK = vocab
     elif vocab <= 32768:
-        BLOCK_N = 4096
+        CHUNK = 4096
     else:
-        BLOCK_N = 4096
-
-    # Ensure BLOCK_N divides N
-    while vocab % BLOCK_N != 0 and BLOCK_N > 64:
-        BLOCK_N //= 2
-
- 

... (964 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -1,17 +1,11 @@
 """Triton softmax with fp32 Kahan-compensated summation.
 
-Multi-block-per-row approach for GPU saturation:
-  Kernel 1: Each chunk computes local-max and Kahan-sum of exp(x - row_max_local).
-            Then atomically max-reduces into a per-row shared max, and a second
-            pass (same kernel, after grid sync) computes Kahan-sum of exp(x-global_max).
-  Kernel 2: Normalize: read x, subtract global max, exp, divide by global sum.
-
-Actually, let's keep it simple and correct:
-  Kernel 1 (max_reduce): multiple blocks per row, each processes a slice,
-    writes per-slice max. Then atomic-max reduces into per-row max.
-  Kernel 2 (sum_write): each block reads row_max, does Kahan-sum of exp(x-row_max),
-    writes normed result.
+Persistent multi-block-per-row approach:
+  Kernel 1 (max):   Multiple blocks per row, each finding local max over a chunk.
+                     Max-reduced into per-row scalar via atomic CAS.
+  Kernel 2 (norm):  One block per row, reads row_max, Kahan-sums exp(x-max), writes output.
 """
+import math
 import torch
 import torch.nn as nn
 import triton
@@ -28,34 +22,35 @@
     row_max_ptr,
     stride_xr,
     N: tl.constexpr,
-    CHUNK: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    NUM_CHUNKS: tl.constexpr,
 ):
     pid = tl.program_id(0)
-    chunks_per_row = tl.num_programs(1) if tl.num_programs(1) > 1 else 1
-    
-    row_id = pid // chunks_per_row
-    chunk_id = pid % chunks_per_row
-    
+    row_id = pid // NUM_CHUNKS
+    chunk_id = pid % NUM_CHUNKS
+
     x_base = X_ptr + row_id * stride_xr
-    
+
     local_max = tl.full([1], -3.4028235e38, dtype=tl.float32)
-    
-    start = chunk_id * CHUNK
-    end = min(start + CHUNK, N)
-    
-    for off in range(start, end, CHUNK):
-        idx = off + tl.arange(0, CHUNK)
+
+    chunk_start = chunk_id * BLOCK_N
+    chunk_end = min(chunk_start + BLOCK_N, N)
+
+    for off in range(chunk_start, chunk_end, BLOCK_N):
+        idx = off + tl.arange(0, BLOCK_N)
         mask = idx < N
         x = tl.load(x_base + idx, mask=mask, other=-3.4028235e38).to(tl.float32)
-        m = tl.max(x, keep_dims=True)
-        local_max = tl.maximum(local_max, m)
-    
+        m = tl.max(x)
+        m_s = m.reshape(())
+        local_max_s = local_max.reshape(())
+        local_max_s = tl.maximum(local_max_s, m_s)
+        local_max = local_max_s.reshape([1])
+
     lm = local_max.reshape(()).to(tl.float32)
-    
-    tl.atomic_max(row_max_ptr + row_id, lm)
+    old = tl.atomic_max(row_max_ptr + row_id, lm.to(tl.float32))
 
 
-@triton.jit 
+@triton.jit
 def _norm_kernel(
     X_ptr,
     Y_ptr,
@@ -63,24 +58,20 @@
     stride_xr,
     stride_yr,
     N: tl.constexpr,
-    CHUNK: tl.constexpr,
+    BLOCK_N: tl.constexpr,
 ):
-    pid = tl.program_id(0)
-    
-    row_id = pid
-    chunk_id = 0  # for now, one block per row
-    
+    row_id = tl.program_id(0)
+
     x_base = X_ptr + row_id * stride_xr
     y_base = Y_ptr + row_id * stride_yr
-    
+
     row_max = tl.load(row_max_ptr + row_id).to(tl.float32)
-    
-    # Kahan sum all chunks
+
     running_sum = tl.full([1], 0.0, dtype=tl.float32)
     comp = tl.full([1], 0.0, dtype=tl.float32)
-    
-    for off in range(0, N, CHUNK):
-        idx = off + tl.arange(0, CHUNK)
+
+    for off in range(0, N, BLOCK_N):
+        idx = off + tl.arange(0, BLOCK_N)
         mask = idx < N
         x = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32)
         e = tl.exp(x - row_max)
@@ -90,11 +81,11 @@
         new_sum = running_sum + y_k
         comp = (new_sum - running_sum) - y_k
         running_sum = new_sum
-    
+
     total = running_sum.reshape(()).to(tl.float32)
-    
-    for off in range(0, N, CHUNK):
-        idx = off + tl.arange(0, CHUNK)
+
+    for off in range(0, N, BLOCK_N):
+        idx = off + tl.arange(0, BLOCK_N)
         mask = idx < N
         x = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32)
         e = tl.exp(x - row_max)
@@ -106,27 +97,35 @@
     assert x.is_cuda and x.dtype == torch.float32 and x.ndim == 2
     batch, vocab = x.shape
     y = torch.empty_like(x)
-    
-    row_max = torch.full((batch,), -float('inf'), dtype=torch.float32, device=x.device)
-    
-    if vocab <= 4096:
-        CHUNK = vocab
-    elif vocab <= 32768:
-        CHUNK = 4096
+
+    row_max = torch.full((batch,), -3.4028235e38, dtype=torch.float32, device=x.device)
+
+    if vocab <= 1024:
+        BLOCK_N = vocab
+    elif vocab <= 4096:
+        BLOCK_N = 2048
+    elif vocab <= 16384:
+        BLOCK_N = 4096
     else:
-        CHUNK = 4096
-    
-    while vocab % CHUNK != 0 and CHUNK > 64:
-        CHUNK //= 2
-    
-    chunks_per_row = (vocab + CHUNK - 1) // CHUNK
-    grid1 = (batch * chunks_per_row,)
-    
-    _max_kernel[grid1](x, row_max, x.stride(0), vocab, CHUNK=CHUNK)
-    
+        BLOCK_N = 4096
+
+    while vocab % BLOCK_N != 0 and BLOCK_N > 32:
+        BLOCK_N //= 2
+
+    num_chunks = (vocab + BLOCK_N - 1) // BLOCK_N
+    grid1 = (batch * num_chunks,)
+
+    _max_kernel[grid1](
+        x, row_max, x.stride(0), vocab,
+        BLOCK_N=BLOCK_N, NUM_CHUNKS=num_chunks,
+    )
+
     grid2 = (batch,)
-    _norm_kernel[grid2](x, y, row_max, x.stride(0), y.stride(0), vocab, CHUNK=CHUNK)
-    
+    _norm_kernel[grid2](
+        x, y, row_max, x.stride(0), y.stride(0), vocab,
+        BLOCK_N=BLOCK_N,
+    )
+
     return y
 
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -98,7 +98,7 @@
     batch, vocab = x.shape
     y = torch.empty_like(x)
 
-    row_max = torch.full((batch,), -3.4028235e38, dtype=torch.float32, device=x.device)
+    row_max = torch.full((batch,), float('-inf'), dtype=torch.float32, device=x.device)
 
     if vocab <= 1024:
         BLOCK_N = vocab

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -31,7 +31,7 @@
 
     x_base = X_ptr + row_id * stride_xr
 
-    local_max = tl.full([1], -3.4028235e38, dtype=tl.float32)
+    local_max = tl.full([1], float('-inf'), dtype=tl.float32)
 
     chunk_start = chunk_id * BLOCK_N
     chunk_end = min(chunk_start + BLOCK_N, N)
@@ -39,7 +39,7 @@
     for off in range(chunk_start, chunk_end, BLOCK_N):
         idx = off + tl.arange(0, BLOCK_N)
         mask = idx < N
-        x = tl.load(x_base + idx, mask=mask, other=-3.4028235e38).to(tl.float32)
+        x = tl.load(x_base + idx, mask=mask, other=float('-inf')).to(tl.float32)
         m = tl.max(x)
         m_s = m.reshape(())
         local_max_s = local_max.reshape(())

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -40,11 +40,8 @@
         idx = off + tl.arange(0, BLOCK_N)
         mask = idx < N
         x = tl.load(x_base + idx, mask=mask, other=float('-inf')).to(tl.float32)
-        m = tl.max(x)
-        m_s = m.reshape(())
-        local_max_s = local_max.reshape(())
-        local_max_s = tl.maximum(local_max_s, m_s)
-        local_max = local_max_s.reshape([1])
+        m = tl.max(x, keep_dims=True)
+        local_max = tl.maximum(local_max, m)
 
     lm = local_max.reshape(()).to(tl.float32)
     old = tl.atomic_max(row_max_ptr + row_id, lm.to(tl.float32))

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/profile_scratch.py
@@ -0,0 +1,51 @@
+"""Quick profiler scratch file."""
+import torch
+import triton
+import triton.language as tl
+
+@triton.jit
+def _max_kernel_v2(
+    X_ptr, row_max_ptr, stride_xr,
+    N: tl.constexpr, BLOCK_N: tl.constexpr, NUM_CHUNKS: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    row_id = pid // NUM_CHUNKS
+    chunk_id = pid % NUM_CHUNKS
+    x_base = X_ptr + row_id * stride_xr
+    local_max = tl.full([1], float('-inf'), dtype=tl.float32)
+    chunk_start = chunk_id * BLOCK_N
+    chunk_end = min(chunk_start + BLOCK_N, N)
+    for off in range(chunk_start, chunk_end, BLOCK_N):
+        idx = off + tl.arange(0, BLOCK_N)
+        mask = idx < N
+        x = tl.load(x_base + idx, mask=mask, other=float('-inf')).to(tl.float32)
+        m = tl.max(x, keep_dims=True)
+        local_max = tl.maximum(local_max, m)
+    lm = local_max.reshape(()).to(tl.float32)
+    tl.atomic_max(row_max_ptr + row_id, lm)
+
+batch, vocab = 8, 32768
+x = torch.randn(batch, vocab, device='cuda', dtype=torch.float32) * 4.0
+row_max = torch.full((batch,), float('-inf'), dtype=torch.float32, device='cuda')
+
+BLOCK_N = 4096
+num_chunks = vocab // BLOCK_N
+grid = (batch * num_chunks,)
+
+# Warm-up
+_max_kernel_v2[grid](x, row_max, x.stride(0), vocab, BLOCK_N=BLOCK_N, NUM_CHUNKS=num_chunks)
+torch.cuda.synchronize()
+
+print(f"Grid: {grid}, BLOCK_N: {BLOCK_N}, num_chunks: {num_chunks}")
+print(f"Total elements: {batch * vocab:,}, bytes: {batch * vocab * 4:,}")
+
+# Profile with torch profiler
+with torch.profiler.profile(
+    activities=[torch.profiler.ProfilerActivity.CUDA],
+    record_shapes=True,
+    with_stack=True,
+) as prof:
+    _max_kernel_v2[grid](x, row_max, x.stride(0), vocab, BLOCK_N=BLOCK_N, NUM_CHUNKS=num_chunks)
+    torch.cuda.synchronize()
+
+print(prof.key_averages().table(sort_by='cuda_time_total', row_limit=15))

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/profile_scratch.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/profile_scratch.py
@@ -1,7 +1,8 @@
-"""Quick profiler scratch file."""
+"""Profile the full softmax kernels."""
 import torch
 import triton
 import triton.language as tl
+import statistics
 
 @triton.jit
 def _max_kernel_v2(
@@ -24,28 +25,85 @@
     lm = local_max.reshape(()).to(tl.float32)
     tl.atomic_max(row_max_ptr + row_id, lm)
 
-batch, vocab = 8, 32768
-x = torch.randn(batch, vocab, device='cuda', dtype=torch.float32) * 4.0
-row_max = torch.full((batch,), float('-inf'), dtype=torch.float32, device='cuda')
+@triton.jit
+def _norm_kernel_v2(
+    X_ptr, Y_ptr, row_max_ptr, stride_xr, stride_yr,
+    N: tl.constexpr, BLOCK_N: tl.constexpr,
+):
+    row_id = tl.program_id(0)
+    x_base = X_ptr + row_id * stride_xr
+    y_base = Y_ptr + row_id * stride_yr
+    row_max = tl.load(row_max_ptr + row_id).to(tl.float32)
 
-BLOCK_N = 4096
-num_chunks = vocab // BLOCK_N
-grid = (batch * num_chunks,)
+    running_sum = tl.full([1], 0.0, dtype=tl.float32)
+    comp = tl.full([1], 0.0, dtype=tl.float32)
 
-# Warm-up
-_max_kernel_v2[grid](x, row_max, x.stride(0), vocab, BLOCK_N=BLOCK_N, NUM_CHUNKS=num_chunks)
-torch.cuda.synchronize()
+    for off in range(0, N, BLOCK_N):
+        idx = off + tl.arange(0, BLOCK_N)
+        mask = idx < N
+        x = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32)
+        e = tl.exp(x - row_max)
+        e_masked = tl.where(mask, e, 0.0)
+        e_sum = tl.sum(e_masked, keep_dims=True).to(tl.float32)
+        y_k = e_sum - comp
+        new_sum = running_sum + y_k
+        comp = (new_sum - running_sum) - y_k
+        running_sum = new_sum
 
-print(f"Grid: {grid}, BLOCK_N: {BLOCK_N}, num_chunks: {num_chunks}")
-print(f"Total elements: {batch * vocab:,}, bytes: {batch * vocab * 4:,}")
+    total = running_sum.reshape(()).to(tl.float32)
 
-# Profile with torch profiler
-with torch.profiler.profile(
-    activities=[torch.profiler.ProfilerActivity.CUDA],
-    record_shapes=True,
-    with_stack=True,
-) as prof:
-    _max_kernel_v2[grid](x, row_max, x.stride(0), vocab, BLOCK_N=BLOCK_N, NUM_CHUNKS=num_chunks)
+    for off in range(0, N, BLOCK_N):
+        idx = off + tl.arange(0, BLOCK_N)
+        mask = idx < N
+        x = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32)
+        e = tl.exp(x - row_max)
+        y = tl.where(mask, e / total, 0.0)
+        tl.store(y_base + idx, y, mask=mask)
+
+def time_kernel(kernel_fn, grid, iters=30, warmup=10, **kwargs):
+    for _ in range(warmup):
+        kernel_fn[grid](**kwargs)
     torch.cuda.synchronize()
+    
+    times = []
+    for _ in range(iters):
+        torch.cuda.synchronize()
+        s = torch.cuda.Event(enable_timing=True)
+        e = torch.cuda.Event(enable_timing=True)
+        s.record()
+        kernel_fn[grid](**kwargs)
+        e.record()
+        torch.cuda.synchronize()
+        times.append(s.elapsed_time(e))
+    return statistics.median(times)
 
-print(prof.key_averages().table(sort_by='cuda_time_total', row_limit=15))
+for batch, vocab in [(32, 4096), (16, 32768), (8, 131072), (4, 262144)]:
+    x = torch.randn(batch, vocab, device='cuda', dtype=torch.float32) * 4.0
+    
+    BLOCK_N = 4096
+    while vocab % BLOCK_N != 0:
+        BLOCK_N //= 2
+    num_chunks = vocab // BLOCK_N
+    
+    y = torch.empty_like(x)
+    row_max = torch.full((batch,), float('-inf'), dtype=torch.float32, device='cuda')
+    
+    grid1 = (batch * num_chunks,)
+    grid2 = (batch,)
+    
+    ms_max = time_kernel(_max_kernel_v2, grid1, x=x, row_max_ptr=row_max, 
+                          stride_xr=x.stride(0), N=vocab, BLOCK_N=BLOCK_N, NUM_CHUNKS=num_chunks)
+    
+    ms_norm = time_kernel(_norm_kernel_v2, grid2, x=x, Y_ptr=y, row_max_ptr=row_max,
+                           stride_xr=x.stride(0), stride_yr=y.stride(0), N=vocab, BLOCK_N=BLOCK_N)
+    
+    total_ms = ms_max + ms_norm
+    input_bytes = batch * vocab * 4
+    output_bytes = batch * vocab * 4
+    total_bytes = input_bytes * 3 + output_bytes  # 3 reads of X, 1 write of Y
+    
+    eff_bw = (input_bytes * 2 + output_bytes) / (total_ms * 1e-3) / 1e9  # counted bytes
+    real_bw = total_bytes / (total_ms * 1e-3) / 1e9  # actual bytes moved
+    
+    print(f"batch={batch:3d} vocab={vocab:7d}: max={ms_max:.3f}ms norm={ms_norm:.3f}ms "
+          f"total={total_ms:.3f}ms counted_bw={eff_bw:.1f}GB/s real_bw={real_bw:.1f}GB/s")

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/profile_scratch.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/profile_scratch.py
@@ -91,10 +91,10 @@
     grid1 = (batch * num_chunks,)
     grid2 = (batch,)
     
-    ms_max = time_kernel(_max_kernel_v2, grid1, x=x, row_max_ptr=row_max, 
+    ms_max = time_kernel(_max_kernel_v2, grid1, X_ptr=x, row_max_ptr=row_max,
                           stride_xr=x.stride(0), N=vocab, BLOCK_N=BLOCK_N, NUM_CHUNKS=num_chunks)
-    
-    ms_norm = time_kernel(_norm_kernel_v2, grid2, x=x, Y_ptr=y, row_max_ptr=row_max,
+
+    ms_norm = time_kernel(_norm_kernel_v2, grid2, X_ptr=x, Y_ptr=y, row_max_ptr=row_max,
                            stride_xr=x.stride(0), stride_yr=y.stride(0), N=vocab, BLOCK_N=BLOCK_N)
     
     total_ms = ms_max + ms_norm

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -1,15 +1,17 @@
-"""Triton softmax with fp32 Kahan-compensated summation.
+"""Triton softmax with online max-reduce + Kahan-compensated summation.
 
-Persistent multi-block-per-row approach:
-  Kernel 1 (max):   Multiple blocks per row, each finding local max over a chunk.
-                     Max-reduced into per-row scalar via atomic CAS.
-  Kernel 2 (norm):  One block per row, reads row_max, Kahan-sums exp(x-max), writes output.
+Single pass reads all elements, maintaining:
+  - running_max: max seen so far
+  - running_sum: Kahan-compensated sum of exp(x - running_max), adjusted when max updates
+  - Then a second pass writes the normalized output.
+
+Net memory traffic: 2 passes (read + write) vs 3 passes for naive approach.
 """
-import math
 import torch
 import torch.nn as nn
 import triton
 import triton.language as tl
+
 
 OP_TYPE = "softmax"
 SUPPORTED_PRECISIONS = ["fp32"]
@@ -17,77 +19,85 @@
 
 
 @triton.jit
-def _max_kernel(
-    X_ptr,
-    row_max_ptr,
-    stride_xr,
-    N: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-    NUM_CHUNKS: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    row_id = pid // NUM_CHUNKS
-    chunk_id = pid % NUM_CHUNKS
-
-    x_base = X_ptr + row_id * stride_xr
-
-    local_max = tl.full([1], float('-inf'), dtype=tl.float32)
-
-    chunk_start = chunk_id * BLOCK_N
-    chunk_end = min(chunk_start + BLOCK_N, N)
-
-    for off in range(chunk_start, chunk_end, BLOCK_N):
-        idx = off + tl.arange(0, BLOCK_N)
-        mask = idx < N
-        x = tl.load(x_base + idx, mask=mask, other=float('-inf')).to(tl.float32)
-        m = tl.max(x, keep_dims=True)
-        local_max = tl.maximum(local_max, m)
-
-    lm = local_max.reshape(()).to(tl.float32)
-    old = tl.atomic_max(row_max_ptr + row_id, lm.to(tl.float32))
-
-
-@triton.jit
-def _norm_kernel(
+def _softmax_kernel_online(
     X_ptr,
     Y_ptr,
-    row_max_ptr,
     stride_xr,
     stride_yr,
     N: tl.constexpr,
     BLOCK_N: tl.constexpr,
 ):
     row_id = tl.program_id(0)
-
     x_base = X_ptr + row_id * stride_xr
     y_base = Y_ptr + row_id * stride_yr
 
-    row_max = tl.load(row_max_ptr + row_id).to(tl.float32)
-
+    # --- Online pass: track running_max and Kahan-sum of exp(x-max) ------
+    running_max = tl.full([1], float('-inf'), dtype=tl.float32)
     running_sum = tl.full([1], 0.0, dtype=tl.float32)
     comp = tl.full([1], 0.0, dtype=tl.float32)
 
     for off in range(0, N, BLOCK_N):
         idx = off + tl.arange(0, BLOCK_N)
         mask = idx < N
-        x = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32)
-        e = tl.exp(x - row_max)
+        x = tl.load(x_base + idx, mask=mask, other=float('-inf')).to(tl.float32)
+
+        # Find max of this chunk
+        chunk_max = tl.max(x, keep_dims=True)
+
+        # New running max
+        new_max = tl.maximum(running_max, chunk_max)
+
+        # Scale the existing sum: exp(x - old_max) -> exp(x - new_max)
+        # factor = exp(old_max - new_max)
+        old_max_val = running_max.reshape(())
+        new_max_val = new_max.reshape(())
+        scale = tl.math.exp(old_max_val - new_max_val)
+
+        # Reshape running_sum to [1] for scalar ops
+        rs_s = running_sum.reshape(())
+        comp_s = comp.reshape(())
+
+        # Scale existing sum and compensation
+        rs_s = rs_s * scale
+        comp_s = comp_s * scale
+
+        # Add new chunk contributions: exp(x - new_max) for this chunk
+        e = tl.exp(x.reshape([-1]) - new_max_val)
         e_masked = tl.where(mask, e, 0.0)
-        e_sum = tl.sum(e_masked, keep_dims=True).to(tl.float32)
-        y_k = e_sum - comp
-        new_sum = running_sum + y_k
-        comp = (new_sum - running_sum) - y_k
-        running_sum = new_sum
+        e_chunk = tl.sum(e_masked, keep_dims=True).to(tl.float32)
+        e_s = e_chunk.reshape(())
 
-    total = running_sum.reshape(()).to(tl.float32)
+        # Kahan add
+        y_k = e_s - comp_s
+        new_sum_s = rs_s + y_k
+        comp_s = (new_sum_s - rs_s) - y_k
+        rs_s = new_sum_s
 
+        running_max = new_max
+        running_sum = rs_s.reshape([1])
+        comp = comp_s.reshape([1])
+
+    final_max = running_max.reshape(()).to(tl.float32)
+    total_sum = running_sum.reshape(()).to(tl.float32)
+
+    # --- Write normalized output ----------------------------------------
     for off in range(0, N, BLOCK_N):
         idx = off + tl.arange(0, BLOCK_N)
         mask = idx < N
         x = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32)
-        e = tl.exp(x - row_max)
-        y = tl.where(mask, e / total, 0.0)
+        e = tl.exp(x - final_max)
+        y = tl.where(mask, e / total_sum, 0.0)
         tl.store(y_base + idx, y, mask=mask)
+
+
+def _pick_block_n(vocab: int) -> int:
+    if vocab <= 1024:
+        return vocab
+    if vocab <= 4096:
+        return 2048
+    if vocab <= 32768:
+        return 4096
+    return 4096
 
 
 def triton_softmax(x: torch.Tensor) -> torch.Tensor:
@@ -95,34 +105,15 @@
     batch, vocab = x.shape
     y = torch.empty_like(x)
 
-    row_max = torch.full((batch,), float('-inf'), dtype=torch.float32, device=x.device)
+    block_n = _pick_block_n(vocab)
+    while vocab % block_n != 0 and block_n > 32:
+        block_n //= 2
 
-    if vocab <= 1024:
-        BLOCK_N = vocab
-    elif vocab <= 4096:
-        BLOCK_N = 2048
-    elif vocab <= 16384:
-        BLOCK_N = 4096
-    else:
-        BLOCK_N = 4096
-
-    while vocab % BLOCK_N != 0 and BLOCK_N > 32:
-        BLOCK_N //= 2
-
-    num_chunks = (vocab + BLOCK_N - 1) // BLOCK_N
-    grid1 = (batch * num_chunks,)
-
-    _max_kernel[grid1](
-        x, row_max, x.stride(0), vocab,
-        BLOCK_N=BLOCK_N, NUM_CHUNKS=num_chunks,
+    grid = (batch,)
+    _softmax_kernel_online[grid](
+        x, y, x.stride(0), y.stride(0), vocab,
+        

... (184 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -62,7 +62,7 @@
         comp_s = comp_s * scale
 
         # Add new chunk contributions: exp(x - new_max) for this chunk
-        e = tl.exp(x.reshape([-1]) - new_max_val)
+        e = tl.exp(x - new_max_val)
         e_masked = tl.where(mask, e, 0.0)
         e_chunk = tl.sum(e_masked, keep_dims=True).to(tl.float32)
         e_s = e_chunk.reshape(())

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -47,38 +47,30 @@
         # New running max
         new_max = tl.maximum(running_max, chunk_max)
 
-        # Scale the existing sum: exp(x - old_max) -> exp(x - new_max)
+        # Scale existing sum: exp(x-old_max) -> exp(x-new_max)
         # factor = exp(old_max - new_max)
-        old_max_val = running_max.reshape(())
-        new_max_val = new_max.reshape(())
-        scale = tl.math.exp(old_max_val - new_max_val)
-
-        # Reshape running_sum to [1] for scalar ops
-        rs_s = running_sum.reshape(())
-        comp_s = comp.reshape(())
+        scale = tl.exp(running_max - new_max)
 
         # Scale existing sum and compensation
-        rs_s = rs_s * scale
-        comp_s = comp_s * scale
+        running_sum = running_sum * scale
+        comp = comp * scale
 
-        # Add new chunk contributions: exp(x - new_max) for this chunk
-        e = tl.exp(x - new_max_val)
+        # Add new chunk contributions: exp(x - new_max)
+        e = tl.exp(x - new_max)
         e_masked = tl.where(mask, e, 0.0)
         e_chunk = tl.sum(e_masked, keep_dims=True).to(tl.float32)
-        e_s = e_chunk.reshape(())
 
         # Kahan add
-        y_k = e_s - comp_s
-        new_sum_s = rs_s + y_k
-        comp_s = (new_sum_s - rs_s) - y_k
-        rs_s = new_sum_s
+        y_k = e_chunk - comp
+        new_sum = running_sum + y_k
+        comp = (new_sum - running_sum) - y_k
+        running_sum = new_sum
 
         running_max = new_max
-        running_sum = rs_s.reshape([1])
-        comp = comp_s.reshape([1])
 
-    final_max = running_max.reshape(()).to(tl.float32)
-    total_sum = running_sum.reshape(()).to(tl.float32)
+    final_max = running_max.broadcast_to([BLOCK_N]).reshape([BLOCK_N])
+    total_sum_vec = running_sum.broadcast_to([BLOCK_N]).reshape([BLOCK_N])
+    total_sum = tl.max(total_sum_vec)
 
     # --- Write normalized output ----------------------------------------
     for off in range(0, N, BLOCK_N):

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -68,9 +68,8 @@
 
         running_max = new_max
 
-    final_max = running_max.broadcast_to([BLOCK_N]).reshape([BLOCK_N])
-    total_sum_vec = running_sum.broadcast_to([BLOCK_N]).reshape([BLOCK_N])
-    total_sum = tl.max(total_sum_vec)
+    final_max = running_max.broadcast_to([BLOCK_N])
+    total_sum_vec = running_sum.broadcast_to([BLOCK_N])
 
     # --- Write normalized output ----------------------------------------
     for off in range(0, N, BLOCK_N):

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -77,7 +77,7 @@
         mask = idx < N
         x = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32)
         e = tl.exp(x - final_max)
-        y = tl.where(mask, e / total_sum, 0.0)
+        y = tl.where(mask, e / total_sum_vec, 0.0)
         tl.store(y_base + idx, y, mask=mask)
 
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -1,110 +1,232 @@
-"""Triton softmax with online max-reduce + Kahan-compensated summation.
-
-Single pass reads all elements, maintaining:
-  - running_max: max seen so far
-  - running_sum: Kahan-compensated sum of exp(x - running_max), adjusted when max updates
-  - Then a second pass writes the normalized output.
-
-Net memory traffic: 2 passes (read + write) vs 3 passes for naive approach.
+"""CUDA C++ softmax with fp32 Kahan-compensated summation.
+
+Two-kernel approach:
+  Kernel 1: Each row splits into blocks; each block finds local max and atomic-max reduces.
+  Kernel 2: Each row (1 block) reads row_max, Kahan-sums exp(x-max), writes normalized output.
 """
 import torch
 import torch.nn as nn
-import triton
-import triton.language as tl
-
+from torch.utils.cpp_extension import load_inline
 
 OP_TYPE = "softmax"
 SUPPORTED_PRECISIONS = ["fp32"]
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
-
-@triton.jit
-def _softmax_kernel_online(
-    X_ptr,
-    Y_ptr,
-    stride_xr,
-    stride_yr,
-    N: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-):
-    row_id = tl.program_id(0)
-    x_base = X_ptr + row_id * stride_xr
-    y_base = Y_ptr + row_id * stride_yr
-
-    # --- Online pass: track running_max and Kahan-sum of exp(x-max) ------
-    running_max = tl.full([1], float('-inf'), dtype=tl.float32)
-    running_sum = tl.full([1], 0.0, dtype=tl.float32)
-    comp = tl.full([1], 0.0, dtype=tl.float32)
-
-    for off in range(0, N, BLOCK_N):
-        idx = off + tl.arange(0, BLOCK_N)
-        mask = idx < N
-        x = tl.load(x_base + idx, mask=mask, other=float('-inf')).to(tl.float32)
-
-        # Find max of this chunk
-        chunk_max = tl.max(x, keep_dims=True)
-
-        # New running max
-        new_max = tl.maximum(running_max, chunk_max)
-
-        # Scale existing sum: exp(x-old_max) -> exp(x-new_max)
-        # factor = exp(old_max - new_max)
-        scale = tl.exp(running_max - new_max)
-
-        # Scale existing sum and compensation
-        running_sum = running_sum * scale
-        comp = comp * scale
-
-        # Add new chunk contributions: exp(x - new_max)
-        e = tl.exp(x - new_max)
-        e_masked = tl.where(mask, e, 0.0)
-        e_chunk = tl.sum(e_masked, keep_dims=True).to(tl.float32)
-
-        # Kahan add
-        y_k = e_chunk - comp
-        new_sum = running_sum + y_k
-        comp = (new_sum - running_sum) - y_k
-        running_sum = new_sum
-
-        running_max = new_max
-
-    final_max = running_max.broadcast_to([BLOCK_N])
-    total_sum_vec = running_sum.broadcast_to([BLOCK_N])
-
-    # --- Write normalized output ----------------------------------------
-    for off in range(0, N, BLOCK_N):
-        idx = off + tl.arange(0, BLOCK_N)
-        mask = idx < N
-        x = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32)
-        e = tl.exp(x - final_max)
-        y = tl.where(mask, e / total_sum_vec, 0.0)
-        tl.store(y_base + idx, y, mask=mask)
-
-
-def _pick_block_n(vocab: int) -> int:
-    if vocab <= 1024:
-        return vocab
-    if vocab <= 4096:
-        return 2048
-    if vocab <= 32768:
-        return 4096
-    return 4096
-
-
-def triton_softmax(x: torch.Tensor) -> torch.Tensor:
+# ---------------------------------------------------------------------------
+# CUDA kernel source
+# ---------------------------------------------------------------------------
+CUDA_SRC = """
+#include <cuda_fp16.h>
+#include <math.h>
+#include <cfloat>
+
+#define WARP_SIZE 32
+
+// --- Warp-level max reduction (shfl) ---
+__device__ __forceinline__ float warp_reduce_max(float val) {
+    val = max(val, __shfl_xor_sync(0xffffffff, val, 16));
+    val = max(val, __shfl_xor_sync(0xffffffff, val, 8));
+    val = max(val, __shfl_xor_sync(0xffffffff, val, 4));
+    val = max(val, __shfl_xor_sync(0xffffffff, val, 2));
+    val = max(val, __shfl_xor_sync(0xffffffff, val, 1));
+    return val;
+}
+
+// --- Warp-level sum reduction (shfl) ---
+__device__ __forceinline__ float warp_reduce_sum(float val) {
+    val += __shfl_xor_sync(0xffffffff, val, 16);
+    val += __shfl_xor_sync(0xffffffff, val, 8);
+    val += __shfl_xor_sync(0xffffffff, val, 4);
+    val += __shfl_xor_sync(0xffffffff, val, 2);
+    val += __shfl_xor_sync(0xffffffff, val, 1);
+    return val;
+}
+
+// --- Kernel 1: per-chunk max, atomic-max into row_max ---
+__global__ void max_kernel(const float* x, float* row_max, int N, int chunk_size) {
+    const int num_chunks = gridDim.x;  // per row: blockDim.y chunks
+    const int row_id = blockIdx.y;
+    const int chunk_id = blockIdx.x;
+
+    const float* row_ptr = x + (long long)row_id * N;
+    const int tid = threadIdx.x;
+
+    // Each thread loads multiple elements
+    float local_max = -FLT_MAX;
+    const int start = chunk_id * chunk_size;
+    const int end = min(start + chunk_size, N);
+
+    for (int i = start + tid; i < end; i += blockDim.x) {
+        local_max = max(local_max, row_ptr[i]);
+    }
+
+    // Warp-level max reduce
+    local_max = warp_reduce_max(local_max);
+
+    // Thread 0 of each warp does atomic max
+    if (tid % WARP_SIZE == 0) {
+        atomicMax(row_max + row_id, __float_as_int(local_max));
+    }
+}
+
+// --- Kernel 2: read row_max, Kahan-sum exp(x-max), write normalized ---
+__global__ void norm_kernel(const float* x, float* y, const float* row_max, int N) {
+    const int row_id = blockIdx.x;
+    const int tid = threadIdx.x;
+
+    const float* row_ptr = x + (long long)row_id * N;
+    float* out_ptr = y + (long long)row_id * N;
+    const float m = row_max[row_id];
+
+    // Pass 1: Kahan sum of exp(x - m)
+    float total = 0.0f;
+    float comp = 0.0f;
+
+    for (int i = tid; i < N; i += blockDim.x) {
+        float val = row_ptr[i] - m;
+        float e = expf(val);
+
+        float yk = e - comp;
+        f

... (4212 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -1,8 +1,8 @@
 """CUDA C++ softmax with fp32 Kahan-compensated summation.
 
-Two-kernel approach:
-  Kernel 1: Each row splits into blocks; each block finds local max and atomic-max reduces.
-  Kernel 2: Each row (1 block) reads row_max, Kahan-sums exp(x-max), writes normalized output.
+Two-kernel approach for GPU saturation:
+  Kernel 1: Multiple blocks per row, each finds local max, atomicCAS-max reduces.
+  Kernel 2: One block per row, reads row_max, Kahan-sums exp(x-max), writes normalized output.
 """
 import torch
 import torch.nn as nn
@@ -12,161 +12,175 @@
 SUPPORTED_PRECISIONS = ["fp32"]
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
-# ---------------------------------------------------------------------------
-# CUDA kernel source
-# ---------------------------------------------------------------------------
-CUDA_SRC = """
-#include <cuda_fp16.h>
+CUDA_SRC = r"""
+#include <cuda_runtime.h>
 #include <math.h>
 #include <cfloat>
 
-#define WARP_SIZE 32
-
-// --- Warp-level max reduction (shfl) ---
-__device__ __forceinline__ float warp_reduce_max(float val) {
-    val = max(val, __shfl_xor_sync(0xffffffff, val, 16));
-    val = max(val, __shfl_xor_sync(0xffffffff, val, 8));
-    val = max(val, __shfl_xor_sync(0xffffffff, val, 4));
-    val = max(val, __shfl_xor_sync(0xffffffff, val, 2));
-    val = max(val, __shfl_xor_sync(0xffffffff, val, 1));
-    return val;
-}
-
-// --- Warp-level sum reduction (shfl) ---
-__device__ __forceinline__ float warp_reduce_sum(float val) {
-    val += __shfl_xor_sync(0xffffffff, val, 16);
-    val += __shfl_xor_sync(0xffffffff, val, 8);
-    val += __shfl_xor_sync(0xffffffff, val, 4);
-    val += __shfl_xor_sync(0xffffffff, val, 2);
-    val += __shfl_xor_sync(0xffffffff, val, 1);
-    return val;
-}
-
-// --- Kernel 1: per-chunk max, atomic-max into row_max ---
-__global__ void max_kernel(const float* x, float* row_max, int N, int chunk_size) {
-    const int num_chunks = gridDim.x;  // per row: blockDim.y chunks
+#define WARP 32
+
+// Atomic max for float using CAS (IEEE 754 bit-cast trick)
+__device__ __forceinline__ void atomicMaxF32(float* addr, float val) {
+    int* iaddr = (int*)addr;
+    int old = *iaddr;
+    while (true) {
+        float old_f = __int_as_float(old);
+        if (val <= old_f) break;
+        int assumed = old;
+        old = atomicCAS(iaddr, assumed, __float_as_int(val));
+        if (old == assumed) break;
+    }
+}
+
+// Warp-level max
+__device__ __forceinline__ float warp_max(float v) {
+    #pragma unroll
+    for (int mask = WARP / 2; mask >= 1; mask >>= 1)
+        v = fmaxf(v, __shfl_xor_sync(0xffffffff, v, mask));
+    return v;
+}
+
+// Warp-level sum
+__device__ __forceinline__ float warp_sum(float v) {
+    #pragma unroll
+    for (int mask = WARP / 2; mask >= 1; mask >>= 1)
+        v += __shfl_xor_sync(0xffffffff, v, mask);
+    return v;
+}
+
+// Kernel 1: find per-row max with multiple blocks per row
+__global__ void max_kernel(const float* x, float* row_max,
+                           int N, int stride_xr) {
     const int row_id = blockIdx.y;
     const int chunk_id = blockIdx.x;
-
-    const float* row_ptr = x + (long long)row_id * N;
+    const int chunks_per_row = gridDim.x;
     const int tid = threadIdx.x;
-
-    // Each thread loads multiple elements
+    const int nthreads = blockDim.x;
+
+    const float* xp = x + (long long)row_id * (long long)stride_xr;
+    const int chunk_len = (N + chunks_per_row - 1) / chunks_per_row;
+    const int start = chunk_id * chunk_len;
+    const int end = min(start + chunk_len, N);
+
     float local_max = -FLT_MAX;
-    const int start = chunk_id * chunk_size;
-    const int end = min(start + chunk_size, N);
-
-    for (int i = start + tid; i < end; i += blockDim.x) {
-        local_max = max(local_max, row_ptr[i]);
-    }
-
-    // Warp-level max reduce
-    local_max = warp_reduce_max(local_max);
-
-    // Thread 0 of each warp does atomic max
-    if (tid % WARP_SIZE == 0) {
-        atomicMax(row_max + row_id, __float_as_int(local_max));
-    }
-}
-
-// --- Kernel 2: read row_max, Kahan-sum exp(x-max), write normalized ---
-__global__ void norm_kernel(const float* x, float* y, const float* row_max, int N) {
+    for (int i = start + tid; i < end; i += nthreads) {
+        local_max = fmaxf(local_max, xp[i]);
+    }
+
+    local_max = warp_max(local_max);
+
+    if (tid % WARP == 0) {
+        atomicMaxF32(row_max + row_id, local_max);
+    }
+}
+
+// Kernel 2: normalize with Kahan-compensated summation
+// Shared memory for inter-warp Kahan sum reduction + final normalization
+__global__ void norm_kernel(const float* x, float* y, const float* row_max,
+                            int N, int stride_x, int stride_y) {
     const int row_id = blockIdx.x;
     const int tid = threadIdx.x;
-
-    const float* row_ptr = x + (long long)row_id * N;
-    float* out_ptr = y + (long long)row_id * N;
+    const int nthreads = blockDim.x;
+
+    const float* xp = x + (long long)row_id * (long long)stride_x;
+    float* yp = y + (long long)row_id * (long long)stride_y;
     const float m = row_max[row_id];
 
-    // Pass 1: Kahan sum of exp(x - m)
-    float total = 0.0f;
-    float comp = 0.0f;
-
-    for (int i = tid; i < N; i += blockDim.x) {
-        float val = row_ptr[i] - m;
-        float e = expf(val);
-
-        float yk = e - comp;
-        float t = total + yk;
-        comp = (t - total) - yk;
-        total = t;
-    }
-
-    // Warp-level sum reduce
-    total = warp_reduce_sum(total);
-
-    // Broadcast total via shared memory
-    __shared__ float shared_total;
-    if (tid % WARP_SIZE == 0) {
-        atomicAdd(&shared_total, total);
+    // --- Pass 1: Kahan sum of exp(x - m) ---
+    // Each warp does its own Kahan sum, then warp sums are combined
+    float w_total = 0.0f;


... (7244 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -2,7 +2,7 @@
 
 Two-kernel approach for GPU saturation:
   Kernel 1: Multiple blocks per row, each finds local max, atomicCAS-max reduces.
-  Kernel 2: One block per row, reads row_max, Kahan-sums exp(x-max), writes normalized output.
+  Kernel 2: One block per row, reads row_max, Kahan-sums exp(x-max), writes output.
 """
 import torch
 import torch.nn as nn
@@ -12,6 +12,7 @@
 SUPPORTED_PRECISIONS = ["fp32"]
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
+# CUDA source: kernels + host launch wrappers (all in .cu so <<<>>> works)
 CUDA_SRC = r"""
 #include <cuda_runtime.h>
 #include <math.h>
@@ -19,7 +20,7 @@
 
 #define WARP 32
 
-// Atomic max for float using CAS (IEEE 754 bit-cast trick)
+// Atomic max for float using CAS
 __device__ __forceinline__ void atomicMaxF32(float* addr, float val) {
     int* iaddr = (int*)addr;
     int old = *iaddr;
@@ -48,7 +49,7 @@
     return v;
 }
 
-// Kernel 1: find per-row max with multiple blocks per row
+// Kernel 1: multi-block max reduction per row
 __global__ void max_kernel(const float* x, float* row_max,
                            int N, int stride_xr) {
     const int row_id = blockIdx.y;
@@ -75,7 +76,6 @@
 }
 
 // Kernel 2: normalize with Kahan-compensated summation
-// Shared memory for inter-warp Kahan sum reduction + final normalization
 __global__ void norm_kernel(const float* x, float* y, const float* row_max,
                             int N, int stride_x, int stride_y) {
     const int row_id = blockIdx.x;
@@ -86,8 +86,7 @@
     float* yp = y + (long long)row_id * (long long)stride_y;
     const float m = row_max[row_id];
 
-    // --- Pass 1: Kahan sum of exp(x - m) ---
-    // Each warp does its own Kahan sum, then warp sums are combined
+    // Pass 1: Kahan sum of exp(x - m)
     float w_total = 0.0f;
     float w_comp = 0.0f;
 
@@ -99,21 +98,15 @@
         w_total = t;
     }
 
-    // Warp-reduce the Kahan sum
     w_total = warp_sum(w_total);
-    // Also reduce compensation across warp
     w_comp = warp_sum(w_comp);
 
-    // Each warp leader writes to shared memory
     extern __shared__ float smem[];
     if (tid % WARP == 0) {
         smem[tid / WARP] = w_total;
-        // We skip Kahan cross-warp; per-warp Kahan + inter-warp plain sum
-        // is sufficient at these tolerance levels
     }
     __syncthreads();
 
-    // Thread 0 sums all warp contributions
     float total_sum = 0.0f;
     if (tid == 0) {
         int nwarps = (nthreads + WARP - 1) / WARP;
@@ -126,58 +119,42 @@
 
     total_sum = smem[0];
 
-    // --- Pass 2: write normalized output ---
+    // Pass 2: write normalized output
     float inv_sum = 1.0f / total_sum;
     for (int i = tid; i < N; i += nthreads) {
         float e = expf(xp[i] - m);
         yp[i] = e * inv_sum;
     }
 }
+
+// Host functions callable from Python
+void softmax_fwd(torch::Tensor x, torch::Tensor y, torch::Tensor row_max,
+                 int N, int stride_xr, int stride_yr, int num_chunks) {
+    int batch = row_max.size(0);
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
+
+    int threads1 = 256;
+    dim3 grid1(num_chunks, batch);
+    dim3 block1(threads1);
+    max_kernel<<<grid1, block1, 0, stream>>>(
+        x.data_ptr<float>(), row_max.data_ptr<float>(), N, stride_xr);
+
+    int threads2 = 256;
+    int smem_size = ((threads2 + WARP - 1) / WARP) * sizeof(float);
+    dim3 grid2(batch);
+    dim3 block2(threads2);
+    norm_kernel<<<grid2, block2, smem_size, stream>>>(
+        x.data_ptr<float>(), y.data_ptr<float>(), row_max.data_ptr<float>(),
+        N, stride_xr, stride_yr);
+}
 """
 
-# C++ wrapper that creates pybind11 functions
-CPP_SRC = """
+CPP_SRC = r"""
 #include <torch/extension.h>
 #include <ATen/cuda/CUDAContext.h>
 
-// Forward declarations of CUDA kernels
-void max_kernel(const float* x, float* row_max, int N, int stride_xr,
-                int num_chunks, int batch, int threads_per_block);
-void norm_kernel(const float* x, float* y, const float* row_max,
-                 int N, int stride_x, int stride_y, int batch, int threads_per_block);
-
-// Launch wrappers
-void launch_max_kernel(torch::Tensor x, torch::Tensor row_max, int N,
-                       int stride_xr, int num_chunks) {
-    at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream();
-    int threads = 256;
-    dim3 grid(num_chunks, row_max.size(0));
-    dim3 block(threads);
-    max_kernel<<<grid, block, 0, stream>>>(
-        x.data_ptr<float>(), row_max.data_ptr<float>(), N, stride_xr);
-}
-
-void launch_norm_kernel(torch::Tensor x, torch::Tensor y, torch::Tensor row_max,
-                        int N, int stride_x, int stride_y) {
-    at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream();
-    int threads = 256;
-    int smem_size = ((threads + 31) / 32) * sizeof(float);
-    dim3 grid(row_max.size(0));
-    dim3 block(threads);
-    norm_kernel<<<grid, block, smem_size, stream>>>(
-        x.data_ptr<float>(), y.data_ptr<float>(), row_max.data_ptr<float>(),
-        N, stride_x, stride_y);
-}
-
 void softmax_fwd(torch::Tensor x, torch::Tensor y, torch::Tensor row_max,
-                 int N, int stride_xr, int stride_yr, int num_chunks) {
-    launch_max_kernel(x, row_max, N, stride_xr, num_chunks);
-    launch_norm_kernel(x, y, row_max, N, stride_xr, stride_yr);
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("softmax_fwd", &softmax_fwd, "CUDA softmax forward pass");
-}
+                 int N, int stride_xr, int stride_yr, int num_chunks);
 """
 
 
@@ -193,12 +170,7 @@
         cpp_sources=CPP_SRC,
         cuda_sources=CUDA_SRC,
         functions=["softmax_fwd"],
-        extra_cuda_cflags=[
-            "-O3",
-            "-arch=sm_120",
-            "-std=c++17",
-            "--use_fast_math",
-        ],
+        extra_cuda_cflags=[

... (382 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -149,13 +149,7 @@
 }
 """
 
-CPP_SRC = r"""
-#include <torch/extension.h>
-#include <ATen/cuda/CUDAContext.h>
-
-void softmax_fwd(torch::Tensor x, torch::Tensor y, torch::Tensor row_max,
-                 int N, int stride_xr, int stride_yr, int num_chunks);
-"""
+CPP_SRC = ""
 
 
 _module = None

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -14,6 +14,8 @@
 
 # CUDA source: kernels + host launch wrappers (all in .cu so <<<>>> works)
 CUDA_SRC = r"""
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
 #include <cuda_runtime.h>
 #include <math.h>
 #include <cfloat>

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -151,7 +151,7 @@
 }
 """
 
-CPP_SRC = ""
+CPP_SRC = "void softmax_fwd(torch::Tensor, torch::Tensor, torch::Tensor, int, int, int, int);"
 
 
 _module = None

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -1,196 +1,118 @@
-"""CUDA C++ softmax with fp32 Kahan-compensated summation.
+"""Triton softmax with online max-reduce + Kahan-compensated summation.
 
-Two-kernel approach for GPU saturation:
-  Kernel 1: Multiple blocks per row, each finds local max, atomicCAS-max reduces.
-  Kernel 2: One block per row, reads row_max, Kahan-sums exp(x-max), writes output.
+Single read pass for max+sum, single write pass for normalization.
+2 memory passes total vs 3 for naive.
 """
 import torch
 import torch.nn as nn
-from torch.utils.cpp_extension import load_inline
+import triton
+import triton.language as tl
+
 
 OP_TYPE = "softmax"
 SUPPORTED_PRECISIONS = ["fp32"]
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
-# CUDA source: kernels + host launch wrappers (all in .cu so <<<>>> works)
-CUDA_SRC = r"""
-#include <torch/extension.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <cuda_runtime.h>
-#include <math.h>
-#include <cfloat>
 
-#define WARP 32
+@triton.jit
+def _softmax_kernel_online(
+    X_ptr,
+    Y_ptr,
+    stride_xr,
+    stride_yr,
+    N,
+    BLOCK_N: tl.constexpr,
+):
+    row_id = tl.program_id(0)
+    x_base = X_ptr + row_id * stride_xr
+    y_base = Y_ptr + row_id * stride_yr
 
-// Atomic max for float using CAS
-__device__ __forceinline__ void atomicMaxF32(float* addr, float val) {
-    int* iaddr = (int*)addr;
-    int old = *iaddr;
-    while (true) {
-        float old_f = __int_as_float(old);
-        if (val <= old_f) break;
-        int assumed = old;
-        old = atomicCAS(iaddr, assumed, __float_as_int(val));
-        if (old == assumed) break;
-    }
-}
+    # --- Online pass: track running_max and Kahan-sum of exp(x-max) ------
+    running_max = tl.full([1], float('-inf'), dtype=tl.float32)
+    running_sum = tl.full([1], 0.0, dtype=tl.float32)
+    comp = tl.full([1], 0.0, dtype=tl.float32)
 
-// Warp-level max
-__device__ __forceinline__ float warp_max(float v) {
-    #pragma unroll
-    for (int mask = WARP / 2; mask >= 1; mask >>= 1)
-        v = fmaxf(v, __shfl_xor_sync(0xffffffff, v, mask));
-    return v;
-}
+    for off in range(0, N, BLOCK_N):
+        idx = off + tl.arange(0, BLOCK_N)
+        mask = idx < N
+        x = tl.load(x_base + idx, mask=mask, other=float('-inf')).to(tl.float32)
 
-// Warp-level sum
-__device__ __forceinline__ float warp_sum(float v) {
-    #pragma unroll
-    for (int mask = WARP / 2; mask >= 1; mask >>= 1)
-        v += __shfl_xor_sync(0xffffffff, v, mask);
-    return v;
-}
+        # Find max of this chunk
+        chunk_max = tl.max(x, keep_dims=True)
 
-// Kernel 1: multi-block max reduction per row
-__global__ void max_kernel(const float* x, float* row_max,
-                           int N, int stride_xr) {
-    const int row_id = blockIdx.y;
-    const int chunk_id = blockIdx.x;
-    const int chunks_per_row = gridDim.x;
-    const int tid = threadIdx.x;
-    const int nthreads = blockDim.x;
+        # New running max
+        new_max = tl.maximum(running_max, chunk_max)
 
-    const float* xp = x + (long long)row_id * (long long)stride_xr;
-    const int chunk_len = (N + chunks_per_row - 1) / chunks_per_row;
-    const int start = chunk_id * chunk_len;
-    const int end = min(start + chunk_len, N);
+        # Scale existing sum by exp(old_max - new_max)
+        scale = tl.exp(running_max - new_max)
 
-    float local_max = -FLT_MAX;
-    for (int i = start + tid; i < end; i += nthreads) {
-        local_max = fmaxf(local_max, xp[i]);
-    }
+        # Scale existing sum and compensation
+        running_sum = running_sum * scale
+        comp = comp * scale
 
-    local_max = warp_max(local_max);
+        # Add new chunk contributions: exp(x - new_max)
+        e = tl.exp(x - new_max)
+        e_masked = tl.where(mask, e, 0.0)
+        e_chunk = tl.sum(e_masked, keep_dims=True).to(tl.float32)
 
-    if (tid % WARP == 0) {
-        atomicMaxF32(row_max + row_id, local_max);
-    }
-}
+        # Kahan add
+        y_k = e_chunk - comp
+        new_sum = running_sum + y_k
+        comp = (new_sum - running_sum) - y_k
+        running_sum = new_sum
 
-// Kernel 2: normalize with Kahan-compensated summation
-__global__ void norm_kernel(const float* x, float* y, const float* row_max,
-                            int N, int stride_x, int stride_y) {
-    const int row_id = blockIdx.x;
-    const int tid = threadIdx.x;
-    const int nthreads = blockDim.x;
+        running_max = new_max
 
-    const float* xp = x + (long long)row_id * (long long)stride_x;
-    float* yp = y + (long long)row_id * (long long)stride_y;
-    const float m = row_max[row_id];
+    # Broadcast [1] scalars to [BLOCK_N] for element-wise ops
+    final_max = running_max.broadcast_to([BLOCK_N])
+    total_sum_vec = running_sum.broadcast_to([BLOCK_N])
 
-    // Pass 1: Kahan sum of exp(x - m)
-    float w_total = 0.0f;
-    float w_comp = 0.0f;
-
-    for (int i = tid; i < N; i += nthreads) {
-        float e = expf(xp[i] - m);
-        float yk = e - w_comp;
-        float t = w_total + yk;
-        w_comp = (t - w_total) - yk;
-        w_total = t;
-    }
-
-    w_total = warp_sum(w_total);
-    w_comp = warp_sum(w_comp);
-
-    extern __shared__ float smem[];
-    if (tid % WARP == 0) {
-        smem[tid / WARP] = w_total;
-    }
-    __syncthreads();
-
-    float total_sum = 0.0f;
-    if (tid == 0) {
-        int nwarps = (nthreads + WARP - 1) / WARP;
-        for (int w = 0; w < nwarps; w++) {
-            total_sum += smem[w];
-        }
-        smem[0] = total_sum;
-    }
-    __syncthreads();
-
-    total_sum = smem[0];
-
-    // Pass 2: write normalized output
-    float inv_sum = 1.0f / total_sum;
-    for (int i = tid; i < N; i += nthreads) {
-        float e = expf(xp[i] - m);
-        yp[i] = e * inv_sum;
-    }
-}
-
-// Host functions callable from Python
-void so

... (3560 more chars)

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/sweep_configs.py
@@ -0,0 +1,124 @@
+"""Sweep over BLOCK_N and num_warps to find optimal config."""
+import statistics
+import torch
+import triton
+import triton.language as tl
+import solution
+
+@triton.jit
+def _softmax_kernel_online(
+    X_ptr, Y_ptr, stride_xr, stride_yr, N, BLOCK_N: tl.constexpr,
+):
+    row_id = tl.program_id(0)
+    x_base = X_ptr + row_id * stride_xr
+    y_base = Y_ptr + row_id * stride_yr
+
+    running_max = tl.full([1], float('-inf'), dtype=tl.float32)
+    running_sum = tl.full([1], 0.0, dtype=tl.float32)
+    comp = tl.full([1], 0.0, dtype=tl.float32)
+
+    for off in range(0, N, BLOCK_N):
+        idx = off + tl.arange(0, BLOCK_N)
+        mask = idx < N
+        x = tl.load(x_base + idx, mask=mask, other=float('-inf')).to(tl.float32)
+        chunk_max = tl.max(x, keep_dims=True)
+        new_max = tl.maximum(running_max, chunk_max)
+        scale = tl.exp(running_max - new_max)
+        running_sum = running_sum * scale
+        comp = comp * scale
+        e = tl.exp(x - new_max)
+        e_masked = tl.where(mask, e, 0.0)
+        e_chunk = tl.sum(e_masked, keep_dims=True).to(tl.float32)
+        y_k = e_chunk - comp
+        new_sum = running_sum + y_k
+        comp = (new_sum - running_sum) - y_k
+        running_sum = new_sum
+        running_max = new_max
+
+    final_max = running_max.broadcast_to([BLOCK_N])
+    total_sum_vec = running_sum.broadcast_to([BLOCK_N])
+
+    for off in range(0, N, BLOCK_N):
+        idx = off + tl.arange(0, BLOCK_N)
+        mask = idx < N
+        x = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32)
+        e = tl.exp(x - final_max)
+        y = tl.where(mask, e / total_sum_vec, 0.0)
+        tl.store(y_base + idx, y, mask=mask)
+
+
+def time_config(batch, vocab, block_n, num_warps, num_stages, iters=10):
+    x = torch.randn(batch, vocab, device='cuda', dtype=torch.float32)
+    y = torch.empty_like(x)
+    
+    while vocab % block_n != 0 and block_n > 32:
+        block_n //= 2
+    
+    for _ in range(3):
+        _softmax_kernel_online[(batch,)](
+            x, y, x.stride(0), y.stride(0), vocab,
+            BLOCK_N=block_n, num_warps=num_warps, num_stages=num_stages,
+        )
+    torch.cuda.synchronize()
+    
+    times = []
+    for _ in range(iters):
+        torch.cuda.synchronize()
+        s = torch.cuda.Event(enable_timing=True)
+        e = torch.cuda.Event(enable_timing=True)
+        s.record()
+        _softmax_kernel_online[(batch,)](
+            x, y, x.stride(0), y.stride(0), vocab,
+            BLOCK_N=block_n, num_warps=num_warps, num_stages=num_stages,
+        )
+        e.record()
+        torch.cuda.synchronize()
+        times.append(s.elapsed_time(e))
+    return statistics.median(times)
+
+
+configs = [
+    (512, 4, 2), (512, 4, 3), (512, 8, 2), (512, 8, 3),
+    (1024, 4, 2), (1024, 4, 3), (1024, 8, 2), (1024, 8, 3),
+    (2048, 4, 2), (2048, 4, 3), (2048, 8, 2), (2048, 8, 3),
+    (4096, 4, 2), (4096, 4, 3), (4096, 8, 2), (4096, 8, 3),
+]
+
+shapes = [
+    (32, 4096, "32x4K"),
+    (16, 32768, "16x32K"),
+    (8, 131072, "8x128K"),
+    (4, 262144, "4x256K"),
+]
+
+for batch, vocab, label in shapes:
+    print(f"\n=== {label} (batch={batch}, vocab={vocab}) ===")
+    best_ms = float('inf')
+    best_cfg = None
+    
+    for block_n_raw, num_warps, num_stages in configs:
+        block_n = min(block_n_raw, vocab)
+        if vocab % block_n != 0:
+            # Find largest divisor of vocab <= block_n_raw
+            candidate = block_n
+            while vocab % candidate != 0 and candidate > 32:
+                candidate //= 2
+            if vocab % candidate != 0:
+                continue
+            block_n = candidate
+        
+        try:
+            ms = time_config(batch, vocab, block_n, num_warps, num_stages)
+        except Exception as e:
+            print(f"  BLK={block_n} warps={num_warps} stages={num_stages}: ERROR ({e})")
+            continue
+        
+        bw = (batch * vocab * 4 * 2) / (ms * 1e-3) / 1e9
+        if ms < best_ms:
+            best_ms = ms
+            best_cfg = (block_n, num_warps, num_stages)
+        
+        print(f"  BLK={block_n:5d} warps={num_warps} stages={num_stages}: {ms:.3f}ms ({bw:.0f} GB/s)")
+    
+    if best_cfg:
+        print(f"  >>> BEST: BLK={best_cfg[0]} warps={best_cfg[1]} stages={best_cfg[2]} @ {best_ms:.3f}ms <<<")

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -78,22 +78,13 @@
 
 
 def _pick_block_n(vocab: int) -> int:
-    """Choose BLOCK_N: trade register pressure vs loop unroll factor.
+    """Choose BLOCK_N based on vocab size for optimal performance.
     
-    Larger BLOCK_N = more ILP per iteration but more registers needed.
-    Optimal depends on the device's register file size and occupancy targets.
+    Empirically determined via sweep: larger blocks + more warps = better.
     """
-    # For Blackwell (sm_120) with large register file:
-    # - BLOCK_N=1024: 4 registers for pointer arithmetic, ~16 for loaded data
-    # - BLOCK_N=2048: similar but bigger load vectors
-    # - BLOCK_N=4096: high register pressure
-    if vocab <= 1024:
+    if vocab <= 2048:
         return vocab
-    if vocab <= 2048:
-        return 1024
-    if vocab <= 8192:
-        return 2048
-    return 2048
+    return min(4096, vocab)
 
 
 def triton_softmax(x: torch.Tensor) -> torch.Tensor:

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -101,8 +101,8 @@
     _softmax_kernel_online[grid](
         x, y, x.stride(0), y.stride(0), vocab,
         BLOCK_N=block_n,
-        num_warps=4,
-        num_stages=3,
+        num_warps=8,
+        num_stages=2,
     )
     return y
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -1,7 +1,10 @@
 """Triton softmax with online max-reduce + Kahan-compensated summation.
 
-Single read pass for max+sum, single write pass for normalization.
-2 memory passes total vs 3 for naive.
+Multi-block-per-row variant for GPU saturation:
+  Each row is split into chunks; multiple programs handle chunks in parallel.
+  Cross-chunk reduction happens through shared (on-chip) atomics + sync.
+
+For small batch sizes, we need multiple programs per row to saturate the GPU.
 """
 import torch
 import torch.nn as nn
@@ -15,7 +18,81 @@
 
 
 @triton.jit
-def _softmax_kernel_online(
+def _softmax_kernel_split(
+    X_ptr,
+    Y_ptr,
+    row_max_ptr,
+    stride_xr,
+    stride_yr,
+    N,
+    BLOCK_N: tl.constexpr,
+    NUM_CHUNKS: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    row_id = pid // NUM_CHUNKS
+    chunk_id = pid % NUM_CHUNKS
+
+    x_base = X_ptr + row_id * stride_xr
+    y_base = Y_ptr + row_id * stride_yr
+
+    # --- Phase 1: Per-block online max+sum ---
+    # Each block handles its chunk online
+    x = tl.load(x_base + chunk_id * BLOCK_N + tl.arange(0, BLOCK_N),
+                mask=(chunk_id * BLOCK_N + tl.arange(0, BLOCK_N)) < N,
+                other=float('-inf')).to(tl.float32)
+
+    chunk_max = tl.max(x, keep_dims=True)
+    chunk_sum = tl.sum(tl.where(
+        (chunk_id * BLOCK_N + tl.arange(0, BLOCK_N)) < N,
+        tl.exp(x - chunk_max),
+        0.0
+    ), keep_dims=True).to(tl.float32)
+
+    # Write partial results to shared workspace
+    # Use atomic max for row_max
+    cm_s = chunk_max.reshape(())
+    # atomic_max on row_max[row_id]
+    tl.atomic_max(row_max_ptr + row_id, cm_s)
+
+
+class Model(nn.Module):
+    def __init__(self, batch: int, vocab: int):
+        super().__init__()
+        self.batch = batch
+        self.vocab = vocab
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Fall back to working single-block kernel for now
+        return _triton_softmax_single(x)
+
+
+def _triton_softmax_single(x: torch.Tensor) -> torch.Tensor:
+    """Proven working single-block online softmax."""
+    assert x.is_cuda and x.dtype == torch.float32 and x.ndim == 2
+    batch, vocab = x.shape
+    y = torch.empty_like(x)
+
+    if vocab <= 2048:
+        block_n = vocab
+        num_warps = 4
+    else:
+        block_n = min(4096, vocab)
+        num_warps = 8
+
+    while vocab % block_n != 0 and block_n > 32:
+        block_n //= 2
+
+    _softmax_kernel_online_v2[(batch,)](
+        x, y, x.stride(0), y.stride(0), vocab,
+        BLOCK_N=block_n,
+        num_warps=num_warps,
+        num_stages=2,
+    )
+    return y
+
+
+@triton.jit
+def _softmax_kernel_online_v2(
     X_ptr,
     Y_ptr,
     stride_xr,
@@ -27,7 +104,6 @@
     x_base = X_ptr + row_id * stride_xr
     y_base = Y_ptr + row_id * stride_yr
 
-    # --- Online pass: track running_max and Kahan-sum of exp(x-max) ------
     running_max = tl.full([1], float('-inf'), dtype=tl.float32)
     running_sum = tl.full([1], 0.0, dtype=tl.float32)
     comp = tl.full([1], 0.0, dtype=tl.float32)
@@ -37,25 +113,17 @@
         mask = idx < N
         x = tl.load(x_base + idx, mask=mask, other=float('-inf')).to(tl.float32)
 
-        # Find max of this chunk
         chunk_max = tl.max(x, keep_dims=True)
-
-        # New running max
         new_max = tl.maximum(running_max, chunk_max)
-
-        # Scale existing sum by exp(old_max - new_max)
         scale = tl.exp(running_max - new_max)
 
-        # Scale existing sum and compensation
         running_sum = running_sum * scale
         comp = comp * scale
 
-        # Add new chunk contributions: exp(x - new_max)
         e = tl.exp(x - new_max)
         e_masked = tl.where(mask, e, 0.0)
         e_chunk = tl.sum(e_masked, keep_dims=True).to(tl.float32)
 
-        # Kahan add
         y_k = e_chunk - comp
         new_sum = running_sum + y_k
         comp = (new_sum - running_sum) - y_k
@@ -63,11 +131,9 @@
 
         running_max = new_max
 
-    # Broadcast [1] scalars to [BLOCK_N] for element-wise ops
     final_max = running_max.broadcast_to([BLOCK_N])
     total_sum_vec = running_sum.broadcast_to([BLOCK_N])
 
-    # --- Write normalized output ----------------------------------------
     for off in range(0, N, BLOCK_N):
         idx = off + tl.arange(0, BLOCK_N)
         mask = idx < N
@@ -75,46 +141,6 @@
         e = tl.exp(x - final_max)
         y = tl.where(mask, e / total_sum_vec, 0.0)
         tl.store(y_base + idx, y, mask=mask)
-
-
-def _pick_block_n(vocab: int) -> int:
-    """Choose BLOCK_N based on vocab size for optimal performance.
-    
-    Empirically determined via sweep: larger blocks + more warps = better.
-    """
-    if vocab <= 2048:
-        return vocab
-    return min(4096, vocab)
-
-
-def triton_softmax(x: torch.Tensor) -> torch.Tensor:
-    assert x.is_cuda and x.dtype == torch.float32 and x.ndim == 2
-    batch, vocab = x.shape
-    y = torch.empty_like(x)
-
-    block_n = _pick_block_n(vocab)
-    # Ensure BLOCK_N divides N
-    while vocab % block_n != 0 and block_n > 32:
-        block_n //= 2
-
-    grid = (batch,)
-    _softmax_kernel_online[grid](
-        x, y, x.stride(0), y.stride(0), vocab,
-        BLOCK_N=block_n,
-        num_warps=8,
-        num_stages=2,
-    )
-    return y
-
-
-class Model(nn.Module):
-    def __init__(self, batch: int, vocab: int):
-        super().__init__()
-        self.batch = batch
-        self.vocab = vocab
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return triton_softmax(x)
 
 
 BATCH = 8

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -1,11 +1,12 @@
 """Triton softmax with online max-reduce + Kahan-compensated summation.
 
-Multi-block-per-row variant for GPU saturation:
-  Each row is split into chunks; multiple programs handle chunks in parallel.
-  Cross-chunk reduction happens through shared (on-chip) atomics + sync.
+Two-phase approach:
+  Phase 1 (online max+sum): 1 block per row reads X once, writes row_max and row_sum.
+  Phase 2 (normalize):       Multiple blocks per row read X, row_max, row_sum and write Y.
 
-For small batch sizes, we need multiple programs per row to saturate the GPU.
+Phase 2 has multiple blocks per row for GPU saturation even at small batch sizes.
 """
+import math
 import torch
 import torch.nn as nn
 import triton
@@ -18,91 +19,17 @@
 
 
 @triton.jit
-def _softmax_kernel_split(
+def _online_max_kernel(
     X_ptr,
-    Y_ptr,
     row_max_ptr,
+    row_sum_ptr,
     stride_xr,
-    stride_yr,
-    N,
-    BLOCK_N: tl.constexpr,
-    NUM_CHUNKS: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    row_id = pid // NUM_CHUNKS
-    chunk_id = pid % NUM_CHUNKS
-
-    x_base = X_ptr + row_id * stride_xr
-    y_base = Y_ptr + row_id * stride_yr
-
-    # --- Phase 1: Per-block online max+sum ---
-    # Each block handles its chunk online
-    x = tl.load(x_base + chunk_id * BLOCK_N + tl.arange(0, BLOCK_N),
-                mask=(chunk_id * BLOCK_N + tl.arange(0, BLOCK_N)) < N,
-                other=float('-inf')).to(tl.float32)
-
-    chunk_max = tl.max(x, keep_dims=True)
-    chunk_sum = tl.sum(tl.where(
-        (chunk_id * BLOCK_N + tl.arange(0, BLOCK_N)) < N,
-        tl.exp(x - chunk_max),
-        0.0
-    ), keep_dims=True).to(tl.float32)
-
-    # Write partial results to shared workspace
-    # Use atomic max for row_max
-    cm_s = chunk_max.reshape(())
-    # atomic_max on row_max[row_id]
-    tl.atomic_max(row_max_ptr + row_id, cm_s)
-
-
-class Model(nn.Module):
-    def __init__(self, batch: int, vocab: int):
-        super().__init__()
-        self.batch = batch
-        self.vocab = vocab
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # Fall back to working single-block kernel for now
-        return _triton_softmax_single(x)
-
-
-def _triton_softmax_single(x: torch.Tensor) -> torch.Tensor:
-    """Proven working single-block online softmax."""
-    assert x.is_cuda and x.dtype == torch.float32 and x.ndim == 2
-    batch, vocab = x.shape
-    y = torch.empty_like(x)
-
-    if vocab <= 2048:
-        block_n = vocab
-        num_warps = 4
-    else:
-        block_n = min(4096, vocab)
-        num_warps = 8
-
-    while vocab % block_n != 0 and block_n > 32:
-        block_n //= 2
-
-    _softmax_kernel_online_v2[(batch,)](
-        x, y, x.stride(0), y.stride(0), vocab,
-        BLOCK_N=block_n,
-        num_warps=num_warps,
-        num_stages=2,
-    )
-    return y
-
-
-@triton.jit
-def _softmax_kernel_online_v2(
-    X_ptr,
-    Y_ptr,
-    stride_xr,
-    stride_yr,
     N,
     BLOCK_N: tl.constexpr,
 ):
+    """Single block per row: online max + Kahan sum of exp(x-max)."""
     row_id = tl.program_id(0)
     x_base = X_ptr + row_id * stride_xr
-    y_base = Y_ptr + row_id * stride_yr
 
     running_max = tl.full([1], float('-inf'), dtype=tl.float32)
     running_sum = tl.full([1], 0.0, dtype=tl.float32)
@@ -131,16 +58,106 @@
 
         running_max = new_max
 
-    final_max = running_max.broadcast_to([BLOCK_N])
-    total_sum_vec = running_sum.broadcast_to([BLOCK_N])
+    # Write results (scalar per row)
+    rm = running_max.reshape(())
+    rs = running_sum.reshape(())
+    tl.store(row_max_ptr + row_id, rm)
+    tl.store(row_sum_ptr + row_id, rs)
 
-    for off in range(0, N, BLOCK_N):
+
+@triton.jit
+def _normalize_kernel(
+    X_ptr,
+    Y_ptr,
+    row_max_ptr,
+    row_sum_ptr,
+    stride_xr,
+    stride_yr,
+    N,
+    BLOCK_N: tl.constexpr,
+    NUM_CHUNKS: tl.constexpr,
+):
+    """Multiple blocks per row for GPU saturation."""
+    pid = tl.program_id(0)
+    row_id = pid // NUM_CHUNKS
+    chunk_id = pid % NUM_CHUNKS
+
+    x_base = X_ptr + row_id * stride_xr
+    y_base = Y_ptr + row_id * stride_yr
+
+    m = tl.load(row_max_ptr + row_id).to(tl.float32)
+    s = tl.load(row_sum_ptr + row_id).to(tl.float32)
+
+    # Broadcast scalars to chunk shape (implicit broadcast)
+    inv_sum = 1.0 / s
+
+    chunk_start = chunk_id * BLOCK_N
+
+    for off in range(chunk_start, min(chunk_start + BLOCK_N, N), BLOCK_N):
         idx = off + tl.arange(0, BLOCK_N)
         mask = idx < N
         x = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32)
-        e = tl.exp(x - final_max)
-        y = tl.where(mask, e / total_sum_vec, 0.0)
+        e = tl.exp(x - m)
+        y = tl.where(mask, e * inv_sum, 0.0)
         tl.store(y_base + idx, y, mask=mask)
+
+
+def triton_softmax(x: torch.Tensor) -> torch.Tensor:
+    assert x.is_cuda and x.dtype == torch.float32 and x.ndim == 2
+    batch, vocab = x.shape
+    y = torch.empty_like(x)
+
+    row_max = torch.empty((batch,), dtype=torch.float32, device=x.device)
+    row_sum = torch.empty((batch,), dtype=torch.float32, device=x.device)
+
+    # Phase 1: online max+sum (1 block per row)
+    if vocab <= 2048:
+        block_n1 = vocab
+    else:
+        block_n1 = min(4096, vocab)
+    while vocab % block_n1 != 0 and block_n1 > 32:
+        block_n1 //= 2
+
+    _online_max_kernel[(batch,)](
+        x, row_max, row_sum, x.stride(0), vocab,
+        BLOCK_N=block_n1, num_warps=8, num_stages=2,
+    )
+
+    # Phase 2: normalize (multiple blocks per row for saturation)
+    if vocab <= 1024:
+        block_n2 = vocab
+        num_chunks = 1
+    elif vocab <= 4096:
+        block_n2 = 1024
+        num_chunks = (vocab + block_n2 - 1) // block_n2
+    elif vocab <= 16384:
+        block_n2 = 2048
+        num_chunks = (vocab + block_n2 - 

... (729 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -1,12 +1,9 @@
 """Triton softmax with online max-reduce + Kahan-compensated summation.
 
-Two-phase approach:
-  Phase 1 (online max+sum): 1 block per row reads X once, writes row_max and row_sum.
-  Phase 2 (normalize):       Multiple blocks per row read X, row_max, row_sum and write Y.
-
-Phase 2 has multiple blocks per row for GPU saturation even at small batch sizes.
+Optimal single-kernel approach tuned for RTX PRO 6000 Blackwell.
+Sweep-validated optimal: BLOCK_N=4096, 8 warps for large shapes;
+BLOCK_N=2048, 4 warps for small shapes.
 """
-import math
 import torch
 import torch.nn as nn
 import triton
@@ -19,18 +16,19 @@
 
 
 @triton.jit
-def _online_max_kernel(
+def _softmax_kernel(
     X_ptr,
-    row_max_ptr,
-    row_sum_ptr,
+    Y_ptr,
     stride_xr,
+    stride_yr,
     N,
     BLOCK_N: tl.constexpr,
 ):
-    """Single block per row: online max + Kahan sum of exp(x-max)."""
     row_id = tl.program_id(0)
     x_base = X_ptr + row_id * stride_xr
+    y_base = Y_ptr + row_id * stride_yr
 
+    # --- Online pass: running max + Kahan-sum -------------------------------
     running_max = tl.full([1], float('-inf'), dtype=tl.float32)
     running_sum = tl.full([1], 0.0, dtype=tl.float32)
     comp = tl.full([1], 0.0, dtype=tl.float32)
@@ -58,47 +56,17 @@
 
         running_max = new_max
 
-    # Write results (scalar per row)
-    rm = running_max.reshape(())
-    rs = running_sum.reshape(())
-    tl.store(row_max_ptr + row_id, rm)
-    tl.store(row_sum_ptr + row_id, rs)
+    # Broadcast [1] to [BLOCK_N] for element-wise ops
+    final_max = running_max.broadcast_to([BLOCK_N])
+    total_sum_vec = running_sum.broadcast_to([BLOCK_N])
 
-
-@triton.jit
-def _normalize_kernel(
-    X_ptr,
-    Y_ptr,
-    row_max_ptr,
-    row_sum_ptr,
-    stride_xr,
-    stride_yr,
-    N,
-    BLOCK_N: tl.constexpr,
-    NUM_CHUNKS: tl.constexpr,
-):
-    """Multiple blocks per row for GPU saturation."""
-    pid = tl.program_id(0)
-    row_id = pid // NUM_CHUNKS
-    chunk_id = pid % NUM_CHUNKS
-
-    x_base = X_ptr + row_id * stride_xr
-    y_base = Y_ptr + row_id * stride_yr
-
-    m = tl.load(row_max_ptr + row_id).to(tl.float32)
-    s = tl.load(row_sum_ptr + row_id).to(tl.float32)
-
-    # Broadcast scalars to chunk shape (implicit broadcast)
-    inv_sum = 1.0 / s
-
-    chunk_start = chunk_id * BLOCK_N
-
-    for off in range(chunk_start, min(chunk_start + BLOCK_N, N), BLOCK_N):
+    # --- Normalize pass: write output ---------------------------------------
+    for off in range(0, N, BLOCK_N):
         idx = off + tl.arange(0, BLOCK_N)
         mask = idx < N
         x = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32)
-        e = tl.exp(x - m)
-        y = tl.where(mask, e * inv_sum, 0.0)
+        e = tl.exp(x - final_max)
+        y = tl.where(mask, e / total_sum_vec, 0.0)
         tl.store(y_base + idx, y, mask=mask)
 
 
@@ -107,46 +75,30 @@
     batch, vocab = x.shape
     y = torch.empty_like(x)
 
-    row_max = torch.empty((batch,), dtype=torch.float32, device=x.device)
-    row_sum = torch.empty((batch,), dtype=torch.float32, device=x.device)
+    # Sweep-validated optimal block sizes
+    if vocab <= 2048:
+        block_n = vocab
+        num_warps = 4
+        num_stages = 3
+    elif vocab <= 4096:
+        block_n = 2048
+        num_warps = 4
+        num_stages = 3
+    else:
+        block_n = min(4096, vocab)
+        num_warps = 8
+        num_stages = 2
 
-    # Phase 1: online max+sum (1 block per row)
-    if vocab <= 2048:
-        block_n1 = vocab
-    else:
-        block_n1 = min(4096, vocab)
-    while vocab % block_n1 != 0 and block_n1 > 32:
-        block_n1 //= 2
+    # Ensure divisibility
+    while vocab % block_n != 0 and block_n > 32:
+        block_n //= 2
 
-    _online_max_kernel[(batch,)](
-        x, row_max, row_sum, x.stride(0), vocab,
-        BLOCK_N=block_n1, num_warps=8, num_stages=2,
+    _softmax_kernel[(batch,)](
+        x, y, x.stride(0), y.stride(0), vocab,
+        BLOCK_N=block_n,
+        num_warps=num_warps,
+        num_stages=num_stages,
     )
-
-    # Phase 2: normalize (multiple blocks per row for saturation)
-    if vocab <= 1024:
-        block_n2 = vocab
-        num_chunks = 1
-    elif vocab <= 4096:
-        block_n2 = 1024
-        num_chunks = (vocab + block_n2 - 1) // block_n2
-    elif vocab <= 16384:
-        block_n2 = 2048
-        num_chunks = (vocab + block_n2 - 1) // block_n2
-    elif vocab <= 65536:
-        block_n2 = 2048
-        num_chunks = (vocab + block_n2 - 1) // block_n2
-    else:
-        block_n2 = 4096
-        num_chunks = (vocab + block_n2 - 1) // block_n2
-
-    grid2 = (batch * num_chunks,)
-    _normalize_kernel[grid2](
-        x, y, row_max, row_sum, x.stride(0), y.stride(0), vocab,
-        BLOCK_N=block_n2, NUM_CHUNKS=num_chunks,
-        num_warps=8, num_stages=2,
-    )
-
     return y
 
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -21,7 +21,7 @@
     Y_ptr,
     stride_xr,
     stride_yr,
-    N,
+    N: tl.constexpr,
     BLOCK_N: tl.constexpr,
 ):
     row_id = tl.program_id(0)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -1,8 +1,7 @@
 """Triton softmax with online max-reduce + Kahan-compensated summation.
 
-Optimal single-kernel approach tuned for RTX PRO 6000 Blackwell.
-Sweep-validated optimal: BLOCK_N=4096, 8 warps for large shapes;
-BLOCK_N=2048, 4 warps for small shapes.
+Strategy: use the fused 2-pass online kernel but with N as constexpr for
+full loop unroll + optimal block/warp/stage config per shape.
 """
 import torch
 import torch.nn as nn
@@ -53,10 +52,8 @@
         new_sum = running_sum + y_k
         comp = (new_sum - running_sum) - y_k
         running_sum = new_sum
-
         running_max = new_max
 
-    # Broadcast [1] to [BLOCK_N] for element-wise ops
     final_max = running_max.broadcast_to([BLOCK_N])
     total_sum_vec = running_sum.broadcast_to([BLOCK_N])
 
@@ -70,26 +67,25 @@
         tl.store(y_base + idx, y, mask=mask)
 
 
+def _get_launch_config(vocab: int) -> tuple[int, int, int]:
+    """Return (block_n, num_warps, num_stages) optimal for this vocab size.
+    
+    Based on empirical sweep on RTX PRO 6000 Blackwell.
+    """
+    if vocab <= 2048:
+        return vocab, 4, 3
+    if vocab <= 4096:
+        return 2048, 4, 3
+    return min(4096, vocab), 8, 2
+
+
 def triton_softmax(x: torch.Tensor) -> torch.Tensor:
     assert x.is_cuda and x.dtype == torch.float32 and x.ndim == 2
     batch, vocab = x.shape
     y = torch.empty_like(x)
 
-    # Sweep-validated optimal block sizes
-    if vocab <= 2048:
-        block_n = vocab
-        num_warps = 4
-        num_stages = 3
-    elif vocab <= 4096:
-        block_n = 2048
-        num_warps = 4
-        num_stages = 3
-    else:
-        block_n = min(4096, vocab)
-        num_warps = 8
-        num_stages = 2
+    block_n, num_warps, num_stages = _get_launch_config(vocab)
 
-    # Ensure divisibility
     while vocab % block_n != 0 and block_n > 32:
         block_n //= 2
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -1,7 +1,7 @@
 """Triton softmax with online max-reduce + Kahan-compensated summation.
 
-Strategy: use the fused 2-pass online kernel but with N as constexpr for
-full loop unroll + optimal block/warp/stage config per shape.
+Persistent kernel variant: all SMs pull rows from an atomic counter,
+giving full GPU saturation even at small batch sizes.
 """
 import torch
 import torch.nn as nn
@@ -15,63 +15,67 @@
 
 
 @triton.jit
-def _softmax_kernel(
+def _softmax_persistent_kernel(
     X_ptr,
     Y_ptr,
     stride_xr,
     stride_yr,
     N: tl.constexpr,
     BLOCK_N: tl.constexpr,
+    BATCH: tl.constexpr,
+    COUNTER_PTR,
 ):
-    row_id = tl.program_id(0)
-    x_base = X_ptr + row_id * stride_xr
-    y_base = Y_ptr + row_id * stride_yr
+    """Persistent kernel: each SM pulls rows from atomic counter."""
+    while True:
+        row_id = tl.atomic_add(COUNTER_PTR, 1)
+        if row_id >= BATCH:
+            return
 
-    # --- Online pass: running max + Kahan-sum -------------------------------
-    running_max = tl.full([1], float('-inf'), dtype=tl.float32)
-    running_sum = tl.full([1], 0.0, dtype=tl.float32)
-    comp = tl.full([1], 0.0, dtype=tl.float32)
+        x_base = X_ptr + row_id * stride_xr
+        y_base = Y_ptr + row_id * stride_yr
 
-    for off in range(0, N, BLOCK_N):
-        idx = off + tl.arange(0, BLOCK_N)
-        mask = idx < N
-        x = tl.load(x_base + idx, mask=mask, other=float('-inf')).to(tl.float32)
+        # --- Online pass: running max + Kahan-sum ---------------------------
+        running_max = tl.full([1], float('-inf'), dtype=tl.float32)
+        running_sum = tl.full([1], 0.0, dtype=tl.float32)
+        comp = tl.full([1], 0.0, dtype=tl.float32)
 
-        chunk_max = tl.max(x, keep_dims=True)
-        new_max = tl.maximum(running_max, chunk_max)
-        scale = tl.exp(running_max - new_max)
+        for off in range(0, N, BLOCK_N):
+            idx = off + tl.arange(0, BLOCK_N)
+            mask = idx < N
+            x = tl.load(x_base + idx, mask=mask, other=float('-inf')).to(tl.float32)
 
-        running_sum = running_sum * scale
-        comp = comp * scale
+            chunk_max = tl.max(x, keep_dims=True)
+            new_max = tl.maximum(running_max, chunk_max)
+            scale = tl.exp(running_max - new_max)
 
-        e = tl.exp(x - new_max)
-        e_masked = tl.where(mask, e, 0.0)
-        e_chunk = tl.sum(e_masked, keep_dims=True).to(tl.float32)
+            running_sum = running_sum * scale
+            comp = comp * scale
 
-        y_k = e_chunk - comp
-        new_sum = running_sum + y_k
-        comp = (new_sum - running_sum) - y_k
-        running_sum = new_sum
-        running_max = new_max
+            e = tl.exp(x - new_max)
+            e_masked = tl.where(mask, e, 0.0)
+            e_chunk = tl.sum(e_masked, keep_dims=True).to(tl.float32)
 
-    final_max = running_max.broadcast_to([BLOCK_N])
-    total_sum_vec = running_sum.broadcast_to([BLOCK_N])
+            y_k = e_chunk - comp
+            new_sum = running_sum + y_k
+            comp = (new_sum - running_sum) - y_k
+            running_sum = new_sum
+            running_max = new_max
 
-    # --- Normalize pass: write output ---------------------------------------
-    for off in range(0, N, BLOCK_N):
-        idx = off + tl.arange(0, BLOCK_N)
-        mask = idx < N
-        x = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32)
-        e = tl.exp(x - final_max)
-        y = tl.where(mask, e / total_sum_vec, 0.0)
-        tl.store(y_base + idx, y, mask=mask)
+        final_max = running_max.broadcast_to([BLOCK_N])
+        total_sum_vec = running_sum.broadcast_to([BLOCK_N])
+
+        # --- Normalize pass: write output -----------------------------------
+        for off in range(0, N, BLOCK_N):
+            idx = off + tl.arange(0, BLOCK_N)
+            mask = idx < N
+            x = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32)
+            e = tl.exp(x - final_max)
+            y = tl.where(mask, e / total_sum_vec, 0.0)
+            tl.store(y_base + idx, y, mask=mask)
 
 
 def _get_launch_config(vocab: int) -> tuple[int, int, int]:
-    """Return (block_n, num_warps, num_stages) optimal for this vocab size.
-    
-    Based on empirical sweep on RTX PRO 6000 Blackwell.
-    """
+    """Return (block_n, num_warps, num_stages) optimal for this vocab size."""
     if vocab <= 2048:
         return vocab, 4, 3
     if vocab <= 4096:
@@ -79,19 +83,30 @@
     return min(4096, vocab), 8, 2
 
 
+# RTX PRO 6000 has ~128 SMs
+_NUM_SMS = 128
+
+
 def triton_softmax(x: torch.Tensor) -> torch.Tensor:
     assert x.is_cuda and x.dtype == torch.float32 and x.ndim == 2
     batch, vocab = x.shape
     y = torch.empty_like(x)
 
+    # Atomic counter on GPU
+    counter = torch.zeros(1, dtype=torch.int32, device=x.device)
+
     block_n, num_warps, num_stages = _get_launch_config(vocab)
-
     while vocab % block_n != 0 and block_n > 32:
         block_n //= 2
 
-    _softmax_kernel[(batch,)](
+    # Use enough blocks to saturate SMs, but not wasteful
+    num_programs = max(batch, _NUM_SMS)
+
+    _softmax_persistent_kernel[(num_programs,)](
         x, y, x.stride(0), y.stride(0), vocab,
         BLOCK_N=block_n,
+        BATCH=batch,
+        COUNTER_PTR=counter,
         num_warps=num_warps,
         num_stages=num_stages,
     )

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -26,52 +26,9 @@
     COUNTER_PTR,
 ):
     """Persistent kernel: each SM pulls rows from atomic counter."""
-    while True:
-        row_id = tl.atomic_add(COUNTER_PTR, 1)
-        if row_id >= BATCH:
-            return
+    pid = tl.program_id(0)
+    row_id = tl.atomic_add(COUNTER_PTR, pid)
 
-        x_base = X_ptr + row_id * stride_xr
-        y_base = Y_ptr + row_id * stride_yr
-
-        # --- Online pass: running max + Kahan-sum ---------------------------
-        running_max = tl.full([1], float('-inf'), dtype=tl.float32)
-        running_sum = tl.full([1], 0.0, dtype=tl.float32)
-        comp = tl.full([1], 0.0, dtype=tl.float32)
-
-        for off in range(0, N, BLOCK_N):
-            idx = off + tl.arange(0, BLOCK_N)
-            mask = idx < N
-            x = tl.load(x_base + idx, mask=mask, other=float('-inf')).to(tl.float32)
-
-            chunk_max = tl.max(x, keep_dims=True)
-            new_max = tl.maximum(running_max, chunk_max)
-            scale = tl.exp(running_max - new_max)
-
-            running_sum = running_sum * scale
-            comp = comp * scale
-
-            e = tl.exp(x - new_max)
-            e_masked = tl.where(mask, e, 0.0)
-            e_chunk = tl.sum(e_masked, keep_dims=True).to(tl.float32)
-
-            y_k = e_chunk - comp
-            new_sum = running_sum + y_k
-            comp = (new_sum - running_sum) - y_k
-            running_sum = new_sum
-            running_max = new_max
-
-        final_max = running_max.broadcast_to([BLOCK_N])
-        total_sum_vec = running_sum.broadcast_to([BLOCK_N])
-
-        # --- Normalize pass: write output -----------------------------------
-        for off in range(0, N, BLOCK_N):
-            idx = off + tl.arange(0, BLOCK_N)
-            mask = idx < N
-            x = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32)
-            e = tl.exp(x - final_max)
-            y = tl.where(mask, e / total_sum_vec, 0.0)
-            tl.store(y_base + idx, y, mask=mask)
 
 
 def _get_launch_config(vocab: int) -> tuple[int, int, int]:

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/04_kahan_softmax/solution.py
@@ -1,7 +1,15 @@
 """Triton softmax with online max-reduce + Kahan-compensated summation.
 
-Persistent kernel variant: all SMs pull rows from an atomic counter,
-giving full GPU saturation even at small batch sizes.
+Optimized for RTX PRO 6000 Blackwell (SM120, GDDR7, 1.8 TB/s).
+
+Strategy: fused 2-pass online kernel with constexpr N for full loop unrolling.
+Sweep-validated optimal config: BLOCK_N=4096, 8 warps for large shapes;
+BLOCK_N=2048, 4 warps for small shapes.
+
+Numerical approach:
+  - Online max tracking with scale adjustment (avoids overflow on extreme logits)
+  - Kahan-compensated summation of exp(x - running_max) for tight fp32 accuracy
+  - Matches fp64 reference within atol=rtol=1e-5 on all shapes including 256K vocab
 """
 import torch
 import torch.nn as nn
@@ -15,55 +23,83 @@
 
 
 @triton.jit
-def _softmax_persistent_kernel(
+def _softmax_kernel(
     X_ptr,
     Y_ptr,
     stride_xr,
     stride_yr,
     N: tl.constexpr,
     BLOCK_N: tl.constexpr,
-    BATCH: tl.constexpr,
-    COUNTER_PTR,
 ):
-    """Persistent kernel: each SM pulls rows from atomic counter."""
-    pid = tl.program_id(0)
-    row_id = tl.atomic_add(COUNTER_PTR, pid)
+    """One-row-per-program online softmax.
+
+    Pass 1: single loop tracking running_max and Kahan-sum of exp(x-max).
+    Pass 2: re-read X, compute exp(x-final_max)/total_sum, write Y.
+
+    2 memory passes (optimal for softmax), full fp32 Kahan-compensated.
+    """
+    row_id = tl.program_id(0)
+    x_base = X_ptr + row_id * stride_xr
+    y_base = Y_ptr + row_id * stride_yr
+
+    running_max = tl.full([1], float('-inf'), dtype=tl.float32)
+    running_sum = tl.full([1], 0.0, dtype=tl.float32)
+    comp = tl.full([1], 0.0, dtype=tl.float32)
+
+    for off in range(0, N, BLOCK_N):
+        idx = off + tl.arange(0, BLOCK_N)
+        mask = idx < N
+        x = tl.load(x_base + idx, mask=mask, other=float('-inf')).to(tl.float32)
+
+        chunk_max = tl.max(x, keep_dims=True)
+        new_max = tl.maximum(running_max, chunk_max)
+        scale = tl.exp(running_max - new_max)
+
+        running_sum = running_sum * scale
+        comp = comp * scale
+
+        e = tl.exp(x - new_max)
+        e_masked = tl.where(mask, e, 0.0)
+        e_chunk = tl.sum(e_masked, keep_dims=True).to(tl.float32)
+
+        y_k = e_chunk - comp
+        new_sum = running_sum + y_k
+        comp = (new_sum - running_sum) - y_k
+        running_sum = new_sum
+        running_max = new_max
+
+    final_max = running_max.broadcast_to([BLOCK_N])
+    total_sum_vec = running_sum.broadcast_to([BLOCK_N])
+
+    for off in range(0, N, BLOCK_N):
+        idx = off + tl.arange(0, BLOCK_N)
+        mask = idx < N
+        x = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32)
+        e = tl.exp(x - final_max)
+        y = tl.where(mask, e / total_sum_vec, 0.0)
+        tl.store(y_base + idx, y, mask=mask)
 
 
-
-def _get_launch_config(vocab: int) -> tuple[int, int, int]:
-    """Return (block_n, num_warps, num_stages) optimal for this vocab size."""
+def get_launch_config(vocab):
+    """Return (BLOCK_N, num_warps, num_stages) validated by sweep."""
     if vocab <= 2048:
         return vocab, 4, 3
     if vocab <= 4096:
         return 2048, 4, 3
-    return min(4096, vocab), 8, 2
+    blk = min(4096, vocab)
+    while vocab % blk != 0 and blk > 32:
+        blk //= 2
+    return blk, 8, 2
 
 
-# RTX PRO 6000 has ~128 SMs
-_NUM_SMS = 128
-
-
-def triton_softmax(x: torch.Tensor) -> torch.Tensor:
+def triton_softmax(x):
     assert x.is_cuda and x.dtype == torch.float32 and x.ndim == 2
     batch, vocab = x.shape
     y = torch.empty_like(x)
-
-    # Atomic counter on GPU
-    counter = torch.zeros(1, dtype=torch.int32, device=x.device)
-
-    block_n, num_warps, num_stages = _get_launch_config(vocab)
-    while vocab % block_n != 0 and block_n > 32:
-        block_n //= 2
-
-    # Use enough blocks to saturate SMs, but not wasteful
-    num_programs = max(batch, _NUM_SMS)
-
-    _softmax_persistent_kernel[(num_programs,)](
+    block_n, num_warps, num_stages = get_launch_config(vocab)
+    _softmax_kernel[(batch,)](
         x, y, x.stride(0), y.stride(0), vocab,
         BLOCK_N=block_n,
-        BATCH=batch,
-        COUNTER_PTR=counter,
         num_warps=num_warps,
         num_stages=num_stages,
     )
@@ -71,12 +107,12 @@
 
 
 class Model(nn.Module):
-    def __init__(self, batch: int, vocab: int):
+    def __init__(self, batch, vocab):
         super().__init__()
         self.batch = batch
         self.vocab = vocab
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x):
         return triton_softmax(x)