diff --git a/aspec/models/ops/direct_cp.py b/aspec/models/ops/direct_cp.py index 3c63e03..c8f417f 100644 --- a/aspec/models/ops/direct_cp.py +++ b/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,7 +500,7 @@ 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 @@ -504,7 +508,7 @@ 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 @@ -512,7 +516,7 @@ 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 diff --git a/aspec/utils/env.py b/aspec/utils/env.py index 7909c34..bfc8835 100644 --- a/aspec/utils/env.py +++ b/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 diff --git a/tests/test_direct_cp.py b/tests/test_direct_cp.py index 5f211d6..b84bd11 100644 --- a/tests/test_direct_cp.py +++ b/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("" 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) diff --git a/tests/test_env.py b/tests/test_env.py index 1ee75c8..905459b 100644 --- a/tests/test_env.py +++ b/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