mxflashattn 1.4.0__py3-none-any.whl

mxflashattn/api.py

@@ -0,0 +1,286 @@
+from __future__ import annotations
+
+import os
+
+import torch
+
+from . import reference
+from .dispatch import dispatch
+from .validation import validate_cumulative_lengths, validate_options, validate_qkv, validate_softmax_scale
+
+
+def _resolved_scale(head_dim: int, softmax_scale: float | None) -> float:
+    return float(softmax_scale) if softmax_scale is not None else float(head_dim**-0.5)
+
+
+def _options(
+    dropout_p: float,
+    window_size: tuple[int, int],
+    softcap: float,
+    alibi_slopes: torch.Tensor | None,
+    deterministic: bool,
+    return_attn_probs: bool,
+) -> None:
+    validate_options(
+        dropout_p=dropout_p,
+        window_size=window_size,
+        softcap=softcap,
+        alibi_slopes=alibi_slopes,
+        deterministic=deterministic,
+        return_attn_probs=return_attn_probs,
+    )
+
+
+def flash_attn_func(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    dropout_p: float = 0.0,
+    softmax_scale: float | None = None,
+    causal: bool = False,
+    window_size: tuple[int, int] = (-1, -1),
+    softcap: float = 0.0,
+    alibi_slopes: torch.Tensor | None = None,
+    deterministic: bool = False,
+    return_attn_probs: bool = False,
+) -> torch.Tensor:
+    validate_qkv(q, k, v)
+    validate_softmax_scale(softmax_scale)
+    _options(dropout_p, window_size, softcap, alibi_slopes, deterministic, return_attn_probs)
+    return dispatch(
+        "flash_attn_func",
+        q,
+        (q, k, v, causal, _resolved_scale(q.shape[-1], softmax_scale)),
+        lambda: reference.attention_batched(q, k, v, softmax_scale=softmax_scale, causal=causal),
+        native_args=(q, k, v),
+        native_kwargs={
+            "dropout_p": dropout_p,
+            "softmax_scale": softmax_scale,
+            "causal": causal,
+            "window_size": window_size,
+            "softcap": softcap,
+            "alibi_slopes": alibi_slopes,
+            "deterministic": deterministic,
+            "return_attn_probs": return_attn_probs,
+        },
+    )
+
+
+def flash_attn_varlen_func(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float = 0.0,
+    softmax_scale: float | None = None,
+    causal: bool = False,
+    window_size: tuple[int, int] = (-1, -1),
+    softcap: float = 0.0,
+    alibi_slopes: torch.Tensor | None = None,
+    deterministic: bool = False,
+    return_attn_probs: bool = False,
+) -> torch.Tensor:
+    validate_softmax_scale(softmax_scale)
+    for name, tensor in (("q", q), ("k", k), ("v", v)):
+        if not isinstance(tensor, torch.Tensor) or tensor.ndim != 3:
+            raise ValueError(f"{name} must have rank 3 with layout [total_tokens, heads, dim]")
+    if k.shape != v.shape:
+        raise ValueError("k and v must have identical shapes")
+    if q.device != k.device or q.device != v.device:
+        raise ValueError("q, k, and v must be on the same device")
+    if q.dtype != k.dtype or q.dtype != v.dtype:
+        raise ValueError("q, k, and v must have the same dtype")
+    if q.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise TypeError("supported dtypes are float16, bfloat16, and float32 for reference execution")
+    if q.shape[1] <= 0 or k.shape[1] <= 0 or q.shape[1] % k.shape[1]:
+        raise ValueError("q heads must be a positive multiple of kv heads")
+    if q.shape[2] not in (64, 128) or k.shape[2] != q.shape[2]:
+        raise ValueError("head dimension must be 64 or 128 and match across q, k, and v")
+    if not isinstance(cu_seqlens_q, torch.Tensor) or not isinstance(cu_seqlens_k, torch.Tensor):
+        raise TypeError("cu_seqlens_q and cu_seqlens_k must be torch.Tensor values")
+    if cu_seqlens_q.device != q.device or cu_seqlens_k.device != q.device:
+        raise ValueError("cu_seqlens_q and cu_seqlens_k must be on the same device as q")
+    _options(dropout_p, window_size, softcap, alibi_slopes, deterministic, return_attn_probs)
+    q_offsets = validate_cumulative_lengths(
+        cu_seqlens_q, total_tokens=q.shape[0], max_seqlen=max_seqlen_q, name="cu_seqlens_q"
+    )
+    k_offsets = validate_cumulative_lengths(
+        cu_seqlens_k,
+        total_tokens=k.shape[0],
+        max_seqlen=max_seqlen_k,
+        name="cu_seqlens_k",
+        require_non_empty=True,
+    )
+    if len(q_offsets) != len(k_offsets):
+        raise ValueError("cu_seqlens_q and cu_seqlens_k must describe the same batch size")
+
+    return dispatch(
+        "flash_attn_varlen_func",
+        q,
+        (
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            causal,
+            _resolved_scale(q.shape[-1], softmax_scale),
+        ),
+        lambda: reference.attention_varlen(
+            q, k, v, q_offsets, k_offsets, softmax_scale=softmax_scale, causal=causal
+        ),
+        native_args=(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k),
+        native_kwargs={
+            "dropout_p": dropout_p,
+            "softmax_scale": softmax_scale,
+            "causal": causal,
+            "window_size": window_size,
+            "softcap": softcap,
+            "alibi_slopes": alibi_slopes,
+            "deterministic": deterministic,
+            "return_attn_probs": return_attn_probs,
+        },
+    )
+
+
+def flash_attn_with_kvcache(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    k: torch.Tensor | None = None,
+    v: torch.Tensor | None = None,
+    cache_seqlens: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    softmax_scale: float | None = None,
+    causal: bool = False,
+    window_size: tuple[int, int] = (-1, -1),
+    softcap: float = 0.0,
+    alibi_slopes: torch.Tensor | None = None,
+    rotary_cos: torch.Tensor | None = None,
+    rotary_sin: torch.Tensor | None = None,
+    return_softmax_lse: bool = False,
+) -> torch.Tensor:
+    validate_softmax_scale(softmax_scale)
+    if rotary_cos is not None or rotary_sin is not None:
+        raise NotImplementedError("rotary embeddings are not supported")
+    if return_softmax_lse:
+        raise NotImplementedError("return_softmax_lse is not supported")
+    _options(0.0, window_size, softcap, alibi_slopes, False, False)
+    if (k is None) != (v is None):
+        raise ValueError("k and v must either both be provided or both be omitted")
+    if q.ndim != 4 or k_cache.ndim != 4 or v_cache.shape != k_cache.shape:
+        raise ValueError("q and cache tensors must have rank 4 and matching k/v cache shapes")
+    if q.device != k_cache.device or q.device != v_cache.device:
+        raise ValueError("q and cache tensors must be on the same device")
+    if q.dtype != k_cache.dtype or q.dtype != v_cache.dtype:
+        raise ValueError("q and cache tensors must have the same dtype")
+    if q.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise TypeError("supported dtypes are float16, bfloat16, and float32 for reference execution")
+    if q.shape[0] != k_cache.shape[0] and block_table is None:
+        raise ValueError("dense cache batch dimension must match q")
+    if q.shape[2] <= 0 or k_cache.shape[2] <= 0:
+        raise ValueError("q and kv cache head counts must be positive")
+    if k_cache.shape[1] <= 0:
+        raise ValueError("cache page size must be positive")
+    if q.shape[3] not in (64, 128) or k_cache.shape[3] != q.shape[3]:
+        raise ValueError("head dimension must be 64 or 128 and match across q and caches")
+    if q.shape[2] % k_cache.shape[2]:
+        raise ValueError("q heads must be divisible by kv cache heads")
+    if (k is not None and (k.shape != v.shape or k.ndim != 4)):
+        raise ValueError("appended k and v must have identical rank-4 shapes")
+    if k is not None and (k.shape[0] != q.shape[0] or k.shape[2:] != k_cache.shape[2:]):
+        raise ValueError("appended k/v must match batch, kv heads, and head dimension")
+    if k is not None and (k.device != q.device or k.dtype != q.dtype or v.device != q.device or v.dtype != q.dtype):
+        raise ValueError("appended k/v must match the query device and dtype")
+    if cache_seqlens is None:
+        raise ValueError("cache_seqlens is required for dense and paged caches")
+    if not isinstance(cache_seqlens, torch.Tensor):
+        raise TypeError("cache_seqlens must be a torch.Tensor")
+    if cache_seqlens.ndim != 1 or cache_seqlens.numel() != q.shape[0]:
+        raise ValueError("cache_seqlens must be one-dimensional with one value per batch item")
+    if cache_seqlens.dtype not in (torch.int32, torch.int64):
+        raise TypeError("cache_seqlens must use int32 or int64")
+    if cache_seqlens.device != q.device:
+        raise ValueError("cache_seqlens must be on the same device as q")
+    if block_table is not None:
+        if not isinstance(block_table, torch.Tensor):
+            raise TypeError("block_table must be a torch.Tensor")
+        if block_table.ndim != 2 or block_table.shape[0] != q.shape[0]:
+            raise ValueError("block_table must have one row per batch item")
+        if block_table.dtype not in (torch.int32, torch.int64):
+            raise TypeError("block_table must use int32 or int64")
+        if block_table.device != q.device:
+            raise ValueError("block_table must be on the same device as q")
+    else:
+        if k_cache.shape[0] != q.shape[0]:
+            raise ValueError("dense cache batch dimension must match q")
+
+    # Deep validation reads cache_seqlens / block_table back to the host, which
+    # costs a device synchronisation on every call. Inside vLLM this runs once
+    # per layer per decode step, so it dominates the end-to-end cost. Cheap
+    # shape/device/dtype checks above always run; the host-syncing checks are
+    # gated and enabled by the test suite via MXFLASHATTN_STRICT_VALIDATION=1.
+    if os.getenv("MXFLASHATTN_STRICT_VALIDATION", "0") == "1":
+        lengths = [int(length) for length in cache_seqlens.detach().cpu().tolist()]
+        if any(length < 0 for length in lengths):
+            raise ValueError("cache sequence lengths must be non-negative")
+        if any(length + (0 if k is None else k.shape[1]) == 0 for length in lengths):
+            raise ValueError("attention requires at least one cached or appended key/value token")
+        if block_table is None:
+            capacity = k_cache.shape[1]
+            if any(length + (0 if k is None else k.shape[1]) > capacity for length in lengths):
+                raise ValueError("cache sequence length exceeds dense cache capacity")
+        else:
+            rows = block_table.detach().cpu().tolist()
+            max_blocks = block_table.shape[1]
+            pages = k_cache.shape[0]
+            for length, row in zip(lengths, rows):
+                used = (length + (0 if k is None else k.shape[1]) + k_cache.shape[1] - 1) // k_cache.shape[1]
+                if used > max_blocks:
+                    raise ValueError("block_table does not contain enough pages for cache sequence")
+                if any(int(page) < 0 or int(page) >= pages for page in row[:used]):
+                    raise ValueError("block_table contains an out-of-range page index")
+
+    return dispatch(
+        "flash_attn_with_kvcache",
+        q,
+        (
+            q,
+            k_cache,
+            v_cache,
+            k,
+            v,
+            cache_seqlens,
+            block_table,
+            _resolved_scale(q.shape[-1], softmax_scale),
+            causal,
+        ),
+        lambda: reference.attention_kvcache(
+            q,
+            k_cache,
+            v_cache,
+            cache_seqlens=cache_seqlens,
+            block_table=block_table,
+            k=k,
+            v=v,
+            softmax_scale=softmax_scale,
+            causal=causal,
+        ),
+        native_args=(q, k_cache, v_cache),
+        native_kwargs={
+            "k": k,
+            "v": v,
+            "cache_seqlens": cache_seqlens,
+            "block_table": block_table,
+            "softmax_scale": softmax_scale,
+            "causal": causal,
+            "window_size": window_size,
+            "softcap": softcap,
+            "alibi_slopes": alibi_slopes,
+        },
+    )

