Add batch inference and LA Flash runtime

README.md

@@ -31,7 +31,8 @@ base_model:
 * 🚀 **Online Demo**: [LocateAnything (Hugging Face Spaces)](https://huggingface.co/spaces/nvidia/LocateAnything)
 * 💻 **GitHub Code**: [NVlabs/Eagle/Embodied](https://github.com/NVlabs/Eagle/tree/main/Embodied)
 * 📄 **Paper**: [arXiv:2605.27365](https://arxiv.org/abs/2605.27365)
-* 
+
+
 # Model Overview
 
 ### Description:
@@ -263,6 +264,37 @@ Test Hardware: H100
 
 We suggest using `max_new_tokens=8192` and `generation_mode="hybrid"` to avoid truncated response and balance speed with accuracy.
 
+### Batch Hybrid Inference
+
+This release includes `batch_infer.py`, `batch_utils`, and `kernel_utils` for
+high-throughput detection and grounding. The `la_flash` backend is a pure
+FlashAttention-varlen sparse range executor: it keeps LocateAnything's hybrid
+MTP decoding path, avoids dense `[B,H,Q,K]` SDPA masks, and does not require a
+custom CUDA extension build.
+
+Use it with:
+
+```bash
+python batch_infer.py \
+  --model . \
+  --attn la_flash \
+  --scheduler pipeline \
+  --batch-size 4 \
+  --image /path/to/image.jpg \
+  --query "person</c>car"
+```
+
+A100 4K probe, real 3840x2160 street image, `query=vehicle`,
+`batch_size=4`, raw PIL input, `in_token_limit=25600`, hybrid MTP inference:
+
+| Backend | Attention Path | Time | Peak Reserved Memory |
+| --- | --- | ---: | ---: |
+| `sdpa` | Dense SDPA masks | 8.2600 s | 35.12 GB |
+| `la_flash` | FlashAttention sparse range plan | 8.0314 s | 11.71 GB |
+
+See `batch_utils/README.md` and `kernel_utils/README.md` for runtime knobs and
+implementation details.
+
 ### Installation
 
 ```bash
@@ -489,9 +521,36 @@ points = LocateAnythingWorker.parse_points(result["answer"], w, h)
 | `slow` | Pure auto-regressive decoding | Slowest | Most robust |
 | `hybrid` (default) | MTP first, falls back to AR on uncertain boxes, switches back after box boundary | Balanced | Best overall |
 
+## Batch Utils and Kernel Utils
+
+This repository also includes optional utilities for high-throughput detection
+runs:
+
+- `batch_infer.py`: JSONL/image-query batch inference CLI.
+- `batch_utils/`: batched hybrid generation runtime. See
+  `batch_utils/README.md`.
+- `kernel_utils/`: LA Flash sparse range utilities. See
+  `kernel_utils/README.md`.
+
+Run a small batch inference job:
+
+```bash
+python batch_infer.py \
+  --model . \
+  --attn la_flash \
+  --scheduler pipeline \
+  --batch-size 4 \
+  --image assets/pointing.png \
+  --query "the object being pointed at"
+```
+
+The batched sparse-plan decode runtime is intended for inference/evaluation and
+does not support the training `labels` path. Training remains on the
+MagiAttention backend.
+
 ## Ethical Considerations:
 NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable development for a wide array of AI applications. When downloaded or used in accordance with our terms of service, developers should work with their internal model team to ensure this model meets requirements for the relevant industry and use case and addresses unforeseen product misuse.
 
 Please make sure you have proper rights and permissions for all input image and video content; if image or video includes people, personal health information, or intellectual property, the image or video generated will not blur or maintain proportions of image subjects included.
 
-Please report model quality, risk, security vulnerabilities or NVIDIA AI Concerns [here](https://app.intigriti.com/programs/nvidia/nvidiavdp/detail).
+Please report model quality, risk, security vulnerabilities or NVIDIA AI Concerns [here](https://app.intigriti.com/programs/nvidia/nvidiavdp/detail).

batch_infer.py

@@ -0,0 +1,129 @@
+#!/usr/bin/env python3
+"""Minimal batch inference CLI for the LocateAnything-3B release code.
+
+Examples:
+  python batch_infer.py --model /path/to/LocateAnything-3B --attn sdpa \
+    --image demo.jpg --query "person</c>car"
+
+  python batch_infer.py --requests requests.jsonl --batch-size 16 --attn la_flash
+
+Each JSONL request should contain {"image": "/path/to.jpg", "query": "person</c>car"}.
+"""
+import argparse
+import json
+import os
+from pathlib import Path
+
+from PIL import Image
+
+
+def _attn_arg(value):
+    mode = (value or "sdpa").strip().lower().replace("-", "_")
+    aliases = {
+        "": "sdpa",
+        "manual": "eager",
+        "torch": "eager",
+        "torch_eager": "eager",
+        "torch_sdpa": "sdpa",
+        "flash": "la_flash",
+        "la_flash": "la_flash",
+        "kernel": "la_flash",
+        "cuda": "la_flash",
+        "range": "la_flash",
+        "range_attention": "la_flash",
+    }
+    mode = aliases.get(mode, mode)
+    if mode not in {"sdpa", "eager", "magi", "la_flash"}:
+        raise argparse.ArgumentTypeError(
+            f"--attn must be one of sdpa, eager, magi, la_flash; got {value!r}"
+        )
+    return mode
+
+
+def _load_requests(args):
+    requests = []
+    if args.requests:
+        with open(args.requests, "r", encoding="utf-8") as f:
+            for line in f:
+                if not line.strip():
+                    continue
+                row = json.loads(line)
+                requests.append((row["image"], row["query"]))
+    if args.image or args.query:
+        if len(args.image or []) != len(args.query or []):
+            raise ValueError("--image and --query must appear the same number of times")
+        requests.extend(zip(args.image, args.query))
+    if not requests:
+        raise ValueError("provide --requests JSONL or at least one --image/--query pair")
+    return requests
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--requests", help="JSONL file with image/query fields")
+    ap.add_argument("--image", action="append", help="Image path; repeat with --query")
+    ap.add_argument("--query", action="append", help="Category query, e.g. person</c>car")
+    ap.add_argument("--model", default=os.environ.get("LA_FLASH_MODEL", "nvidia/LocateAnything-3B"))
+    ap.add_argument("--attn", type=_attn_arg, default=os.environ.get("LA_FLASH_ATTN", "sdpa"),
+                    help="LLM attention backend: sdpa, eager, magi, or la_flash")
+    ap.add_argument("--vision-attn", default=os.environ.get("LA_FLASH_VISION_ATTN", "auto"),
+                    choices=["auto", "flash_attention_2", "sdpa", "eager"])
+    ap.add_argument("--batch-size", type=int, default=1)
+    ap.add_argument("--scheduler", default=os.environ.get("LA_FLASH_HYBRID_SCHEDULER", "eager"),
+                    choices=["eager", "hold_ar", "ar_first", "pipeline", "adaptive"])
+    ap.add_argument("--group-size", type=int, default=int(os.environ.get("LA_FLASH_HYBRID_GROUP_SIZE", "0")))
+    ap.add_argument("--max-new-tokens", type=int, default=2048)
+    ap.add_argument("--temperature", type=float, default=0.7)
+    ap.add_argument("--top-p", type=float, default=0.9)
+    ap.add_argument("--top-k", type=int, default=0)
+    ap.add_argument("--repetition-penalty", type=float, default=1.1)
+    ap.add_argument("--strict-attn", action="store_true",
+                    help="Fail instead of falling back to SDPA if magi/la_flash is unavailable")
+    ap.add_argument("--out", default="", help="Optional output JSONL path; stdout if omitted")
+    args = ap.parse_args()
+    args.attn = _attn_arg(args.attn)
+
+    os.environ["LA_FLASH_MODEL"] = args.model
+    os.environ["LA_FLASH_ATTN"] = args.attn
+    os.environ["LA_FLASH_VISION_ATTN"] = args.vision_attn
+    os.environ["LA_FLASH_HYBRID_SCHEDULER"] = args.scheduler
+    os.environ["LA_FLASH_HYBRID_GROUP_SIZE"] = str(args.group_size)
+    if args.strict_attn:
+        os.environ["LA_FLASH_STRICT_ATTN"] = "1"
+
+    from batch_utils import generate_batch_hybrid, get_last_hybrid_stats, load
+    from batch_utils.hybrid_runtime import load_pil
+
+    requests = _load_requests(args)
+    load()
+
+    writer = open(args.out, "w", encoding="utf-8") if args.out else None
+    try:
+        for start in range(0, len(requests), max(1, args.batch_size)):
+            chunk = requests[start:start + max(1, args.batch_size)]
+            pairs = [(load_pil(image), query) for image, query in chunk]
+            texts = generate_batch_hybrid(
+                pairs,
+                temperature=args.temperature,
+                top_p=None if args.top_p < 0 else args.top_p,
+                top_k=None if args.top_k <= 0 else args.top_k,
+                repetition_penalty=args.repetition_penalty,
+                max_new_tokens=args.max_new_tokens,
+                scheduler=args.scheduler,
+                group_size=args.group_size,
+            )
+            stats = get_last_hybrid_stats()
+            for (image, query), text in zip(chunk, texts):
+                row = {"image": str(Path(image)), "query": query, "raw_response": text, "stats": stats}
+                line = json.dumps(row, ensure_ascii=False)
+                if writer:
+                    writer.write(line + "\n")
+                else:
+                    print(line, flush=True)
+    finally:
+        if writer:
+            writer.close()
+
+
+if __name__ == "__main__":
+    main()

batch_utils/README.md

@@ -0,0 +1,53 @@
+# 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.

batch_utils/__init__.py

@@ -0,0 +1,12 @@
+"""Hybrid batched inference helpers for NVIDIA LocateAnything-3B."""
+from .hybrid_runtime import load, load_pil
+from .engine_hybrid import generate_batch_hybrid, generate_batch_grouped_hybrid, get_last_hybrid_stats
+
+__all__ = [
+    "load",
+    "load_pil",
+    "generate_batch_hybrid",
+    "generate_batch_grouped_hybrid",
+    "get_last_hybrid_stats",
+]
+__version__ = "0.1.0"

batch_utils/engine_hybrid.py

@@ -0,0 +1,1357 @@
+"""Batched hybrid-mode generation for LocateAnything-3B.
+
+This module keeps the stock hybrid state machine:
+
+  MTP -> error_box -> AR
+  AR  -> box_end_ar -> MTP
+
+Rows in a batch may be in different modes. The decode loop therefore stores
+per-row KV caches, packs rows with the same mode for one forward call, then
+unpacks the clean KV back per row.
+"""
+import copy
+import importlib
+import os
+
+import torch
+
+
+from .hybrid_runtime import (
+    ATTN_MODE,
+    AR_BATCH_SAN,
+    BATCH_SAN,
+    DEV,
+    N_FUTURE,
+    _encode_images,
+    _helpers,
+    _pad_generated,
+    _set_llm_mode,
+    _tokenize,
+    _tokenize_cached_image,
+    build_magi_scheduler_ranges,
+    language_model_forward,
+    load,
+    sample_next_tokens_batched,
+    sample_tokens_batched,
+)
+
+
+README_MAX_NEW_TOKENS = 2048
+README_TEMPERATURE = 0.7
+README_TOP_P = 0.9
+README_REPETITION_PENALTY = 1.1
+
+_LAST_HYBRID_STATS = None
+
+
+def _row_len(kv):
+    return kv[0][0].shape[2]
+
+
+def _pack_stock_kv_rows(kv_rows, rows, dev):
+    """Left-pad per-row real-token KV caches for stock-style decoding."""
+    lengths = [0 if kv_rows[r] is None else _row_len(kv_rows[r]) for r in rows]
+    kmax = max(lengths) if lengths else 0
+    if kmax == 0:
+        return None, torch.zeros((len(rows), 0), dtype=torch.long, device=dev), lengths, 0
+
+    ref = next(kv_rows[r] for r in rows if kv_rows[r] is not None)
+    packed = []
+    for layer in range(len(ref)):
+        ref_k, ref_v = ref[layer]
+        ks, vs = [], []
+        for r, length in zip(rows, lengths):
+            if length == 0:
+                k = ref_k.new_zeros((1, ref_k.shape[1], kmax, ref_k.shape[3]))
+                v = ref_v.new_zeros((1, ref_v.shape[1], kmax, ref_v.shape[3]))
+            else:
+                k, v = kv_rows[r][layer]
+                if length < kmax:
+                    pad_shape = (1, k.shape[1], kmax - length, k.shape[3])
+                    k = torch.cat([k.new_zeros(pad_shape), k], dim=2)
+                    v = torch.cat([v.new_zeros(pad_shape), v], dim=2)
+            ks.append(k)
+            vs.append(v)
+        packed.append((torch.cat(ks, dim=0), torch.cat(vs, dim=0)))
+
+    kvalid = torch.zeros((len(rows), kmax), dtype=torch.long, device=dev)
+    for i, length in enumerate(lengths):
+        if length:
+            kvalid[i, kmax - length :] = 1
+    return tuple(packed), kvalid, lengths, kmax
+
+
+def _unpack_stock_after_forward(out_kv, local_row, old_len, uncached_len, kmax, umax):
+    """Keep old real KV plus the right-aligned uncached real tokens; drop pads/window."""
+    out = []
+    u0 = kmax + (umax - uncached_len)
+    u1 = kmax + umax
+    for k, v in out_kv:
+        parts_k, parts_v = [], []
+        if old_len:
+            parts_k.append(k[local_row : local_row + 1, :, kmax - old_len : kmax, :])
+            parts_v.append(v[local_row : local_row + 1, :, kmax - old_len : kmax, :])
+        if uncached_len:
+            parts_k.append(k[local_row : local_row + 1, :, u0:u1, :])
+            parts_v.append(v[local_row : local_row + 1, :, u0:u1, :])
+        out.append((torch.cat(parts_k, dim=2).contiguous(),
+                    torch.cat(parts_v, dim=2).contiguous()))
+    return tuple(out)
+
+
+def _mk_generate_kwargs(temperature, top_p, top_k, repetition_penalty, row_temp=None):
+    t = temperature if row_temp is None else row_temp
+    gk = {"repetition_penalty": repetition_penalty, "generation_mode": "hybrid"}
+    if t and t > 0:
+        gk["temperature"] = t
+    if top_p is not None:
+        gk["top_p"] = top_p
+    if top_k is not None:
+        gk["top_k"] = top_k
+    return gk
+
+
+def _classify_ar_token(token_val, tids):
+    if token_val == tids["box_end_token_id"]:
+        return "box_end_ar"
+    if tids["coord_start_token_id"] <= token_val <= tids["coord_end_token_id"]:
+        return "coord_ar"
+    if token_val == tids["none_token_id"]:
+        return "coord_ar"
+    return "im_end"
+
+
+def _env_flag(name, default=False):
+    val = os.environ.get(name)
+    if val is None:
+        return default
+    return val.lower() not in {"0", "false", "no", "off", ""}
+
+
+def _env_int(name, default):
+    val = os.environ.get(name)
+    if val is None or val == "":
+        return default
+    return int(val)
+
+
+def _kv_pack_token_budget():
+    return max(0, _env_int("LA_FLASH_KV_PACK_TOKEN_BUDGET", 0))
+
+
+def _debug_enabled(debug):
+    return _env_flag("LA_FLASH_DEBUG", False) if debug is None else bool(debug)
+
+
+def _new_hybrid_stats(total_rows, scheduler, group_size, hold_max_steps, adaptive_hold_mtp_max=0):
+    return {
+        "scheduler": scheduler,
+        "requested_group_size": int(group_size or 0),
+        "hold_max_steps": int(hold_max_steps),
+        "adaptive_hold_mtp_max": int(adaptive_hold_mtp_max),
+        "input_batches": 1,
+        "input_rows": int(total_rows),
+        "groups": 0,
+        "group_sizes": [],
+        "decode_loops": 0,
+        "mixed_mode_cycles": 0,
+        "eager_mtp_then_ar_cycles": 0,
+        "ar_first_cycles": 0,
+        "pipeline_ar_after_mtp_cycles": 0,
+        "adaptive_hold_cycles": 0,
+        "adaptive_ar_first_cycles": 0,
+        "hold_ar_steps": 0,
+        "hold_ar_held_mtp_rows": 0,
+        "hold_ar_limit_mtp_forwards": 0,
+        "mtp_forwards": 0,
+        "ar_forwards": 0,
+        "mtp_forward_rows": 0,
+        "ar_forward_rows": 0,
+        "mtp_forward_query_tokens": 0,
+        "ar_forward_query_tokens": 0,
+        "max_mtp_forward_rows": 0,
+        "max_ar_forward_rows": 0,
+        "mtp_max_uncached_len": 0,
+        "ar_max_uncached_len": 0,
+        "mtp_forward_row_hist": {},
+        "ar_forward_row_hist": {},
+        "prompt_prefill_mode": _hybrid_prefill_mode(),
+        "prompt_prefill_forwards": 0,
+        "prompt_prefill_forward_rows": 0,
+        "prompt_prefill_forward_query_tokens": 0,
+        "prompt_prefill_real_tokens": 0,
+        "prompt_prefill_shared_groups": 0,
+        "prompt_prefill_shared_rows": 0,
+        "prompt_prefill_shared_saved_tokens": 0,
+        "kv_bucket_splits": 0,
+        "kv_bucket_groups": 0,
+        "kv_bucket_max_packed_tokens": 0,
+    }
+
+
+def _set_last_hybrid_stats(stats):
+    global _LAST_HYBRID_STATS
+    _LAST_HYBRID_STATS = copy.deepcopy(stats) if stats is not None else None
+
+
+def get_last_hybrid_stats():
+    """Return scheduler/forward statistics from the most recent hybrid batch."""
+    return copy.deepcopy(_LAST_HYBRID_STATS)
+
+
+def _record_group_stats(stats, bsz):
+    if stats is None:
+        return
+    stats["groups"] += 1
+    stats["group_sizes"].append(int(bsz))
+
+
+def _bump_hist(hist, val):
+    key = str(int(val))
+    hist[key] = int(hist.get(key, 0)) + 1
+
+
+def _record_forward_stats(stats, kind, rows, q_len, uncached_lens):
+    if stats is None:
+        return
+    prefix = "mtp" if kind == "mtp" else "ar"
+    nrows = int(len(rows))
+    q_len = int(q_len)
+    stats[f"{prefix}_forwards"] += 1
+    stats[f"{prefix}_forward_rows"] += nrows
+    stats[f"{prefix}_forward_query_tokens"] += nrows * q_len
+    stats[f"max_{prefix}_forward_rows"] = max(stats[f"max_{prefix}_forward_rows"], nrows)
+    stats[f"{prefix}_max_uncached_len"] = max(
+        stats[f"{prefix}_max_uncached_len"],
+        max((int(x) for x in uncached_lens), default=0),
+    )
+    _bump_hist(stats[f"{prefix}_forward_row_hist"], nrows)
+
+
+def _record_prefill_stats(stats, rows, q_len, real_tokens, shared_groups=0, shared_rows=0, saved_tokens=0):
+    if stats is None:
+        return
+    nrows = int(rows)
+    stats["prompt_prefill_forwards"] += 1
+    stats["prompt_prefill_forward_rows"] += nrows
+    stats["prompt_prefill_forward_query_tokens"] += nrows * int(q_len)
+    stats["prompt_prefill_real_tokens"] += int(real_tokens)
+    stats["prompt_prefill_shared_groups"] += int(shared_groups)
+    stats["prompt_prefill_shared_rows"] += int(shared_rows)
+    stats["prompt_prefill_shared_saved_tokens"] += int(saved_tokens)
+
+
+def _split_rows_by_kv_budget(rows, kv_rows):
+    """Keep dense left-padded KV packs bounded when a few rows become long tails."""
+    budget = _kv_pack_token_budget()
+    if budget <= 0 or len(rows) <= 1:
+        return [rows]
+    lengths = [0 if kv_rows[r] is None else _row_len(kv_rows[r]) for r in rows]
+    if not lengths or max(lengths) * len(rows) <= budget:
+        return [rows]
+
+    groups = []
+    current = []
+    current_max = 0
+    for row, length in sorted(zip(rows, lengths), key=lambda item: item[1]):
+        next_max = max(current_max, int(length))
+        if current and next_max * (len(current) + 1) > budget:
+            groups.append(current)
+            current = [row]
+            current_max = int(length)
+        else:
+            current.append(row)
+            current_max = next_max
+    if current:
+        groups.append(current)
+    return groups or [rows]
+
+
+def _record_kv_bucket_stats(stats, groups, kv_rows):
+    if stats is None:
+        return
+    max_packed = 0
+    for group in groups:
+        if not group:
+            continue
+        kmax = max((0 if kv_rows[r] is None else _row_len(kv_rows[r])) for r in group)
+        max_packed = max(max_packed, int(kmax) * len(group))
+    stats["kv_bucket_max_packed_tokens"] = max(stats["kv_bucket_max_packed_tokens"], max_packed)
+    if len(groups) > 1:
+        stats["kv_bucket_splits"] += 1
+        stats["kv_bucket_groups"] += len(groups)
+
+
+def _hybrid_scheduler(scheduler):
+    val = os.environ.get("LA_FLASH_HYBRID_SCHEDULER", "eager") if scheduler is None else scheduler
+    val = str(val).strip().lower()
+    aliases = {
+        "": "eager",
+        "default": "eager",
+        "normal": "eager",
+        "hold": "hold_ar",
+        "hold-ar": "hold_ar",
+        "hold_mtp": "hold_ar",
+        "hold-mtp": "hold_ar",
+        "repair_first": "ar_first",
+        "repair-first": "ar_first",
+        "ar-first": "ar_first",
+    }
+    val = aliases.get(val, val)
+    if val not in {"eager", "hold_ar", "ar_first", "pipeline", "adaptive"}:
+        raise ValueError("scheduler must be one of: eager, hold_ar, ar_first, pipeline, adaptive")
+    return val
+
+
+def _hybrid_group_size(group_size):
+    if group_size is None:
+        return max(0, _env_int("LA_FLASH_HYBRID_GROUP_SIZE", 0))
+    return max(0, int(group_size))
+
+
+def _hybrid_prefill_mode():
+    val = os.environ.get("LA_FLASH_HYBRID_PREFILL", "shared").strip().lower()
+    aliases = {
+        "0": "none",
+        "false": "none",
+        "off": "none",
+        "legacy": "none",
+        "1": "per_row",
+        "true": "per_row",
+        "on": "per_row",
+        "single": "per_row",
+        "row": "per_row",
+        "rows": "per_row",
+        "batched": "batch",
+        "prefix": "shared",
+        "shared_prefix": "shared",
+        "shared-image": "shared",
+        "shared_image": "shared",
+        "vision": "shared",
+    }
+    val = aliases.get(val, val)
+    if val not in {"none", "per_row", "batch", "shared"}:
+        raise ValueError("LA_FLASH_HYBRID_PREFILL must be one of none, per_row, batch, shared")
+    return val
+
+
+def _tolist(t):
+    return t.detach().cpu().tolist()
+
+
+def _safe_decode_rows(tok, input_ids):
+    rows = []
+    for row in _tolist(input_ids):
+        try:
+            rows.append(tok.decode(torch.tensor(row), skip_special_tokens=False))
+        except Exception:
+            rows.append("<decode failed>")
+    return rows
+
+
+def _safe_decode_row(tok, row):
+    try:
+        return tok.decode(torch.tensor(row), skip_special_tokens=False)
+    except Exception:
+        return "<decode failed>"
+
+
+def _effective_allowed_mask(mask2d, q_len, past_len, mtp_window=False):
+    """Readable 1/0 q-by-k mask derived from the 2D key-valid mask.
+
+    This mirrors the model path at a high level:
+    causal + padding columns, then the MTP window update
+    attn[-block:, -block:] = visible and attn[-block:, -block-1] = masked.
+    """
+    rows = []
+    key_valid = mask2d.detach().cpu().bool()
+    total_len = int(key_valid.numel())
+    for qi in range(q_len):
+        q_abs = past_len + qi
+        row = []
+        for ki in range(total_len):
+            row.append(1 if bool(key_valid[ki]) and ki <= q_abs else 0)
+        rows.append(row)
+
+    if mtp_window and q_len >= N_FUTURE and total_len >= N_FUTURE:
+        q0 = q_len - N_FUTURE
+        k0 = total_len - N_FUTURE
+        for qi in range(q0, q_len):
+            for ki in range(k0, total_len):
+                rows[qi][ki] = 1
+            if k0 - 1 >= 0:
+                rows[qi][k0 - 1] = 0
+    return rows
+
+
+def _tail_matrix(mat, rows=None, cols=None):
+    if rows is not None:
+        mat = mat[-rows:]
+    if cols is not None:
+        mat = [row[-cols:] for row in mat]
+    return mat
+
+
+def _format_01_matrix(mat):
+    return "\n".join("  " + " ".join(str(int(v)) for v in row) for row in mat)
+
+
+def _safe_sdpa_mask_enabled():
+    return _env_flag("LA_FLASH_SDPA_SAFE_4D_MASK", True)
+
+
+def _build_safe_sdpa_visible_mask(attention_mask_2d, input_ids, past_len, mtp_window=False):
+    """Build a 4D 1/0 visible mask, with harmless visibility for all-masked pad queries.
+
+    The remote Qwen2 SDPA path uses a 2D key-valid mask and can create fully
+    masked query rows for left-padded, no-cache prefill. Those rows can produce
+    NaNs inside SDPA and later contaminate real tokens through masked K columns.
+    This 4D mask keeps real-token visibility identical, and only gives otherwise
+    all-masked query rows one valid fallback key so their activations stay finite.
+    """
+    bsz, q_len = int(input_ids.shape[0]), int(input_ids.shape[1])
+    key_len = int(attention_mask_2d.shape[1])
+    dev = input_ids.device
+    key_valid = attention_mask_2d.to(dtype=torch.bool, device=dev)
+    key_idx = torch.arange(key_len, device=dev).view(1, 1, key_len)
+    q_abs = (past_len + torch.arange(q_len, device=dev)).view(1, q_len, 1)
+    visible = key_valid[:, None, :] & (key_idx <= q_abs)
+
+    if mtp_window and q_len >= N_FUTURE and key_len >= N_FUTURE:
+        k0 = key_len - N_FUTURE
+        visible[:, -N_FUTURE:, k0:key_len] = key_valid[:, None, k0:key_len]
+        blocked_k = k0 - 1
+        if blocked_k >= 0:
+            visible[:, -N_FUTURE:, blocked_k] = False
+
+    row_has_key = visible.any(dim=-1)
+    fallback_rows = int((~row_has_key).sum().item())
+    if fallback_rows:
+        for b in range(bsz):
+            valid = torch.nonzero(key_valid[b], as_tuple=False).flatten()
+            fallback = int(valid[0].item()) if valid.numel() else 0
+            missing = torch.nonzero(~row_has_key[b], as_tuple=False).flatten()
+            if missing.numel():
+                visible[b, missing, fallback] = True
+
+    mask = visible[:, None, :, :].to(dtype=torch.bfloat16)
+    try:
+        mask._la_flash_visible_mask = True
+    except Exception:
+        pass
+    return mask, fallback_rows
+
+
+def _mask_desc(mask):
+    if mask is None:
+        return "none"
+    if isinstance(mask, dict):
+        return "magi_ranges"
+    if hasattr(mask, "dim"):
+        return "4d_safe_sdpa" if mask.dim() == 4 else "2d_key_valid"
+    return type(mask).__name__
+
+
+def _forward_attention_mask(model, input_ids, attention_mask_2d, past_len, mtp_window=False, range_plan=False):
+    llm = model.language_model.model
+    if getattr(model, "_la_flash_requested_attn", ATTN_MODE) in {"magi", "la_flash"}:
+        range_plan = build_magi_scheduler_ranges(
+            model, attention_mask_2d, input_ids, past_len, mtp_window=mtp_window)
+        if range_plan is not None:
+            return range_plan, 0
+    needs_safe_pad = (
+        past_len == 0
+        and attention_mask_2d is not None
+        and attention_mask_2d.dim() == 2
+        and input_ids.shape[0] > 1
+    )
+    if (
+        getattr(llm, "_attn_implementation", None) == "sdpa"
+        and _safe_sdpa_mask_enabled()
+        and needs_safe_pad
+        and attention_mask_2d is not None
+        and attention_mask_2d.dim() == 2
+    ):
+        return _build_safe_sdpa_visible_mask(attention_mask_2d, input_ids, past_len, mtp_window)
+    return attention_mask_2d, 0
+
+
+def _actual_sdpa_allowed_masks(model, input_ids, attention_mask, past_len):
+    """Recreate the remote Qwen2 SDPA 4D additive mask and return a 0/1 view."""
+    llm = model.language_model.model
+    mod = importlib.import_module(type(llm).__module__)
+    bsz, q_len = int(input_ids.shape[0]), int(input_ids.shape[1])
+    dummy = torch.empty(
+        (bsz, q_len, 1),
+        dtype=torch.bfloat16,
+        device=input_ids.device,
+    )
+    mask4 = mod._prepare_4d_causal_attention_mask(
+        attention_mask,
+        (bsz, q_len),
+        dummy,
+        past_len,
+        sliding_window=getattr(llm.config, "sliding_window", None),
+    )
+    remote_ar_decode = q_len == 1 or (
+        input_ids is not None and int(input_ids[0, -1].item()) != int(llm.text_mask_token_id)
+    )
+    if not remote_ar_decode and mask4 is not None and mask4.dim() == 4:
+        rows = []
+        for b in range(bsz):
+            rows.append(
+                mod.update_causal_mask_for_one_gen_window_2d(
+                    input_ids[b],
+                    mask4[b][0].clone(),
+                    block_size=int(llm.block_size),
+                    use_cache=True,
+                    causal_attn=bool(getattr(llm, "causal_attn", False)),
+                ).unsqueeze(0)
+            )
+        mask4 = torch.stack(rows, dim=0)
+    allowed = (mask4[:, 0] >= 0).to(torch.int8).detach().cpu().tolist()
+    return allowed, tuple(mask4.shape), remote_ar_decode
+
+
+def _debug_magi_ranges(q_len, past_len, mtp_window=False):
+    kv_len = past_len + q_len
+    ar_decode = not mtp_window
+    if ar_decode:
+        return {
+            "q_ranges": [[0, q_len]],
+            "k_ranges": [[0, kv_len]],
+            "attn_type_map": ["CAUSAL"],
+        }
+
+    block = N_FUTURE
+    if not (0 < block <= q_len <= kv_len):
+        return {"error": f"invalid magi MTP shape: block={block}, q_len={q_len}, kv_len={kv_len}"}
+
+    prefix_len = kv_len - block
+    blocked_k = prefix_len - 1
+    q_ranges, k_ranges, attn_types = [], [], []
+    if q_len == kv_len:
+        if prefix_len > 0:
+            q_ranges.append([0, prefix_len])
+            k_ranges.append([0, prefix_len])
+            attn_types.append("CAUSAL")
+        if prefix_len > 0 and blocked_k > 0:
+            q_ranges.append([prefix_len, kv_len])
+            k_ranges.append([0, blocked_k])
+            attn_types.append("FULL")
+        q_ranges.append([prefix_len, kv_len])
+        k_ranges.append([prefix_len, kv_len])
+        attn_types.append("FULL")
+    else:
+        recompute = q_len - block
+        q_global_start = kv_len - q_len
+        for i in range(recompute):
+            g = q_global_start + i
+            q_ranges.append([i, i + 1])
+            k_ranges.append([0, g + 1])
+            attn_types.append("FULL")
+        q_win = [recompute, q_len]
+        if blocked_k > 0:
+            q_ranges.append(q_win)
+            k_ranges.append([0, blocked_k])
+            attn_types.append("FULL")
+        q_ranges.append(q_win)
+        k_ranges.append([prefix_len, kv_len])
+        attn_types.append("FULL")
+
+    return {"q_ranges": q_ranges, "k_ranges": k_ranges, "attn_type_map": attn_types}
+
+
+def _print_debug_forward(label, model, tok, input_ids, attention_mask, position_ids,
+                         past_len, mtp_window=False, extra=None, attention_impl="sdpa"):
+    print(f"\n========== LA Flash DEBUG {label} ==========", flush=True)
+    if extra:
+        for k, v in extra.items():
+            print(f"{k}: {v}", flush=True)
+    tail = int(os.environ.get("LA_FLASH_DEBUG_TAIL", "15"))
+    bsz, q_len = int(input_ids.shape[0]), int(input_ids.shape[1])
+    key_len = int(attention_mask.shape[1])
+    q_tail, k_tail = min(tail, q_len), min(tail, key_len)
+    print(
+        "shapes: "
+        f"input_ids={tuple(input_ids.shape)} "
+        f"position_ids={tuple(position_ids.shape)} "
+        f"attention_mask_key_valid={tuple(attention_mask.shape)} "
+        f"mask_2d_q_by_k=({bsz}, {q_len}, {key_len}) "
+        f"mask_2d_tail=({bsz}, {q_tail}, {k_tail}) "
+        f"past_len={past_len} q_len={q_len} "
+        f"mtp_window={mtp_window} ar_decode={not mtp_window}",
+        flush=True,
+    )
+    print(f"dtypes/devices: input_ids={input_ids.dtype}@{input_ids.device} position_ids={position_ids.dtype}@{position_ids.device} attention_mask={attention_mask.dtype}@{attention_mask.device}", flush=True)
+    print(f"attention_impl={attention_impl}", flush=True)
+    input_rows = _tolist(input_ids)
+    pos_rows = _tolist(position_ids)
+    print(f"tail_window_last={tail}", flush=True)
+    print(f"input_ids_tail.shape=({bsz}, {q_tail})", flush=True)
+    print(f"position_ids_tail.shape=({bsz}, {q_tail})", flush=True)
+    actual_sdpa = None
+    if attention_impl in {"sdpa", "eager", "la_flash"}:
+        try:
+            actual_sdpa = _actual_sdpa_allowed_masks(model, input_ids, attention_mask, past_len)
+            print(
+                f"actual_sdpa_4d_mask_shape={actual_sdpa[1]} "
+                f"remote_ar_decode={actual_sdpa[2]}",
+                flush=True,
+            )
+        except Exception as e:
+            print(f"actual_sdpa_4d_mask_debug_failed={type(e).__name__}: {e}", flush=True)
+
+    for b in range(input_ids.shape[0]):
+        ids_tail = input_rows[b][-tail:]
+        pos_tail = pos_rows[b][-tail:]
+        allowed = _effective_allowed_mask(attention_mask[b], input_ids.shape[1], past_len, mtp_window)
+        q_tail = min(tail, len(allowed))
+        k_tail = min(tail, len(allowed[0]) if allowed else 0)
+        allowed_tail = _tail_matrix(allowed, rows=q_tail, cols=k_tail)
+        print(f"batch_row={b} ar_decode={not mtp_window}", flush=True)
+        print(f"input_ids_tail[-{tail}:]: {ids_tail}", flush=True)
+        print(f"decoded_tail[-{tail}:]: {_safe_decode_row(tok, ids_tail)}", flush=True)
+        print(f"position_ids_tail[-{tail}:]: {pos_tail}", flush=True)
+        print(f"expected_mask_2d_tail[-{q_tail}:,-{k_tail}:].shape=({q_tail}, {k_tail})", flush=True)
+        print(_format_01_matrix(allowed_tail), flush=True)
+        if actual_sdpa is not None:
+            actual = actual_sdpa[0][b]
+            actual_tail = _tail_matrix(actual, rows=q_tail, cols=k_tail)
+            mismatch = sum(
+                int(allowed[qi][ki] != actual[qi][ki])
+                for qi in range(len(allowed))
+                for ki in range(len(allowed[qi]))
+            )
+            print(
+                f"actual_sdpa_mask_2d_tail[-{q_tail}:,-{k_tail}:].shape=({q_tail}, {k_tail})",
+                flush=True,
+            )
+            print(_format_01_matrix(actual_tail), flush=True)
+            print(f"expected_vs_actual_sdpa_mismatch_count={mismatch}", flush=True)
+
+    if _env_flag("LA_FLASH_DEBUG_FULL_MASK", False):
+        masks = [
+            _effective_allowed_mask(attention_mask[b], input_ids.shape[1], past_len, mtp_window)
+            for b in range(input_ids.shape[0])
+        ]
+        print("effective_allowed_mask_q_by_k_FULL:", masks, flush=True)
+    if attention_impl == "magi":
+        if bsz == 1:
+            print(
+                "magi_ranges:",
+                _debug_magi_ranges(input_ids.shape[1], past_len, mtp_window),
+                flush=True,
+            )
+        else:
+            print(
+                "magi_ranges: built once per forward from the batched scheduler mask",
+                flush=True,
+            )
+            print(
+                "magi_ranges_single_row_template:",
+                _debug_magi_ranges(input_ids.shape[1], past_len, mtp_window),
+                flush=True,
+            )
+
+
+def _common_prefix_len(prompt_ids, rows):
+    if not rows:
+        return 0
+    first = prompt_ids[rows[0]]
+    max_len = min(int(prompt_ids[r].numel()) for r in rows)
+    prefix_len = 0
+    for idx in range(max_len):
+        val = int(first[idx].item())
+        if all(int(prompt_ids[r][idx].item()) == val for r in rows[1:]):
+            prefix_len += 1
+        else:
+            break
+    return prefix_len
+
+
+def _prefill_shared_prefix_kv_rows(model, prompt_ids, vit_list, img_tok, pad, dev, stats=None, debug=False):
+    """Cache one common prompt prefix per image-feature group.
+
+    Multi-category split repeats the same image feature tensor for each
+    category prompt.  Token ids are identical through the image tokens and the
+    fixed prompt prefix, so we prefill that shared prefix once and let each
+    category row forward only its text suffix.
+    """
+    bsz = len(prompt_ids)
+    kv_rows = [None] * bsz
+    cached_lens = [0] * bsz
+    groups = {}
+    for row, vit in enumerate(vit_list):
+        groups.setdefault(id(vit), []).append(row)
+
+    items = []
+    min_prefix_len = max(1, _env_int("LA_FLASH_SHARED_PREFILL_MIN_PREFIX", 64))
+    for rows in groups.values():
+        if len(rows) < 2:
+            continue
+        prefix_len = _common_prefix_len(prompt_ids, rows)
+        if prefix_len < min_prefix_len:
+            continue
+        prefix_ids = prompt_ids[rows[0]][:prefix_len]
+        image_token_count = int((prefix_ids == img_tok).sum().item())
+        if image_token_count != int(vit_list[rows[0]].shape[0]):
+            if debug:
+                print(
+                    "LA Flash shared prefill skip group: "
+                    f"rows={rows} prefix_len={prefix_len} "
+                    f"image_tokens={image_token_count} visual_rows={int(vit_list[rows[0]].shape[0])}",
+                    flush=True,
+                )
+            continue
+        items.append((rows, prefix_ids, vit_list[rows[0]]))
+
+    if not items:
+        return kv_rows, cached_lens
+
+    lengths = [int(ids.numel()) for _rows, ids, _vit in items]
+    pmax = max(lengths)
+    input_ids = torch.full((len(items), pmax), pad, dtype=torch.long, device=dev)
+    amask = torch.zeros((len(items), pmax), dtype=torch.long, device=dev)
+    pos = torch.ones((len(items), pmax), dtype=torch.long, device=dev)
+    for item_idx, (_rows, ids, _vit) in enumerate(items):
+        length = lengths[item_idx]
+        left = pmax - length
+        input_ids[item_idx, left:] = ids.to(dev)
+        amask[item_idx, left:] = 1
+        pos[item_idx, left:] = torch.arange(length, dtype=torch.long, device=dev)
+
+    visual_features = torch.cat([vit for _rows, _ids, vit in items], dim=0)
+    assert int((input_ids == img_tok).sum().item()) == visual_features.shape[0], \
+        "shared-prefix image-token count != supplied visual_features rows"
+
+    if debug:
+        group_sizes = [len(rows) for rows, _ids, _vit in items]
+        print(
+            "LA Flash hybrid shared prompt prefill "
+            f"groups={len(items)} group_sizes={group_sizes} prefix_lens={lengths}",
+            flush=True,
+        )
+
+    forward_mask, fallback_rows = _forward_attention_mask(
+        model, input_ids, amask, 0, mtp_window=False)
+    if debug and fallback_rows:
+        print(
+            "LA Flash hybrid shared prefill safe SDPA fallback "
+            f"query_rows={fallback_rows}",
+            flush=True,
+        )
+    forward_kwargs = dict(
+        input_ids=input_ids,
+        visual_features=visual_features,
+        image_token_index=img_tok,
+        attention_mask=forward_mask,
+        position_ids=pos,
+        past_key_values=None,
+        use_cache=True,
+    )
+    if isinstance(forward_mask, dict):
+        out = language_model_forward(model, **forward_kwargs, return_logits=False)
+    else:
+        out = model.language_model.model(**forward_kwargs)
+
+    real_tokens = sum(lengths)
+    shared_rows = sum(len(rows) for rows, _ids, _vit in items)
+    saved_tokens = sum((len(rows) - 1) * length for (rows, _ids, _vit), length in zip(items, lengths))
+    _record_prefill_stats(
+        stats,
+        rows=len(items),
+        q_len=pmax,
+        real_tokens=real_tokens,
+        shared_groups=len(items),
+        shared_rows=shared_rows,
+        saved_tokens=saved_tokens,
+    )
+
+    for item_idx, (rows, _ids, _vit) in enumerate(items):
+        prefix_len = lengths[item_idx]
+        prefix_kv = _unpack_stock_after_forward(out.past_key_values, item_idx, 0, prefix_len, 0, pmax)
+        for row in rows:
+            kv_rows[row] = prefix_kv
+            cached_lens[row] = prefix_len
+
+    return kv_rows, cached_lens
+
+
+@torch.no_grad()
+def _prefill_prompt_kv_rows(model, prompt_ids, vit_list, img_tok, pad, dev, mode, debug=False, stats=None):
+    """Return per-row prompt KV caches and cached lengths.
+
+    ``mode='none'`` preserves the legacy stock-like first MTP forward where the
+    whole prompt and the 6-token MTP window are forwarded together.  The split
+    prefill modes keep prompt KV clean before the scheduler batches only short
+    suffix/window forwards, which avoids ragged prompt+window masking in the
+    first decode step.
+    """
+    bsz = len(prompt_ids)
+    lengths = [int(p.numel()) for p in prompt_ids]
+    if mode == "none":
+        return [None] * bsz, [0] * bsz
+
+    base = model.language_model.model
+    if debug:
+        print(f"LA Flash hybrid prompt prefill mode={mode} rows={bsz} lengths={lengths}", flush=True)
+
+    if mode == "shared":
+        return _prefill_shared_prefix_kv_rows(
+            model, prompt_ids, vit_list, img_tok, pad, dev, stats=stats, debug=debug)
+
+    if mode == "per_row":
+        kv_rows = []
+        for b, ids in enumerate(prompt_ids):
+            ids = ids.to(dev).unsqueeze(0)
+            pos = torch.arange(ids.shape[1], dtype=torch.long, device=dev).unsqueeze(0)
+            out = base(
+                input_ids=ids,
+                visual_features=vit_list[b],
+                image_token_index=img_tok,
+                attention_mask=None,
+                position_ids=pos,
+                past_key_values=None,
+                use_cache=True,
+            )
+            kv_rows.append(out.past_key_values)
+            _record_prefill_stats(stats, rows=1, q_len=ids.shape[1], real_tokens=ids.shape[1])
+        return kv_rows, lengths
+
+    pmax = max(lengths)
+    input_ids = torch.full((bsz, pmax), pad, dtype=torch.long, device=dev)
+    amask = torch.zeros((bsz, pmax), dtype=torch.long, device=dev)
+    pos = torch.ones((bsz, pmax), dtype=torch.long, device=dev)
+    for b, ids in enumerate(prompt_ids):
+        left = pmax - lengths[b]
+        input_ids[b, left:] = ids.to(dev)
+        amask[b, left:] = 1
+        pos[b, left:] = torch.arange(lengths[b], dtype=torch.long, device=dev)
+
+    visual_features = torch.cat(vit_list, dim=0)
+    assert int((input_ids == img_tok).sum().item()) == visual_features.shape[0], \
+        "image-token count != supplied visual_features rows"
+    forward_mask, fallback_rows = _forward_attention_mask(
+        model, input_ids, amask, 0, mtp_window=False)
+    if debug and fallback_rows:
+        print(
+            "LA Flash hybrid batch prefill safe SDPA fallback "
+            f"query_rows={fallback_rows}",
+            flush=True,
+        )
+    forward_kwargs = dict(
+        input_ids=input_ids,
+        visual_features=visual_features,
+        image_token_index=img_tok,
+        attention_mask=forward_mask,
+        position_ids=pos,
+        past_key_values=None,
+        use_cache=True,
+    )
+    if isinstance(forward_mask, dict):
+        out = language_model_forward(model, **forward_kwargs, return_logits=False)
+    else:
+        out = base(**forward_kwargs)
+    _record_prefill_stats(stats, rows=bsz, q_len=pmax, real_tokens=sum(lengths))
+    kv_rows = [
+        _unpack_stock_after_forward(out.past_key_values, b, 0, lengths[b], 0, pmax)
+        for b in range(bsz)
+    ]
+    return kv_rows, lengths
+
+
+@torch.no_grad()
+def generate_batch_hybrid(pairs, temperature=README_TEMPERATURE, top_p=README_TOP_P, top_k=None,
+                          repetition_penalty=README_REPETITION_PENALTY,
+                          max_new_tokens=README_MAX_NEW_TOKENS, temps=None,
+                          debug=None, scheduler=None, group_size=None,
+                          vision_features=None, _stats=None):
+    """Batched stock-style LocateAnything-3B hybrid generation.
+
+    This mirrors ``model.generate(..., generation_mode='hybrid')``: each row
+    owns a full ``generated`` token stream plus a KV cache truncated to real
+    generated tokens before sampling.  MTP forwards
+    ``generated[cached_len:] + duplicate-last + mask*5``; AR forwards
+    ``generated[cached_len:]``.
+    """
+    tok, _, model = load()
+    san, hpat = _helpers()
+    tids = model.token_ids
+    img_tok = model.config.image_token_index
+    mask_tok = tids["default_mask_token_id"]
+    im_end = tids["im_end_token_id"]
+    pad = tok.pad_token_id if tok.pad_token_id is not None else im_end
+    dev = DEV
+
+    if not pairs:
+        return []
+    if temps is not None and len(temps) != len(pairs):
+        raise ValueError("temps must have the same length as pairs")
+    if vision_features is not None and len(vision_features) != len(pairs):
+        raise ValueError("vision_features must have the same length as pairs")
+    debug = _debug_enabled(debug)
+    scheduler = _hybrid_scheduler(scheduler)
+    group_size = _hybrid_group_size(group_size)
+    requested_attn = getattr(model, "_la_flash_requested_attn", ATTN_MODE)
+    use_magi = requested_attn == "magi"
+    prefill_mode = _hybrid_prefill_mode()
+    hold_max_steps = max(0, _env_int("LA_FLASH_HYBRID_HOLD_MAX_STEPS", 5))
+    adaptive_hold_mtp_max = max(0, _env_int("LA_FLASH_HYBRID_ADAPTIVE_HOLD_MTP_MAX", 3))
+    top_level_stats = _stats is None
+    if top_level_stats:
+        _stats = _new_hybrid_stats(
+            len(pairs), scheduler, group_size, hold_max_steps, adaptive_hold_mtp_max)
+        if os.environ.get("LA_FLASH_PLAN_STATS", "0") == "1":
+            model._la_flash_sparse_plan_stats = None
+    if group_size and len(pairs) > group_size:
+        outs = []
+        if debug:
+            print(
+                f"LA Flash hybrid grouped scheduling: total_rows={len(pairs)} "
+                f"group_size={group_size} scheduler={scheduler} hold_max_steps={hold_max_steps} "
+                f"adaptive_hold_mtp_max={adaptive_hold_mtp_max}",
+                flush=True,
+            )
+        for start in range(0, len(pairs), group_size):
+            end = min(start + group_size, len(pairs))
+            chunk_temps = temps[start:end] if temps is not None else None
+            chunk_vision_features = (
+                vision_features[start:end] if vision_features is not None else None
+            )
+            if debug:
+                print(f"LA Flash hybrid group rows=[{start}:{end}]", flush=True)
+            outs.extend(generate_batch_hybrid(
+                pairs[start:end],
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+                repetition_penalty=repetition_penalty,
+                max_new_tokens=max_new_tokens,
+                temps=chunk_temps,
+                debug=debug,
+                scheduler=scheduler,
+                group_size=0,
+                vision_features=chunk_vision_features,
+                _stats=_stats,
+            ))
+        if top_level_stats:
+            _set_last_hybrid_stats(_stats)
+        return outs
+
+    use_cached_tokenize = (
+        vision_features is not None
+        and os.environ.get("LA_FLASH_CACHE_TOKENIZE", "1") != "0"
+    )
+    if use_cached_tokenize:
+        try:
+            prompt_ids = [
+                _tokenize_cached_image(q, int(v.shape[0]), im=im)
+                for (im, q), v in zip(pairs, vision_features)
+            ]
+        except Exception as exc:
+            if os.environ.get("LA_FLASH_CACHE_TOKENIZE_STRICT", "0") == "1":
+                raise
+            if debug:
+                print(f"LA Flash cached tokenize fallback: {exc}", flush=True)
+            prompt_ids = [_tokenize(im, q) for im, q in pairs]
+    else:
+        prompt_ids = [_tokenize(im, q) for im, q in pairs]
+    vit_list = (
+        list(vision_features)
+        if vision_features is not None
+        else _encode_images([im for im, _ in pairs])
+    )
+    lengths = [int(p.numel()) for p in prompt_ids]
+    bsz = len(pairs)
+    _record_group_stats(_stats, bsz)
+
+    _set_llm_mode(model, requested_attn)
+
+    modes = ["mtp"] * bsz
+    finished = [False] * bsz
+    gen_ids = [[] for _ in range(bsz)]
+    full_ids = [list(ids.detach().cpu().tolist()) for ids in prompt_ids]
+    kv_rows, cached_lens = _prefill_prompt_kv_rows(
+        model, prompt_ids, vit_list, img_tok, pad, dev, prefill_mode, debug=debug, stats=_stats)
+    total_limits = [lengths[b] + max_new_tokens for b in range(bsz)]
+
+    row_temps = [float(temperature or 0.0)] * bsz if temps is None else [float(t or 0.0) for t in temps]
+
+    def run_ar(ar_rows, step_idx):
+        row_groups = _split_rows_by_kv_budget(ar_rows, kv_rows)
+        _record_kv_bucket_stats(_stats, row_groups, kv_rows)
+        for row_group in row_groups:
+            _step_stock_ar_rows(
+                model, san, tids, prompt_ids, kv_rows, row_group,
+                cached_lens, full_ids, gen_ids, modes, finished, total_limits,
+                pad, img_tok, row_temps, temperature, top_p, top_k,
+                repetition_penalty, dev, tok, debug, step_idx, use_magi, _stats,
+            )
+
+    def run_mtp(mtp_rows, step_idx):
+        if any(cached_lens[r] == 0 for r in mtp_rows) and any(cached_lens[r] > 0 for r in mtp_rows):
+            first_rows = [r for r in mtp_rows if cached_lens[r] == 0]
+            cached_rows = [r for r in mtp_rows if cached_lens[r] > 0]
+            if first_rows:
+                run_mtp(first_rows, step_idx)
+            if cached_rows:
+                run_mtp(cached_rows, step_idx)
+            return
+        row_groups = _split_rows_by_kv_budget(mtp_rows, kv_rows)
+        _record_kv_bucket_stats(_stats, row_groups, kv_rows)
+        if len(row_groups) > 1:
+            for row_group in row_groups:
+                run_mtp(row_group, step_idx)
+            return
+        _step_stock_mtp_rows(
+            model, san, hpat, tids, prompt_ids, kv_rows, mtp_rows,
+            cached_lens, full_ids, gen_ids, modes, finished, total_limits,
+            vit_list, pad, mask_tok, img_tok, row_temps, top_p, top_k,
+            repetition_penalty, dev, tok, debug, step_idx, use_magi, _stats,
+        )
+
+    def live_rows(mode):
+        return [b for b in range(bsz) if not finished[b] and modes[b] == mode]
+
+    step = 0
+    hold_steps = 0
+    while not all(finished) and step <= max_new_tokens:
+        step += 1
+        if _stats is not None:
+            _stats["decode_loops"] += 1
+        if scheduler == "hold_ar" and hold_max_steps > 0:
+            ar_rows = live_rows("ar")
+            mtp_rows = live_rows("mtp")
+            if ar_rows and mtp_rows and _stats is not None:
+                _stats["mixed_mode_cycles"] += 1
+            if ar_rows and (hold_steps < hold_max_steps or not mtp_rows):
+                if mtp_rows and _stats is not None:
+                    _stats["hold_ar_steps"] += 1
+                    _stats["hold_ar_held_mtp_rows"] += len(mtp_rows)
+                run_ar(ar_rows, step)
+                hold_steps += 1
+                continue
+            if mtp_rows:
+                if ar_rows and _stats is not None:
+                    _stats["hold_ar_limit_mtp_forwards"] += 1
+                run_mtp(mtp_rows, step)
+                hold_steps = 0
+            continue
+
+        if scheduler in {"ar_first", "pipeline", "adaptive"}:
+            ar_rows_at_loop_start = live_rows("ar")
+            mtp_rows_at_loop_start = live_rows("mtp")
+            mixed = bool(ar_rows_at_loop_start and mtp_rows_at_loop_start)
+            if mixed and _stats is not None:
+                _stats["mixed_mode_cycles"] += 1
+
+            if scheduler == "adaptive" and mixed and hold_max_steps > 0:
+                should_hold = len(mtp_rows_at_loop_start) <= adaptive_hold_mtp_max
+                if should_hold and hold_steps < hold_max_steps:
+                    if _stats is not None:
+                        _stats["adaptive_hold_cycles"] += 1
+                        _stats["hold_ar_steps"] += 1
+                        _stats["hold_ar_held_mtp_rows"] += len(mtp_rows_at_loop_start)
+                    run_ar(ar_rows_at_loop_start, step)
+                    hold_steps += 1
+                    continue
+
+            if ar_rows_at_loop_start:
+                if mixed and _stats is not None:
+                    if scheduler == "adaptive":
+                        _stats["adaptive_ar_first_cycles"] += 1
+                    else:
+                        _stats["ar_first_cycles"] += 1
+                run_ar(ar_rows_at_loop_start, step)
+
+            mtp_rows = live_rows("mtp")
+            if mtp_rows:
+                run_mtp(mtp_rows, step)
+                hold_steps = 0
+
+            if scheduler == "pipeline" and mtp_rows:
+                old_ar = set(ar_rows_at_loop_start)
+                new_ar_rows = [b for b in live_rows("ar") if b not in old_ar]
+                if new_ar_rows:
+                    if _stats is not None:
+                        _stats["pipeline_ar_after_mtp_cycles"] += 1
+                    run_ar(new_ar_rows, step)
+            continue
+
+        mtp_rows = live_rows("mtp")
+        ar_rows_at_loop_start = live_rows("ar")
+        if mtp_rows and ar_rows_at_loop_start and _stats is not None:
+            _stats["mixed_mode_cycles"] += 1
+        if mtp_rows:
+            run_mtp(mtp_rows, step)
+
+        ar_rows = [b for b in range(bsz) if not finished[b] and modes[b] == "ar"]
+        if mtp_rows and ar_rows and _stats is not None:
+            _stats["eager_mtp_then_ar_cycles"] += 1
+        if ar_rows:
+            run_ar(ar_rows, step)
+
+    outs = [
+        tok.decode(torch.tensor(gen_ids[b], dtype=torch.long, device=dev),
+                   skip_special_tokens=False) if gen_ids[b] else ""
+        for b in range(bsz)
+    ]
+    if top_level_stats:
+        if os.environ.get("LA_FLASH_PLAN_STATS", "0") == "1":
+            _stats["sparse_plan_stats"] = copy.deepcopy(
+                getattr(model, "_la_flash_sparse_plan_stats", None) or {}
+            )
+        _set_last_hybrid_stats(_stats)
+    return outs
+
+
+@torch.no_grad()
+def _step_stock_mtp_rows(model, san, hpat, tids, prompt_ids, kv_rows, rows,
+                         cached_lens, full_ids, gen_ids, modes, finished, total_limits,
+                         vit_list, pad, mask_tok, img_tok, row_temps, top_p, top_k,
+                         repetition_penalty, dev, tok, debug, step_idx, use_magi, stats=None):
+    kv, kvalid, old_lens, kmax = _pack_stock_kv_rows(kv_rows, rows, dev)
+    uncached_lens = [len(full_ids[r]) - cached_lens[r] for r in rows]
+    umax = max(uncached_lens)
+    seq_len = umax + N_FUTURE
+    _record_forward_stats(stats, "mtp", rows, seq_len, uncached_lens)
+
+    suf_ids = torch.full((len(rows), seq_len), pad, dtype=torch.long, device=dev)
+    suf_pos = torch.ones((len(rows), seq_len), dtype=torch.long, device=dev)
+    q_valid = torch.zeros((len(rows), seq_len), dtype=torch.long, device=dev)
+
+    for i, r in enumerate(rows):
+        uncached = full_ids[r][cached_lens[r] :]
+        left = umax - len(uncached)
+        if uncached:
+            suf_ids[i, left : left + len(uncached)] = torch.tensor(uncached, dtype=torch.long, device=dev)
+            suf_pos[i, left : left + len(uncached)] = torch.arange(
+                cached_lens[r], len(full_ids[r]), dtype=torch.long, device=dev)
+            q_valid[i, left : left + len(uncached)] = 1
+
+        rep = full_ids[r][-1]
+        cur_len = len(full_ids[r])
+        suf_ids[i, umax] = rep
+        suf_pos[i, umax] = cur_len - 1
+        q_valid[i, umax] = 1
+        for j in range(1, N_FUTURE):
+            suf_ids[i, umax + j] = mask_tok
+            suf_pos[i, umax + j] = cur_len + (j - 1)
+            q_valid[i, umax + j] = 1
+
+    full_mask = torch.cat([kvalid, q_valid], dim=1)
+
+    if debug:
+        forward_mask, fallback_rows = _forward_attention_mask(
+            model, suf_ids, full_mask, kmax, mtp_window=True, range_plan=True)
+        _print_debug_forward(
+            f"MTP step={step_idx}",
+            model,
+            tok,
+            suf_ids,
+            full_mask,
+            suf_pos,
+            past_len=kmax,
+            mtp_window=True,
+            extra={
+                "global_rows": rows,
+                "old_kv_lens": old_lens,
+                "cached_lens": [cached_lens[r] for r in rows],
+                "full_lens": [len(full_ids[r]) for r in rows],
+                "uncached_lens": uncached_lens,
+                "forward_attention_mask": _mask_desc(forward_mask),
+                "safe_sdpa_fallback_query_rows": fallback_rows,
+            },
+            attention_impl="magi" if use_magi else ATTN_MODE,
+        )
+    else:
+        forward_mask, _ = _forward_attention_mask(
+            model, suf_ids, full_mask, kmax, mtp_window=True, range_plan=True)
+    first_rows = [r for r in rows if cached_lens[r] == 0]
+    visual_features = None
+    if first_rows:
+        if first_rows != rows:
+            raise RuntimeError("mixed first/non-first MTP rows are not supported")
+        visual_features = torch.cat([vit_list[r] for r in rows], dim=0)
+        assert int((suf_ids == img_tok).sum().item()) == visual_features.shape[0], \
+            "image-token count != supplied visual_features rows"
+    out = language_model_forward(
+        model, input_ids=suf_ids, attention_mask=forward_mask,
+        position_ids=suf_pos, past_key_values=kv, use_cache=True,
+        visual_features=visual_features,
+        image_token_index=img_tok if visual_features is not None else None,
+        logits_slice=slice(-N_FUTURE, None))
+
+    for i, r in enumerate(rows):
+        kv_rows[r] = _unpack_stock_after_forward(
+            out.past_key_values, i, old_lens[i], uncached_lens[i], kmax, umax)
+        cached_lens[r] = len(full_ids[r])
+
+    wlogits = out.logits[:, -N_FUTURE:, :]
+    local_prompts = [prompt_ids[r] for r in rows]
+    local_gen = [gen_ids[r] for r in rows]
+    gen_pad = _pad_generated(local_prompts, local_gen, img_tok, dev)
+    per_row_temp = torch.tensor([row_temps[r] for r in rows], dtype=torch.float32, device=dev)
+
+    if BATCH_SAN:
+        x0_all, boxes_all = sample_tokens_batched(
+            wlogits, gen_pad, tids, per_row_temp,
+            repetition_penalty=repetition_penalty, top_p=top_p, top_k=top_k,
+            keep_k_avg=4, generation_mode="hybrid")
+
+    for i, r in enumerate(rows):
+        if finished[r]:
+            continue
+        if BATCH_SAN:
+            x0b, boxb = x0_all[i], boxes_all[i]
+        else:
+            gk = _mk_generate_kwargs(row_temps[r], top_p, top_k, repetition_penalty)
+            _, _, x0, box_avg = san(wlogits[i : i + 1], gen_pad[i : i + 1], tids, keep_k=5, **gk)
+            x0b, boxb = x0[0], box_avg[0]
+        nt = x0b if bool((boxb == 0).all()) else boxb
+        op = hpat(nt, tids, "hybrid")
+
+        toks = [int(t) for t in op["tokens"]]
+        for t in toks:
+            gen_ids[r].append(t)
+            full_ids[r].append(t)
+
+        if op["type"] == "im_end":
+            finished[r] = True
+        elif op["type"] == "error_box":
+            modes[r] = "ar"
+        if len(full_ids[r]) >= total_limits[r]:
+            finished[r] = True
+
+
+@torch.no_grad()
+def _step_stock_ar_rows(model, san, tids, prompt_ids, kv_rows, rows,
+                        cached_lens, full_ids, gen_ids, modes, finished, total_limits,
+                        pad, img_tok, row_temps, temperature, top_p, top_k,
+                        repetition_penalty, dev, tok, debug, step_idx, use_magi, stats=None):
+    kv, kvalid, old_lens, kmax = _pack_stock_kv_rows(kv_rows, rows, dev)
+    uncached_lens = [len(full_ids[r]) - cached_lens[r] for r in rows]
+    if any(n <= 0 for n in uncached_lens):
+        raise RuntimeError(f"AR rows have no uncached tokens: {rows}")
+    umax = max(uncached_lens)
+    _record_forward_stats(stats, "ar", rows, umax, uncached_lens)
+
+    suf_ids = torch.full((len(rows), umax), pad, dtype=torch.long, device=dev)
+    suf_pos = torch.ones((len(rows), umax), dtype=torch.long, device=dev)
+    q_valid = torch.zeros((len(rows), umax), dtype=torch.long, device=dev)
+
+    for i, r in enumerate(rows):
+        uncached = full_ids[r][cached_lens[r] :]
+        left = umax - len(uncached)
+        suf_ids[i, left:] = torch.tensor(uncached, dtype=torch.long, device=dev)
+        suf_pos[i, left:] = torch.arange(cached_lens[r], len(full_ids[r]), dtype=torch.long, device=dev)
+        q_valid[i, left:] = 1
+
+    full_mask = torch.cat([kvalid, q_valid], dim=1)
+    if debug:
+        forward_mask, fallback_rows = _forward_attention_mask(
+            model, suf_ids, full_mask, kmax, mtp_window=False, range_plan=True)
+        _print_debug_forward(
+            f"AR step={step_idx}",
+            model,
+            tok,
+            suf_ids,
+            full_mask,
+            suf_pos,
+            past_len=kmax,
+            mtp_window=False,
+            extra={
+                "global_rows": rows,
+                "old_kv_lens": old_lens,
+                "cached_lens": [cached_lens[r] for r in rows],
+                "full_lens": [len(full_ids[r]) for r in rows],
+                "uncached_lens": uncached_lens,
+                "forward_attention_mask": _mask_desc(forward_mask),
+                "safe_sdpa_fallback_query_rows": fallback_rows,
+            },
+            attention_impl="magi" if use_magi else ATTN_MODE,
+        )
+    else:
+        forward_mask, _ = _forward_attention_mask(
+            model, suf_ids, full_mask, kmax, mtp_window=False, range_plan=True)
+
+    out = language_model_forward(
+        model, input_ids=suf_ids, attention_mask=forward_mask,
+        position_ids=suf_pos, past_key_values=kv, use_cache=True,
+        logits_slice=slice(-1, None))
+
+    for i, r in enumerate(rows):
+        kv_rows[r] = _unpack_stock_after_forward(
+            out.past_key_values, i, old_lens[i], uncached_lens[i], kmax, umax)
+        cached_lens[r] = len(full_ids[r])
+
+    if AR_BATCH_SAN:
+        local_prompts = [prompt_ids[r] for r in rows]
+        local_gen = [gen_ids[r] for r in rows]
+        gen_pad = _pad_generated(local_prompts, local_gen, img_tok, dev)
+        per_row_temp = torch.tensor([row_temps[r] for r in rows], dtype=torch.float32, device=dev)
+        x0_all = sample_next_tokens_batched(
+            out.logits[:, -1:, :],
+            gen_pad,
+            per_row_temp,
+            repetition_penalty=repetition_penalty,
+            top_p=top_p,
+            top_k=top_k,
+        )
+
+    for i, r in enumerate(rows):
+        if AR_BATCH_SAN:
+            token_val = int(x0_all[i, 0].item())
+        else:
+            logits = out.logits[i : i + 1, -1:, :]
+            gen_pad = _pad_generated([prompt_ids[r]], [gen_ids[r]], img_tok, dev)
+            gk = _mk_generate_kwargs(temperature, top_p, top_k, repetition_penalty, row_temp=row_temps[r])
+            _, _, x0, _ = san(logits, gen_pad, tids, **gk)
+            token_val = int(x0[0, 0].item())
+        out_type = _classify_ar_token(token_val, tids)
+
+        gen_ids[r].append(token_val)
+        full_ids[r].append(token_val)
+
+        if out_type == "im_end":
+            finished[r] = True
+        elif out_type == "box_end_ar":
+            modes[r] = "mtp"
+
+        if len(full_ids[r]) >= total_limits[r]:
+            finished[r] = True
+
+
+def generate_batch_grouped_hybrid(groups, temperature=README_TEMPERATURE, top_p=README_TOP_P,
+                                  top_k=None, repetition_penalty=README_REPETITION_PENALTY,
+                                  max_new_tokens=README_MAX_NEW_TOKENS, temps=None,
+                                  debug=None, scheduler=None, group_size=None,
+                                  vision_features=None):
+    """Hybrid grouped API shape.
+
+    This preserves grouped return shape, but intentionally uses the generic
+    hybrid decoder rather than the fast engine's shared-prefix optimization.
+    """
+    flat = []
+    flat_vision_features = [] if vision_features is not None else None
+    counts = []
+    for group_idx, (im, queries) in enumerate(groups):
+        counts.append(len(queries))
+        flat.extend((im, q) for q in queries)
+        if flat_vision_features is not None:
+            flat_vision_features.extend([vision_features[group_idx]] * len(queries))
+
+    outs = generate_batch_hybrid(
+        flat, temperature=temperature, top_p=top_p, top_k=top_k,
+        repetition_penalty=repetition_penalty, max_new_tokens=max_new_tokens,
+        temps=temps, debug=debug, scheduler=scheduler, group_size=group_size,
+        vision_features=flat_vision_features)
+    res, offset = [], 0
+    for n in counts:
+        res.append(outs[offset : offset + n])
+        offset += n
+    return res
+
+
+__all__ = ["generate_batch_hybrid", "generate_batch_grouped_hybrid", "get_last_hybrid_stats"]

batch_utils/hybrid_runtime.py

@@ -0,0 +1,1842 @@
+"""Internal runtime support for the LocateAnything-3B hybrid batch decoder.
+
+This file keeps only the model-loading, tokenization, image-encoding, stock
+processor, and sample-token helpers that ``engine_hybrid.py`` needs.
+
+Important env knobs:
+  LA_FLASH_MODEL          HF repo id / local path of the model (default nvidia/LocateAnything-3B)
+  HF_HUB_OFFLINE=1    read the local HF cache only (no network); unset -> download on first use
+  LA_FLASH_ATTN           sdpa, eager, magi, or la_flash; la_flash uses FlashAttention sparse ranges
+  LA_FLASH_STRICT_ATTN    1 -> fail if the requested backend is unavailable;
+                      default 0 falls back to sdpa
+  LA_FLASH_VISION_ATTN    auto, flash_attention_2, sdpa, or eager (default auto)
+  LA_FLASH_HYBRID_PREFILL shared, none, per_row, or batch prompt KV prefill (default shared)
+  MTP_BATCH_VISION    0 -> per-image vision encode (default 1: batched when flash is present)
+  LA_FLASH_VISION_ENCODE_BATCH_SIZE
+                      max images per MoonViT encode micro-batch (default 8; <=0 disables limit)
+  MTP_BATCH_SAN       0 -> per-row logits/sample pipeline (default 1: batched over [B,6,V])
+  AR_BATCH_SAN        0 -> per-row AR sample pipeline (default 1: batched over [B,1,V])
+"""
+import inspect
+import os, warnings, importlib, torch
+from types import SimpleNamespace
+import numpy as np
+from transformers import AutoModel, AutoTokenizer, AutoProcessor
+
+
+# By default let transformers fetch the model on first use; set HF_HUB_OFFLINE=1 yourself
+# to read the local HF cache only (e.g. air-gapped / already-downloaded runs).
+MODEL = os.environ.get("LA_FLASH_MODEL", "nvidia/LocateAnything-3B")
+
+
+LLM_ATTN_MODES = ("sdpa", "eager", "magi", "la_flash")
+VISION_ATTN_MODES = ("auto", "flash_attention_2", "sdpa", "eager")
+
+
+def _normalize_attn_mode(value):
+    mode = (value or "sdpa").strip().lower().replace("-", "_")
+    aliases = {
+        "": "sdpa",
+        "manual": "eager",
+        "torch": "eager",
+        "torch_eager": "eager",
+        "torch_sdpa": "sdpa",
+        "scaled_dot_product_attention": "sdpa",
+        "flash": "la_flash",
+        "la_flash": "la_flash",
+        "kernel": "la_flash",
+        "cuda": "la_flash",
+        "range": "la_flash",
+        "range_attention": "la_flash",
+        "flex_flash": "magi",
+        "flex_flash_attention": "magi",
+        "flex_flash_attn": "magi",
+    }
+    mode = aliases.get(mode, mode)
+    if mode not in LLM_ATTN_MODES:
+        raise ValueError(
+            f"LA_FLASH_ATTN must be one of {', '.join(LLM_ATTN_MODES)}; got {value!r}"
+        )
+    return mode
+
+
+def _normalize_vision_attn_mode(value):
+    mode = (value or "auto").strip().lower().replace("-", "_")
+    aliases = {
+        "": "auto",
+        "flash": "flash_attention_2",
+        "flash_attention2": "flash_attention_2",
+        "fa2": "flash_attention_2",
+        "manual": "eager",
+    }
+    mode = aliases.get(mode, mode)
+    if mode not in VISION_ATTN_MODES:
+        raise ValueError(
+            f"LA_FLASH_VISION_ATTN must be one of {', '.join(VISION_ATTN_MODES)}; got {value!r}"
+        )
+    return mode
+
+
+ATTN_MODE = _normalize_attn_mode(os.environ.get("LA_FLASH_ATTN", "sdpa"))
+REMOTE_ATTN_MODE = "sdpa" if ATTN_MODE in {"la_flash", "magi"} else ATTN_MODE
+VISION_ATTN_MODE = _normalize_vision_attn_mode(os.environ.get("LA_FLASH_VISION_ATTN", "auto"))
+MAX_DIM = 1024
+DEV, DT = "cuda", torch.bfloat16
+N_FUTURE = 6                                            # = config.block_size (MTP window)
+_PROMPT = "Locate all the instances that matches the following description: "
+
+
+def _env_flag(name, default=False):
+    val = os.environ.get(name)
+    if val is None:
+        return default
+    return val.strip().lower() not in {"0", "false", "no", "off"}
+
+
+def _env_int(name):
+    val = os.environ.get(name)
+    if val is None or val.strip() == "":
+        return None
+    return int(val)
+
+
+def _strict_attn():
+    return _env_flag("LA_FLASH_STRICT_ATTN", False)
+
+
+def _fallback_to_sdpa(model, requested, reason):
+    if requested == "sdpa":
+        raise RuntimeError(f"LA_FLASH_ATTN=sdpa failed: {reason}") from reason
+    message = f"LA_FLASH_ATTN={requested} is unavailable; falling back to sdpa. Reason: {reason}"
+    if _strict_attn():
+        raise RuntimeError(message) from reason
+    warnings.warn(message)
+    _set_llm_mode(model, "sdpa")
+    model._la_flash_requested_attn_original = requested
+    model._la_flash_attn_fallback_reason = str(reason)
+    return "sdpa"
+
+
+# Optional compile for the shared Qwen2 core. This is off by default because the
+# hybrid scheduler already varies query/cache shapes and first-call compile cost is high.
+MTP_COMPILE = os.environ.get("MTP_COMPILE", "0") == "1"
+
+# Batch the MoonViT vision encode across a micro-batch's images: pack N images into ONE
+# extract_feature. With flash present, MoonViT's varlen cu_seqlens path is block-diagonal per
+# image and equivalent to per-image encode.
+# Without flash, sdpa builds a dense [1,S,S] mask -> O(S^2) N^2 -> per-image fallback (auto, see
+# _vision_is_flash). Default ON; set MTP_BATCH_VISION=0 to force per-image.
+BATCH_VISION = os.environ.get("MTP_BATCH_VISION", "1") == "1"
+_vision_encode_batch_size = _env_int("LA_FLASH_VISION_ENCODE_BATCH_SIZE")
+VISION_ENCODE_BATCH_SIZE = 8 if _vision_encode_batch_size is None else max(0, _vision_encode_batch_size)
+
+# Batch the per-row box-decode (sample_tokens): run the row-independent logits pipeline
+# (rep-penalty / per-row temperature / top_p / top_k / softmax / sample) ONCE over the whole
+# [B,6,V] step instead of B times on [1,6,V]; only the variable-length box assembly stays per-row.
+# Greedy is BIT-IDENTICAL to the per-row san (argmax, no RNG). Default ON; MTP_BATCH_SAN=0 -> per-row.
+BATCH_SAN = os.environ.get("MTP_BATCH_SAN", "1") == "1"
+
+# Batch the AR repair sampler over [B,1,V]. This shares the exact filtering
+# helpers with MTP batching but skips box/ref decoding, so it only replaces the
+# repeated stock one-token sample calls.  Sampling itself stays row-ordered by
+# default to preserve the stock RNG consumption pattern for AR repair.
+AR_BATCH_SAN = os.environ.get("AR_BATCH_SAN", "1") == "1"
+
+_tok = _proc = _model = None
+
+def _magi_diag():
+    lines = []
+    try:
+        import magi_attention
+        lines.append(f"magi_attention: OK file={getattr(magi_attention, '__file__', None)}")
+        lines.append(f"magi_attention.__version__={getattr(magi_attention, '__version__', '<missing>')}")
+    except Exception as e:
+        lines.append(f"magi_attention: FAIL {type(e).__name__}: {e}")
+        return "\n".join(lines)
+    try:
+        from magi_attention.functional.flex_flash_attn import flex_flash_attn_func
+        lines.append(f"magi_attention.functional.flex_flash_attn: OK func={flex_flash_attn_func}")
+    except Exception as e:
+        lines.append(f"magi_attention.functional.flex_flash_attn: FAIL {type(e).__name__}: {e}")
+    return "\n".join(lines)
+
+def _remote_magi_diag(model=None):
+    lines = []
+    try:
+        if model is not None:
+            mod = importlib.import_module(type(model.language_model.model).__module__)
+        else:
+            # Best effort: if the dynamic module is not imported yet this may fail;
+            # the post-load diagnostic below will still work.
+            mod = importlib.import_module("transformers_modules.LocateAnything-3B.modeling_qwen2")
+        lines.append(f"remote_qwen2_module={getattr(mod, '__file__', None)}")
+        lines.append(f"remote_qwen2._MAGI_AVAILABLE={getattr(mod, '_MAGI_AVAILABLE', '<missing>')!r}")
+        lines.append(f"remote_qwen2.flex_flash_attn_func={getattr(mod, 'flex_flash_attn_func', '<missing>')}")
+    except Exception as e:
+        lines.append(f"remote_qwen2: diagnostic failed {type(e).__name__}: {e}")
+    return "\n".join(lines)
+
+def _attn_class_diag(model):
+    try:
+        llm = model.language_model.model
+        classes = [type(layer.self_attn).__name__ for layer in llm.layers[:4]]
+        return (
+            f"llm._attn_implementation={getattr(llm, '_attn_implementation', None)!r}\n"
+            f"config._attn_implementation={getattr(llm.config, '_attn_implementation', None)!r}\n"
+            f"first_attn_classes={classes}"
+        )
+    except Exception as e:
+        return f"attention class diagnostic failed {type(e).__name__}: {e}"
+
+
+def _set_vision_attention_mode(model):
+    """Match HF's MoonViT policy: prefer flash_attention_2, then sdpa, then eager."""
+    vm = getattr(model, "vision_model", None)
+    if vm is None:
+        return None
+    mod = importlib.import_module(type(vm).__module__)
+    funcs = getattr(mod, "VL_VISION_ATTENTION_FUNCTIONS", {})
+    has_flash = getattr(mod, "flash_attn_varlen_func", None) is not None
+    requested = VISION_ATTN_MODE
+
+    if requested == "auto":
+        candidates = ("flash_attention_2", "sdpa", "eager")
+    else:
+        candidates = (requested, "flash_attention_2", "sdpa", "eager")
+
+    chosen = None
+    for candidate in candidates:
+        if candidate == "flash_attention_2" and not has_flash:
+            continue
+        if candidate in funcs:
+            chosen = candidate
+            break
+    if chosen is None:
+        raise RuntimeError("MoonViT has no supported attention implementation.")
+
+    if requested == "flash_attention_2" and chosen != "flash_attention_2":
+        warnings.warn("LA_FLASH_VISION_ATTN=flash_attention_2 requested but flash-attn is unavailable; "
+                      f"using {chosen}.")
+    elif requested not in {"auto", chosen}:
+        warnings.warn(f"LA_FLASH_VISION_ATTN={requested} is unavailable; using {chosen}.")
+
+    if hasattr(model.config, "vision_config"):
+        model.config.vision_config._attn_implementation = chosen
+    try:
+        vm.config._attn_implementation = chosen
+    except Exception:
+        pass
+    try:
+        for block in vm.encoder.blocks:
+            block.attn_implementation = chosen
+    except Exception as exc:
+        raise RuntimeError("Failed to configure MoonViT attention implementation.") from exc
+    model._la_flash_vision_attn = chosen
+    return chosen
+
+
+def load():
+    """Lazy model load with HF remote-code semantics plus release backends.
+
+    The text decoder is pinned to one of sdpa/eager/magi/la_flash. MoonViT is
+    configured independently and follows the HF policy: flash_attention_2 when
+    flash-attn is importable, otherwise sdpa, otherwise eager.
+    """
+    global _tok, _proc, _model
+    if _model is None:
+        _tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
+        _proc = AutoProcessor.from_pretrained(MODEL, trust_remote_code=True)
+        attn_impl = REMOTE_ATTN_MODE
+        if ATTN_MODE == "magi" and os.environ.get("LA_FLASH_DEBUG", "0") != "0":
+            print("LA Flash magi pre-load diagnostic:", flush=True)
+            print(_magi_diag(), flush=True)
+        _model = AutoModel.from_pretrained(MODEL, torch_dtype=DT, trust_remote_code=True,
+                                           attn_implementation=attn_impl).to(DEV).eval()
+        _set_vision_attention_mode(_model)
+        actual_attn = getattr(_model.language_model.model, "_attn_implementation", None)
+        if ATTN_MODE == "magi" and os.environ.get("LA_FLASH_DEBUG", "0") != "0":
+            print("LA Flash magi post-load diagnostic:", flush=True)
+            print(_remote_magi_diag(_model), flush=True)
+            print(_attn_class_diag(_model), flush=True)
+        if ATTN_MODE == "magi":
+            try:
+                qwen2_mod = importlib.import_module(type(_model.language_model.model).__module__)
+                if not getattr(qwen2_mod, "_MAGI_AVAILABLE", False):
+                    raise RuntimeError(
+                        "remote module reports _MAGI_AVAILABLE=False.\n"
+                        f"{_remote_magi_diag(_model)}\n{_magi_diag()}"
+                    )
+                first_attn = type(_model.language_model.model.layers[0].self_attn).__name__
+                if actual_attn != "sdpa" or first_attn != "_BatchedMagiAttention":
+                    _set_llm_mode(_model, "magi")
+                    actual_attn = getattr(_model.language_model.model, "_attn_implementation", None)
+                    first_attn = type(_model.language_model.model.layers[0].self_attn).__name__
+                    if os.environ.get("LA_FLASH_DEBUG", "0") != "0":
+                        print("LA Flash magi post-swap diagnostic:", flush=True)
+                        print(_attn_class_diag(_model), flush=True)
+                if actual_attn != "sdpa" or first_attn != "_BatchedMagiAttention":
+                    raise RuntimeError(
+                        "batched magi attention did not activate. "
+                        f"actual_attn={actual_attn!r}; first_attn={first_attn!r}; "
+                        f"{_remote_magi_diag(_model)}; {_attn_class_diag(_model)}"
+                    )
+                _model._la_flash_requested_attn = "magi"
+            except Exception as exc:
+                _fallback_to_sdpa(_model, "magi", exc)
+        else:
+            try:
+                _set_llm_mode(_model, ATTN_MODE)   # decode-safe mask plumbing for sdpa/eager/la_flash
+            except Exception as exc:
+                _fallback_to_sdpa(_model, ATTN_MODE, exc)
+        if MTP_COMPILE:
+            _maybe_compile(_model)
+    return _tok, _proc, _model
+
+
+def _maybe_compile(model):
+    """Compile the shared Qwen2Model core (base.forward). It backs BOTH prefill (called directly)
+    and decode (language_model.forward -> self.model). lm_head + MoonViT left eager. dynamic=True
+    so the varying decode S/kvlen don't trigger a recompile storm. No-op + warning if triton is
+    missing (inductor needs it on GPU). First call pays the compile cost (~42s warm / ~187s cold)."""
+    try:
+        import triton  # noqa: F401
+    except Exception:
+        warnings.warn("MTP_COMPILE set but triton is unavailable; running without torch.compile.")
+        return
+    import torch._dynamo as _dyn
+    _dyn.config.cache_size_limit = max(_dyn.config.cache_size_limit, 64)
+    base = model.language_model.model
+    if not getattr(base, "_mtp_compiled", False):
+        base.forward = torch.compile(base.forward, dynamic=True)
+        base._mtp_compiled = True
+
+
+def build_batched_magi_attention_class(mod):
+    """Build a Qwen2 attention subclass backed by Magi's flex_flash_attn.
+
+    The official LocateAnything ``Qwen2MagiAttention`` asserts ``bsz == 1`` and
+    relies on ``Qwen2Model._attn_implementation == "magi"`` to build a single
+    sample range plan.  For release batch inference the hybrid scheduler passes
+    a batched Magi range plan directly to this layer; a 4D-mask conversion path
+    remains as a compatibility fallback.
+    """
+    flex_flash_attn_func = getattr(mod, "flex_flash_attn_func", None)
+    if flex_flash_attn_func is None:
+        try:
+            from magi_attention.functional.flex_flash_attn import flex_flash_attn_func
+        except Exception as exc:
+            raise RuntimeError(
+                "LA_FLASH_ATTN=magi requires "
+                "magi_attention.functional.flex_flash_attn.flex_flash_attn_func."
+            ) from exc
+
+    FULL, CAUSAL = 0, 1
+    causal_plan_cache = {}
+    try:
+        magi_params = set(inspect.signature(flex_flash_attn_func).parameters)
+    except (TypeError, ValueError):
+        magi_params = set()
+    supports_disable_fwd_atomic = "disable_fwd_atomic_reduction" in magi_params
+
+    def _disjoint_q_ranges(q_ranges):
+        seen = set()
+        for start, end in q_ranges:
+            key = (int(start), int(end))
+            if key in seen:
+                return False
+            seen.add(key)
+        return True
+
+    def _plan_disjoint_q_ranges(plan):
+        cached = plan.get("_la_flash_disjoint_q_ranges")
+        if cached is not None:
+            return bool(cached)
+        q_ranges = plan["q_ranges"].detach().to(device="cpu", dtype=torch.int32).tolist()
+        disjoint = _disjoint_q_ranges(q_ranges)
+        try:
+            plan["_la_flash_disjoint_q_ranges"] = disjoint
+        except Exception:
+            pass
+        return disjoint
+
+    def _tensor_plan(q_ranges, k_ranges, types, device):
+        return {
+            "q_ranges": torch.tensor(q_ranges, dtype=torch.int32, device=device).contiguous(),
+            "k_ranges": torch.tensor(k_ranges, dtype=torch.int32, device=device).contiguous(),
+            "attn_type_map": torch.tensor(types, dtype=torch.int32, device=device).contiguous(),
+            "_la_flash_disjoint_q_ranges": _disjoint_q_ranges(q_ranges),
+        }
+
+    def _offset_plan(plan, q_offset, k_offset):
+        return (
+            (plan["q_ranges"] + int(q_offset)).tolist(),
+            (plan["k_ranges"] + int(k_offset)).tolist(),
+            plan["attn_type_map"].tolist(),
+        )
+
+    def _causal_plan(bsz, q_len, kv_seq_len, device):
+        key = (int(bsz), int(q_len), int(kv_seq_len), device.type, device.index)
+        cached = causal_plan_cache.get(key)
+        if cached is not None:
+            return cached
+        q_ranges, k_ranges, types = [], [], []
+        for b in range(int(bsz)):
+            q_base = b * int(q_len)
+            k_base = b * int(kv_seq_len)
+            q_ranges.append([q_base, q_base + int(q_len)])
+            k_ranges.append([k_base, k_base + int(kv_seq_len)])
+            types.append(CAUSAL)
+        plan = _tensor_plan(q_ranges, k_ranges, types, device)
+        plan.update(
+            {
+                "flash_cu_seqlens_q": torch.arange(
+                    0,
+                    (int(bsz) + 1) * int(q_len),
+                    int(q_len),
+                    dtype=torch.int32,
+                    device=device,
+                ),
+                "flash_cu_seqlens_k": torch.arange(
+                    0,
+                    (int(bsz) + 1) * int(kv_seq_len),
+                    int(kv_seq_len),
+                    dtype=torch.int32,
+                    device=device,
+                ),
+                "flash_causal": True,
+            }
+        )
+        causal_plan_cache[key] = plan
+        return plan
+
+    def _row_segments(row):
+        idx = np.flatnonzero(row)
+        if idx.size == 0:
+            return ((0, 1),)
+        split = np.flatnonzero(np.diff(idx) > 1) + 1
+        starts = np.concatenate((idx[:1], idx[split]))
+        ends = np.concatenate((idx[split - 1], idx[-1:])) + 1
+        return tuple((int(s), int(e)) for s, e in zip(starts, ends))
+
+    def _visible_from_4d_mask(attention_mask, kv_seq_len):
+        mask = attention_mask[:, :, :, :kv_seq_len]
+        if mask.dtype == torch.bool:
+            return mask[:, 0].detach().to(device="cpu", dtype=torch.bool).contiguous()
+        mask_cpu = mask[:, 0].detach().to(device="cpu").contiguous()
+        if getattr(attention_mask, "_la_flash_visible_mask", False):
+            return (mask_cpu > 0).to(dtype=torch.bool)
+
+        max_value = float(mask_cpu.max().item()) if mask_cpu.numel() else 0.0
+        min_value = float(mask_cpu.min().item()) if mask_cpu.numel() else 0.0
+        if max_value > 0.0 and min_value >= 0.0:
+            return (mask_cpu > 0).to(dtype=torch.bool)
+        return (mask_cpu >= 0).to(dtype=torch.bool)
+
+    def _plan_from_visible_mask(attention_mask, bsz, q_len, kv_seq_len, device):
+        cache_key = (int(bsz), int(q_len), int(kv_seq_len), device.type, device.index)
+        cached = getattr(attention_mask, "_la_flash_magi_plan", None)
+        if cached is not None and cached[0] == cache_key:
+            return cached[1]
+
+        visible = _visible_from_4d_mask(attention_mask, int(kv_seq_len)).numpy()
+        q_ranges, k_ranges, types = [], [], []
+        for b in range(int(bsz)):
+            q_base = b * int(q_len)
+            k_base = b * int(kv_seq_len)
+            run_start = 0
+            run_segments = _row_segments(visible[b, 0])
+            for q in range(1, int(q_len)):
+                segments = _row_segments(visible[b, q])
+                if segments == run_segments:
+                    continue
+                for start, end in run_segments:
+                    q_ranges.append([q_base + run_start, q_base + q])
+                    k_ranges.append([k_base + start, k_base + end])
+                    types.append(FULL)
+                run_start = q
+                run_segments = segments
+            for start, end in run_segments:
+                q_ranges.append([q_base + run_start, q_base + int(q_len)])
+                k_ranges.append([k_base + start, k_base + end])
+                types.append(FULL)
+
+        plan = _tensor_plan(q_ranges, k_ranges, types, device)
+        try:
+            attention_mask._la_flash_magi_plan = (cache_key, plan)
+        except Exception:
+            pass
+        return plan
+
+    def _plan_from_magi_dict(attention_mask, bsz, q_len, kv_seq_len, device):
+        if int(bsz) == 1:
+            return attention_mask
+        q_ranges, k_ranges, types = [], [], []
+        for b in range(int(bsz)):
+            qs, ks, ts = _offset_plan(
+                attention_mask,
+                q_offset=b * int(q_len),
+                k_offset=b * int(kv_seq_len),
+            )
+            q_ranges.extend(qs)
+            k_ranges.extend(ks)
+            types.extend(ts)
+        return _tensor_plan(q_ranges, k_ranges, types, device)
+
+    def _magi_plan(attention_mask, bsz, q_len, kv_seq_len, device):
+        if isinstance(attention_mask, dict):
+            if attention_mask.get("_la_flash_batched", False):
+                return attention_mask
+            return _plan_from_magi_dict(attention_mask, bsz, q_len, kv_seq_len, device)
+        if attention_mask is None:
+            return _causal_plan(bsz, q_len, kv_seq_len, device)
+        if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+            raise ValueError(
+                f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, "
+                f"but is {attention_mask.size()}"
+            )
+        return _plan_from_visible_mask(attention_mask, bsz, q_len, kv_seq_len, device)
+
+    class _BatchedMagiAttention(mod.Qwen2Attention):
+        """MagiAttention path with true batch inference via packed token ranges."""
+
+        def forward(
+            self,
+            hidden_states: torch.Tensor,
+            attention_mask=None,
+            position_ids=None,
+            past_key_value=None,
+            output_attentions=False,
+            use_cache=False,
+            **kwargs,
+        ):
+            if output_attentions:
+                raise NotImplementedError("MagiAttention does not support output_attentions=True")
+
+            bsz, q_len, _ = hidden_states.size()
+            query_states = self.q_proj(hidden_states)
+            key_states = self.k_proj(hidden_states)
+            value_states = self.v_proj(hidden_states)
+
+            query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+            key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+            value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+            kv_seq_len = key_states.shape[-2]
+            if past_key_value is not None:
+                if self.layer_idx is None:
+                    raise ValueError(
+                        f"The cache structure has changed since version v4.36. If you are using "
+                        f"{self.__class__.__name__} for auto-regressive decoding with k/v caching, "
+                        "please initialize the attention class with a layer index."
+                    )
+                kv_seq_len += past_key_value.get_seq_length(self.layer_idx)
+
+            cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+            query_states, key_states = mod.apply_rotary_pos_emb(
+                query_states, key_states, cos, sin, position_ids)
+
+            if past_key_value is not None:
+                cache_kwargs = {"sin": sin, "cos": cos}
+                key_states, value_states = past_key_value.update(
+                    key_states, value_states, self.layer_idx, cache_kwargs)
+
+            kv_seq_len = key_states.shape[-2]
+            plan = _magi_plan(attention_mask, bsz, q_len, kv_seq_len, query_states.device)
+            magi_extra_kwargs = {}
+            if supports_disable_fwd_atomic:
+                magi_extra_kwargs["disable_fwd_atomic_reduction"] = (
+                    (not self.training) and _plan_disjoint_q_ranges(plan)
+                )
+
+            query_states = query_states.transpose(1, 2).reshape(
+                bsz * q_len, self.num_heads, self.head_dim).contiguous()
+            key_states = key_states.transpose(1, 2).reshape(
+                bsz * kv_seq_len, self.num_key_value_heads, self.head_dim).contiguous()
+            value_states = value_states.transpose(1, 2).reshape(
+                bsz * kv_seq_len, self.num_key_value_heads, self.head_dim).contiguous()
+
+            attn_output, _ = flex_flash_attn_func(
+                query_states,
+                key_states,
+                value_states,
+                q_ranges=plan["q_ranges"],
+                k_ranges=plan["k_ranges"],
+                attn_type_map=plan["attn_type_map"],
+                softmax_scale=getattr(self, "softmax_scale", self.head_dim ** -0.5),
+                softcap=0.0,
+                deterministic=False,
+                **magi_extra_kwargs,
+            )
+            attn_output = attn_output.view(bsz, q_len, self.hidden_size)
+            attn_output = self.o_proj(attn_output)
+            return attn_output, None, past_key_value
+
+    return _BatchedMagiAttention
+
+
+def build_la_flash_attention_class(mod):
+    """Build a Qwen2 attention subclass backed by LA Flash sparse ranges."""
+    try:
+        from kernel_utils import is_available, range_attention
+    except Exception as exc:
+        raise RuntimeError(
+            "LA_FLASH_ATTN=la_flash requires kernel_utils and FlashAttention."
+        ) from exc
+    if not is_available():
+        raise RuntimeError(
+            "LA_FLASH_ATTN=la_flash requires flash_attn.flash_attn_varlen_func."
+        )
+
+    FULL, CAUSAL = 0, 1
+    causal_plan_cache = {}
+
+    def _tensor_plan(q_ranges, k_ranges, types, device):
+        max_q_len = max((int(end) - int(start) for start, end in q_ranges), default=0)
+        max_k_len = max((int(end) - int(start) for start, end in k_ranges), default=0)
+        plan = {
+            "q_ranges": torch.tensor(q_ranges, dtype=torch.int32, device=device).contiguous(),
+            "k_ranges": torch.tensor(k_ranges, dtype=torch.int32, device=device).contiguous(),
+            "attn_type_map": torch.tensor(types, dtype=torch.int32, device=device).contiguous(),
+            "max_q_len": max_q_len,
+            "max_k_len": max_k_len,
+        }
+        plan.update(_la_flash_group_plan_tensors(q_ranges, types, device))
+        return plan
+
+    def _offset_plan(plan, q_offset, k_offset):
+        return (
+            (plan["q_ranges"] + int(q_offset)).tolist(),
+            (plan["k_ranges"] + int(k_offset)).tolist(),
+            plan["attn_type_map"].tolist(),
+        )
+
+    def _causal_plan(bsz, q_len, kv_seq_len, device):
+        key = (int(bsz), int(q_len), int(kv_seq_len), device.type, device.index)
+        cached = causal_plan_cache.get(key)
+        if cached is not None:
+            return cached
+        q_ranges, k_ranges, types = [], [], []
+        for b in range(int(bsz)):
+            q_base = b * int(q_len)
+            k_base = b * int(kv_seq_len)
+            q_ranges.append([q_base, q_base + int(q_len)])
+            k_ranges.append([k_base, k_base + int(kv_seq_len)])
+            types.append(CAUSAL)
+        plan = _tensor_plan(q_ranges, k_ranges, types, device)
+        plan.update(
+            {
+                "flash_cu_seqlens_q": torch.arange(
+                    0,
+                    (int(bsz) + 1) * int(q_len),
+                    int(q_len),
+                    dtype=torch.int32,
+                    device=device,
+                ),
+                "flash_cu_seqlens_k": torch.arange(
+                    0,
+                    (int(bsz) + 1) * int(kv_seq_len),
+                    int(kv_seq_len),
+                    dtype=torch.int32,
+                    device=device,
+                ),
+                "flash_causal": True,
+            }
+        )
+        causal_plan_cache[key] = plan
+        return plan
+
+    def _row_segments(row):
+        idx = np.flatnonzero(row)
+        if idx.size == 0:
+            return ((0, 1),)
+        split = np.flatnonzero(np.diff(idx) > 1) + 1
+        starts = np.concatenate((idx[:1], idx[split]))
+        ends = np.concatenate((idx[split - 1], idx[-1:])) + 1
+        return tuple((int(s), int(e)) for s, e in zip(starts, ends))
+
+    def _visible_from_4d_mask(attention_mask, kv_seq_len):
+        mask = attention_mask[:, :, :, :kv_seq_len]
+        if mask.dtype == torch.bool:
+            return mask[:, 0].detach().to(device="cpu", dtype=torch.bool).contiguous()
+        mask_cpu = mask[:, 0].detach().to(device="cpu").contiguous()
+        if getattr(attention_mask, "_la_flash_visible_mask", False):
+            return (mask_cpu > 0).to(dtype=torch.bool)
+
+        max_value = float(mask_cpu.max().item()) if mask_cpu.numel() else 0.0
+        min_value = float(mask_cpu.min().item()) if mask_cpu.numel() else 0.0
+        if max_value > 0.0 and min_value >= 0.0:
+            return (mask_cpu > 0).to(dtype=torch.bool)
+        return (mask_cpu >= 0).to(dtype=torch.bool)
+
+    def _prefix_len(row):
+        idx = np.flatnonzero(row)
+        if idx.size == 0:
+            return None
+        end = int(idx[-1]) + 1
+        if not bool(row[:end].all()) or bool(row[end:].any()):
+            return None
+        return end
+
+    def _causal_plan_from_visible(visible, bsz, q_len, kv_seq_len, device):
+        q_ranges, k_ranges, types = [], [], []
+        packed_flash = True
+        for b in range(int(bsz)):
+            first_len = _prefix_len(visible[b, 0])
+            if first_len is None:
+                return None
+            valid_len = int(first_len) + int(q_len) - 1
+            if valid_len < int(q_len) or valid_len > int(kv_seq_len):
+                return None
+            for q in range(int(q_len)):
+                row_len = _prefix_len(visible[b, q])
+                expected = valid_len - int(q_len) + q + 1
+                if row_len != expected:
+                    return None
+            q_base = b * int(q_len)
+            k_base = b * int(kv_seq_len)
+            q_ranges.append([q_base, q_base + int(q_len)])
+            k_ranges.append([k_base, k_base + valid_len])
+            types.append(CAUSAL)
+            packed_flash = packed_flash and valid_len == int(kv_seq_len)
+
+        plan = _tensor_plan(q_ranges, k_ranges, types, device)
+        plan["_la_flash_disjoint_q_ranges"] = True
+        if packed_flash:
+            plan.update(
+                {
+                    "flash_cu_seqlens_q": torch.arange(
+                        0,
+                        (int(bsz) + 1) * int(q_len),
+                        int(q_len),
+                        dtype=torch.int32,
+                        device=device,
+                    ),
+                    "flash_cu_seqlens_k": torch.arange(
+                        0,
+                        (int(bsz) + 1) * int(kv_seq_len),
+                        int(kv_seq_len),
+                        dtype=torch.int32,
+                        device=device,
+                    ),
+                    "flash_causal": True,
+                }
+            )
+        return plan
+
+    def _plan_from_visible_mask(attention_mask, bsz, q_len, kv_seq_len, device):
+        cache_key = (int(bsz), int(q_len), int(kv_seq_len), device.type, device.index, "la_flash")
+        cached = getattr(attention_mask, "_la_flash_range_plan", None)
+        if cached is not None and cached[0] == cache_key:
+            return cached[1]
+
+        visible = _visible_from_4d_mask(attention_mask, int(kv_seq_len)).numpy()
+        plan = _causal_plan_from_visible(visible, bsz, q_len, kv_seq_len, device)
+        if plan is not None:
+            try:
+                attention_mask._la_flash_range_plan = (cache_key, plan)
+            except Exception:
+                pass
+            return plan
+
+        q_ranges, k_ranges, types = [], [], []
+        for b in range(int(bsz)):
+            q_base = b * int(q_len)
+            k_base = b * int(kv_seq_len)
+            run_start = 0
+            run_segments = _row_segments(visible[b, 0])
+            for q in range(1, int(q_len)):
+                segments = _row_segments(visible[b, q])
+                if segments == run_segments:
+                    continue
+                for start, end in run_segments:
+                    q_ranges.append([q_base + run_start, q_base + q])
+                    k_ranges.append([k_base + start, k_base + end])
+                    types.append(FULL)
+                run_start = q
+                run_segments = segments
+            for start, end in run_segments:
+                q_ranges.append([q_base + run_start, q_base + int(q_len)])
+                k_ranges.append([k_base + start, k_base + end])
+                types.append(FULL)
+
+        plan = _tensor_plan(q_ranges, k_ranges, types, device)
+        try:
+            attention_mask._la_flash_range_plan = (cache_key, plan)
+        except Exception:
+            pass
+        return plan
+
+    def _plan_from_magi_dict(attention_mask, bsz, q_len, kv_seq_len, device):
+        if int(bsz) == 1:
+            return attention_mask
+        q_ranges, k_ranges, types = [], [], []
+        for b in range(int(bsz)):
+            qs, ks, ts = _offset_plan(
+                attention_mask,
+                q_offset=b * int(q_len),
+                k_offset=b * int(kv_seq_len),
+            )
+            q_ranges.extend(qs)
+            k_ranges.extend(ks)
+            types.extend(ts)
+        return _tensor_plan(q_ranges, k_ranges, types, device)
+
+    def _range_plan(attention_mask, bsz, q_len, kv_seq_len, device):
+        if isinstance(attention_mask, dict):
+            if attention_mask.get("_la_flash_batched", False):
+                return attention_mask
+            return _plan_from_magi_dict(attention_mask, bsz, q_len, kv_seq_len, device)
+        if attention_mask is None:
+            return _causal_plan(bsz, q_len, kv_seq_len, device)
+        if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+            raise ValueError(
+                f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, "
+                f"but is {attention_mask.size()}"
+            )
+        return _plan_from_visible_mask(attention_mask, bsz, q_len, kv_seq_len, device)
+
+    class _LaFlashAttention(mod.Qwen2Attention):
+        """Range-plan attention path backed by FlashAttention sparse ranges."""
+
+        def forward(
+            self,
+            hidden_states: torch.Tensor,
+            attention_mask=None,
+            position_ids=None,
+            past_key_value=None,
+            output_attentions=False,
+            use_cache=False,
+            **kwargs,
+        ):
+            if output_attentions:
+                raise NotImplementedError("LA Flash attention does not support output_attentions=True")
+
+            bsz, q_len, _ = hidden_states.size()
+            query_states = self.q_proj(hidden_states)
+            key_states = self.k_proj(hidden_states)
+            value_states = self.v_proj(hidden_states)
+
+            query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+            key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+            value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+            kv_seq_len = key_states.shape[-2]
+            if past_key_value is not None:
+                if self.layer_idx is None:
+                    raise ValueError(
+                        f"The cache structure has changed since version v4.36. If you are using "
+                        f"{self.__class__.__name__} for auto-regressive decoding with k/v caching, "
+                        "please initialize the attention class with a layer index."
+                    )
+                kv_seq_len += past_key_value.get_seq_length(self.layer_idx)
+
+            cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+            query_states, key_states = mod.apply_rotary_pos_emb(
+                query_states, key_states, cos, sin, position_ids)
+
+            if past_key_value is not None:
+                cache_kwargs = {"sin": sin, "cos": cos}
+                key_states, value_states = past_key_value.update(
+                    key_states, value_states, self.layer_idx, cache_kwargs)
+
+            kv_seq_len = key_states.shape[-2]
+            dense_backend = os.environ.get("LA_FLASH_DENSE_BACKEND", "sdpa").strip().lower()
+            if dense_backend == "sdpa" and not isinstance(attention_mask, dict):
+                dense_key_states = mod.repeat_kv(key_states, self.num_key_value_groups)
+                dense_value_states = mod.repeat_kv(value_states, self.num_key_value_groups)
+                if attention_mask is not None:
+                    if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                        raise ValueError(
+                            f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, "
+                            f"but is {attention_mask.size()}"
+                        )
+                    query_for_sdpa = query_states.contiguous()
+                    key_for_sdpa = dense_key_states.contiguous()
+                    value_for_sdpa = dense_value_states.contiguous()
+                    is_causal = False
+                elif past_key_value is None:
+                    query_for_sdpa = query_states
+                    key_for_sdpa = dense_key_states
+                    value_for_sdpa = dense_value_states
+                    is_causal = bool(self.is_causal and q_len > 1)
+                else:
+                    query_for_sdpa = key_for_sdpa = value_for_sdpa = None
+                    is_causal = False
+                if query_for_sdpa is not None:
+                    attn_output = torch.nn.functional.scaled_dot_product_attention(
+                        query_for_sdpa,
+                        key_for_sdpa,
+                        value_for_sdpa,
+                        attn_mask=attention_mask,
+                        dropout_p=self.attention_dropout if self.training else 0.0,
+                        is_causal=is_causal,
+                    )
+                    attn_output = attn_output.transpose(1, 2).contiguous()
+                    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+                    attn_output = self.o_proj(attn_output)
+                    return attn_output, None, past_key_value
+
+            plan = _range_plan(attention_mask, bsz, q_len, kv_seq_len, query_states.device)
+
+            query_states = query_states.transpose(1, 2).reshape(
+                bsz * q_len, self.num_heads, self.head_dim).contiguous()
+            key_states = key_states.transpose(1, 2).reshape(
+                bsz * kv_seq_len, self.num_key_value_heads, self.head_dim).contiguous()
+            value_states = value_states.transpose(1, 2).reshape(
+                bsz * kv_seq_len, self.num_key_value_heads, self.head_dim).contiguous()
+
+            attn_output = range_attention(
+                query_states,
+                key_states,
+                value_states,
+                plan["q_ranges"],
+                plan["k_ranges"],
+                plan["attn_type_map"],
+                getattr(self, "softmax_scale", self.head_dim ** -0.5),
+                segment_offsets=plan.get("segment_offsets"),
+                group_q_ranges=plan.get("group_q_ranges"),
+                group_attn_type_map=plan.get("group_attn_type_map"),
+                max_q_len=plan.get("max_q_len"),
+                max_k_len=plan.get("max_k_len"),
+                flash_cu_seqlens_q=plan.get("flash_cu_seqlens_q"),
+                flash_cu_seqlens_k=plan.get("flash_cu_seqlens_k"),
+                flash_causal=plan.get("flash_causal"),
+                disjoint_q_ranges=plan.get("_la_flash_disjoint_q_ranges"),
+            )
+            attn_output = attn_output.view(bsz, q_len, self.hidden_size)
+            attn_output = self.o_proj(attn_output)
+            return attn_output, None, past_key_value
+
+    return _LaFlashAttention
+
+
+def _is_magi_plan(obj):
+    return isinstance(obj, dict) and {
+        "q_ranges",
+        "k_ranges",
+        "attn_type_map",
+    }.issubset(obj.keys())
+
+
+def _la_flash_group_plan_tensors(q_ranges, types, device):
+    """Group consecutive Magi range entries that share the same query span.
+
+    Magi-style plans may represent one query span with multiple disjoint key
+    spans.  LA Flash consumes those as one FlashAttention-backed softmax group.
+    """
+    if not q_ranges:
+        return {
+            "group_q_ranges": torch.empty((0, 2), dtype=torch.int32, device=device),
+            "segment_offsets": torch.zeros((1,), dtype=torch.int32, device=device),
+            "group_attn_type_map": torch.empty((0,), dtype=torch.int32, device=device),
+        }
+
+    grouped_q, grouped_types, offsets = [], [], [0]
+    last_q = None
+    last_type = None
+    for idx, (q_range, attn_type) in enumerate(zip(q_ranges, types)):
+        key = (int(q_range[0]), int(q_range[1]))
+        attn_type = int(attn_type)
+        if last_q is None:
+            grouped_q.append([key[0], key[1]])
+            grouped_types.append(attn_type)
+            last_q = key
+            last_type = attn_type
+            continue
+        if key == last_q and attn_type == last_type:
+            continue
+        offsets.append(idx)
+        grouped_q.append([key[0], key[1]])
+        grouped_types.append(attn_type)
+        last_q = key
+        last_type = attn_type
+    offsets.append(len(q_ranges))
+
+    return {
+        "group_q_ranges": torch.tensor(grouped_q, dtype=torch.int32, device=device).contiguous(),
+        "segment_offsets": torch.tensor(offsets, dtype=torch.int32, device=device).contiguous(),
+        "group_attn_type_map": torch.tensor(grouped_types, dtype=torch.int32, device=device).contiguous(),
+        "max_q_len": max((end - start for start, end in grouped_q), default=0),
+    }
+
+
+def _record_sparse_plan_stats(model, q_ranges, k_ranges, types):
+    if os.environ.get("LA_FLASH_PLAN_STATS", "0") != "1":
+        return
+    stats = getattr(model, "_la_flash_sparse_plan_stats", None)
+    if stats is None:
+        stats = {
+            "calls": 0,
+            "ranges": 0,
+            "q_tokens": 0,
+            "k_tokens": 0,
+            "max_q_len": 0,
+            "max_k_len": 0,
+            "full_ranges": 0,
+            "causal_ranges": 0,
+            "other_ranges": 0,
+        }
+        model._la_flash_sparse_plan_stats = stats
+    stats["calls"] += 1
+    stats["ranges"] += len(q_ranges)
+    for (q_start, q_end), (k_start, k_end), attn_type in zip(q_ranges, k_ranges, types):
+        q_len = int(q_end) - int(q_start)
+        k_len = int(k_end) - int(k_start)
+        stats["q_tokens"] += q_len
+        stats["k_tokens"] += k_len
+        stats["max_q_len"] = max(stats["max_q_len"], q_len)
+        stats["max_k_len"] = max(stats["max_k_len"], k_len)
+        attn_type = int(attn_type)
+        if attn_type == 0:
+            stats["full_ranges"] += 1
+        elif attn_type == 1:
+            stats["causal_ranges"] += 1
+        else:
+            stats["other_ranges"] += 1
+
+
+def build_magi_scheduler_ranges(model, attention_mask_2d, input_ids, past_len, mtp_window=False):
+    """Build batched Magi ranges directly from the hybrid scheduler mask.
+
+    The official Qwen2 SDPA dispatcher may optimize an all-valid 2D mask to
+    ``None`` before decoder layers see it. That is correct for plain causal
+    attention but loses LocateAnything's MTP generation-window rule. Building
+    ranges here keeps Magi batch inference exact and avoids per-layer dense
+    mask conversion.
+    """
+    requested_attn = getattr(model, "_la_flash_requested_attn", ATTN_MODE)
+    if requested_attn not in {"magi", "la_flash"}:
+        return None
+    if attention_mask_2d is None or not hasattr(attention_mask_2d, "dim") or attention_mask_2d.dim() != 2:
+        return None
+
+    bsz, q_len = int(input_ids.shape[0]), int(input_ids.shape[1])
+    key_len = int(attention_mask_2d.shape[1])
+    dev = input_ids.device
+    llm = model.language_model.model
+    block = int(getattr(llm, "block_size", N_FUTURE))
+    causal_attn = bool(getattr(llm, "causal_attn", False))
+    use_mtp_window = bool(mtp_window and q_len >= block and key_len >= block)
+    q0 = max(0, q_len - block)
+    k0 = max(0, key_len - block)
+    blocked_k = k0 - 1
+    past_len = int(past_len)
+
+    key_valid = attention_mask_2d.detach().to(device="cpu", dtype=torch.bool).contiguous().numpy()
+    key_idx = np.arange(key_len)
+    q_ranges, k_ranges, types = [], [], []
+    if not use_mtp_window:
+        causal_q_ranges, causal_k_ranges, causal_types = [], [], []
+        causal_fast_path = True
+        packed_flash = True
+        for b in range(bsz):
+            valid = np.flatnonzero(key_valid[b])
+            if valid.size == 0:
+                causal_fast_path = False
+                break
+            valid_len = int(valid[-1]) + 1
+            if valid_len < q_len or not bool(key_valid[b, :valid_len].all()) or bool(key_valid[b, valid_len:].any()):
+                causal_fast_path = False
+                break
+            packed_flash = packed_flash and valid_len == key_len
+            q_base = b * q_len
+            k_base = b * key_len
+            causal_q_ranges.append([q_base, q_base + q_len])
+            causal_k_ranges.append([k_base, k_base + valid_len])
+            causal_types.append(1)
+        if causal_fast_path:
+            plan = {
+                "q_ranges": torch.tensor(causal_q_ranges, dtype=torch.int32, device=dev).contiguous(),
+                "k_ranges": torch.tensor(causal_k_ranges, dtype=torch.int32, device=dev).contiguous(),
+                "attn_type_map": torch.tensor(causal_types, dtype=torch.int32, device=dev).contiguous(),
+                "max_q_len": q_len,
+                "max_k_len": max((end - start for start, end in causal_k_ranges), default=0),
+                "_la_flash_batched": True,
+                "_la_flash_disjoint_q_ranges": True,
+            }
+            if packed_flash:
+                plan.update(
+                    {
+                        "flash_cu_seqlens_q": torch.arange(
+                            0,
+                            (bsz + 1) * q_len,
+                            q_len,
+                            dtype=torch.int32,
+                            device=dev,
+                        ),
+                        "flash_cu_seqlens_k": torch.arange(
+                            0,
+                            (bsz + 1) * key_len,
+                            key_len,
+                            dtype=torch.int32,
+                            device=dev,
+                        ),
+                        "flash_causal": True,
+                    }
+                )
+            plan.update(_la_flash_group_plan_tensors(causal_q_ranges, causal_types, dev))
+            _record_sparse_plan_stats(model, causal_q_ranges, causal_k_ranges, causal_types)
+            return plan
+
+    def row_segments(row):
+        idx = np.flatnonzero(row)
+        if idx.size == 0:
+            return ((0, 1),)
+        split = np.flatnonzero(np.diff(idx) > 1) + 1
+        starts = np.concatenate((idx[:1], idx[split]))
+        ends = np.concatenate((idx[split - 1], idx[-1:])) + 1
+        return tuple((int(s), int(e)) for s, e in zip(starts, ends))
+
+    for b in range(bsz):
+        q_base = b * q_len
+        k_base = b * key_len
+        run_start = 0
+        run_segments = None
+        if use_mtp_window and not causal_attn:
+            prefix_q_len = q0
+            prefix_k_end = past_len + prefix_q_len
+            prefix_ok = (
+                prefix_q_len > 0
+                and prefix_k_end <= key_len
+                and bool(key_valid[b, :prefix_k_end].all())
+            )
+            window_prefix_ok = blocked_k <= 0 or bool(key_valid[b, :blocked_k].all())
+            window_ok = bool(key_valid[b, k0:key_len].all())
+            if prefix_ok:
+                q_ranges.append([q_base, q_base + prefix_q_len])
+                k_ranges.append([k_base, k_base + prefix_k_end])
+                types.append(1)
+                run_start = prefix_q_len
+            if run_start == prefix_q_len and prefix_q_len < q_len and window_prefix_ok and window_ok:
+                if blocked_k > 0:
+                    q_ranges.append([q_base + prefix_q_len, q_base + q_len])
+                    k_ranges.append([k_base, k_base + blocked_k])
+                    types.append(0)
+                q_ranges.append([q_base + prefix_q_len, q_base + q_len])
+                k_ranges.append([k_base + k0, k_base + key_len])
+                types.append(0)
+                continue
+
+        for q in range(run_start, q_len):
+            visible = key_valid[b] & (key_idx <= q + past_len)
+            if use_mtp_window and q >= q0:
+                if not causal_attn:
+                    visible = visible.copy()
+                    visible[k0:key_len] = key_valid[b, k0:key_len]
+                if blocked_k >= 0:
+                    if visible.base is None:
+                        visible[blocked_k] = False
+                    else:
+                        visible = visible.copy()
+                        visible[blocked_k] = False
+            segments = row_segments(visible)
+            if run_segments is None:
+                run_segments = segments
+                continue
+            if segments == run_segments:
+                continue
+            for start, end in run_segments:
+                q_ranges.append([q_base + run_start, q_base + q])
+                k_ranges.append([k_base + start, k_base + end])
+                types.append(0)
+            run_start = q
+            run_segments = segments
+        for start, end in run_segments:
+            q_ranges.append([q_base + run_start, q_base + q_len])
+            k_ranges.append([k_base + start, k_base + end])
+            types.append(0)
+
+    seen_q_ranges = set()
+    disjoint_q_ranges = True
+    for start, end in q_ranges:
+        key = (int(start), int(end))
+        if key in seen_q_ranges:
+            disjoint_q_ranges = False
+            break
+        seen_q_ranges.add(key)
+
+    plan = {
+        "q_ranges": torch.tensor(q_ranges, dtype=torch.int32, device=dev).contiguous(),
+        "k_ranges": torch.tensor(k_ranges, dtype=torch.int32, device=dev).contiguous(),
+        "attn_type_map": torch.tensor(types, dtype=torch.int32, device=dev).contiguous(),
+        "max_q_len": max((end - start for start, end in q_ranges), default=0),
+        "max_k_len": max((end - start for start, end in k_ranges), default=0),
+        "_la_flash_batched": True,
+        "_la_flash_disjoint_q_ranges": disjoint_q_ranges,
+    }
+    plan.update(_la_flash_group_plan_tensors(q_ranges, types, dev))
+    _record_sparse_plan_stats(model, q_ranges, k_ranges, types)
+    return plan
+
+
+def _direct_base_forward(
+    base,
+    input_ids=None,
+    visual_features=None,
+    image_token_index=None,
+    attention_mask=None,
+    position_ids=None,
+    past_key_values=None,
+    inputs_embeds=None,
+    use_cache=None,
+    output_attentions=None,
+    output_hidden_states=None,
+    return_dict=None,
+):
+    mod = importlib.import_module(type(base).__module__)
+    output_attentions = output_attentions if output_attentions is not None else base.config.output_attentions
+    output_hidden_states = (
+        output_hidden_states if output_hidden_states is not None else base.config.output_hidden_states
+    )
+    use_cache = use_cache if use_cache is not None else base.config.use_cache
+
+    if input_ids is not None and inputs_embeds is not None:
+        raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+    if input_ids is not None:
+        batch_size, seq_length = input_ids.shape
+    elif inputs_embeds is not None:
+        batch_size, seq_length, _ = inputs_embeds.shape
+    else:
+        raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
+
+    past_key_values_length = 0
+    use_legacy_cache = False
+    if use_cache:
+        Cache = getattr(mod, "Cache")
+        DynamicCache = getattr(mod, "DynamicCache")
+        use_legacy_cache = not isinstance(past_key_values, Cache)
+        if use_legacy_cache:
+            if past_key_values is None:
+                past_key_values = DynamicCache()
+            else:
+                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+        past_key_values_length = past_key_values.get_seq_length()
+
+    if position_ids is None:
+        dev = input_ids.device if input_ids is not None else inputs_embeds.device
+        position_ids = torch.arange(
+            past_key_values_length,
+            seq_length + past_key_values_length,
+            dtype=torch.long,
+            device=dev,
+        ).unsqueeze(0).view(-1, seq_length)
+    else:
+        position_ids = position_ids.view(-1, seq_length).long()
+
+    if inputs_embeds is None:
+        inputs_embeds = base.image_processing(input_ids, visual_features, image_token_index)
+
+    hidden_states = inputs_embeds
+    all_hidden_states = () if output_hidden_states else None
+    all_self_attns = () if output_attentions else None
+    next_decoder_cache = None
+
+    for decoder_layer in base.layers:
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+        layer_outputs = decoder_layer(
+            hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_values,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+        )
+        hidden_states = layer_outputs[0]
+        if use_cache:
+            next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+        if output_attentions:
+            all_self_attns += (layer_outputs[1],)
+
+    hidden_states = base.norm(hidden_states)
+    if output_hidden_states:
+        all_hidden_states += (hidden_states,)
+    next_cache = None
+    if use_cache:
+        next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
+    return SimpleNamespace(
+        last_hidden_state=hidden_states,
+        past_key_values=next_cache,
+        hidden_states=all_hidden_states,
+        attentions=all_self_attns,
+    )
+
+
+def language_model_forward(model, **kwargs):
+    """Forward through the text LM, bypassing official dense-mask prep for sparse plans."""
+    lm = model.language_model
+    return_logits = kwargs.pop("return_logits", True)
+    logits_slice = kwargs.pop("logits_slice", None)
+    attention_mask = kwargs.get("attention_mask")
+    use_direct_sparse = (
+        getattr(model, "_la_flash_requested_attn", ATTN_MODE) in {"magi", "la_flash"}
+        and _is_magi_plan(attention_mask)
+    )
+    if not use_direct_sparse:
+        return lm(**kwargs)
+
+    labels = kwargs.pop("labels", None)
+    if labels is not None:
+        raise NotImplementedError("labels are not supported in the direct sparse-plan decode forward")
+    output_attentions = kwargs.get("output_attentions", None)
+    output_hidden_states = kwargs.get("output_hidden_states", None)
+    base_out = _direct_base_forward(lm.model, **kwargs)
+    logits = None
+    if return_logits:
+        hidden_states = base_out.last_hidden_state
+        if logits_slice is not None:
+            hidden_states = hidden_states[:, logits_slice, :]
+        logits = lm.lm_head(hidden_states).float()
+    return SimpleNamespace(
+        logits=logits,
+        past_key_values=base_out.past_key_values,
+        hidden_states=base_out.hidden_states if output_hidden_states else None,
+        attentions=base_out.attentions if output_attentions else None,
+    )
+
+
+_EagerCls = _SdpaCls = _LaFlashCls = _MagiCls = None
+def _attn_classes(mode=None):
+    """Attention classes from the dynamic Qwen2 remote module.
+
+    The official Qwen2Model mask dispatcher only implements ``sdpa`` and
+    single-row ``magi``.  Eager, LA Flash, and batched Magi inference
+    therefore swap the layer class while keeping the model's mask dispatcher
+    pinned to ``sdpa``.
+    """
+    global _EagerCls, _SdpaCls, _LaFlashCls, _MagiCls
+    mode = _normalize_attn_mode(mode) if mode is not None else None
+    if _SdpaCls is None:
+        mod = importlib.import_module(type(_model.language_model.model).__module__)
+        _EagerCls = mod.Qwen2Attention
+        _SdpaCls = mod.Qwen2SdpaAttention
+    else:
+        mod = importlib.import_module(type(_model.language_model.model).__module__)
+    if (mode is None or mode == "la_flash") and _LaFlashCls is None:
+        _LaFlashCls = build_la_flash_attention_class(mod)
+    if (mode is None or mode == "magi") and _MagiCls is None:
+        _MagiCls = build_batched_magi_attention_class(mod) if getattr(mod, "_MAGI_AVAILABLE", False) else None
+    return _EagerCls, _SdpaCls, _LaFlashCls, _MagiCls
+
+def _set_llm_mode(model, mode):
+    """Swap every Qwen2 decoder layer's attention class.
+
+    Release backends keep ``Qwen2Model._attn_implementation='sdpa'`` so the
+    official Qwen2 mask dispatcher stays available for dense-mask modes. The
+    local ``la_flash`` and batched ``magi`` wrappers can also consume scheduler-built
+    sparse plans directly, avoiding repeated per-layer dense mask conversion.
+    """
+    mode = _normalize_attn_mode(mode)
+    eager, sdpa, la_flash, magi = _attn_classes(mode)
+    impl = "sdpa"
+    if mode == "sdpa":
+        cls = sdpa
+    elif mode == "eager":
+        cls = eager
+    elif mode == "la_flash":
+        cls = la_flash
+    elif mode == "magi":
+        if magi is None:
+            raise RuntimeError("MagiAttention is unavailable in the current Python environment.")
+        cls = magi
+    else:
+        raise ValueError(f"unknown LLM attention mode: {mode}")
+    llm = model.language_model.model
+    for lyr in llm.layers:
+        lyr.self_attn.__class__ = cls
+        if mode == "magi":
+            lyr.self_attn.softmax_scale = lyr.self_attn.head_dim ** -0.5
+    llm._attn_implementation = impl
+    llm.config._attn_implementation = llm._attn_implementation
+    if hasattr(model.config, "text_config"):
+        model.config.text_config._attn_implementation = llm._attn_implementation
+    model.config._attn_implementation = llm._attn_implementation
+    model._la_flash_requested_attn = mode
+
+_st = _hp = None
+def _helpers():
+    """The model's own sample_tokens / handle_pattern (the exact box decoders)."""
+    global _st, _hp
+    if _st is None:
+        m = importlib.import_module(type(load()[2]).__module__)
+        _st, _hp = m.sample_tokens, m.handle_pattern
+    return _st, _hp
+
+
+_gu = None
+def _gen_utils():
+    """The model's generate_utils module (apply_repetition_penalty / top_p_logits / top_k_logits /
+    decode_bbox_avg / decode_ref / dists) -- the pieces sample_tokens_batched reuses verbatim."""
+    global _gu
+    if _gu is None:
+        m = importlib.import_module(type(load()[2]).__module__)
+        _gu = importlib.import_module(m.sample_tokens.__module__)
+    return _gu
+
+
+def _env_float(name, default):
+    val = os.environ.get(name)
+    if val is None or val.strip() == "":
+        return float(default)
+    return float(val)
+
+
+def _coord_fallback_mode():
+    mode = os.environ.get("LA_FLASH_COORD_FALLBACK_MODE", "legacy").strip().lower().replace("-", "_")
+    aliases = {
+        "": "legacy",
+        "official": "legacy",
+        "range": "legacy",
+        "spread": "legacy",
+        "none": "off",
+        "disable": "off",
+        "disabled": "off",
+        "entropy_variance": "uncertainty",
+        "entropy_var": "uncertainty",
+        "ent_var": "uncertainty",
+        "entropy_std": "uncertainty",
+    }
+    mode = aliases.get(mode, mode)
+    if mode not in {"legacy", "uncertainty", "off"}:
+        raise ValueError(
+            "LA_FLASH_COORD_FALLBACK_MODE must be one of legacy, uncertainty, off"
+        )
+    return mode
+
+
+def _coord_uncertainty_threshold(coord_start_token_id, coord_end_token_id):
+    """Return the coord uncertainty threshold in raw coord-token units.
+
+    Backward-compatible behavior:
+    - LA_FLASH_COORD_UNCERTAINTY_THRESH > 1 is treated as raw coord-token RMSE.
+    - LA_FLASH_COORD_UNCERTAINTY_THRESH <= 1 is treated as normalized by coord span.
+    - LA_FLASH_COORD_UNCERTAINTY_NORM_THRESH is an explicit normalized override.
+    """
+    coord_span = max(float(coord_end_token_id - coord_start_token_id + 1), 1.0)
+    norm_val = os.environ.get("LA_FLASH_COORD_UNCERTAINTY_NORM_THRESH")
+    if norm_val is not None and norm_val.strip() != "":
+        return float(norm_val) * coord_span
+
+    val = os.environ.get("LA_FLASH_COORD_UNCERTAINTY_THRESH")
+    if val is None or val.strip() == "":
+        return 20.0
+    threshold = float(val)
+    if 0.0 < threshold <= 1.0:
+        return threshold * coord_span
+    return threshold
+
+
+def _decode_bbox_with_uncertainty(logits, probs, token_ids, keep_k=4, generation_mode="hybrid"):
+    """Decode an MTP box with configurable coord uncertainty fallback.
+
+    The default mode is the official LocateAnything rule.  ``uncertainty`` keeps
+    the same frame checks and top-k coord selection, but uses one scalar
+    criterion per coordinate: the posterior RMSE of committing to the current
+    MAP coordinate among valid coord candidates.  This is the Bayes risk under
+    squared coordinate error, so probabilities and token distances are folded
+    into one threshold in coordinate-token units.
+    """
+    gu = _gen_utils()
+    mode = _coord_fallback_mode()
+    if mode == "legacy" or generation_mode != "hybrid":
+        return gu.decode_bbox_avg(logits, probs, token_ids, keep_k=keep_k, generation_mode=generation_mode)
+
+    coord_start_token_id = token_ids["coord_start_token_id"]
+    coord_end_token_id = token_ids["coord_end_token_id"]
+    box_start_token_id = token_ids["box_start_token_id"]
+    box_end_token_id = token_ids["box_end_token_id"]
+    none_token_id = token_ids["none_token_id"]
+    null_token_id = token_ids["null_token_id"]
+    device = logits.device
+
+    box_type = gu.is_valid_box_frame(
+        probs,
+        token_ids,
+        start_thresh=_env_float("LA_FLASH_COORD_BOX_START_THRESH", 0.7),
+        end_thresh=_env_float("LA_FLASH_COORD_BOX_END_THRESH", 0.2),
+        topk=keep_k,
+    )
+    if box_type == "empty_box":
+        return torch.tensor([
+            box_start_token_id,
+            none_token_id,
+            box_end_token_id,
+            null_token_id,
+            null_token_id,
+            null_token_id,
+        ], dtype=torch.long, device=device)
+    if box_type == "illegal_box":
+        return None
+
+    pos_probs, pos_ids = torch.topk(probs[1:5], k=keep_k, dim=-1)
+    valid = (pos_ids >= coord_start_token_id) & (pos_ids <= coord_end_token_id)
+    has_valid = valid.any(dim=-1)
+    if not has_valid.all():
+        return None
+
+    first_valid_idx = valid.long().argmax(dim=-1, keepdim=True)
+    first_valid_ids = pos_ids.gather(-1, first_valid_idx).squeeze(-1)
+    if mode == "off":
+        final_coords = first_valid_ids
+    else:
+        valid_counts = valid.sum(dim=-1)
+        valid_probs = torch.where(valid, pos_probs, torch.zeros_like(pos_probs))
+        valid_mass = valid_probs.sum(dim=-1).clamp_min(1e-12)
+        weights = valid_probs / valid_mass.unsqueeze(-1)
+        coord_values = (pos_ids - coord_start_token_id).to(dtype=torch.float32)
+        map_coord = (first_valid_ids - coord_start_token_id).to(dtype=torch.float32)
+        uncertainty = (weights * (coord_values - map_coord.unsqueeze(-1)).pow(2)).sum(dim=-1).sqrt()
+        is_abnormal = (
+            (valid_counts > 1)
+            & (uncertainty > _coord_uncertainty_threshold(coord_start_token_id, coord_end_token_id))
+        )
+        final_coords = torch.where(is_abnormal, torch.tensor(0, device=device), first_valid_ids)
+
+    start_t = torch.tensor([box_start_token_id], dtype=final_coords.dtype, device=device)
+    end_t = torch.tensor([box_end_token_id], dtype=final_coords.dtype, device=device)
+    return torch.cat([start_t, final_coords, end_t])
+
+
+def _apply_repetition_penalty_lowmem(logits, generated, repetition_penalty):
+    """Apply the stock repetition penalty without allocating a [B, S, V] mask."""
+    if repetition_penalty == 1.0:
+        return logits
+    _, _, vocab_size = logits.shape
+    for row in range(logits.shape[0]):
+        valid_tokens = generated[row].unique()
+        valid_tokens = valid_tokens[(valid_tokens >= 0) & (valid_tokens < vocab_size)]
+        if valid_tokens.numel() == 0:
+            continue
+        row_logits = logits[row, :, valid_tokens]
+        logits[row, :, valid_tokens] = torch.where(
+            row_logits > 0,
+            row_logits / repetition_penalty,
+            row_logits * repetition_penalty,
+        )
+    return logits
+
+
+def _finite_logit_bounds(dtype):
+    finfo = torch.finfo(dtype)
+    return finfo.min, finfo.max
+
+
+def _finite_logits(logits):
+    if not logits.dtype.is_floating_point:
+        logits = logits.float()
+    min_val, max_val = _finite_logit_bounds(logits.dtype)
+    return torch.nan_to_num(logits, nan=min_val, posinf=max_val, neginf=min_val)
+
+
+def _finite_logits_(logits):
+    if not logits.dtype.is_floating_point:
+        return logits.float()
+    min_val, max_val = _finite_logit_bounds(logits.dtype)
+    return logits.nan_to_num_(nan=min_val, posinf=max_val, neginf=min_val)
+
+
+def _top_p_logits_slice_(logits, top_p):
+    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+    cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
+    sorted_indices_to_remove = cumulative_probs > top_p
+    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+    sorted_indices_to_remove[..., 0] = False
+
+    remove = torch.zeros_like(logits, dtype=torch.bool, device=logits.device)
+    remove.scatter_(-1, sorted_indices, sorted_indices_to_remove)
+    logits.masked_fill_(remove, torch.finfo(logits.dtype).min)
+    return logits
+
+
+def _top_p_logits_(logits, top_p):
+    """In-place nucleus filtering with bounded sort workspace.
+
+    The MTP sampler uses logits shaped ``[B, 6, V]``.  Top-p is independent for
+    each row and each future position, so filtering one position at a time keeps
+    the expensive sorted-index workspace at ``[B, V]`` instead of ``[B, 6, V]``.
+    """
+    if logits.dim() == 3 and logits.shape[1] > 1:
+        for pos in range(logits.shape[1]):
+            _top_p_logits_slice_(logits[:, pos, :], top_p)
+        return logits
+    return _top_p_logits_slice_(logits, top_p)
+
+
+def _top_k_logits_(logits, top_k):
+    """In-place top-k filtering mirroring generate_utils.top_k_logits."""
+    top_k = min(int(top_k), logits.size(-1))
+    threshold = torch.topk(logits, top_k)[0][..., -1, None]
+    logits.masked_fill_(logits < threshold, torch.finfo(logits.dtype).min)
+    return logits
+
+
+def _safe_probs(filtered_logits):
+    """Softmax with CUDA-multinomial-safe cleanup and row-wise argmax fallback."""
+    filtered_logits = _finite_logits(filtered_logits)
+    probs = torch.softmax(filtered_logits, dim=-1, dtype=torch.float32)
+    probs = torch.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0).clamp_min_(0.0)
+    row_sum = probs.sum(dim=-1, keepdim=True)
+    bad = (~torch.isfinite(row_sum)) | (row_sum <= 0)
+    if bool(bad.any().item()):
+        fallback = torch.zeros_like(probs)
+        fallback.scatter_(-1, filtered_logits.argmax(dim=-1, keepdim=True), 1.0)
+        probs = torch.where(bad, fallback, probs)
+        row_sum = probs.sum(dim=-1, keepdim=True)
+    return probs / row_sum.clamp_min(1.0e-20)
+
+
+def _sample_top_p_sorted_tokens(logits, top_p):
+    """Sample from top-p filtered logits without scattering back to vocab order."""
+    sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+    cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
+    remove = cumulative_probs > top_p
+    remove[..., 1:] = remove[..., :-1].clone()
+    remove[..., 0] = False
+    sorted_logits.masked_fill_(remove, torch.finfo(sorted_logits.dtype).min)
+    sorted_probs = _safe_probs(sorted_logits)
+    sample_idx = sorted_probs.argmax(dim=-1)
+    try:
+        sample_idx = torch.distributions.Categorical(probs=sorted_probs).sample()
+    except Exception:
+        pass
+    return sorted_indices.gather(-1, sample_idx.unsqueeze(-1)).squeeze(-1)
+
+
+@torch.no_grad()
+def sample_tokens_batched(logits, generated, token_ids, per_row_temp,
+                          repetition_penalty=1.0, top_p=None, top_k=None,
+                          keep_k_avg=4, generation_mode='fast'):
+    """Batched fork of generate_utils.sample_tokens for the MTP window [B,6,V]. The logits pipeline
+    (rep-penalty / per-row temperature / top_p / top_k / softmax / sample) is ROW-INDEPENDENT, so run
+    it ONCE over the whole batch instead of B times on [1,6,V] (the per-row san defeats batching by
+    slicing wlogits[b:b+1]). Only the variable-length box ASSEMBLY (decode_bbox_avg -> ragged shapes,
+    where sample_tokens' final torch.stack throws) stays per-row, returned as a LIST.
+
+    Equivalence to per-row san: every pipeline op reduces on dim=-1 only (never crosses the row dim),
+    so row b's processed logits/probs are bit-identical to slicing first -> greedy (per_row_temp==0,
+    argmax branch, no RNG) is BIT-EXACT. Under sampling, one batched Categorical changes the global
+    RNG consumption order vs B per-row draws -> box-size jitter (blessed; greedy is the exact gate).
+    apply_repetition_penalty already loops per-row internally, so passing the full [B,M] `generated`
+    is row-correct. keep_k_avg/generation_mode mirror sample_tokens' decode_bbox_avg call EXACTLY
+    (note: the per-row san passes keep_k=5 but decode_bbox_avg reads keep_k_avg, default 4 -- so 5 is
+    a no-op there; we replicate keep_k_avg=4). Returns (x0[B,6], boxes: list of B 1-D LongTensors)."""
+    gu = _gen_utils()
+    B, S, V = logits.shape                                        # S = N_FUTURE = 6
+    if repetition_penalty != 1.0:
+        logits = _apply_repetition_penalty_lowmem(logits, generated, repetition_penalty)
+    t = per_row_temp.to(dtype=logits.dtype).view(B, 1, 1)
+    sample_rows = per_row_temp > 0
+    if bool(sample_rows.all().item()):
+        logits.div_(t.clamp(min=1e-8))
+    elif bool(sample_rows.any().item()):
+        idx = sample_rows.nonzero(as_tuple=True)[0]
+        logits[idx].div_(t[idx].clamp(min=1e-8))
+    logits = _finite_logits_(logits)
+    if top_p is not None and top_p < 1:
+        logits = _top_p_logits_(logits, top_p)
+    if top_k is not None and top_k > 0:
+        logits = _top_k_logits_(logits, top_k)
+    probs = _safe_probs(logits)
+    x0 = probs.argmax(dim=-1)                                     # [B,6]; greedy rows are final here
+    samp = per_row_temp > 0
+    if bool(samp.any()):                                         # sampling rows: ONE batched Categorical draw
+        idx = samp.nonzero(as_tuple=True)[0]
+        try:
+            x0[idx] = gu.dists.Categorical(probs=probs[idx]).sample()
+        except Exception:
+            pass                                                # keep argmax (matches san's except: probs.max)
+    boxes = []
+    fallback = torch.zeros(1, dtype=x0.dtype, device=x0.device)
+    for b in range(B):                                          # variable-length box assembly (per-row, exact)
+        db = _decode_bbox_with_uncertainty(
+            logits[b], probs[b], token_ids,
+            keep_k=keep_k_avg, generation_mode=generation_mode)
+        if db is not None:
+            boxes.append(db)
+        else:
+            ref = gu.decode_ref(logits[b], probs[b], token_ids)
+            if ref is None:
+                boxes.append(fallback)
+            elif torch.is_tensor(ref):
+                boxes.append(ref.to(dtype=x0.dtype, device=x0.device))
+            else:
+                boxes.append(torch.tensor(ref, dtype=x0.dtype, device=x0.device))
+    return x0, boxes
+
+
+@torch.no_grad()
+def sample_next_tokens_batched(logits, generated, per_row_temp,
+                               repetition_penalty=1.0, top_p=None, top_k=None):
+    """Batched one-token sampler for AR repair rows.
+
+    This mirrors the row-independent part of ``sample_tokens`` for logits shaped
+    ``[B,1,V]``.  It intentionally does not run bbox/ref assembly because AR mode
+    only needs the next token before the state machine classifies it.
+    """
+    gu = _gen_utils()
+    if logits.dim() != 3 or logits.shape[1] != 1:
+        raise ValueError(f"AR batched sampler expects logits [B,1,V], got {tuple(logits.shape)}")
+    B = int(logits.shape[0])
+    if repetition_penalty != 1.0:
+        logits = _apply_repetition_penalty_lowmem(logits, generated, repetition_penalty)
+    t = per_row_temp.to(dtype=logits.dtype).view(B, 1, 1)
+    sample_rows = per_row_temp > 0
+    if bool(sample_rows.all().item()):
+        logits.div_(t.clamp(min=1e-8))
+    elif bool(sample_rows.any().item()):
+        idx = sample_rows.nonzero(as_tuple=True)[0]
+        logits[idx].div_(t[idx].clamp(min=1e-8))
+    logits = _finite_logits_(logits)
+    sorted_top_p = os.environ.get("AR_SORTED_TOPP", "0") == "1"
+    default_top_p = sorted_top_p and top_p is not None and top_p < 1 and (top_k is None or top_k <= 0)
+    if default_top_p and bool(sample_rows.all().item()):
+        return _sample_top_p_sorted_tokens(logits, top_p)
+    if top_p is not None and top_p < 1:
+        logits = _top_p_logits_(logits, top_p)
+    if top_k is not None and top_k > 0:
+        logits = _top_k_logits_(logits, top_k)
+    probs = _safe_probs(logits)
+    x0 = probs.argmax(dim=-1)
+    if bool(sample_rows.any().item()):
+        # Keep row-ordered sampling as the release default.  A single batched
+        # Categorical is faster, but it consumes RNG differently from stock AR
+        # repair and can alter default-temperature termination behavior.
+        for row in sample_rows.nonzero(as_tuple=True)[0].tolist():
+            try:
+                x0[row : row + 1] = gu.dists.Categorical(probs=probs[row : row + 1]).sample()
+            except Exception:
+                pass
+    return x0
+
+
+def load_pil(p):
+    from PIL import Image
+    im = Image.open(p).convert("RGB"); w, h = im.size
+    if max(w, h) > MAX_DIM:
+        s = MAX_DIM / max(w, h); im = im.resize((max(1, round(w*s)), max(1, round(h*s))), Image.LANCZOS)
+    return im
+
+def _preproc_one(im):
+    """CPU-side processor for one image -> (pixel_values[bf16], grid[int32]). Split out of
+    _encode_image so _encode_images can batch the GPU encode while preprocessing stays per-image."""
+    tok, proc, model = load()
+    msg = [{"role": "user", "content": [{"type": "image", "image": im}, {"type": "text", "text": "x"}]}]
+    text = proc.py_apply_chat_template(msg, tokenize=False, add_generation_prompt=True)
+    imgs, vids = proc.process_vision_info(msg)
+    inp = proc(text=[text], images=imgs, videos=vids, return_tensors="pt").to(DEV)
+    grid = inp.get("image_grid_hws")
+    if isinstance(grid, np.ndarray): grid = torch.from_numpy(grid).to(DEV, dtype=torch.int32)
+    return inp["pixel_values"].to(DT), grid
+
+
+def _vision_is_flash():
+    """True iff MoonViT will actually run flash_attn_varlen (so cross-image packing is
+    block-diagonal = exact AND a win). If the vision blocks are on sdpa/eager, OR the flash
+    wheel is absent (multihead_attention falls back to the dense-mask sdpa path), packing is
+    O(S^2) N^2 -> caller must stay per-image."""
+    vm = load()[2].vision_model
+    mod = importlib.import_module(type(vm).__module__)
+    if getattr(mod, "flash_attn_varlen_func", None) is None:
+        return False
+    try:
+        return vm.encoder.blocks[0].attn_implementation == "flash_attention_2"
+    except Exception:
+        return False
+
+
+@torch.no_grad()
+def _encode_images(ims):
+    """N images -> list of [n_img_tokens, C] mlp1-projected visual_features, one per image
+    (row-order). Drop-in for [_encode_image(im) for im in ims].
+
+    With flash present (_vision_is_flash) and N>1, packs images into
+    extract_feature micro-batches: MoonViT's varlen cu_seqlens path is
+    block-diagonal by image. Without flash, the dense SDPA fallback would scale
+    with the packed total sequence length, so this function falls back to
+    per-image encode. MTP_BATCH_VISION=0 also forces per-image encode."""
+    tok, proc, model = load()
+    pvs, grids = [], []
+    for im in ims:
+        pv, g = _preproc_one(im)
+        pvs.append(pv); grids.append(g)
+    if BATCH_VISION and len(ims) > 1 and _vision_is_flash():
+        if VISION_ENCODE_BATCH_SIZE <= 0 or VISION_ENCODE_BATCH_SIZE >= len(ims):
+            vit_list = model.extract_feature(torch.cat(pvs, dim=0), torch.cat(grids, dim=0))
+        else:
+            vit_list = []
+            for start in range(0, len(ims), VISION_ENCODE_BATCH_SIZE):
+                end = min(start + VISION_ENCODE_BATCH_SIZE, len(ims))
+                vit_list.extend(
+                    model.extract_feature(
+                        torch.cat(pvs[start:end], dim=0),
+                        torch.cat(grids[start:end], dim=0),
+                    )
+                )
+        return [model.mlp1(v) for v in vit_list]        # one [P_i, C] per image (patch_merger split)
+    return [model.mlp1(torch.cat(model.extract_feature(pv, g), dim=0))
+            for pv, g in zip(pvs, grids)]                # per-image (flash absent / N==1 / forced off)
+
+
+@torch.no_grad()
+def _encode_image(im):
+    """Single-image convenience wrapper (single-image callers); = _encode_images([im])[0]
+    (takes the per-image path inside _encode_images, so bit-identical to the original)."""
+    return _encode_images([im])[0]
+
+@torch.no_grad()
+def _tokenize(im, query):
+    """1-D prompt token ids for (image, query). Uses the model's own chat template."""
+    tok, proc, model = load()
+    msg = [{"role": "user", "content": [{"type": "image", "image": im},
+                                        {"type": "text", "text": _PROMPT + query + "."}]}]
+    text = proc.py_apply_chat_template(msg, tokenize=False, add_generation_prompt=True)
+    imgs, vids = proc.process_vision_info(msg)
+    return proc(text=[text], images=imgs, videos=vids, return_tensors="pt").to(DEV)["input_ids"][0]
+
+
+@torch.no_grad()
+def _tokenize_cached_image(query, image_token_count, im=None):
+    """Tokenize a prompt when the image token count is already known.
+
+    This keeps the processor's chat template, but directly expands ``<image-1>``
+    from the cached visual feature length.  It avoids re-running the CPU image
+    processor for every category prompt that shares the same image.
+    """
+    tok, proc, model = load()
+    msg = [{"role": "user", "content": [{"type": "image", "image": im},
+                                        {"type": "text", "text": _PROMPT + query + "."}]}]
+    text = proc.py_apply_chat_template(msg, tokenize=False, add_generation_prompt=True)
+    placeholder = f"<{getattr(proc, 'image_placeholder', 'image')}-1>"
+    image_token = getattr(proc, "image_token", "<IMG_CONTEXT>")
+    image_start = getattr(proc, "image_start_token", "<img>")
+    image_end = getattr(proc, "image_end_token", "</img>")
+    replacement = f"<image 1>{image_start}{image_token * int(image_token_count)}{image_end}"
+    if placeholder not in text:
+        raise ValueError(f"cached image placeholder {placeholder!r} was not found in chat template")
+    text = text.replace(placeholder, replacement, 1)
+    return tok([text], return_tensors="pt").to(DEV)["input_ids"][0]
+
+
+def _proc_full(im, query):
+    """Full processor dict (input_ids, attention_mask, pixel_values, image_grid_hws) —
+    used by the bench to drive the STOCK generate for the equivalence check."""
+    tok, proc, model = load()
+    msg = [{"role": "user", "content": [{"type": "image", "image": im},
+                                        {"type": "text", "text": _PROMPT + query + "."}]}]
+    text = proc.py_apply_chat_template(msg, tokenize=False, add_generation_prompt=True)
+    imgs, vids = proc.process_vision_info(msg)
+    inp = proc(text=[text], images=imgs, videos=vids, return_tensors="pt").to(DEV)
+    grid = inp.get("image_grid_hws")
+    if isinstance(grid, np.ndarray): grid = torch.from_numpy(grid).to(DEV, dtype=torch.int32)
+    inp["image_grid_hws"] = grid
+    return inp
+
+def _pad_generated(prompt_ids, gen_ids, img_tok, dev):
+    """Per-row [prompt + accepted] left-padded with the image token (already in every
+    prompt -> .unique() unchanged -> repetition penalty identical to single-run)."""
+    rows = [list(prompt_ids[b].tolist()) + gen_ids[b] for b in range(len(prompt_ids))]
+    M = max(len(r) for r in rows)
+    out = torch.full((len(rows), M), img_tok, dtype=torch.long, device=dev)
+    for b, r in enumerate(rows):
+        out[b, M - len(r):] = torch.tensor(r, dtype=torch.long, device=dev)
+    return out

kernel_utils/README.md

@@ -0,0 +1,53 @@
+# 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. |
+
+The old local CUDA-only encodings (`100 + block_size` and `10000 + ...`) were
+removed from the release.
+
+## 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.

kernel_utils/__init__.py

@@ -0,0 +1,5 @@
+"""FlashAttention sparse range utilities for LocateAnything batch inference."""
+
+from .range_attention import range_attention, is_available
+
+__all__ = ["range_attention", "is_available"]

kernel_utils/range_attention.py

@@ -0,0 +1,394 @@
+"""Sparse LocateAnything attention implemented with FlashAttention varlen.
+
+The public API accepts flattened query/key/value tensors:
+
+  q: [total_q, num_q_heads, head_dim]
+  k: [total_k, num_kv_heads, head_dim]
+  v: [total_k, num_kv_heads, head_dim]
+
+and a Magi-style range plan:
+
+  q_ranges: [num_ranges, 2]
+  k_ranges: [num_key_segments, 2]
+  segment_offsets: [num_query_groups + 1]
+  attn_type_map:
+    0 = full attention over the listed key segment(s)
+    1 = bottom-right causal attention
+
+For LocateAnything hybrid MTP decode, batch_utils represents the window as a
+causal prefix plus full-attention sparse window segments.  This module packs
+those visible KV segments and calls FlashAttention varlen, avoiding dense masks.
+"""
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+import torch
+
+
+_FLASH_ATTN_VARLEN = None
+_FLASH_ATTN_ERROR: Optional[BaseException] = None
+
+
+def _env_enabled(name: str, default: str = "auto") -> bool:
+    value = os.environ.get(name, default).strip().lower()
+    return value in {"", "auto", "1", "on", "true", "yes", "force"}
+
+
+def is_available() -> bool:
+    try:
+        _load_flash_attn_varlen()
+        return True
+    except Exception:
+        return False
+
+
+def _flash_fastpath_enabled() -> bool:
+    return _env_enabled("LA_FLASH_FASTPATH", "auto")
+
+
+def _flash_segment_fastpath_enabled() -> bool:
+    return _env_enabled("LA_FLASH_SEGMENT_FASTPATH", "auto")
+
+
+def _load_flash_attn_varlen():
+    global _FLASH_ATTN_VARLEN, _FLASH_ATTN_ERROR
+    if _FLASH_ATTN_VARLEN is not None:
+        return _FLASH_ATTN_VARLEN
+    if _FLASH_ATTN_ERROR is not None:
+        raise _FLASH_ATTN_ERROR
+    try:
+        from flash_attn import flash_attn_varlen_func
+
+        _FLASH_ATTN_VARLEN = flash_attn_varlen_func
+        return _FLASH_ATTN_VARLEN
+    except BaseException as exc:
+        _FLASH_ATTN_ERROR = exc
+        raise
+
+
+def _coalesce_query_groups(q_ranges, k_ranges, attn_type_map):
+    """Group consecutive entries that share the same query span and mask type."""
+    if q_ranges.numel() == 0:
+        segment_offsets = torch.zeros((1,), dtype=torch.int32, device=q_ranges.device)
+        return q_ranges, k_ranges, segment_offsets, attn_type_map, 0, 0
+
+    q_cpu = q_ranges.detach().to(device="cpu", dtype=torch.int32).contiguous()
+    t_cpu = attn_type_map.detach().to(device="cpu", dtype=torch.int32).contiguous()
+    grouped_q = []
+    grouped_t = []
+    offsets = [0]
+    max_q_len = 0
+    last_q = None
+    last_t = None
+    for idx, (qr, attn_type) in enumerate(zip(q_cpu.tolist(), t_cpu.tolist())):
+        key = (int(qr[0]), int(qr[1]))
+        attn_type = int(attn_type)
+        if attn_type not in (0, 1):
+            raise RuntimeError(
+                "LA Flash path only supports attn_type 0/1. "
+                f"Got attn_type={attn_type}; regenerate the plan without the old local kernel encodings."
+            )
+        if last_q is None:
+            grouped_q.append([key[0], key[1]])
+            grouped_t.append(attn_type)
+            max_q_len = max(max_q_len, key[1] - key[0])
+            last_q = key
+            last_t = attn_type
+            continue
+        if key == last_q and attn_type == last_t:
+            continue
+        offsets.append(idx)
+        grouped_q.append([key[0], key[1]])
+        grouped_t.append(attn_type)
+        max_q_len = max(max_q_len, key[1] - key[0])
+        last_q = key
+        last_t = attn_type
+    offsets.append(int(q_ranges.shape[0]))
+
+    k_cpu = k_ranges.detach().to(device="cpu", dtype=torch.int32).contiguous()
+    max_k_len = max((int(end) - int(start) for start, end in k_cpu.tolist()), default=0)
+
+    return (
+        torch.tensor(grouped_q, dtype=torch.int32, device=q_ranges.device).contiguous(),
+        k_ranges,
+        torch.tensor(offsets, dtype=torch.int32, device=q_ranges.device).contiguous(),
+        torch.tensor(grouped_t, dtype=torch.int32, device=q_ranges.device).contiguous(),
+        int(max_q_len),
+        int(max_k_len),
+    )
+
+
+def _flash_lse_to_tq_h(lse, total_q, q_lengths=None):
+    if lse is None:
+        return None
+    if lse.dim() != 2:
+        if lse.dim() == 3 and q_lengths is not None and lse.shape[0] == len(q_lengths):
+            chunks = []
+            for idx, q_len in enumerate(q_lengths):
+                q_len = int(q_len)
+                if lse.shape[1] == 0 or q_len > lse.shape[2]:
+                    return None
+                chunks.append(lse[idx, :, :q_len].transpose(0, 1).contiguous())
+            merged = torch.cat(chunks, dim=0).float()
+            return merged if merged.shape[0] == total_q else None
+        return None
+    if lse.shape[0] == total_q:
+        return lse.float()
+    if lse.shape[1] == total_q:
+        return lse.transpose(0, 1).contiguous().float()
+    return None
+
+
+def _make_cu_seqlens(lengths, device):
+    return torch.tensor([0] + list(torch.tensor(lengths).cumsum(0).tolist()), device=device, dtype=torch.int32)
+
+
+def _try_flash_segment_merge(
+    q,
+    k,
+    v,
+    k_ranges,
+    segment_offsets,
+    group_q_ranges,
+    group_attn_type_map,
+    softmax_scale,
+):
+    if not _flash_segment_fastpath_enabled():
+        return None
+    if q.dtype not in (torch.float16, torch.bfloat16) or k.dtype != q.dtype or v.dtype != q.dtype:
+        return None
+    if group_q_ranges is None or segment_offsets is None or group_attn_type_map is None:
+        return None
+
+    flash_attn_varlen = _load_flash_attn_varlen()
+    gq_cpu = group_q_ranges.detach().to(device="cpu", dtype=torch.int32).contiguous()
+    kr_cpu = k_ranges.detach().to(device="cpu", dtype=torch.int32).contiguous()
+    seg_cpu = segment_offsets.detach().to(device="cpu", dtype=torch.int32).contiguous()
+    type_cpu = group_attn_type_map.detach().to(device="cpu", dtype=torch.int32).contiguous()
+
+    groups = []
+    max_segments = 0
+    for group_idx, (q_start, q_end) in enumerate(gq_cpu.tolist()):
+        attn_type = int(type_cpu[group_idx].item())
+        if attn_type not in (0, 1):
+            return None
+        seg_start = int(seg_cpu[group_idx].item())
+        seg_end = int(seg_cpu[group_idx + 1].item())
+        if seg_end <= seg_start or q_end <= q_start:
+            return None
+        segments = kr_cpu[seg_start:seg_end].tolist()
+        max_segments = max(max_segments, len(segments))
+        groups.append((int(q_start), int(q_end), attn_type, [(int(a), int(b)) for a, b in segments]))
+
+    if not groups or max_segments == 0:
+        return None
+
+    can_pack_full_groups = all(attn_type == 0 or len(segments) == 1 for _, _, attn_type, segments in groups)
+    if can_pack_full_groups:
+        merged = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=q.dtype)
+        covered = torch.zeros((q.shape[0],), device=q.device, dtype=torch.bool)
+        for attn_type in (0, 1):
+            q_slices = []
+            k_slices = []
+            v_slices = []
+            q_lengths = []
+            k_lengths = []
+            targets = []
+            for q_start, q_end, group_type, segments in groups:
+                if group_type != attn_type:
+                    continue
+                q_slices.append(q[q_start:q_end])
+                if attn_type == 0 and len(segments) > 1:
+                    k_slices.append(torch.cat([k[start:end] for start, end in segments], dim=0))
+                    v_slices.append(torch.cat([v[start:end] for start, end in segments], dim=0))
+                    k_lengths.append(sum(end - start for start, end in segments))
+                else:
+                    k_start, k_end = segments[0]
+                    k_slices.append(k[k_start:k_end])
+                    v_slices.append(v[k_start:k_end])
+                    k_lengths.append(k_end - k_start)
+                q_lengths.append(q_end - q_start)
+                targets.append((q_start, q_end))
+            if not q_slices:
+                continue
+
+            out_pass = flash_attn_varlen(
+                torch.cat(q_slices, dim=0).contiguous(),
+                torch.cat(k_slices, dim=0).contiguous(),
+                torch.cat(v_slices, dim=0).contiguous(),
+                _make_cu_seqlens(q_lengths, q.device),
+                _make_cu_seqlens(k_lengths, q.device),
+                int(max(q_lengths)),
+                int(max(k_lengths)),
+                dropout_p=0.0,
+                softmax_scale=float(softmax_scale),
+                causal=bool(attn_type == 1),
+            )
+            if isinstance(out_pass, tuple):
+                out_pass = out_pass[0]
+
+            cursor = 0
+            for q_start, q_end in targets:
+                q_len = q_end - q_start
+                merged[q_start:q_end] = out_pass[cursor:cursor + q_len]
+                covered[q_start:q_end] = True
+                cursor += q_len
+
+        if bool(covered.all().item()):
+            return merged
+
+    merged = torch.zeros((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
+    merged_lse = torch.full((q.shape[0], q.shape[1]), -float("inf"), device=q.device, dtype=torch.float32)
+    covered = torch.zeros((q.shape[0],), device=q.device, dtype=torch.bool)
+
+    for segment_idx in range(max_segments):
+        for attn_type in (0, 1):
+            q_slices = []
+            k_slices = []
+            v_slices = []
+            q_lengths = []
+            k_lengths = []
+            targets = []
+            for q_start, q_end, group_type, segments in groups:
+                if group_type != attn_type or segment_idx >= len(segments):
+                    continue
+                k_start, k_end = segments[segment_idx]
+                if k_end <= k_start:
+                    continue
+                q_slices.append(q[q_start:q_end])
+                k_slices.append(k[k_start:k_end])
+                v_slices.append(v[k_start:k_end])
+                q_lengths.append(q_end - q_start)
+                k_lengths.append(k_end - k_start)
+                targets.append((q_start, q_end))
+            if not q_slices:
+                continue
+
+            result = flash_attn_varlen(
+                torch.cat(q_slices, dim=0).contiguous(),
+                torch.cat(k_slices, dim=0).contiguous(),
+                torch.cat(v_slices, dim=0).contiguous(),
+                _make_cu_seqlens(q_lengths, q.device),
+                _make_cu_seqlens(k_lengths, q.device),
+                int(max(q_lengths)),
+                int(max(k_lengths)),
+                dropout_p=0.0,
+                softmax_scale=float(softmax_scale),
+                causal=bool(attn_type == 1),
+                return_attn_probs=True,
+            )
+            if not isinstance(result, tuple) or len(result) < 2:
+                return None
+            out_pass = result[0]
+            lse_pass = _flash_lse_to_tq_h(result[1], out_pass.shape[0], q_lengths)
+            if lse_pass is None:
+                return None
+
+            cursor = 0
+            for q_start, q_end in targets:
+                q_len = q_end - q_start
+                out_seg = out_pass[cursor:cursor + q_len].float()
+                lse_seg = lse_pass[cursor:cursor + q_len]
+                old_lse = merged_lse[q_start:q_end]
+                new_lse = torch.maximum(old_lse, lse_seg)
+                old_w = torch.exp(old_lse - new_lse)
+                seg_w = torch.exp(lse_seg - new_lse)
+                denom = (old_w + seg_w).clamp_min(1e-20)
+                merged[q_start:q_end] = (
+                    merged[q_start:q_end] * old_w.unsqueeze(-1)
+                    + out_seg * seg_w.unsqueeze(-1)
+                ) / denom.unsqueeze(-1)
+                merged_lse[q_start:q_end] = new_lse + torch.log(denom)
+                covered[q_start:q_end] = True
+                cursor += q_len
+
+    if not bool(covered.all().item()):
+        return None
+    return merged.to(dtype=q.dtype)
+
+
+def range_attention(
+    q,
+    k,
+    v,
+    q_ranges,
+    k_ranges,
+    attn_type_map,
+    softmax_scale: float,
+    *,
+    segment_offsets=None,
+    group_q_ranges=None,
+    group_attn_type_map=None,
+    max_q_len=None,
+    max_k_len=None,
+    flash_cu_seqlens_q=None,
+    flash_cu_seqlens_k=None,
+    flash_causal=None,
+    disjoint_q_ranges=None,
+):
+    """Run sparse range attention through FlashAttention varlen."""
+    del disjoint_q_ranges
+    if not q.is_cuda:
+        raise RuntimeError("LA Flash range_attention requires CUDA tensors")
+    if segment_offsets is None or group_q_ranges is None or group_attn_type_map is None:
+        (
+            group_q_ranges,
+            k_ranges,
+            segment_offsets,
+            group_attn_type_map,
+            computed_max_q_len,
+            computed_max_k_len,
+        ) = _coalesce_query_groups(q_ranges, k_ranges, attn_type_map)
+        if max_q_len is None:
+            max_q_len = computed_max_q_len
+        if max_k_len is None:
+            max_k_len = computed_max_k_len
+    elif max_q_len is None:
+        lengths = (group_q_ranges[:, 1] - group_q_ranges[:, 0]).detach().to(device="cpu")
+        max_q_len = int(lengths.max().item()) if lengths.numel() else 0
+    if max_k_len is None:
+        k_lengths = (k_ranges[:, 1] - k_ranges[:, 0]).detach().to(device="cpu")
+        max_k_len = int(k_lengths.max().item()) if k_lengths.numel() else 0
+
+    if (
+        flash_cu_seqlens_q is not None
+        and flash_cu_seqlens_k is not None
+        and flash_causal is not None
+        and _flash_fastpath_enabled()
+        and q.dtype in (torch.float16, torch.bfloat16)
+        and k.dtype == q.dtype
+        and v.dtype == q.dtype
+    ):
+        flash_attn_varlen = _load_flash_attn_varlen()
+        return flash_attn_varlen(
+            q.contiguous(),
+            k.contiguous(),
+            v.contiguous(),
+            flash_cu_seqlens_q.contiguous().to(device=q.device, dtype=torch.int32),
+            flash_cu_seqlens_k.contiguous().to(device=q.device, dtype=torch.int32),
+            int(max_q_len),
+            int(max_k_len),
+            dropout_p=0.0,
+            softmax_scale=float(softmax_scale),
+            causal=bool(flash_causal),
+        )
+
+    segment_out = _try_flash_segment_merge(
+        q,
+        k,
+        v,
+        k_ranges,
+        segment_offsets,
+        group_q_ranges,
+        group_attn_type_map,
+        softmax_scale,
+    )
+    if segment_out is not None:
+        return segment_out
+
+    raise RuntimeError(
+        "LA Flash could not express this range plan with FlashAttention varlen. "
+        "Only attn_type 0/1 range plans are supported in the release path."
+    )

pyproject.toml

@@ -0,0 +1,18 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "locateanything-la-flash"
+version = "0.1.0"
+description = "LocateAnything batch utils with LA Flash inference backend."
+requires-python = ">=3.10"
+dependencies = [
+  "numpy",
+  "pillow",
+  "torch",
+  "transformers",
+]
+
+[tool.setuptools]
+packages = ["batch_utils", "kernel_utils"]