fix(fa4): reuse compiled masks across variable-length batches

aspec/models/ops/direct_cp.py

@@ -7,6 +7,8 @@
 
 from __future__ import annotations
 
+import hashlib
+import inspect
 from dataclasses import dataclass, field
 from functools import lru_cache
 from typing import Optional
@@ -477,16 +479,18 @@ def _apply_block_keep(
 
 @lru_cache(maxsize=None)
 def _build_fa4_mask(
-    context_length: int,
     block_size: int,
-    num_blocks: int,
-    query_length: int,
-    key_length: int,
     packed: bool,
     context_window: Optional[int] = None,
     query_window: Optional[int] = None,
     draft_is_causal: bool = False,
 ):
+    """Specialize only on mask semantics; row geometry stays on the device.
+
+    FA4 hashes modifier closures. Capturing concrete row lengths creates fresh
+    forward/backward kernels throughout a variable-length corpus. Auxiliary
+    tensor dimensions are dynamic in FA4, so they also provide safe bounds.
+    """
     import cutlass
     import cutlass.cute as cute
     from flash_attn.cute import utils
@@ -496,23 +500,23 @@ def read_block_aux(aux, batch, index: "cute.TensorSSA") -> "cute.TensorSSA":
         index_fragment = cute.make_rmem_tensor(1, cutlass.Int32)
         index_fragment.store(index)
         value_fragment = cute.make_rmem_tensor(1, cutlass.Int32)
-        value_fragment[0] = aux[batch[0], cutlass.min(index_fragment[0], num_blocks - 1)]
+        value_fragment[0] = aux[batch[0], cutlass.min(index_fragment[0], aux.shape[1] - 1)]
         return value_fragment.load()
 
     @cute.jit
     def read_query_aux(aux, batch, index: "cute.TensorSSA") -> "cute.TensorSSA":
         index_fragment = cute.make_rmem_tensor(1, cutlass.Int32)
         index_fragment.store(index)
         value_fragment = cute.make_rmem_tensor(1, cutlass.Int32)
-        value_fragment[0] = aux[batch[0], cutlass.min(index_fragment[0], query_length - 1)]
+        value_fragment[0] = aux[batch[0], cutlass.min(index_fragment[0], aux.shape[1] - 1)]
         return value_fragment.load()
 
     @cute.jit
     def read_key_aux(aux, batch, index: "cute.TensorSSA") -> "cute.TensorSSA":
         index_fragment = cute.make_rmem_tensor(1, cutlass.Int32)
         index_fragment.store(index)
         value_fragment = cute.make_rmem_tensor(1, cutlass.Int32)
-        value_fragment[0] = aux[batch[0], cutlass.min(index_fragment[0], key_length - 1)]
+        value_fragment[0] = aux[batch[0], cutlass.min(index_fragment[0], aux.shape[1] - 1)]
         return value_fragment.load()
 
     @cute.jit
@@ -528,8 +532,9 @@ def mask_mod(batch, head, query_index, key_index, seqlen_info, aux_tensors):
             key_token_keep,
             anchor_doc_ids,
             key_context_doc_ids,
+            context_length,
         ) = aux_tensors
-        context_len = utils.scalar_to_ssa(context_length, cutlass.Int32)
+        context_len = utils.scalar_to_ssa(context_length[0], cutlass.Int32)
         zero = utils.scalar_to_ssa(0, cutlass.Int32)
         query_block = read_query_aux(query_blocks, batch, query_index)
         anchor = read_block_aux(anchors, batch, query_block)
@@ -566,6 +571,14 @@ def mask_mod(batch, head, query_index, key_index, seqlen_info, aux_tensors):
             & (key_valid > zero)
         )
 
+    # Nested cute.jit helpers otherwise contribute process-specific function
+    # addresses to FA4's closure hash. Bind every remaining specialization and
+    # the complete factory source so disk artifacts are reusable across ranks
+    # while code or semantics changes still invalidate the key.
+    identity = (block_size, packed, context_window, query_window, draft_is_causal)
+    mask_mod.__cute_hash__ = hashlib.sha256(
+        (inspect.getsource(_build_fa4_mask) + repr(identity)).encode()
+    ).hexdigest()
     return mask_mod
 
 