mxflashattn/dispatch.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+import importlib
+import importlib.metadata
+import os
+import threading
+import warnings
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Any, Callable
+
+import torch
+
+
+class BackendUnavailableError(RuntimeError):
+    """Raised when native execution is unavailable and fallback was not enabled."""
+
+
+@dataclass(frozen=True)
+class DispatchInfo:
+    backend: str
+    operation: str
+    fallback_reason: str | None = None
+
+
+_state = threading.local()
+
+
+def get_last_dispatch_info() -> DispatchInfo | None:
+    return getattr(_state, "dispatch_info", None)
+
+
+def _native_extension() -> Any | None:
+    try:
+        return importlib.import_module("mxflashattn._C")
+    except (ImportError, OSError):
+        return None
+
+
+@lru_cache(maxsize=1)
+def _metax_flash_attn_module() -> Any | None:
+    try:
+        module = importlib.import_module("flash_attn")
+    except (ImportError, OSError):
+        return None
+    try:
+        version = importlib.metadata.version("flash-attn")
+    except importlib.metadata.PackageNotFoundError:
+        version = str(getattr(module, "__version__", ""))
+    if "metax" not in version.casefold():
+        return None
+    return module
+
+
+def _flash_attn_function(operation: str) -> Callable[..., torch.Tensor] | None:
+    if os.getenv("MXFLASHATTN_BACKEND", "auto").strip().casefold() == "aten":
+        return None
+    module = _metax_flash_attn_module()
+    if module is None:
+        return None
+    function = getattr(module, operation, None)
+    return function if callable(function) else None
+
+
+def native_unavailable_reason(
+    tensor: torch.Tensor,
+    operation: str,
+    *,
+    require_aten: bool = False,
+) -> str | None:
+    if tensor.device.type not in {"cuda", "maca", "musa"}:
+        return "native_backend_requires_accelerator_tensor"
+    if tensor.dtype not in (torch.float16, torch.bfloat16):
+        return "native_backend_supports_only_float16_and_bfloat16"
+    if _flash_attn_function(operation) is not None and not require_aten:
+        return None
+    extension = _native_extension()
+    if extension is None:
+        if require_aten:
+            return "zero_softmax_scale_requires_aten_extension"
+        return "native_extension_not_built_or_not_importable"
+    if not bool(extension.is_available()):
+        return "native_extension_reports_unavailable_runtime"
+    if not hasattr(extension, operation):
+        return f"native_extension_does_not_implement_{operation}"
+    return None
+
+
+def dispatch(
+    operation: str,
+    tensor: torch.Tensor,
+    extension_args: tuple[Any, ...],
+    reference: Callable[[], torch.Tensor],
+    *,
+    native_args: tuple[Any, ...],
+    native_kwargs: dict[str, Any],
+) -> torch.Tensor:
+    """Dispatch vendor-direct first, then ATen correctness, then opt-in reference.
+
+    The vendor wheel is the performance path. Any vendor import, signature, or
+    runtime failure is recorded and routed to the ATen extension when it is
+    available; it is never silently relabeled as a native FlashAttention path.
+    """
+    _state.dispatch_info = None
+    require_aten = native_kwargs.get("softmax_scale") == 0.0
+    vendor = None if require_aten else _flash_attn_function(operation)
+    # Preserve the actionable contract reason before attempting fallbacks. In
+    # particular, CPU reference tests and callers should see that native
+    # execution requires an accelerator rather than a generic import failure.
+    failure_reason: str | None = native_unavailable_reason(
+        tensor, operation, require_aten=require_aten
+    )
+
+    if vendor is not None:
+        try:
+            result = vendor(*native_args, **native_kwargs)
+            if operation == "flash_attn_with_kvcache" and native_kwargs.get("k") is not None:
+                # Documented contract: appending k/v advances the caller's
+                # cache_seqlens in place. Only this branch mutates it, so callers
+                # that pass k=None (the vLLM decode path) need no defensive copy.
+                native_kwargs["cache_seqlens"].add_(native_kwargs["k"].shape[1])
+            _state.dispatch_info = DispatchInfo(backend="vendor_direct", operation=operation)
+            return result
+        except (Exception, RuntimeError) as exc:
+            failure_reason = f"vendor_call_failed:{type(exc).__name__}:{exc}"
+
+    extension = _native_extension()
+    extension_reason = ""
+    if extension is not None:
+        try:
+            available = bool(extension.is_available())
+        except Exception as exc:
+            available = False
+            extension_reason = f"aten_availability_check_failed:{type(exc).__name__}:{exc}"
+        if available and hasattr(extension, operation):
+            try:
+                result = getattr(extension, operation)(*extension_args)
+                _state.dispatch_info = DispatchInfo(
+                    backend="mxmac_aten_correctness",
+                    operation=operation,
+                    fallback_reason=failure_reason,
+                )
+                return result
+            except Exception as exc:
+                extension_reason = f"aten_call_failed:{type(exc).__name__}:{exc}"
+        elif extension is None:
+            extension_reason = "aten_extension_not_importable"
+        elif not hasattr(extension, operation):
+            extension_reason = f"aten_extension_does_not_implement_{operation}"
+        elif not extension_reason:
+            extension_reason = "aten_extension_reports_unavailable_runtime"
+
+    reason = ";".join(part for part in (failure_reason, extension_reason) if part)
+    if not reason:
+        reason = "vendor_unavailable_or_parameter_incompatible"
+    if os.getenv("MXFLASHATTN_ALLOW_FALLBACK", "").strip() != "1":
+        raise BackendUnavailableError(
+            f"MXFlashAttn dispatch failed ({reason}); set MXFLASHATTN_ALLOW_FALLBACK=1 "
+            "to explicitly allow the PyTorch reference path"
+        )
+
+    _state.dispatch_info = DispatchInfo(
+        backend="reference", operation=operation, fallback_reason=reason
+    )
+    warnings.warn(
+        f"MXFlashAttn is using the PyTorch reference path: {reason}",
+        RuntimeWarning,
+        stacklevel=3,
+    )
+    return reference()
+
+

