More tests for FlashInfer SDPA

Author: mseeger

keys_values/attention/flashinfer_wrapper.py

@@ -160,6 +160,7 @@ def triton_score_sum(
     scale: float,
     n_kv_heads: int,
     group_size: int,
+    causal_masking: bool = True,
 ) -> torch.Tensor:
     """Compute attention weight sums using Triton (no V needed).
 
@@ -174,6 +175,8 @@ def triton_score_sum(
         scale: softmax scale factor (1/sqrt(head_size))
         n_kv_heads: number of KV heads
         group_size: GQA group size (n_head // n_kv_heads)
+        causal_masking: Whether to use causal attention mask or not. Defaults
+            to `True`
 
     Returns:
         W: [batch, n_kv_heads, kv_len] (fp32) attention weight sums
@@ -258,7 +261,7 @@ def triton_score_sum(
         BLOCK_Q=BLOCK_Q,
         HEAD_DIM=head_size,
         GROUP_SIZE=group_size,
-        HAS_CAUSAL=True,
+        HAS_CAUSAL=causal_masking,
         num_warps=NUM_WARPS,
         num_stages=NUM_STAGES,
     )

keys_values/attention/sdpa_wrapper.py

@@ -330,13 +330,13 @@ def _reorder(
 
 def reorder_key_value(
     key: torch.Tensor,
-    value: torch.Tensor,
+    value: Optional[torch.Tensor],
     token_positions: torch.Tensor,
     input_pos: int,
     q_len: int,
     sort_if_3d: bool = True,
     check_token_pos: bool = False,
-) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor]]:
+) -> Tuple[torch.Tensor, Optional[torch.Tensor], Dict[str, torch.Tensor]]:
     """
     Reorder `key, value` tensors using permutations (for each b, h) which, if
     applied to `token_positions`, place `input_pos:(input_pos + q_len)` at the
@@ -364,7 +364,7 @@ def reorder_key_value(
         extra_info = dict(index_gat=index_gat, index_scat=index_scat)
     return (
         reorder_buffer_given_extra_info(key, **extra_info),
-        reorder_buffer_given_extra_info(value, **extra_info),
+        None if value is None else reorder_buffer_given_extra_info(value, **extra_info),
         extra_info,
     )
 

keys_values/scripts/collect_eval_results.py

@@ -63,14 +63,20 @@ def main(
 
 if __name__ == "__main__":
     base_path = Path.home() / "out/finetune/neurips_exp/lora/qwen3_4b"
+    dataset_size = "64k"
     datasets = [
-        "helmet_nq_32k",
-        #        "helmet_trivia_qa_32k",
-        #        "helmet_hotpot_qa_32k",
-        #        "helmet_pop_qa_32k",
+        f"helmet_nq_{dataset_size}",
+        f"helmet_trivia_qa_{dataset_size}",
+        f"helmet_hotpot_qa_{dataset_size}",
+        f"helmet_pop_qa_{dataset_size}",
     ]
     cases = [
-        "lr_4gpu_lpc2_avg1_lr5",
+        "lr_4gpu_cs2048_lr5",
+        "h2o_4gpu_cs2048_lr5",
+        "slr_4gpu_cs2048_lr5",
+        #    "qh2o_4gpu_cs2048_lr5",
+        #    "h2onorm_4gpu_cs2048_lr5",
+        #    "qh2onorm_4gpu_cs2048_lr5",
     ]
     model_type = "lora"
     for dataset, case in product(datasets, cases):