Fix Blelloch exclusive scan: avoid numerical instability

Changed Blelloch scan to compute exclusive prefix directly instead of
converting from inclusive, avoiding division by lambda^n which causes
overflow when lambda is small.

Implementation:
1. Compute inclusive prefix using standard up-sweep + down-sweep
2. Convert to exclusive via simple rank shift: each rank i receives
   inclusive[i-1] from rank i-1, rank 0 gets zero

This matches the pattern used in lasp_naive where the ring naturally
produces exclusive prefix, avoiding the numerical issues of computing
1/lambda^n which overflows to infinity when s >= 1.0.

Fixes NaN gradients in backward pass.

Author: Hoang Phan

lasp/lasp_blelloch.py

@@ -161,30 +161,8 @@ def forward(ctx, q, k, v, s, KV, DKV):
             )
 
             # Blelloch scan: O(log P) tree communication
-            # IMPORTANT: Blelloch returns INCLUSIVE prefix (includes current rank)
-            # but LASP needs EXCLUSIVE prefix (only previous ranks)
-            KV_prefix_inclusive = scanner.scan(local_kv)
-
-            # Convert inclusive to exclusive
-            # For the LASP associative operation (λ^C, KV), we have:
-            #   inclusive[i] = λ^(C*i)*KV[0] + ... + λ^C*KV[i-1] + KV[i]
-            #   exclusive[i] = λ^(C*(i-1))*KV[0] + ... + KV[i-1]
-            #
-            # To convert: exclusive = λ^(-C) * (inclusive - KV[i])
-            #
-            # NOTE: Create new tensor instead of modifying KV with .copy_()
-            # This avoids modifying input buffers which can cause issues
-            if rank > 0:
-                # Compute λ^(-C) = 1 / λ^C
-                lambda_C_inv = 1.0 / lambda_decay ** n
-                # Expand to match tensor dimensions [h] → [b, h, d, e]
-                lambda_C_inv_expanded = lambda_C_inv.view(1, h, 1, 1).expand(b, h, d, e)
-                # exclusive = λ^(-C) * (inclusive - local)
-                KV_prefix = lambda_C_inv_expanded * (KV_prefix_inclusive - local_kv)
-            else:
-                # Rank 0 has no previous ranks, so prefix is zero
-                # Use KV which is already zeroed
-                KV_prefix = KV
+            # Returns EXCLUSIVE prefix (only previous ranks, not including current)
+            KV_prefix = scanner.scan(local_kv)
 
         # ===== STEP 4: Inter-chunk attention using fused kernel =====
         # This is the key improvement: use _fwd_none_diag_kernel instead of torch.matmul
@@ -318,25 +296,8 @@ def backward(ctx, do):
             )
 
             # Reverse scan for gradients
-            # IMPORTANT: Blelloch returns INCLUSIVE suffix (includes current rank)
-            # but LASP needs EXCLUSIVE suffix (only future ranks)
-            DKV_suffix_inclusive = scanner.scan(local_dkv)
-
-            # Convert inclusive to exclusive
-            # Same logic as forward: exclusive = λ^(-C) * (inclusive - local)
-            # NOTE: Create new tensor instead of modifying DKV with .copy_()
-            # This avoids modifying saved tensors which can cause CUDA errors
-            if rank < world_size - 1:
-                # Compute λ^(-C) = 1 / λ^C
-                lambda_C_inv = 1.0 / lambda_decay ** n
-                # Expand to match tensor dimensions [h] → [b, h, d, e]
-                lambda_C_inv_expanded = lambda_C_inv.view(1, h, 1, 1).expand(b, h, d, e)
-                # exclusive = λ^(-C) * (inclusive - local)
-                DKV_suffix = lambda_C_inv_expanded * (DKV_suffix_inclusive - local_dkv)
-            else:
-                # Last rank (which is rank 0 in forward) has no future ranks
-                # Return zero suffix (use DKV which is already zeroed)
-                DKV_suffix = DKV
+            # Returns EXCLUSIVE suffix (only future ranks, not including current)
+            DKV_suffix = scanner.scan(local_dkv)
 
         # ===== STEP 4: Inter-chunk gradient contribution using fused kernel =====
         with torch.cuda.device(q.device.index):