@@ -591,11 +604,7 @@ def _fa4_partial(
 
     batch, heads, queries, _ = query.shape
     mask_mod = _build_fa4_mask(
-        metadata.context_length,
         metadata.block_size,
-        metadata.num_blocks,
-        queries,
-        key.shape[2],
         metadata.context_doc_ids is not None,
         metadata.context_window,
         metadata.query_window,
@@ -670,6 +679,13 @@ def _fa4_partial(
         ),
         dim=1,
     ).contiguous()
+    geometry_key = ("mask_context_length", query.device, metadata.context_length)
+    context_length_tensor = metadata._fa4_cache.get(geometry_key)
+    if context_length_tensor is None:
+        context_length_tensor = torch.full(
+            (1,), metadata.context_length, dtype=torch.int32, device=query.device
+        )
+        metadata._fa4_cache[geometry_key] = context_length_tensor
     aux_tensors = [
         metadata.anchor_positions.to(torch.int32).contiguous(),
         metadata.block_keep_mask.to(torch.int32).contiguous(),
@@ -680,6 +696,7 @@ def _fa4_partial(
         key_token_keep.contiguous(),
         anchor_doc_ids.contiguous(),
         key_context_doc_ids,
+        context_length_tensor,
     ]
     from aspec.models.ops.fa4_sparse import fixed_mask_aux
 

aspec/utils/env.py

@@ -35,6 +35,8 @@
     "ASPEC_LOG_LEVEL",
     "TP_SOCKET_IFNAME",
     "CUTE_DSL_CACHE_DIR",
+    "FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED",
+    "FLASH_ATTENTION_CUTE_DSL_CACHE_DIR",
     "ASPEC_FLASH_ATTN_OPT_LEVEL",
     "ASPEC_RUN_ID",
     "ASPEC_ATTEMPT_ID",
@@ -81,6 +83,7 @@ def get_aspec_env_vars(args=None) -> dict[str, str]:
     if getattr(args, "cluster_mode", "single_node") == "static_multi_node":
         # Ray is started from each node's own venv/environment. A driver's
         # library/cache paths must not overwrite another node's local layout.
-        for key in ("LD_LIBRARY_PATH", "PYTHONPATH", "HF_HOME", "CUTE_DSL_CACHE_DIR", "ASPEC_LOG_DIR"):
+        for key in ("LD_LIBRARY_PATH", "PYTHONPATH", "HF_HOME", "CUTE_DSL_CACHE_DIR",
+                    "FLASH_ATTENTION_CUTE_DSL_CACHE_DIR", "ASPEC_LOG_DIR"):
             env.pop(key, None)
     return env

tests/test_direct_cp.py

@@ -17,6 +17,7 @@
 
 # Adapted from AngelSpec tests/test_direct_cp.py (Apache-2.0).
 
+import inspect
 import math
 import os
 
@@ -352,3 +353,84 @@ def test_direct_cp_sparse_tiles_cover_visible_tokens_and_only_mark_full_tiles():
 
     assert torch.all(expected_any <= (partial | full))
     assert torch.all(full <= expected_all)
+
+
+@pytest.fixture(scope="module")
+def cold_dispatch_caches():
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9:
+        pytest.skip("requires SM90+ CUDA")
+    from aspec.compat.fa4_cutlass import (
+        configure_fa4_cutlass_compilation,
+        install_ampere_helpers_shim,
+    )
+
+    install_ampere_helpers_shim()
+    configure_fa4_cutlass_compilation()
+    import cutlass.cute as cute
+    from flash_attn.cute import cache_utils, interface
+
+    # Start cold, then retain entries across semantic variants as well as
+    # lengths. A bad hash that aliases two different masks must fail numerics.
+    with pytest.MonkeyPatch.context() as patch:
+        patch.setattr(interface._flash_attn_fwd, "compile_cache", cache_utils.JITCache())
+        patch.setattr(interface._flash_attn_bwd, "compile_cache", cache_utils.JITCache())
+        yield cute
+
+
+@pytest.mark.skipif(
+    os.getenv("ASPEC_RUN_GPU_INTEGRATION") != "1",
+    reason="set ASPEC_RUN_GPU_INTEGRATION=1 to exercise real FA4 compilation",
+)
+@pytest.mark.parametrize("context_window,query_window,causal", [
+    (None, None, False), (97, None, False), (None, 3, False), (None, 3, True),
+])
+def test_variable_lengths_reuse_mask_kernels_and_match_math(monkeypatch, cold_dispatch_caches, context_window, query_window, causal):
+    cute = cold_dispatch_caches
+    original_compile = cute.compile
+    compiled_families = set()
+
+    def observe_compile(*args, **kwargs):
+        name = type(args[0]).__name__
+        if name.startswith(("FlashAttentionForwardSm", "FlashAttentionBackwardSm")):
+            caller = inspect.currentframe().f_back.f_locals
+            # FA4 deliberately separates singleton block strides in backward.
+            # Preserve its ABI/tile families, but never create another kernel
+            # merely because the modifier captured a different row length.
+            mask_hash = caller["mask_mod_hash"]
+            family = (name, tuple("<mask>" if isinstance(value, str) and value == mask_hash else value
+                                  for value in caller["compile_key"]))
+            assert family not in compiled_families, "row geometry recompiled an existing FA4 dispatch family"
+            compiled_families.add(family)
+        return original_compile(*args, **kwargs)
+
+    monkeypatch.setattr(cute, "compile", observe_compile)
+    for index, (context, blocks) in enumerate([(129, 4), (257, 9), (641, 13), (257, 9)]):
+        torch.manual_seed(31 + index)
+        queries = blocks * 7
+        shapes = [(1, 4, queries, 128), (1, 1, context, 128), (1, 1, context, 128),
+                  (1, 1, queries, 128), (1, 1, queries, 128)]
+        actual_inputs = [torch.randn(*shape, device="cuda", dtype=torch.bfloat16, requires_grad=True) for shape in shapes]
+        reference_inputs = [value.detach().clone().requires_grad_() for value in actual_inputs]
+        anchors = torch.randint(context, (1, blocks), device="cuda").sort().values
+        anchors[0, 0], anchors[0, -1] = 0, context - 1
+        keep = torch.ones_like(anchors, dtype=torch.bool)
+        keep[0, -2] = False
+
+        def metadata(backend):
+            return DirectCPMetadata(
+                anchor_positions=anchors, block_keep_mask=keep,
+                context_position_ids=torch.arange(context, device="cuda")[None],
+                block_size=7, cp_rank=0, cp_size=1, group=None, sequence_length=context,
+                backend=backend, fa4_sparse=True, context_window=context_window,
+                query_window=query_window, draft_is_causal=causal,
+            )
+
+        actual, actual_lse = direct_cp_attention(*actual_inputs, metadata("fa4"))
+        reference, reference_lse = direct_cp_attention(*reference_inputs, metadata("math"))
+        torch.testing.assert_close(actual, reference, rtol=5e-2, atol=5e-2)
+        torch.testing.assert_close(actual_lse, reference_lse, rtol=5e-2, atol=5e-2)
+        upstream = torch.randn_like(actual)
+        (actual * upstream).sum().backward()
+        (reference * upstream).sum().backward()
+        for value, expected in zip(actual_inputs, reference_inputs, strict=True):
+            torch.testing.assert_close(value.grad, expected.grad, rtol=8e-2, atol=8e-2)

tests/test_env.py

@@ -66,3 +66,16 @@ def test_aspec_actor_env_forwards_explicit_pythonpath(monkeypatch):
     monkeypatch.setenv("PYTHONPATH", python_path)
 
     assert get_aspec_env_vars()["PYTHONPATH"] == python_path
+
+
+def test_fa4_disk_cache_policy_reaches_actor_processes(monkeypatch):
+    from types import SimpleNamespace
+
+    monkeypatch.setenv("FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED", "1")
+    monkeypatch.setenv("FLASH_ATTENTION_CUTE_DSL_CACHE_DIR", "/task-local/compiled-fa4")
+    env = get_aspec_env_vars()
+    assert env["FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED"] == "1"
+    assert env["FLASH_ATTENTION_CUTE_DSL_CACHE_DIR"] == "/task-local/compiled-fa4"
+    node_local = get_aspec_env_vars(SimpleNamespace(cluster_mode="static_multi_node"))
+    assert node_local["FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED"] == "1"
+    assert "FLASH_ATTENTION_CUTE_DSL_CACHE_DIR" not in node_local