mxflashattn/policy.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+
+Phase = Literal["prefill", "decode"]
+Backend = Literal["vendor", "m1_prefill_chunked", "m1_split_kv"]
+
+
+@dataclass(frozen=True)
+class WorkloadShape:
+    """Shape metadata used by the v1.4 candidate policy.
+
+    This object contains no tensors and performs no device queries. It is
+    deliberately suitable for local tests and for later integration with a
+    vLLM adapter.
+    """
+
+    phase: Phase
+    batch_size: int
+    query_len: int
+    kv_len: int
+    head_dim: int
+    gqa_ratio: int
+
+    def __post_init__(self) -> None:
+        if self.phase not in ("prefill", "decode"):
+            raise ValueError("phase must be 'prefill' or 'decode'")
+        if self.batch_size <= 0 or self.query_len <= 0 or self.kv_len <= 0:
+            raise ValueError("batch_size, query_len, and kv_len must be positive")
+        if self.head_dim <= 0 or self.gqa_ratio <= 0:
+            raise ValueError("head_dim and gqa_ratio must be positive")
+
+
+@dataclass(frozen=True)
+class DispatchPolicyConfig:
+    """Candidate thresholds, not production calibration.
+
+    The values are intentionally disabled by default. They must be calibrated
+    on the target C500 environment before being used for runtime dispatch.
+    """
+
+    enabled: bool = False
+    calibrated: bool = False
+    prefill_batch_min: int = 4
+    prefill_kv_len_min: int = 512
+    decode_batch_min: int = 16
+    decode_kv_len_min: int = 256
+
+
+@dataclass(frozen=True)
+class DispatchDecision:
+    backend: Backend
+    reason: str
+    eligible: bool
+    policy_enabled: bool
+    thresholds_calibrated: bool
+
+
+class AdaptiveDispatchPolicy:
+    """Safe, tensor-free candidate policy for later runtime integration."""
+
+    def __init__(self, config: DispatchPolicyConfig | None = None) -> None:
+        self.config = config or DispatchPolicyConfig()
+
+    def select(self, shape: WorkloadShape) -> DispatchDecision:
+        config = self.config
+        if not config.enabled:
+            return DispatchDecision(
+                backend="vendor",
+                reason="policy_disabled",
+                eligible=False,
+                policy_enabled=False,
+                thresholds_calibrated=config.calibrated,
+            )
+        if not config.calibrated:
+            return DispatchDecision(
+                backend="vendor",
+                reason="thresholds_not_calibrated",
+                eligible=False,
+                policy_enabled=True,
+                thresholds_calibrated=False,
+            )
+
+        if shape.phase == "prefill":
+            eligible = (
+                shape.batch_size >= config.prefill_batch_min
+                and shape.kv_len >= config.prefill_kv_len_min
+            )
+            return DispatchDecision(
+                backend="m1_prefill_chunked" if eligible else "vendor",
+                reason="prefill_threshold_match" if eligible else "prefill_below_threshold",
+                eligible=eligible,
+                policy_enabled=True,
+                thresholds_calibrated=True,
+            )
+
+        eligible = (
+            shape.batch_size >= config.decode_batch_min
+            and shape.kv_len >= config.decode_kv_len_min
+            and shape.gqa_ratio in (2, 4, 8)
+            and shape.head_dim in (64, 128)
+        )
+        return DispatchDecision(
+            backend="m1_split_kv" if eligible else "vendor",
+            reason="decode_threshold_match" if eligible else "decode_shape_or_threshold_miss",
+            eligible=eligible,
+            policy_enabled=True,
+            thresholds_calibrated=True,
+        )

