File size: 3,350 Bytes
8d21b77 c32291c 8d21b77 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | # LA Flash Utils
This folder contains the sparse attention utilities used by
`LA_FLASH_ATTN=la_flash`. The release path is implemented with
FlashAttention varlen over LocateAnything range plans. It does not include or
build a local C++/CUDA extension.
## Features
- Supports batched LocateAnything hybrid MTP inference on A100, RTX 4090, and H100.
- Consumes Magi-style `q_ranges`, `k_ranges`, `segment_offsets`, and
`attn_type_map` plans generated by `batch_utils.hybrid_runtime`.
- Uses FlashAttention varlen for packed causal/full plans.
- Packs LocateAnything MTP full-window key segments before calling
FlashAttention, avoiding dense `[B,H,Q,K]` masks.
- Supports log-sum-exp merging for compatible non-packed multi-segment plans.
## Attention Types
The release path intentionally supports only FlashAttention-compatible plan
types:
| Value | Meaning |
| --- | --- |
| `0` | Full attention over the listed key segment or packed key segments. |
| `1` | Bottom-right causal attention. |
## How It Works
`batch_utils.hybrid_runtime` builds sparse range plans for the text decoder.
Each plan describes which query token intervals attend to which key/value token
intervals. `kernel_utils.range_attention` executes those plans with
FlashAttention instead of materializing dense SDPA masks.
The runtime follows three paths:
- **Packed simple plans:** when each query range maps to one contiguous
key/value range, LA Flash flattens the selected ranges, builds FlashAttention
`cu_seqlens_q` / `cu_seqlens_k`, and calls `flash_attn_varlen_func` directly.
- **Packed MTP full-window plans:** for hybrid MTP decode, multiple full
key/value windows for the same query block are concatenated into one packed
key/value sequence before the FlashAttention call. This keeps the sparse
memory profile without constructing a `[B,H,Q,K]` attention mask.
- **Compatible multi-segment plans:** when a query range attends to multiple
segments that cannot be packed as one sequence, each segment is evaluated with
FlashAttention and the partial outputs are merged with the standard
log-sum-exp softmax composition.
The output tensor shape and dtype match the decoder attention output expected
by the model. This path is inference-oriented and depends on FlashAttention's
forward kernels; it is not a custom autograd training backend.
## Runtime Knobs
| Variable | Default | Meaning |
| --- | --- | --- |
| `LA_FLASH_ATTN` | `sdpa` | Set to `la_flash` to enable this backend through `batch_utils`. |
| `LA_FLASH_FASTPATH` | `auto` | Use FlashAttention varlen for packed simple plans. |
| `LA_FLASH_SEGMENT_FASTPATH` | `auto` | Use FlashAttention varlen for multi-segment sparse plans. Full segments are packed first; other compatible segments use LSE merging. |
| `LA_FLASH_PLAN_STATS` | `0` | Record sparse plan statistics in inference summaries. |
## Notes
Dense prefill and stock worker-style generation should keep
`LA_FLASH_DENSE_BACKEND=sdpa`; LA Flash is used for sparse range plans
produced by `batch_utils`.
This package is for inference and evaluation. Training remains on the
MagiAttention backend; the batched sparse-plan decode runtime does not support
the `labels` training path.
## Source Layout
- `range_attention.py`: FlashAttention varlen dispatch, sparse KV packing, LSE
merge fallback, and availability checks.
|