nimo1234/LocateAnything-3B
011
1"""Internal runtime support for the LocateAnything-3B hybrid batch decoder.2 3This file keeps only the model-loading, tokenization, image-encoding, stock4processor, and sample-token helpers that ``engine_hybrid.py`` needs.5 6Important env knobs:7 LA_FLASH_MODEL HF repo id / local path of the model (default nvidia/LocateAnything-3B)8 HF_HUB_OFFLINE=1 read the local HF cache only (no network); unset -> download on first use9 LA_FLASH_ATTN sdpa, eager, magi, or la_flash; la_flash uses FlashAttention sparse ranges10 LA_FLASH_STRICT_ATTN 1 -> fail if the requested backend is unavailable;11 default 0 falls back to sdpa12 LA_FLASH_VISION_ATTN auto, flash_attention_2, sdpa, or eager (default auto)13 LA_FLASH_HYBRID_PREFILL shared, none, per_row, or batch prompt KV prefill (default shared)14 MTP_BATCH_VISION 0 -> per-image vision encode (default 1: batched when flash is present)15 LA_FLASH_VISION_ENCODE_BATCH_SIZE16 max images per MoonViT encode micro-batch (default 8; <=0 disables limit)17 MTP_BATCH_SAN 0 -> per-row logits/sample pipeline (default 1: batched over [B,6,V])18 AR_BATCH_SAN 0 -> per-row AR sample pipeline (default 1: batched over [B,1,V])19"""20import inspect21import os, warnings, importlib, torch22from types import SimpleNamespace23import numpy as np24from transformers import AutoModel, AutoTokenizer, AutoProcessor25 26 27# By default let transformers fetch the model on first use; set HF_HUB_OFFLINE=1 yourself28# to read the local HF cache only (e.g. air-gapped / already-downloaded runs).29MODEL = os.environ.get("LA_FLASH_MODEL", "nvidia/LocateAnything-3B")30 31 32LLM_ATTN_MODES = ("sdpa", "eager", "magi", "la_flash")33VISION_ATTN_MODES = ("auto", "flash_attention_2", "sdpa", "eager")34 35 36def _normalize_attn_mode(value):37 mode = (value or "sdpa").strip().lower().replace("-", "_")38 aliases = {39 "": "sdpa",40 "manual": "eager",41 "torch": "eager",42 "torch_eager": "eager",43 "torch_sdpa": "sdpa",44 "scaled_dot_product_attention": "sdpa",45 "flash": "la_flash",46 "la_flash": "la_flash",47 "kernel": "la_flash",48 "cuda": "la_flash",49 "range": "la_flash",50 "range_attention": "la_flash",51 "flex_flash": "magi",52 "flex_flash_attention": "magi",53 "flex_flash_attn": "magi",54 }55 mode = aliases.get(mode, mode)56 if mode not in LLM_ATTN_MODES:57 raise ValueError(58 f"LA_FLASH_ATTN must be one of {', '.join(LLM_ATTN_MODES)}; got {value!r}"59 )60 return mode61 62 63def _normalize_vision_attn_mode(value):64 mode = (value or "auto").strip().lower().replace("-", "_")65 aliases = {66 "": "auto",67 "flash": "flash_attention_2",68 "flash_attention2": "flash_attention_2",69 "fa2": "flash_attention_2",70 "manual": "eager",71 }72 mode = aliases.get(mode, mode)73 if mode not in VISION_ATTN_MODES:74 raise ValueError(75 f"LA_FLASH_VISION_ATTN must be one of {', '.join(VISION_ATTN_MODES)}; got {value!r}"76 )77 return mode78 79 80ATTN_MODE = _normalize_attn_mode(os.environ.get("LA_FLASH_ATTN", "sdpa"))81REMOTE_ATTN_MODE = "sdpa" if ATTN_MODE in {"la_flash", "magi"} else ATTN_MODE82VISION_ATTN_MODE = _normalize_vision_attn_mode(os.environ.get("LA_FLASH_VISION_ATTN", "auto"))83MAX_DIM = 102484DEV, DT = "cuda", torch.bfloat1685N_FUTURE = 6 # = config.block_size (MTP window)86_PROMPT = "Locate all the instances that matches the following description: "87 88 89def _env_flag(name, default=False):90 val = os.environ.get(name)91 if val is None:92 return default93 return val.strip().lower() not in {"0", "false", "no", "off"}94 95 96def _env_int(name):97 val = os.environ.get(name)98 if val is None or val.strip() == "":99 return None100 return int(val)101 102 103def _strict_attn():104 return _env_flag("LA_FLASH_STRICT_ATTN", False)105 106 107def _fallback_to_sdpa(model, requested, reason):108 if requested == "sdpa":109 raise RuntimeError(f"LA_FLASH_ATTN=sdpa failed: {reason}") from reason110 message = f"LA_FLASH_ATTN={requested} is unavailable; falling back to sdpa. Reason: {reason}"111 if _strict_attn():112 raise RuntimeError(message) from reason113 warnings.warn(message)114 _set_llm_mode(model, "sdpa")115 model._la_flash_requested_attn_original = requested116 model._la_flash_attn_fallback_reason = str(reason)117 return "sdpa"118 119 120# Optional compile for the shared Qwen2 core. This is off by default because the121# hybrid scheduler already varies query/cache shapes and first-call compile cost is high.122MTP_COMPILE = os.environ.get("MTP_COMPILE", "0") == "1"123 124# Batch the MoonViT vision encode across a micro-batch's images: pack N images into ONE125# extract_feature. With flash present, MoonViT's varlen cu_seqlens path is block-diagonal per126# image and equivalent to per-image encode.127# Without flash, sdpa builds a dense [1,S,S] mask -> O(S^2) N^2 -> per-image fallback (auto, see128# _vision_is_flash). Default ON; set MTP_BATCH_VISION=0 to force per-image.129BATCH_VISION = os.environ.get("MTP_BATCH_VISION", "1") == "1"130_vision_encode_batch_size = _env_int("LA_FLASH_VISION_ENCODE_BATCH_SIZE")131VISION_ENCODE_BATCH_SIZE = 8 if _vision_encode_batch_size is None else max(0, _vision_encode_batch_size)132 133# Batch the per-row box-decode (sample_tokens): run the row-independent logits pipeline134# (rep-penalty / per-row temperature / top_p / top_k / softmax / sample) ONCE over the whole135# [B,6,V] step instead of B times on [1,6,V]; only the variable-length box assembly stays per-row.136# Greedy is BIT-IDENTICAL to the per-row san (argmax, no RNG). Default ON; MTP_BATCH_SAN=0 -> per-row.137BATCH_SAN = os.environ.get("MTP_BATCH_SAN", "1") == "1"138 139# Batch the AR repair sampler over [B,1,V]. This shares the exact filtering140# helpers with MTP batching but skips box/ref decoding, so it only replaces the141# repeated stock one-token sample calls. Sampling itself stays row-ordered by142# default to preserve the stock RNG consumption pattern for AR repair.143AR_BATCH_SAN = os.environ.get("AR_BATCH_SAN", "1") == "1"144 145_tok = _proc = _model = None146 147def _magi_diag():148 lines = []149 try:150 import magi_attention151 lines.append(f"magi_attention: OK file={getattr(magi_attention, '__file__', None)}")152 lines.append(f"magi_attention.__version__={getattr(magi_attention, '__version__', '<missing>')}")153 except Exception as e:154 lines.append(f"magi_attention: FAIL {type(e).__name__}: {e}")155 return "\n".join(lines)156 try:157 from magi_attention.functional.flex_flash_attn import flex_flash_attn_func158 lines.append(f"magi_attention.functional.flex_flash_attn: OK func={flex_flash_attn_func}")159 except Exception as e:160 lines.append(f"magi_attention.functional.flex_flash_attn: FAIL {type(e).__name__}: {e}")161 return "\n".join(lines)162 163def _remote_magi_diag(model=None):164 lines = []165 try:166 if model is not None:167 mod = importlib.import_module(type(model.language_model.model).__module__)168 else:169 # Best effort: if the dynamic module is not imported yet this may fail;170 # the post-load diagnostic below will still work.171 mod = importlib.import_module("transformers_modules.LocateAnything-3B.modeling_qwen2")172 lines.append(f"remote_qwen2_module={getattr(mod, '__file__', None)}")173 lines.append(f"remote_qwen2._MAGI_AVAILABLE={getattr(mod, '_MAGI_AVAILABLE', '<missing>')!r}")174 lines.append(f"remote_qwen2.flex_flash_attn_func={getattr(mod, 'flex_flash_attn_func', '<missing>')}")175 except Exception as e:176 lines.append(f"remote_qwen2: diagnostic failed {type(e).__name__}: {e}")177 return "\n".join(lines)178 179def _attn_class_diag(model):180 try:181 llm = model.language_model.model182 classes = [type(layer.self_attn).__name__ for layer in llm.layers[:4]]183 return (184 f"llm._attn_implementation={getattr(llm, '_attn_implementation', None)!r}\n"185 f"config._attn_implementation={getattr(llm.config, '_attn_implementation', None)!r}\n"186 f"first_attn_classes={classes}"187 )188 except Exception as e:189 return f"attention class diagnostic failed {type(e).__name__}: {e}"190 191 192def _set_vision_attention_mode(model):193 """Match HF's MoonViT policy: prefer flash_attention_2, then sdpa, then eager."""194 vm = getattr(model, "vision_model", None)195 if vm is None:196 return None197 mod = importlib.import_module(type(vm).__module__)198 funcs = getattr(mod, "VL_VISION_ATTENTION_FUNCTIONS", {})199 has_flash = getattr(mod, "flash_attn_varlen_func", None) is not None200 requested = VISION_ATTN_MODE201 202 if requested == "auto":203 candidates = ("flash_attention_2", "sdpa", "eager")204 else:205 candidates = (requested, "flash_attention_2", "sdpa", "eager")206 207 chosen = None208 for candidate in candidates:209 if candidate == "flash_attention_2" and not has_flash:210 continue211 if candidate in funcs:212 chosen = candidate213 break214 if chosen is None:215 raise RuntimeError("MoonViT has no supported attention implementation.")216 217 if requested == "flash_attention_2" and chosen != "flash_attention_2":218 warnings.warn("LA_FLASH_VISION_ATTN=flash_attention_2 requested but flash-attn is unavailable; "219 f"using {chosen}.")220 elif requested not in {"auto", chosen}:221 warnings.warn(f"LA_FLASH_VISION_ATTN={requested} is unavailable; using {chosen}.")222 223 if hasattr(model.config, "vision_config"):224 model.config.vision_config._attn_implementation = chosen225 try:226 vm.config._attn_implementation = chosen227 except Exception:228 pass229 try:230 for block in vm.encoder.blocks:231 block.attn_implementation = chosen232 except Exception as exc:233 raise RuntimeError("Failed to configure MoonViT attention implementation.") from exc234 model._la_flash_vision_attn = chosen235 return chosen236 237 238def load():239 """Lazy model load with HF remote-code semantics plus release backends.240 241 The text decoder is pinned to one of sdpa/eager/magi/la_flash. MoonViT is242 configured independently and follows the HF policy: flash_attention_2 when243 flash-attn is importable, otherwise sdpa, otherwise eager.244 """245 global _tok, _proc, _model246 if _model is None:247 _tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)248 _proc = AutoProcessor.from_pretrained(MODEL, trust_remote_code=True)249 attn_impl = REMOTE_ATTN_MODE250 if ATTN_MODE == "magi" and os.environ.get("LA_FLASH_DEBUG", "0") != "0":251 print("LA Flash magi pre-load diagnostic:", flush=True)252 print(_magi_diag(), flush=True)253 _model = AutoModel.from_pretrained(MODEL, torch_dtype=DT, trust_remote_code=True,254 attn_implementation=attn_impl).to(DEV).eval()255 _set_vision_attention_mode(_model)256 actual_attn = getattr(_model.language_model.model, "_attn_implementation", None)257 if ATTN_MODE == "magi" and os.environ.get("LA_FLASH_DEBUG", "0") != "0":258 print("LA Flash magi post-load diagnostic:", flush=True)259 print(_remote_magi_diag(_model), flush=True)260 print(_attn_class_diag(_model), flush=True)261 if ATTN_MODE == "magi":262 try:263 qwen2_mod = importlib.import_module(type(_model.language_model.model).__module__)264 if not getattr(qwen2_mod, "_MAGI_AVAILABLE", False):265 raise RuntimeError(266 "remote module reports _MAGI_AVAILABLE=False.\n"267 f"{_remote_magi_diag(_model)}\n{_magi_diag()}"268 )269 first_attn = type(_model.language_model.model.layers[0].self_attn).__name__270 if actual_attn != "sdpa" or first_attn != "_BatchedMagiAttention":271 _set_llm_mode(_model, "magi")272 actual_attn = getattr(_model.language_model.model, "_attn_implementation", None)273 first_attn = type(_model.language_model.model.layers[0].self_attn).__name__274 if os.environ.get("LA_FLASH_DEBUG", "0") != "0":275 print("LA Flash magi post-swap diagnostic:", flush=True)276 print(_attn_class_diag(_model), flush=True)277 if actual_attn != "sdpa" or first_attn != "_BatchedMagiAttention":278 raise RuntimeError(279 "batched magi attention did not activate. "280 f"actual_attn={actual_attn!r}; first_attn={first_attn!r}; "281 f"{_remote_magi_diag(_model)}; {_attn_class_diag(_model)}"282 )283 _model._la_flash_requested_attn = "magi"284 except Exception as exc:285 _fallback_to_sdpa(_model, "magi", exc)286 else:287 try:288 _set_llm_mode(_model, ATTN_MODE) # decode-safe mask plumbing for sdpa/eager/la_flash289 except Exception as exc:290 _fallback_to_sdpa(_model, ATTN_MODE, exc)291 if MTP_COMPILE:292 _maybe_compile(_model)293 return _tok, _proc, _model294 295 296def _maybe_compile(model):297 """Compile the shared Qwen2Model core (base.forward). It backs BOTH prefill (called directly)298 and decode (language_model.forward -> self.model). lm_head + MoonViT left eager. dynamic=True299 so the varying decode S/kvlen don't trigger a recompile storm. No-op + warning if triton is300 missing (inductor needs it on GPU). First call pays the compile cost (~42s warm / ~187s cold)."""301 try:302 import triton # noqa: F401303 except Exception:304 warnings.warn("MTP_COMPILE set but triton is unavailable; running without torch.compile.")305 return306 import torch._dynamo as _dyn307 _dyn.config.cache_size_limit = max(_dyn.config.cache_size_limit, 64)308 base = model.language_model.model309 if not getattr(base, "_mtp_compiled", False):310 base.forward = torch.compile(base.forward, dynamic=True)311 base._mtp_compiled = True312 313 314def build_batched_magi_attention_class(mod):315 """Build a Qwen2 attention subclass backed by Magi's flex_flash_attn.316 317 The official LocateAnything ``Qwen2MagiAttention`` asserts ``bsz == 1`` and318 relies on ``Qwen2Model._attn_implementation == "magi"`` to build a single319 sample range plan. For release batch inference the hybrid scheduler passes320 a batched Magi range plan directly to this layer; a 4D-mask conversion path321 remains as a compatibility fallback.322 """323 flex_flash_attn_func = getattr(mod, "flex_flash_attn_func", None)324 if flex_flash_attn_func is None:325 try:326 from magi_attention.functional.flex_flash_attn import flex_flash_attn_func327 except Exception as exc:328 raise RuntimeError(329 "LA_FLASH_ATTN=magi requires "330 "magi_attention.functional.flex_flash_attn.flex_flash_attn_func."331 ) from exc332 333 FULL, CAUSAL = 0, 1334 causal_plan_cache = {}335 try:336 magi_params = set(inspect.signature(flex_flash_attn_func).parameters)337 except (TypeError, ValueError):338 magi_params = set()339 supports_disable_fwd_atomic = "disable_fwd_atomic_reduction" in magi_params340 341 def _disjoint_q_ranges(q_ranges):342 seen = set()343 for start, end in q_ranges:344 key = (int(start), int(end))345 if key in seen:346 return False347 seen.add(key)348 return True349 350 def _plan_disjoint_q_ranges(plan):351 cached = plan.get("_la_flash_disjoint_q_ranges")352 if cached is not None:353 return bool(cached)354 q_ranges = plan["q_ranges"].detach().to(device="cpu", dtype=torch.int32).tolist()355 disjoint = _disjoint_q_ranges(q_ranges)356 try:357 plan["_la_flash_disjoint_q_ranges"] = disjoint358 except Exception:359 pass360 return disjoint361 362 def _tensor_plan(q_ranges, k_ranges, types, device):363 return {364 "q_ranges": torch.tensor(q_ranges, dtype=torch.int32, device=device).contiguous(),365 "k_ranges": torch.tensor(k_ranges, dtype=torch.int32, device=device).contiguous(),366 "attn_type_map": torch.tensor(types, dtype=torch.int32, device=device).contiguous(),367 "_la_flash_disjoint_q_ranges": _disjoint_q_ranges(q_ranges),368 }369 370 def _offset_plan(plan, q_offset, k_offset):371 return (372 (plan["q_ranges"] + int(q_offset)).tolist(),373 (plan["k_ranges"] + int(k_offset)).tolist(),374 plan["attn_type_map"].tolist(),375 )376 377 def _causal_plan(bsz, q_len, kv_seq_len, device):378 key = (int(bsz), int(q_len), int(kv_seq_len), device.type, device.index)379 cached = causal_plan_cache.get(key)380 if cached is not None:381 return cached382 q_ranges, k_ranges, types = [], [], []383 for b in range(int(bsz)):384 q_base = b * int(q_len)385 k_base = b * int(kv_seq_len)386 q_ranges.append([q_base, q_base + int(q_len)])387 k_ranges.append([k_base, k_base + int(kv_seq_len)])388 types.append(CAUSAL)389 plan = _tensor_plan(q_ranges, k_ranges, types, device)390 plan.update(391 {392 "flash_cu_seqlens_q": torch.arange(393 0,394 (int(bsz) + 1) * int(q_len),395 int(q_len),396 dtype=torch.int32,397 device=device,398 ),399 "flash_cu_seqlens_k": torch.arange(400 0,401 (int(bsz) + 1) * int(kv_seq_len),402 int(kv_seq_len),403 dtype=torch.int32,404 device=device,405 ),406 "flash_causal": True,407 }408 )409 causal_plan_cache[key] = plan410 return plan411 412 def _row_segments(row):413 idx = np.flatnonzero(row)414 if idx.size == 0:415 return ((0, 1),)416 split = np.flatnonzero(np.diff(idx) > 1) + 1417 starts = np.concatenate((idx[:1], idx[split]))418 ends = np.concatenate((idx[split - 1], idx[-1:])) + 1419 return tuple((int(s), int(e)) for s, e in zip(starts, ends))420 421 def _visible_from_4d_mask(attention_mask, kv_seq_len):422 mask = attention_mask[:, :, :, :kv_seq_len]423 if mask.dtype == torch.bool:424 return mask[:, 0].detach().to(device="cpu", dtype=torch.bool).contiguous()425 mask_cpu = mask[:, 0].detach().to(device="cpu").contiguous()426 if getattr(attention_mask, "_la_flash_visible_mask", False):427 return (mask_cpu > 0).to(dtype=torch.bool)428 429 max_value = float(mask_cpu.max().item()) if mask_cpu.numel() else 0.0430 min_value = float(mask_cpu.min().item()) if mask_cpu.numel() else 0.0431 if max_value > 0.0 and min_value >= 0.0:432 return (mask_cpu > 0).to(dtype=torch.bool)433 return (mask_cpu >= 0).to(dtype=torch.bool)434 435 def _plan_from_visible_mask(attention_mask, bsz, q_len, kv_seq_len, device):436 cache_key = (int(bsz), int(q_len), int(kv_seq_len), device.type, device.index)437 cached = getattr(attention_mask, "_la_flash_magi_plan", None)438 if cached is not None and cached[0] == cache_key:439 return cached[1]440 441 visible = _visible_from_4d_mask(attention_mask, int(kv_seq_len)).numpy()442 q_ranges, k_ranges, types = [], [], []443 for b in range(int(bsz)):444 q_base = b * int(q_len)445 k_base = b * int(kv_seq_len)446 run_start = 0447 run_segments = _row_segments(visible[b, 0])448 for q in range(1, int(q_len)):449 segments = _row_segments(visible[b, q])450 if segments == run_segments:451 continue452 for start, end in run_segments:453 q_ranges.append([q_base + run_start, q_base + q])454 k_ranges.append([k_base + start, k_base + end])455 types.append(FULL)456 run_start = q457 run_segments = segments458 for start, end in run_segments:459 q_ranges.append([q_base + run_start, q_base + int(q_len)])460 k_ranges.append([k_base + start, k_base + end])461 types.append(FULL)462 463 plan = _tensor_plan(q_ranges, k_ranges, types, device)464 try:465 attention_mask._la_flash_magi_plan = (cache_key, plan)466 except Exception:467 pass468 return plan469 470 def _plan_from_magi_dict(attention_mask, bsz, q_len, kv_seq_len, device):471 if int(bsz) == 1:472 return attention_mask473 q_ranges, k_ranges, types = [], [], []474 for b in range(int(bsz)):475 qs, ks, ts = _offset_plan(476 attention_mask,477 q_offset=b * int(q_len),478 k_offset=b * int(kv_seq_len),479 )480 q_ranges.extend(qs)481 k_ranges.extend(ks)482 types.extend(ts)483 return _tensor_plan(q_ranges, k_ranges, types, device)484 485 def _magi_plan(attention_mask, bsz, q_len, kv_seq_len, device):486 if isinstance(attention_mask, dict):487 if attention_mask.get("_la_flash_batched", False):488 return attention_mask489 return _plan_from_magi_dict(attention_mask, bsz, q_len, kv_seq_len, device)490 if attention_mask is None:491 return _causal_plan(bsz, q_len, kv_seq_len, device)492 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):493 raise ValueError(494 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, "495 f"but is {attention_mask.size()}"496 )497 return _plan_from_visible_mask(attention_mask, bsz, q_len, kv_seq_len, device)498 499 class _BatchedMagiAttention(mod.Qwen2Attention):500 """MagiAttention path with true batch inference via packed token ranges."""501 502 def forward(503 self,504 hidden_states: torch.Tensor,505 attention_mask=None,506 position_ids=None,507 past_key_value=None,508 output_attentions=False,509 use_cache=False,510 **kwargs,511 ):512 if output_attentions:513 raise NotImplementedError("MagiAttention does not support output_attentions=True")514 515 bsz, q_len, _ = hidden_states.size()516 query_states = self.q_proj(hidden_states)517 key_states = self.k_proj(hidden_states)518 value_states = self.v_proj(hidden_states)519 520 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)521 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)522 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)523 524 kv_seq_len = key_states.shape[-2]525 if past_key_value is not None:526 if self.layer_idx is None:527 raise ValueError(528 f"The cache structure has changed since version v4.36. If you are using "529 f"{self.__class__.__name__} for auto-regressive decoding with k/v caching, "530 "please initialize the attention class with a layer index."531 )532 kv_seq_len += past_key_value.get_seq_length(self.layer_idx)533 534 cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)535 query_states, key_states = mod.apply_rotary_pos_emb(536 query_states, key_states, cos, sin, position_ids)537 538 if past_key_value is not None:539 cache_kwargs = {"sin": sin, "cos": cos}540 key_states, value_states = past_key_value.update(541 key_states, value_states, self.layer_idx, cache_kwargs)542 543 kv_seq_len = key_states.shape[-2]544 plan = _magi_plan(attention_mask, bsz, q_len, kv_seq_len, query_states.device)545 magi_extra_kwargs = {}546 if supports_disable_fwd_atomic:547 magi_extra_kwargs["disable_fwd_atomic_reduction"] = (548 (not self.training) and _plan_disjoint_q_ranges(plan)549 )550 551 query_states = query_states.transpose(1, 2).reshape(552 bsz * q_len, self.num_heads, self.head_dim).contiguous()553 key_states = key_states.transpose(1, 2).reshape(554 bsz * kv_seq_len, self.num_key_value_heads, self.head_dim).contiguous()555 value_states = value_states.transpose(1, 2).reshape(556 bsz * kv_seq_len, self.num_key_value_heads, self.head_dim).contiguous()557 558 attn_output, _ = flex_flash_attn_func(559 query_states,560 key_states,561 value_states,562 q_ranges=plan["q_ranges"],563 k_ranges=plan["k_ranges"],564 attn_type_map=plan["attn_type_map"],565 softmax_scale=getattr(self, "softmax_scale", self.head_dim ** -0.5),566 softcap=0.0,567 deterministic=False,568 **magi_extra_kwargs,569 )570 attn_output = attn_output.view(bsz, q_len, self.hidden_size)571 attn_output = self.o_proj(attn_output)572 return attn_output, None, past_key_value573 574 return _BatchedMagiAttention575 576 577def build_la_flash_attention_class(mod):578 """Build a Qwen2 attention subclass backed by LA Flash sparse ranges."""579 try:580 from kernel_utils import is_available, range_attention581 except Exception as exc:582 raise RuntimeError(583 "LA_FLASH_ATTN=la_flash requires kernel_utils and FlashAttention."584 ) from exc585 if not is_available():586 raise RuntimeError(587 "LA_FLASH_ATTN=la_flash requires flash_attn.flash_attn_varlen_func."588 )589 590 FULL, CAUSAL = 0, 1591 causal_plan_cache = {}592 593 def _tensor_plan(q_ranges, k_ranges, types, device):594 max_q_len = max((int(end) - int(start) for start, end in q_ranges), default=0)595 max_k_len = max((int(end) - int(start) for start, end in k_ranges), default=0)596 plan = {597 "q_ranges": torch.tensor(q_ranges, dtype=torch.int32, device=device).contiguous(),598 "k_ranges": torch.tensor(k_ranges, dtype=torch.int32, device=device).contiguous(),599 "attn_type_map": torch.tensor(types, dtype=torch.int32, device=device).contiguous(),600 "max_q_len": max_q_len,601 "max_k_len": max_k_len,602 }603 plan.update(_la_flash_group_plan_tensors(q_ranges, types, device))604 return plan605 606 def _offset_plan(plan, q_offset, k_offset):607 return (608 (plan["q_ranges"] + int(q_offset)).tolist(),609 (plan["k_ranges"] + int(k_offset)).tolist(),610 plan["attn_type_map"].tolist(),611 )612 613 def _causal_plan(bsz, q_len, kv_seq_len, device):614 key = (int(bsz), int(q_len), int(kv_seq_len), device.type, device.index)615 cached = causal_plan_cache.get(key)616 if cached is not None:617 return cached618 q_ranges, k_ranges, types = [], [], []619 for b in range(int(bsz)):620 q_base = b * int(q_len)621 k_base = b * int(kv_seq_len)622 q_ranges.append([q_base, q_base + int(q_len)])623 k_ranges.append([k_base, k_base + int(kv_seq_len)])624 types.append(CAUSAL)625 plan = _tensor_plan(q_ranges, k_ranges, types, device)626 plan.update(627 {628 "flash_cu_seqlens_q": torch.arange(629 0,630 (int(bsz) + 1) * int(q_len),631 int(q_len),632 dtype=torch.int32,633 device=device,634 ),635 "flash_cu_seqlens_k": torch.arange(636 0,637 (int(bsz) + 1) * int(kv_seq_len),638 int(kv_seq_len),639 dtype=torch.int32,640 device=device,641 ),642 "flash_causal": True,643 }644 )645 causal_plan_cache[key] = plan646 return plan647 648 def _row_segments(row):649 idx = np.flatnonzero(row)650 if idx.size == 0:651 return ((0, 1),)652 split = np.flatnonzero(np.diff(idx) > 1) + 1653 starts = np.concatenate((idx[:1], idx[split]))654 ends = np.concatenate((idx[split - 1], idx[-1:])) + 1655 return tuple((int(s), int(e)) for s, e in zip(starts, ends))656 657 def _visible_from_4d_mask(attention_mask, kv_seq_len):658 mask = attention_mask[:, :, :, :kv_seq_len]659 if mask.dtype == torch.bool:660 return mask[:, 0].detach().to(device="cpu", dtype=torch.bool).contiguous()661 mask_cpu = mask[:, 0].detach().to(device="cpu").contiguous()662 if getattr(attention_mask, "_la_flash_visible_mask", False):663 return (mask_cpu > 0).to(dtype=torch.bool)664 665 max_value = float(mask_cpu.max().item()) if mask_cpu.numel() else 0.0666 min_value = float(mask_cpu.min().item()) if mask_cpu.numel() else 0.0667 if max_value > 0.0 and min_value >= 0.0:668 return (mask_cpu > 0).to(dtype=torch.bool)669 return (mask_cpu >= 0).to(dtype=torch.bool)670 671 def _prefix_len(row):672 idx = np.flatnonzero(row)673 if idx.size == 0:674 return None675 end = int(idx[-1]) + 1676 if not bool(row[:end].all()) or bool(row[end:].any()):677 return None678 return end679 680 def _causal_plan_from_visible(visible, bsz, q_len, kv_seq_len, device):681 q_ranges, k_ranges, types = [], [], []682 packed_flash = True683 for b in range(int(bsz)):684 first_len = _prefix_len(visible[b, 0])685 if first_len is None:686 return None687 valid_len = int(first_len) + int(q_len) - 1688 if valid_len < int(q_len) or valid_len > int(kv_seq_len):689 return None690 for q in range(int(q_len)):691 row_len = _prefix_len(visible[b, q])692 expected = valid_len - int(q_len) + q + 1693 if row_len != expected:694 return None695 q_base = b * int(q_len)696 k_base = b * int(kv_seq_len)697 q_ranges.append([q_base, q_base + int(q_len)])698 k_ranges.append([k_base, k_base + valid_len])699 types.append(CAUSAL)700 packed_flash = packed_flash and valid_len == int(kv_seq_len)701 702 plan = _tensor_plan(q_ranges, k_ranges, types, device)703 plan["_la_flash_disjoint_q_ranges"] = True704 if packed_flash:705 plan.update(706 {707 "flash_cu_seqlens_q": torch.arange(708 0,709 (int(bsz) + 1) * int(q_len),710 int(q_len),711 dtype=torch.int32,712 device=device,713 ),714 "flash_cu_seqlens_k": torch.arange(715 0,716 (int(bsz) + 1) * int(kv_seq_len),717 int(kv_seq_len),718 dtype=torch.int32,719 device=device,720 ),721 "flash_causal": True,722 }723 )724 return plan725 726 def _plan_from_visible_mask(attention_mask, bsz, q_len, kv_seq_len, device):727 cache_key = (int(bsz), int(q_len), int(kv_seq_len), device.type, device.index, "la_flash")728 cached = getattr(attention_mask, "_la_flash_range_plan", None)729 if cached is not None and cached[0] == cache_key:730 return cached[1]731 732 visible = _visible_from_4d_mask(attention_mask, int(kv_seq_len)).numpy()733 plan = _causal_plan_from_visible(visible, bsz, q_len, kv_seq_len, device)734 if plan is not None:735 try:736 attention_mask._la_flash_range_plan = (cache_key, plan)737 except Exception:738 pass739 return plan740 741 q_ranges, k_ranges, types = [], [], []742 for b in range(int(bsz)):743 q_base = b * int(q_len)744 k_base = b * int(kv_seq_len)745 run_start = 0746 run_segments = _row_segments(visible[b, 0])747 for q in range(1, int(q_len)):748 segments = _row_segments(visible[b, q])749 if segments == run_segments:750 continue751 for start, end in run_segments:752 q_ranges.append([q_base + run_start, q_base + q])753 k_ranges.append([k_base + start, k_base + end])754 types.append(FULL)755 run_start = q756 run_segments = segments757 for start, end in run_segments:758 q_ranges.append([q_base + run_start, q_base + int(q_len)])759 k_ranges.append([k_base + start, k_base + end])760 types.append(FULL)761 762 plan = _tensor_plan(q_ranges, k_ranges, types, device)763 try:764 attention_mask._la_flash_range_plan = (cache_key, plan)765 except Exception:766 pass767 return plan768 769 def _plan_from_magi_dict(attention_mask, bsz, q_len, kv_seq_len, device):770 if int(bsz) == 1:771 return attention_mask772 q_ranges, k_ranges, types = [], [], []773 for b in range(int(bsz)):774 qs, ks, ts = _offset_plan(775 attention_mask,776 q_offset=b * int(q_len),777 k_offset=b * int(kv_seq_len),778 )779 q_ranges.extend(qs)780 k_ranges.extend(ks)781 types.extend(ts)782 return _tensor_plan(q_ranges, k_ranges, types, device)783 784 def _range_plan(attention_mask, bsz, q_len, kv_seq_len, device):785 if isinstance(attention_mask, dict):786 if attention_mask.get("_la_flash_batched", False):787 return attention_mask788 return _plan_from_magi_dict(attention_mask, bsz, q_len, kv_seq_len, device)789 if attention_mask is None:790 return _causal_plan(bsz, q_len, kv_seq_len, device)791 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):792 raise ValueError(793 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, "794 f"but is {attention_mask.size()}"795 )796 return _plan_from_visible_mask(attention_mask, bsz, q_len, kv_seq_len, device)797 798 class _LaFlashAttention(mod.Qwen2Attention):799 """Range-plan attention path backed by FlashAttention sparse ranges."""800 801 def forward(802 self,803 hidden_states: torch.Tensor,804 attention_mask=None,805 position_ids=None,806 past_key_value=None,807 output_attentions=False,808 use_cache=False,809 **kwargs,810 ):811 if output_attentions:812 raise NotImplementedError("LA Flash attention does not support output_attentions=True")813 814 bsz, q_len, _ = hidden_states.size()815 query_states = self.q_proj(hidden_states)816 key_states = self.k_proj(hidden_states)817 value_states = self.v_proj(hidden_states)818 819 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)820 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)821 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)822 823 kv_seq_len = key_states.shape[-2]824 if past_key_value is not None:825 if self.layer_idx is None:826 raise ValueError(827 f"The cache structure has changed since version v4.36. If you are using "828 f"{self.__class__.__name__} for auto-regressive decoding with k/v caching, "829 "please initialize the attention class with a layer index."830 )831 kv_seq_len += past_key_value.get_seq_length(self.layer_idx)832 833 cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)834 query_states, key_states = mod.apply_rotary_pos_emb(835 query_states, key_states, cos, sin, position_ids)836 837 if past_key_value is not None:838 cache_kwargs = {"sin": sin, "cos": cos}839 key_states, value_states = past_key_value.update(840 key_states, value_states, self.layer_idx, cache_kwargs)841 842 kv_seq_len = key_states.shape[-2]843 dense_backend = os.environ.get("LA_FLASH_DENSE_BACKEND", "sdpa").strip().lower()844 if dense_backend == "sdpa" and not isinstance(attention_mask, dict):845 dense_key_states = mod.repeat_kv(key_states, self.num_key_value_groups)846 dense_value_states = mod.repeat_kv(value_states, self.num_key_value_groups)847 if attention_mask is not None:848 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):849 raise ValueError(850 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, "851 f"but is {attention_mask.size()}"852 )853 query_for_sdpa = query_states.contiguous()854 key_for_sdpa = dense_key_states.contiguous()855 value_for_sdpa = dense_value_states.contiguous()856 is_causal = False857 elif past_key_value is None:858 query_for_sdpa = query_states859 key_for_sdpa = dense_key_states860 value_for_sdpa = dense_value_states861 is_causal = bool(self.is_causal and q_len > 1)862 else:863 query_for_sdpa = key_for_sdpa = value_for_sdpa = None864 is_causal = False865 if query_for_sdpa is not None:866 attn_output = torch.nn.functional.scaled_dot_product_attention(867 query_for_sdpa,868 key_for_sdpa,869 value_for_sdpa,870 attn_mask=attention_mask,871 dropout_p=self.attention_dropout if self.training else 0.0,872 is_causal=is_causal,873 )874 attn_output = attn_output.transpose(1, 2).contiguous()875 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)876 attn_output = self.o_proj(attn_output)877 return attn_output, None, past_key_value878 879 plan = _range_plan(attention_mask, bsz, q_len, kv_seq_len, query_states.device)880 881 query_states = query_states.transpose(1, 2).reshape(882 bsz * q_len, self.num_heads, self.head_dim).contiguous()883 key_states = key_states.transpose(1, 2).reshape(884 bsz * kv_seq_len, self.num_key_value_heads, self.head_dim).contiguous()885 value_states = value_states.transpose(1, 2).reshape(886 bsz * kv_seq_len, self.num_key_value_heads, self.head_dim).contiguous()887 888 attn_output = range_attention(889 query_states,890 key_states,891 value_states,892 plan["q_ranges"],893 plan["k_ranges"],894 plan["attn_type_map"],895 getattr(self, "softmax_scale", self.head_dim ** -0.5),896 segment_offsets=plan.get("segment_offsets"),897 group_q_ranges=plan.get("group_q_ranges"),898 group_attn_type_map=plan.get("group_attn_type_map"),899 max_q_len=plan.get("max_q_len"),900 max_k_len=plan.get("max_k_len"),901 flash_cu_seqlens_q=plan.get("flash_cu_seqlens_q"),902 flash_cu_seqlens_k=plan.get("flash_cu_seqlens_k"),903 flash_causal=plan.get("flash_causal"),904 disjoint_q_ranges=plan.get("_la_flash_disjoint_q_ranges"),905 )906 attn_output = attn_output.view(bsz, q_len, self.hidden_size)907 attn_output = self.o_proj(attn_output)908 return attn_output, None, past_key_value909 910 return _LaFlashAttention911 912 913def _is_magi_plan(obj):914 return isinstance(obj, dict) and {915 "q_ranges",916 "k_ranges",917 "attn_type_map",918 }.issubset(obj.keys())919 920 921def _la_flash_group_plan_tensors(q_ranges, types, device):922 """Group consecutive Magi range entries that share the same query span.923 924 Magi-style plans may represent one query span with multiple disjoint key925 spans. LA Flash consumes those as one FlashAttention-backed softmax group.926 """927 if not q_ranges:928 return {929 "group_q_ranges": torch.empty((0, 2), dtype=torch.int32, device=device),930 "segment_offsets": torch.zeros((1,), dtype=torch.int32, device=device),931 "group_attn_type_map": torch.empty((0,), dtype=torch.int32, device=device),932 }933 934 grouped_q, grouped_types, offsets = [], [], [0]935 last_q = None936 last_type = None937 for idx, (q_range, attn_type) in enumerate(zip(q_ranges, types)):938 key = (int(q_range[0]), int(q_range[1]))939 attn_type = int(attn_type)940 if last_q is None:941 grouped_q.append([key[0], key[1]])942 grouped_types.append(attn_type)943 last_q = key944 last_type = attn_type945 continue946 if key == last_q and attn_type == last_type:947 continue948 offsets.append(idx)949 grouped_q.append([key[0], key[1]])950 grouped_types.append(attn_type)951 last_q = key952 last_type = attn_type953 offsets.append(len(q_ranges))954 955 return {956 "group_q_ranges": torch.tensor(grouped_q, dtype=torch.int32, device=device).contiguous(),957 "segment_offsets": torch.tensor(offsets, dtype=torch.int32, device=device).contiguous(),958 "group_attn_type_map": torch.tensor(grouped_types, dtype=torch.int32, device=device).contiguous(),959 "max_q_len": max((end - start for start, end in grouped_q), default=0),960 }961 962 963def _record_sparse_plan_stats(model, q_ranges, k_ranges, types):964 if os.environ.get("LA_FLASH_PLAN_STATS", "0") != "1":965 return966 stats = getattr(model, "_la_flash_sparse_plan_stats", None)967 if stats is None:968 stats = {969 "calls": 0,970 "ranges": 0,971 "q_tokens": 0,972 "k_tokens": 0,973 "max_q_len": 0,974 "max_k_len": 0,975 "full_ranges": 0,976 "causal_ranges": 0,977 "other_ranges": 0,978 }979 model._la_flash_sparse_plan_stats = stats980 stats["calls"] += 1981 stats["ranges"] += len(q_ranges)982 for (q_start, q_end), (k_start, k_end), attn_type in zip(q_ranges, k_ranges, types):983 q_len = int(q_end) - int(q_start)984 k_len = int(k_end) - int(k_start)985 stats["q_tokens"] += q_len986 stats["k_tokens"] += k_len987 stats["max_q_len"] = max(stats["max_q_len"], q_len)988 stats["max_k_len"] = max(stats["max_k_len"], k_len)989 attn_type = int(attn_type)990 if attn_type == 0:991 stats["full_ranges"] += 1992 elif attn_type == 1:993 stats["causal_ranges"] += 1994 else:995 stats["other_ranges"] += 1996 997 998def build_magi_scheduler_ranges(model, attention_mask_2d, input_ids, past_len, mtp_window=False):999 """Build batched Magi ranges directly from the hybrid scheduler mask.1000 1001 The official Qwen2 SDPA dispatcher may optimize an all-valid 2D mask to1002 ``None`` before decoder layers see it. That is correct for plain causal1003 attention but loses LocateAnything's MTP generation-window rule. Building1004 ranges here keeps Magi batch inference exact and avoids per-layer dense1005 mask conversion.1006 """1007 requested_attn = getattr(model, "_la_flash_requested_attn", ATTN_MODE)1008 if requested_attn not in {"magi", "la_flash"}:1009 return None1010 if attention_mask_2d is None or not hasattr(attention_mask_2d, "dim") or attention_mask_2d.dim() != 2:1011 return None1012 1013 bsz, q_len = int(input_ids.shape[0]), int(input_ids.shape[1])1014 key_len = int(attention_mask_2d.shape[1])1015 dev = input_ids.device1016 llm = model.language_model.model1017 block = int(getattr(llm, "block_size", N_FUTURE))1018 causal_attn = bool(getattr(llm, "causal_attn", False))1019 use_mtp_window = bool(mtp_window and q_len >= block and key_len >= block)1020 q0 = max(0, q_len - block)1021 k0 = max(0, key_len - block)1022 blocked_k = k0 - 11023 past_len = int(past_len)1024 1025 key_valid = attention_mask_2d.detach().to(device="cpu", dtype=torch.bool).contiguous().numpy()1026 key_idx = np.arange(key_len)1027 q_ranges, k_ranges, types = [], [], []1028 if not use_mtp_window:1029 causal_q_ranges, causal_k_ranges, causal_types = [], [], []1030 causal_fast_path = True1031 packed_flash = True1032 for b in range(bsz):1033 valid = np.flatnonzero(key_valid[b])1034 if valid.size == 0:1035 causal_fast_path = False1036 break1037 valid_len = int(valid[-1]) + 11038 if valid_len < q_len or not bool(key_valid[b, :valid_len].all()) or bool(key_valid[b, valid_len:].any()):1039 causal_fast_path = False1040 break1041 packed_flash = packed_flash and valid_len == key_len1042 q_base = b * q_len1043 k_base = b * key_len1044 causal_q_ranges.append([q_base, q_base + q_len])1045 causal_k_ranges.append([k_base, k_base + valid_len])1046 causal_types.append(1)1047 if causal_fast_path:1048 plan = {1049 "q_ranges": torch.tensor(causal_q_ranges, dtype=torch.int32, device=dev).contiguous(),1050 "k_ranges": torch.tensor(causal_k_ranges, dtype=torch.int32, device=dev).contiguous(),1051 "attn_type_map": torch.tensor(causal_types, dtype=torch.int32, device=dev).contiguous(),1052 "max_q_len": q_len,1053 "max_k_len": max((end - start for start, end in causal_k_ranges), default=0),1054 "_la_flash_batched": True,1055 "_la_flash_disjoint_q_ranges": True,1056 }1057 if packed_flash:1058 plan.update(1059 {1060 "flash_cu_seqlens_q": torch.arange(1061 0,1062 (bsz + 1) * q_len,1063 q_len,1064 dtype=torch.int32,1065 device=dev,1066 ),1067 "flash_cu_seqlens_k": torch.arange(1068 0,1069 (bsz + 1) * key_len,1070 key_len,1071 dtype=torch.int32,1072 device=dev,1073 ),1074 "flash_causal": True,1075 }1076 )1077 plan.update(_la_flash_group_plan_tensors(causal_q_ranges, causal_types, dev))1078 _record_sparse_plan_stats(model, causal_q_ranges, causal_k_ranges, causal_types)1079 return plan1080 1081 def row_segments(row):1082 idx = np.flatnonzero(row)1083 if idx.size == 0:1084 return ((0, 1),)1085 split = np.flatnonzero(np.diff(idx) > 1) + 11086 starts = np.concatenate((idx[:1], idx[split]))1087 ends = np.concatenate((idx[split - 1], idx[-1:])) + 11088 return tuple((int(s), int(e)) for s, e in zip(starts, ends))1089 1090 for b in range(bsz):1091 q_base = b * q_len1092 k_base = b * key_len1093 run_start = 01094 run_segments = None1095 if use_mtp_window and not causal_attn:1096 prefix_q_len = q01097 prefix_k_end = past_len + prefix_q_len1098 prefix_ok = (1099 prefix_q_len > 01100 and prefix_k_end <= key_len1101 and bool(key_valid[b, :prefix_k_end].all())1102 )1103 window_prefix_ok = blocked_k <= 0 or bool(key_valid[b, :blocked_k].all())1104 window_ok = bool(key_valid[b, k0:key_len].all())1105 if prefix_ok:1106 q_ranges.append([q_base, q_base + prefix_q_len])1107 k_ranges.append([k_base, k_base + prefix_k_end])1108 types.append(1)1109 run_start = prefix_q_len1110 if run_start == prefix_q_len and prefix_q_len < q_len and window_prefix_ok and window_ok:1111 if blocked_k > 0:1112 q_ranges.append([q_base + prefix_q_len, q_base + q_len])1113 k_ranges.append([k_base, k_base + blocked_k])1114 types.append(0)1115 q_ranges.append([q_base + prefix_q_len, q_base + q_len])1116 k_ranges.append([k_base + k0, k_base + key_len])1117 types.append(0)1118 continue1119 1120 for q in range(run_start, q_len):1121 visible = key_valid[b] & (key_idx <= q + past_len)1122 if use_mtp_window and q >= q0:1123 if not causal_attn:1124 visible = visible.copy()1125 visible[k0:key_len] = key_valid[b, k0:key_len]1126 if blocked_k >= 0:1127 if visible.base is None:1128 visible[blocked_k] = False1129 else:1130 visible = visible.copy()1131 visible[blocked_k] = False1132 segments = row_segments(visible)1133 if run_segments is None:1134 run_segments = segments1135 continue1136 if segments == run_segments:1137 continue1138 for start, end in run_segments:1139 q_ranges.append([q_base + run_start, q_base + q])1140 k_ranges.append([k_base + start, k_base + end])1141 types.append(0)1142 run_start = q1143 run_segments = segments1144 for start, end in run_segments:1145 q_ranges.append([q_base + run_start, q_base + q_len])1146 k_ranges.append([k_base + start, k_base + end])1147 types.append(0)1148 1149 seen_q_ranges = set()1150 disjoint_q_ranges = True1151 for start, end in q_ranges:1152 key = (int(start), int(end))1153 if key in seen_q_ranges:1154 disjoint_q_ranges = False1155 break1156 seen_q_ranges.add(key)1157 1158 plan = {1159 "q_ranges": torch.tensor(q_ranges, dtype=torch.int32, device=dev).contiguous(),1160 "k_ranges": torch.tensor(k_ranges, dtype=torch.int32, device=dev).contiguous(),1161 "attn_type_map": torch.tensor(types, dtype=torch.int32, device=dev).contiguous(),1162 "max_q_len": max((end - start for start, end in q_ranges), default=0),1163 "max_k_len": max((end - start for start, end in k_ranges), default=0),1164 "_la_flash_batched": True,1165 "_la_flash_disjoint_q_ranges": disjoint_q_ranges,1166 }1167 plan.update(_la_flash_group_plan_tensors(q_ranges, types, dev))1168 _record_sparse_plan_stats(model, q_ranges, k_ranges, types)1169 return plan1170 1171 1172def _direct_base_forward(1173 base,1174 input_ids=None,1175 visual_features=None,1176 image_token_index=None,1177 attention_mask=None,1178 position_ids=None,1179 past_key_values=None,1180 inputs_embeds=None,1181 use_cache=None,1182 output_attentions=None,1183 output_hidden_states=None,1184 return_dict=None,1185):1186 mod = importlib.import_module(type(base).__module__)1187 output_attentions = output_attentions if output_attentions is not None else base.config.output_attentions1188 output_hidden_states = (1189 output_hidden_states if output_hidden_states is not None else base.config.output_hidden_states1190 )1191 use_cache = use_cache if use_cache is not None else base.config.use_cache1192 1193 if input_ids is not None and inputs_embeds is not None:1194 raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")1195 if input_ids is not None:1196 batch_size, seq_length = input_ids.shape1197 elif inputs_embeds is not None:1198 batch_size, seq_length, _ = inputs_embeds.shape1199 else:1200 raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")