mxflashattn/prefill.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+import math
+
+import torch
+
+from .reference import attention_batched
+from .validation import validate_qkv, validate_softmax_scale
+
+
+def prefill_chunked_reference(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    chunk_size: int = 256,
+    softmax_scale: float | None = None,
+    causal: bool = False,
+) -> torch.Tensor:
+    """Correctness-first chunked attention prototype for Prefill.
+
+    The tensors use the public layout ``[batch, sequence, heads, head_dim]``.
+    KV is scanned in chunks and combined with an online softmax accumulator.
+    This is a reference implementation for algorithm validation, not a
+    C500-optimized kernel and not part of the production dispatch path.
+    """
+    validate_qkv(q, k, v)
+    validate_softmax_scale(softmax_scale)
+    if not isinstance(chunk_size, int) or chunk_size <= 0:
+        raise ValueError("chunk_size must be a positive integer")
+
+    batch, query_len, query_heads, head_dim = q.shape
+    key_len = k.shape[1]
+    scale = float(softmax_scale) if softmax_scale is not None else 1.0 / math.sqrt(head_dim)
+    groups = query_heads // k.shape[2]
+
+    # Use float accumulators to make the prototype's numerical behavior
+    # comparable across float16, bfloat16, and float32 inputs.
+    qf = q.float().transpose(1, 2)
+    kf = k.float().transpose(1, 2).repeat_interleave(groups, dim=1)
+    vf = v.float().transpose(1, 2).repeat_interleave(groups, dim=1)
+    output = torch.zeros((batch, query_heads, query_len, head_dim), device=q.device, dtype=torch.float32)
+    normalizer = torch.zeros((batch, query_heads, query_len, 1), device=q.device, dtype=torch.float32)
+    running_max = torch.full_like(normalizer, float("-inf"))
+
+    query_positions = torch.arange(query_len, device=q.device) + (key_len - query_len)
+    for start in range(0, key_len, chunk_size):
+        end = min(start + chunk_size, key_len)
+        scores = torch.matmul(qf, kf[:, :, start:end].transpose(-2, -1)) * scale
+
+        if causal:
+            key_positions = torch.arange(start, end, device=q.device)
+            blocked = key_positions.unsqueeze(0) > query_positions.unsqueeze(1)
+            scores = scores.masked_fill(blocked.unsqueeze(0).unsqueeze(0), float("-inf"))
+
+        chunk_max = scores.amax(dim=-1, keepdim=True)
+        old_valid = torch.isfinite(running_max)
+        chunk_valid = torch.isfinite(chunk_max)
+        new_max = torch.maximum(running_max, chunk_max)
+
+        old_factor = torch.where(
+            old_valid,
+            torch.exp(running_max - new_max),
+            torch.zeros_like(running_max),
+        )
+        exp_scores = torch.where(
+            torch.isfinite(scores),
+            torch.exp(scores - new_max),
+            torch.zeros_like(scores),
+        )
+        normalizer = normalizer * old_factor + exp_scores.sum(dim=-1, keepdim=True)
+        output = output * old_factor
+        output = output + torch.matmul(exp_scores, vf[:, :, start:end])
+        running_max = torch.where(chunk_valid | old_valid, new_max, running_max)
+
+    output = torch.where(
+        normalizer > 0,
+        output / normalizer.clamp_min(torch.finfo(output.dtype).tiny),
+        torch.zeros_like(output),
+    )
+    return output.transpose(1, 2).to(dtype=q.dtype)
+
+
+def prefill_chunked_reference_fallback(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    chunk_size: int = 256,
+    softmax_scale: float | None = None,
+    causal: bool = False,
+) -> torch.Tensor:
+    """Convenience wrapper used by local experiments.
+
+    The fallback provides a single switch for comparing the prototype with the
+    existing readable reference implementation. It is intentionally not used
+    by the runtime dispatcher.
+    """
+    if chunk_size >= k.shape[1]:
+        return attention_batched(q, k, v, softmax_scale=softmax_scale, causal=causal)
+    return prefill_chunked_reference(
+        q,
+        k,
+        v,
+        chunk_size=chunk_size,
+        softmax_scale=softmax_scale,
+        causal=causal,
+    )