lasp/utils/blelloch_ops.py

@@ -239,22 +239,23 @@ def combine(
 
     def scan(self, local_value: torch.Tensor) -> torch.Tensor:
         """
-        Perform parallel prefix scan on local KV contribution.
+        Perform parallel EXCLUSIVE prefix scan on local KV contribution.
 
         Args:
             local_value: Local KV state b[rank] (shape: [b, h, d, e])
 
         Returns:
-            prefix_sum: KV[0:rank+1] - prefix sum up to this rank
+            exclusive_prefix: KV[0:rank] - prefix sum excluding current rank
+                              (rank 0 gets zero, rank i gets sum from ranks 0 to i-1)
         """
         if self.world_size == 1:
-            # Single GPU: no communication needed
-            return local_value
+            # Single GPU: exclusive prefix is zero (no previous ranks)
+            return torch.zeros_like(local_value)
 
         b, h, d, e = local_value.shape
 
         # ============ UP-SWEEP PHASE ============
-        # Build tree bottom-up, accumulating partial sums
+        # Build tree bottom-up, accumulating partial sums (inclusive)
 
         current_value = local_value.clone()
         tree_values = [current_value]  # Store for down-sweep
@@ -283,9 +284,9 @@ def scan(self, local_value: torch.Tensor) -> torch.Tensor:
                 tree_values.append(current_value)
 
         # ============ DOWN-SWEEP PHASE ============
-        # Distribute prefix sums top-down
+        # Distribute inclusive prefix sums top-down
 
-        prefix_sum = None
+        inclusive_prefix = None
 
         for level in range(self.num_levels - 1, -1, -1):
             partner = self.get_partner_rank(level, 'down')
@@ -304,19 +305,37 @@ def scan(self, local_value: torch.Tensor) -> torch.Tensor:
                 distance = abs(self.scan_rank - partner)
                 # Use the tree value stored during up-sweep
                 tree_idx = min(level, len(tree_values) - 1)
-                prefix_sum = self.combine(left_prefix, tree_values[tree_idx], distance)
+                inclusive_prefix = self.combine(left_prefix, tree_values[tree_idx], distance)
 
             elif self.is_sender(level, 'down') and partner < self.world_size:
                 # Send to right child (convert to global rank)
                 global_partner = self.local_to_global_rank(partner)
-                send_value = prefix_sum if prefix_sum is not None else tree_values[min(level, len(tree_values) - 1)]
+                send_value = inclusive_prefix if inclusive_prefix is not None else tree_values[min(level, len(tree_values) - 1)]
                 dist.send(tensor=send_value.contiguous(), dst=global_partner, group=self.group)
 
-        # Rank 0 has no left prefix, uses its accumulated tree value
-        if prefix_sum is None:
-            prefix_sum = tree_values[-1] if len(tree_values) > 1 else local_value
+        # Compute inclusive prefix for this rank
+        if inclusive_prefix is None:
+            inclusive_prefix = tree_values[-1] if len(tree_values) > 1 else local_value
 
-        return prefix_sum
+        # ============ CONVERT TO EXCLUSIVE ============
+        # Simple approach: rank i sends inclusive[i] to rank i+1
+        # Rank 0 returns zero, rank i returns inclusive[i-1]
+
+        exclusive_prefix = torch.zeros_like(local_value)
+
+        if self.scan_rank > 0:
+            # Receive from left neighbor (scan_rank - 1)
+            left_neighbor = self.scan_rank - 1
+            global_left = self.local_to_global_rank(left_neighbor)
+            dist.recv(tensor=exclusive_prefix, src=global_left, group=self.group)
+
+        if self.scan_rank < self.world_size - 1:
+            # Send to right neighbor (scan_rank + 1)
+            right_neighbor = self.scan_rank + 1
+            global_right = self.local_to_global_rank(right_neighbor)
+            dist.send(tensor=inclusive_prefix.contiguous(), dst=global_right, group=self.group)
+
+        return exclusive_prefix
 
 
 def safe_decay_power(base: float, exponent: int, use_log_space: bool = True) -> float: