diff --git a/tests/unit/test_d5c_multitile.py b/tests/unit/test_d5c_multitile.py
index b0a39544..c0e01bcc 100644
--- a/tests/unit/test_d5c_multitile.py
+++ b/tests/unit/test_d5c_multitile.py
@@ -51,48 +51,40 @@ def run_segment(q, k_seg, v_seg, kernel, compiled, stream,
     """Run one 128-token segment of FMHA, return normalized O and LSE."""
     m = q.shape[0]
     hd = v_seg.shape[1]
-    pv_n_tile = kernel.pv_n_tile
 
     # Allocate per-segment outputs
     o_seg = torch.zeros(m, hd, dtype=torch.bfloat16, device='cuda')
     lse_seg = torch.zeros(m, dtype=torch.float32, device='cuda')
+    rs_seg = torch.zeros(m, dtype=torch.float32, device='cuda')
 
     def to_cute(t):
         return ct.from_dlpack(t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(t))
 
-    for nt in range(kernel.n_pv_tiles):
-        v_start = nt * pv_n_tile
-        v_end = v_start + pv_n_tile
-        v_tile = v_seg[:, v_start:v_end].contiguous()
-        c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
-        lse_tile = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
-        rs_tile = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
+    c_tile = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
+    lse_tile = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
+    rs_tile = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
 
-        mQ = to_cute(q.unsqueeze(-1))
-        mK = to_cute(k_seg.unsqueeze(-1))
-        mV = to_cute(v_tile.unsqueeze(-1))
-        mC = to_cute(c_tile)
-        mLSE = to_cute(lse_tile)
-        mRS = to_cute(rs_tile)
+    mQ = to_cute(q.unsqueeze(-1))
+    mK = to_cute(k_seg.unsqueeze(-1))
+    mV = to_cute(v_seg.unsqueeze(-1))  # (n_k, hd, 1)
+    mC = to_cute(c_tile)
+    mLSE = to_cute(lse_tile)
+    mRS = to_cute(rs_tile)
 
-        if sink_bias is not None:
-            mSB = to_cute(sink_bias)
-            compiled(mQ, mK, mV, mC, stream, mLSE,
-                     swa_len=swa_len, sink_bias=mSB, row_sums=mRS)
-        else:
-            compiled(mQ, mK, mV, mC, stream, mLSE,
-                     swa_len=swa_len, row_sums=mRS)
+    if sink_bias is not None:
+        mSB = to_cute(sink_bias)
+        compiled(mQ, mK, mV, mC, stream, mLSE,
+                 swa_len=swa_len, sink_bias=mSB, row_sums=mRS)
+    else:
+        compiled(mQ, mK, mV, mC, stream, mLSE,
+                 swa_len=swa_len, row_sums=mRS)
 
-        torch.cuda.synchronize()
-        o_seg[:, v_start:v_end] = c_tile[:, :, 0]
-        if nt == 0:
-            lse_seg = lse_tile[:, 0, 0].clone()
-            rs_seg = rs_tile[:, 0, 0].clone()
-        # Note: LSE and row_sum are the same across PV tiles (same softmax)
+    torch.cuda.synchronize()
+    rs = rs_tile[:, 0, 0].float()
+    o_norm = c_tile[:, :, 0].float() / rs.unsqueeze(1).clamp(min=1e-30)
+    lse_val = lse_tile[:, 0, 0].clone()
 
-    # Normalize using row_sum
-    o_norm = o_seg.float() / rs_seg.unsqueeze(1).clamp(min=1e-30)
-    return o_norm, lse_seg
+    return o_norm, lse_val
 
 
 def python_kv_merge(segment_results):