kernel_utils/README.md

# 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.