Fix: pass scaling, dropout to add_mask

  1. apply scaling when we use exp-weights
  2. remove previous mask before we do top-k / top-p

Author: apd10

sparse_attention_hub/metric_logging/logger.py

@@ -190,7 +190,7 @@ def flush(self) -> None:
             return
 
         # Get current timestamp for filename
-        filename = f"micro_metrics.jsonl"
+        filename = "micro_metrics.jsonl"
         filepath = os.path.join(self.log_path, filename)
 
         # Write events to file

sparse_attention_hub/sparse_attention/research_attention/base.py

@@ -105,6 +105,8 @@ def custom_attention(
                 queries=queries,
                 values=values,
                 attention_mask=attention_mask,
+                scaling=scaling,
+                dropout=dropout,
                 sparse_meta_data=sparse_meta_data,
                 previous_mask=sparse_attention_mask,
                 **kwargs,

sparse_attention_hub/sparse_attention/research_attention/maskers/base.py

@@ -160,6 +160,8 @@ def add_mask(
         queries: torch.Tensor,
         values: torch.Tensor,
         attention_mask: torch.Tensor,
+        scaling: float,
+        dropout: float,
         sparse_meta_data: Dict[Any, Any],  # want to keep it general here.
         previous_mask: Mask,
         **kwargs: Dict[str, Any],

sparse_attention_hub/sparse_attention/research_attention/maskers/fixed/implementations/basic_fixed.py

@@ -39,6 +39,8 @@ def add_mask(
         queries: torch.Tensor,
         values: torch.Tensor,
         attention_mask: torch.Tensor,
+        scaling: float,
+        dropout: float,
         sparse_meta_data: Dict[Any, Any],
         previous_mask: Mask,
         **kwargs: Dict[str, Any],
@@ -142,6 +144,8 @@ def add_mask(
         queries: torch.Tensor,
         values: torch.Tensor,
         attention_mask: torch.Tensor,
+        scaling: float,
+        dropout: float,
         sparse_meta_data: Dict[Any, Any],
         previous_mask: Mask,
         **kwargs: Dict[str, Any],
@@ -182,6 +186,8 @@ def add_mask(
         queries: torch.Tensor,
         values: torch.Tensor,
         attention_mask: torch.Tensor,
+        scaling: float,
+        dropout: float,
         sparse_meta_data: Dict[Any, Any],
         previous_mask: Mask,
         **kwargs: Dict[str, Any],

sparse_attention_hub/sparse_attention/research_attention/maskers/fixed/implementations/double_sparsity_top_k.py

@@ -53,6 +53,8 @@ def add_mask(
         queries: torch.Tensor,
         values: torch.Tensor,
         attention_mask: torch.Tensor,
+        scaling: float,
+        dropout: float,
         sparse_meta_data: Dict[Any, Any],
         previous_mask: Mask,
         **kwargs: Dict[str, Any],

sparse_attention_hub/sparse_attention/research_attention/maskers/fixed/implementations/hashattention_top_k.py

@@ -57,6 +57,8 @@ def add_mask(
         queries: torch.Tensor,
         values: torch.Tensor,
         attention_mask: torch.Tensor,
+        scaling: float,
+        dropout: float,
         sparse_meta_data: Dict[str, Dict[int, Optional[torch.Tensor]]],
         previous_mask: Mask,
         **kwargs: Dict[str, Any],
@@ -85,6 +87,8 @@ def add_mask(
             effective_heavy_size,
             keys,
             queries,
+            attention_mask,
+            previous_mask.get_dense_mask(),
             sparse_meta_data,
             previous_mask,
             layer_idx,
@@ -143,14 +147,21 @@ def _create_hash_topk_mask(
         heavy_size: int,
         keys: torch.Tensor,
         queries: torch.Tensor,
+        attention_mask: torch.Tensor,
         sparse_meta_data: Dict[str, Dict[int, Optional[torch.Tensor]]],
         previous_mask: Mask,
         layer_idx: int,
         **kwargs: Dict[str, Any],
     ) -> Mask:
         """Create hash attention top-K mask using hash-based scoring."""
         scores: torch.Tensor = self._compute_hashattention_score(
-            queries, keys, sparse_meta_data, layer_idx, **kwargs
+            queries,
+            keys,
+            attention_mask,
+            previous_mask.get_dense_mask(),
+            sparse_meta_data,
+            layer_idx,
+            **kwargs,
         )
         top_k_indices: torch.Tensor = self._get_topk_indices_from_inactive_positions(
             scores, previous_mask, heavy_size
@@ -303,6 +314,8 @@ def _compute_hashattention_score(
         self,
         queries: torch.Tensor,
         keys: torch.Tensor,
+        attention_mask: torch.Tensor,
+        previous_dense_mask: torch.Tensor,
         sparse_meta_data: Dict[str, Dict[int, Optional[torch.Tensor]]],
         layer_idx: int,
         **kwargs: Dict[str, Any],
@@ -319,7 +332,13 @@ def _compute_hashattention_score(
         )
 
         # (B, H, #queries, hat_bits) x (B, H, hat_bits, #keys) -> (B, H, #queries, #keys)
-        return torch.matmul(query_signatures, key_signatures.transpose(-2, -1))
+        scores: torch.Tensor = torch.matmul(
+            query_signatures, key_signatures.transpose(-2, -1)
+        )
+        if attention_mask is not None:
+            scores = scores + attention_mask[:, :, :, : keys.shape[-2]]
+        scores[previous_dense_mask != 0] = torch.finfo(scores.dtype).min
+        return scores
 
     @classmethod
     def create_from_config(cls, config: MaskerConfig) -> "HashAttentionTopKMasker":

sparse_attention_hub/sparse_attention/research_attention/maskers/fixed/implementations/oracle_top_k.py

@@ -43,6 +43,8 @@ def add_mask(
         queries: torch.Tensor,
         values: torch.Tensor,
         attention_mask: torch.Tensor,
+        scaling: float,
+        dropout: float,
         sparse_meta_data: Dict[Any, Any],
         previous_mask: Mask,
         **kwargs: Dict[str, Any],
@@ -64,7 +66,12 @@ def add_mask(
 
         # Create oracle top-K mask
         oracle_mask: Mask = self._create_oracle_topk_mask(
-            tensor_dims, effective_heavy_size, keys, queries, previous_mask
+            tensor_dims,
+            effective_heavy_size,
+            keys,
+            queries,
+            attention_mask,
+            previous_mask,
         )
         return previous_mask.merge_mask(oracle_mask, inplace=False)
 
@@ -84,11 +91,12 @@ def _create_oracle_topk_mask(
         heavy_size: int,
         keys: torch.Tensor,
         queries: torch.Tensor,
+        attention_mask: torch.Tensor,
         previous_mask: Mask,
     ) -> Mask:
         """Create oracle top-K mask using raw attention scores."""
         raw_attention_scores: torch.Tensor = self._compute_raw_attention_scores(
-            keys, queries
+            keys, queries, attention_mask, previous_mask.get_dense_mask()
         )
         top_k_indices: torch.Tensor = self._get_topk_indices_from_inactive_positions(
             raw_attention_scores, previous_mask, heavy_size
@@ -98,12 +106,20 @@ def _create_oracle_topk_mask(
         )
 
     def _compute_raw_attention_scores(
-        self, keys: torch.Tensor, queries: torch.Tensor
+        self,
+        keys: torch.Tensor,
+        queries: torch.Tensor,
+        attention_mask: torch.Tensor,
+        previous_dense_mask: torch.Tensor,
     ) -> torch.Tensor:
         """Compute raw attention scores using query-key dot product."""
         ngroups = _get_num_key_value_groups(queries, keys)
         keys = repeat_kv(keys, ngroups)
-        return torch.matmul(queries, keys.transpose(-2, -1))
+        scores: torch.Tensor = torch.matmul(queries, keys.transpose(-2, -1))
+        if attention_mask is not None:
+            scores = scores + attention_mask[:, :, :, : keys.shape[-2]]
+        scores[previous_dense_mask != 0] = torch.finfo(scores.dtype).min
+        return scores
 
     @classmethod
     def create_from_config(cls, config: MaskerConfig) -> "OracleTopK":

sparse_attention_hub/sparse_attention/research_attention/maskers/fixed/implementations/oracle_top_p.py

@@ -44,6 +44,8 @@ def add_mask(
         queries: torch.Tensor,
         values: torch.Tensor,
         attention_mask: torch.Tensor,
+        scaling: float,
+        dropout: float,
         sparse_meta_data: Dict[Any, Any],
         previous_mask: Mask,
         **kwargs: Dict[str, Any],
@@ -62,7 +64,7 @@ def add_mask(
 
         # Create oracle top-P attention mask
         oracle_mask: Mask = self._create_oracle_top_p_mask(
-            tensor_dims, keys, queries, previous_mask
+            tensor_dims, keys, queries, previous_mask, attention_mask, scaling
         )
         return previous_mask.merge_mask(oracle_mask, inplace=False)
 
@@ -71,13 +73,25 @@ def _should_use_full_attention(self, dims: AttentionTensorDimensions) -> bool:
         effective_size: int = int(self.top_p * dims.seq_len_keys)
         return dims.seq_len_keys <= effective_size
 
-    def _compute_attention_scores(
-        self, keys: torch.Tensor, queries: torch.Tensor
+    def _compute_exp_attention_scores(
+        self,
+        keys: torch.Tensor,
+        queries: torch.Tensor,
+        previous_dense_mask: torch.Tensor,
+        attention_mask: torch.Tensor,
+        scaling: float,
     ) -> torch.Tensor:
         """Compute exp(attention scores) between queries and keys."""
         ngroups = _get_num_key_value_groups(queries, keys)
         keys = repeat_kv(keys, ngroups)
-        raw_attention_scores = queries @ keys.transpose(-2, -1)
+        raw_attention_scores = torch.matmul(queries, keys.transpose(2, 3)) * scaling
+        if attention_mask is not None:
+            raw_attention_scores = (
+                raw_attention_scores + attention_mask[:, :, :, : keys.shape[-2]]
+            )
+        raw_attention_scores[previous_dense_mask != 0] = torch.finfo(
+            raw_attention_scores.dtype
+        ).min
         _max_attention_score = raw_attention_scores.max(dim=-1, keepdim=True)[0]
         adjusted = torch.exp(raw_attention_scores - _max_attention_score)
         return adjusted
@@ -101,42 +115,33 @@ def _compute_top_p_thresholds(
 
         # Find positions where normalized_cumsum >= top_p
         threshold_positions = torch.searchsorted(
-            normalized_cumsum, top_p_tensor, side="left"
-        )
-
-        # Prepare indices for advanced indexing (shape-agnostic)
-        leading_shape = scores.shape[:-1]
-        idx_grids = torch.meshgrid(
-            *[torch.arange(s, device=scores.device) for s in leading_shape],
-            indexing="ij",
+            normalized_cumsum, top_p_tensor, side="right"
         )
-        thresholds = sorted_scores[idx_grids + (threshold_positions.squeeze(-1),)]
-
-        # Add trailing singleton dimension for broadcasting
-        return thresholds.unsqueeze(-1)
+        thresholds = torch.gather(sorted_scores, dim=-1, index=threshold_positions)
+        return thresholds
 
     def _create_oracle_top_p_mask(
         self,
         dims: AttentionTensorDimensions,
         keys: torch.Tensor,
         queries: torch.Tensor,
         previous_mask: Mask,
+        attention_mask: torch.Tensor,
+        scaling: float,
     ) -> Mask:
         """Create oracle top-P attention mask using vectorized computation."""
-        # Get attention scores
-        scores: torch.Tensor = self._compute_attention_scores(keys, queries)
-        # Get previous dense mask and mask out already active positions
+        # Get attention scores after masking out already active positions
         previous_dense_mask: torch.Tensor = previous_mask.get_dense_mask()
-        masked_scores: torch.Tensor = scores.clone()
-        masked_scores[previous_dense_mask != 0] = float("-inf")
+        scores: torch.Tensor = self._compute_exp_attention_scores(
+            keys, queries, previous_dense_mask, attention_mask, scaling
+        )
 
         # Compute thresholds using vectorized operations
-        thresholds: torch.Tensor = self._compute_top_p_thresholds(
-            masked_scores, self.top_p
-        )
+        thresholds: torch.Tensor = self._compute_top_p_thresholds(scores, self.top_p)
+        thresholds = thresholds.to(queries.dtype)
 
         # Create dense mask: scores >= thresholds
-        dense_mask: torch.Tensor = masked_scores >= thresholds
+        dense_mask: torch.Tensor = scores >= thresholds
 
         # Create mask object
         mask_shape: tuple = (

sparse_attention_hub/sparse_attention/research_attention/maskers/fixed/implementations/pq_top_k.py

@@ -39,6 +39,8 @@ def add_mask(
         queries: torch.Tensor,
         values: torch.Tensor,
         attention_mask: torch.Tensor,
+        scaling: float,
+        dropout: float,
         sparse_meta_data: Dict[Any, Any],
         previous_mask: Mask,
         **kwargs: Dict[str, Any],

sparse_attention_hub/sparse_attention/research_attention/maskers/sampling/implementations/adaptive_sampling.py

@@ -144,12 +144,18 @@ def __init__(self, config: AdaptiveSamplingMaskerConfig) -> None:
         self.delta_ppf = float(norm.ppf(1 - self.delta))
 
     def _compute_exp_attention_scores(
-        self, queries: torch.Tensor, keys: torch.Tensor
+        self,
+        queries: torch.Tensor,
+        keys: torch.Tensor,
+        scaling: float,
+        attention_mask: torch.Tensor,
     ) -> torch.Tensor:
         """Compute exponential attention scores with numerical stability."""
         ngroups = _get_num_key_value_groups(queries, keys)
         keys = repeat_kv(keys, ngroups)
-        raw_scores = torch.matmul(queries, keys.transpose(-2, -1))
+        raw_scores = torch.matmul(queries, keys.transpose(-2, -1)) * scaling
+        if attention_mask is not None:
+            raw_scores = raw_scores + attention_mask[:, :, :, : keys.shape[-2]]
         max_scores = torch.max(raw_scores, dim=-1, keepdim=True)[0]
         return torch.exp(raw_scores - max_scores)
 
@@ -244,6 +250,8 @@ def add_mask(
         queries: torch.Tensor,
         values: torch.Tensor,
         attention_mask: torch.Tensor,
+        scaling: float,
+        dropout: float,
         sparse_meta_data: Dict[Any, Any],
         previous_mask: Mask,
         **kwargs: Dict[str, Any],
@@ -280,8 +288,10 @@ def add_mask(
             dims.seq_len_queries,
             dims.seq_len_keys,
         )
-
-        expwts = self._compute_exp_attention_scores(queries, keys)
+        # Compute attention scores after removing attention_mask
+        expwts = self._compute_exp_attention_scores(
+            queries, keys, scaling, attention_mask
+        )
         static_denominator = apply_inv_mask_sum(expwts, previous_mask)
 
         # Get sampling parameters