From cea79502ba65784f741bbab0ef372e89b3e92d2c Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Fri, 25 Sep 2026 00:24:58 +0800 Subject: [PATCH 1/3] fix(hi3): fail fast when the AR KV cache cannot be built `HunyuanImage3ARStep._build_kv_cache` returned None when the upstream `HunyuanStaticCache` class was missing, the transformer had no config, or the constructor raised. Upstream `HunyuanImage3ForCausalMM` never creates a cache on its own, so a None cache made prefill correct but every decode step attend only to its own single token, silently producing context-free rollouts. Build the cache directly and let the original error propagate. --- unirl/models/hunyuan_image3/ar.py | 25 ++++++++----------------- 1 file changed, 8 insertions(+), 17 deletions(-) diff --git a/unirl/models/hunyuan_image3/ar.py b/unirl/models/hunyuan_image3/ar.py index 000e78b55..94a306b62 100644 --- a/unirl/models/hunyuan_image3/ar.py +++ b/unirl/models/hunyuan_image3/ar.py @@ -219,23 +219,14 @@ def _build_kv_cache(transformer, *, batch_size: int, max_cache_len: int): """Pre-build a ``HunyuanStaticCache`` for the AR loop.""" import sys as _sys - upstream_mod = _sys.modules.get(type(transformer).__module__) - cache_cls = getattr(upstream_mod, "HunyuanStaticCache", None) - if cache_cls is None: - return None - config = getattr(transformer, "config", None) - if config is None: - return None - try: - return cache_cls( - config=config, - batch_size=batch_size, - max_cache_len=max_cache_len, - dtype=torch.bfloat16, - dynamic=True, - ) - except Exception: # noqa: BLE001 -- fall back to HF default cache - return None + cache_cls = _sys.modules[type(transformer).__module__].HunyuanStaticCache + return cache_cls( + config=transformer.config, + batch_size=batch_size, + max_cache_len=max_cache_len, + dtype=torch.bfloat16, + dynamic=True, + ) class HunyuanImage3ARStage(ARStage[HunyuanImage3ARConditions]): From ebaac30661b2a599860645ba5f94e748900b610c Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Fri, 25 Sep 2026 01:32:35 +0800 Subject: [PATCH 2/3] refactor(hi3): import sys at module level in AR stage --- unirl/models/hunyuan_image3/ar.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/unirl/models/hunyuan_image3/ar.py b/unirl/models/hunyuan_image3/ar.py index 94a306b62..88379d050 100644 --- a/unirl/models/hunyuan_image3/ar.py +++ b/unirl/models/hunyuan_image3/ar.py @@ -2,6 +2,7 @@ from __future__ import annotations +import sys from dataclasses import dataclass from dataclasses import field as dc_field from typing import Any, Dict, List, Optional, Tuple @@ -217,9 +218,7 @@ def step( @staticmethod def _build_kv_cache(transformer, *, batch_size: int, max_cache_len: int): """Pre-build a ``HunyuanStaticCache`` for the AR loop.""" - import sys as _sys - - cache_cls = _sys.modules[type(transformer).__module__].HunyuanStaticCache + cache_cls = sys.modules[type(transformer).__module__].HunyuanStaticCache return cache_cls( config=transformer.config, batch_size=batch_size, From 14c63f3cd923bebf56d99682346cff25ecd443f1 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Fri, 25 Sep 2026 01:34:23 +0800 Subject: [PATCH 3/3] refactor(hi3): inline AR KV cache construction into init_state --- unirl/models/hunyuan_image3/ar.py | 21 +++++++-------------- 1 file changed, 7 insertions(+), 14 deletions(-) diff --git a/unirl/models/hunyuan_image3/ar.py b/unirl/models/hunyuan_image3/ar.py index 88379d050..ef36235b3 100644 --- a/unirl/models/hunyuan_image3/ar.py +++ b/unirl/models/hunyuan_image3/ar.py @@ -100,8 +100,13 @@ def init_state( batch_size = int(input_ids.shape[0]) prompt_len = int(input_ids.shape[1]) - past_kv_initial = self._build_kv_cache( - transformer, batch_size=batch_size, max_cache_len=prompt_len + int(max_new_tokens) + cache_cls = sys.modules[type(transformer).__module__].HunyuanStaticCache + past_kv_initial = cache_cls( + config=transformer.config, + batch_size=batch_size, + max_cache_len=prompt_len + int(max_new_tokens), + dtype=torch.bfloat16, + dynamic=True, ) cond_vit = conditions.cond_vit @@ -215,18 +220,6 @@ def step( return token_id, log_prob, state - @staticmethod - def _build_kv_cache(transformer, *, batch_size: int, max_cache_len: int): - """Pre-build a ``HunyuanStaticCache`` for the AR loop.""" - cache_cls = sys.modules[type(transformer).__module__].HunyuanStaticCache - return cache_cls( - config=transformer.config, - batch_size=batch_size, - max_cache_len=max_cache_len, - dtype=torch.bfloat16, - dynamic=True, - ) - class HunyuanImage3ARStage(ARStage[HunyuanImage3ARConditions]): """Rollout-level AR stage: ``HunyuanImage3ARConditions → TextSegment``."""