diff --git a/tests/unit/test_d5c_fused.py b/tests/unit/test_d5c_fused.py
index a9335f93..e623bad0 100644
--- a/tests/unit/test_d5c_fused.py
+++ b/tests/unit/test_d5c_fused.py
@@ -91,7 +91,7 @@ def reference_sink_merge(q, k_comp, v_comp, k_swa, v_swa,
     lse_swa = torch.logsumexp(attn_swa, dim=-1, keepdim=True)
 
     # Sink merge (normalized formula)
-    exp_sink = attn_sink.exp()
+    exp_sink = math.exp(attn_sink)
     numerator = lse_comp.exp() * o_norm_comp + exp_sink * lse_swa.exp() * o_norm_swa
     denominator = (lse_comp.exp() + exp_sink * lse_swa.exp()).clamp(min=1e-30)
     o = numerator / denominator
@@ -117,7 +117,8 @@ def test_d5c_combined():
     v_swa = torch.randn(n_swa, hd, dtype=torch.bfloat16, device='cuda')
 
     # Sink weight (in log domain, per-head). n_h=1 so it's a scalar.
-    attn_sink = torch.tensor([0.5], dtype=torch.float32, device='cuda')
+    attn_sink_val = 0.5
+    attn_sink = torch.tensor([attn_sink_val], dtype=torch.float32, device='cuda')
 
     # === FP32 References ===
     qf = q[:, :, 0]
@@ -125,13 +126,13 @@ def test_d5c_combined():
     # Reference 1: Combined softmax with sink bias (our kernel's approach)
     ref_combined = reference_combined_attention(
         qf, k_comp[:, :, 0], v_comp, k_swa[:, :, 0], v_swa,
-        attn_sink[0].item(), scale, swa_len
+        attn_sink_val, scale, swa_len
     )
 
     # Reference 2: Separate attention + sink merge (original D5b formula)
     ref_merge = reference_sink_merge(
         qf, k_comp[:, :, 0], v_comp, k_swa[:, :, 0], v_swa,
-        attn_sink[0].item(), scale, swa_len
+        attn_sink_val, scale, swa_len
     )
 
     # Verify the two references agree
@@ -222,12 +223,13 @@ def test_d5c_with_causal():
     k_swa = torch.randn(n_swa, hd, 1, dtype=torch.bfloat16, device='cuda')
     v_swa = torch.randn(n_swa, hd, dtype=torch.bfloat16, device='cuda')
 
-    attn_sink = torch.tensor([0.3], dtype=torch.float32, device='cuda')
+    attn_sink_val = 0.3
+    attn_sink = torch.tensor([attn_sink_val], dtype=torch.float32, device='cuda')
     qf = q[:, :, 0]
 
     ref = reference_combined_attention(
         qf, k_comp[:, :, 0], v_comp, k_swa[:, :, 0], v_swa,
-        attn_sink[0].item(), scale, swa_len, is_causal=True
+        attn_sink_val, scale, swa_len, is_causal=True
     )
 
     # Concatenate KV