File size: 2,094 Bytes
c338b8f | 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 | # Batch Utils
`batch_utils` contains the optional batched hybrid generation runtime for
LocateAnything. It keeps the model loading, tokenization, image feature caching,
sampling, and scheduler code used by `batch_infer.py` and the detection
experiments.
## Runtime Modes
- `LA_FLASH_ATTN=sdpa`: stock PyTorch SDPA path.
- `LA_FLASH_ATTN=eager`: eager attention path for debugging.
- `LA_FLASH_ATTN=magi`: MagiAttention path when MagiAttention is installed.
- `LA_FLASH_ATTN=la_flash`: LA Flash sparse range backend
from `kernel_utils`.
## Common Knobs
| Variable | Default | Meaning |
| --- | --- | --- |
| `LA_FLASH_MODEL` | `nvidia/LocateAnything-3B` | HF model id or local model directory. |
| `LA_FLASH_ATTN` | `sdpa` | LLM attention backend. |
| `LA_FLASH_VISION_ATTN` | `auto` | Vision encoder attention: `auto`, `flash_attention_2`, `sdpa`, or `eager`. |
| `LA_FLASH_STRICT_ATTN` | `0` | Set `1` to fail instead of falling back to SDPA. |
| `LA_FLASH_HYBRID_SCHEDULER` | `eager` | Hybrid decode scheduler. |
| `LA_FLASH_HYBRID_GROUP_SIZE` | `0` | Scheduler group size; `0` lets the runtime decide. |
| `LA_FLASH_VISION_ENCODE_BATCH_SIZE` | `8` | Maximum images per MoonViT encode micro-batch. |
| `LA_FLASH_KV_PACK_TOKEN_BUDGET` | `0` | Optional KV packing memory cap for long-tail batches. |
| `LA_FLASH_DENSE_BACKEND` | `sdpa` | Dense worker/prefill attention backend. Keep this as `sdpa`; LA Flash is used for sparse range plans. |
| `LA_FLASH_SEGMENT_FASTPATH` | `auto` | Sparse MTP decode uses FlashAttention varlen multi-segment merge by default. |
## CLI Example
```bash
python batch_infer.py \
--model nvidia/LocateAnything-3B \
--attn la_flash \
--scheduler pipeline \
--batch-size 4 \
--image /path/to/image.jpg \
--query "person</c>car"
```
For JSONL input, each row should contain:
```json
{"image": "/path/to/image.jpg", "query": "person</c>car"}
```
## Training Boundary
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.
|