batch_utils/README.md

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