fix(hi3): fail fast when the AR KV cache cannot be built by CjhHa1 · Pull Request #508 · Tencent-Hunyuan/UniRL

CjhHa1 · 2026-09-24T16:25:21Z

Summary

HunyuanImage3ARStep._build_kv_cache returned None in three cases: the upstream HunyuanStaticCache class could not be found, the transformer had no config, or the cache constructor raised (swallowed by except Exception). The AR loop then kept running with past_key_values=None.

That fallback does not degrade gracefully. Upstream HunyuanImage3ForCausalMM never creates a cache itself: attention only writes KV when past_key_value is not None, and HunyuanImage3Model.forward hands the incoming cache straight back out. So with None:

prefill sees the full prompt and the first token is correct;
every decode step feeds only the new token (prepare_inputs_for_generation gathers input_ids down to [B, 1] at position_ids), with no KV, so it attends only to itself.

The result is context-free generations with valid shapes and no error. The most realistic trigger is a transformers upgrade that changes the StaticCache constructor signature that HunyuanStaticCache inherits.

This PR builds the cache directly in init_state and lets the original error propagate. With the fallbacks gone, _build_kv_cache was a single-caller wrapper around one constructor call, so it is inlined.

Related Issue

N/A. Found while reviewing #479.

Test Plan

Behavioral check against a stub upstream module:
- normal construction returns a HunyuanStaticCache with the requested batch_size / max_cache_len / dynamic=True;
- a raising constructor propagates its TypeError;
- a module without HunyuanStaticCache raises AttributeError: module ... has no attribute 'HunyuanStaticCache'.
pre-commit run --files unirl/models/hunyuan_image3/ar.py: all hooks pass.
Not run: real-checkpoint AR rollout. The success path constructs the same object with the same arguments as before.

Compatibility / Risk

The success path is unchanged.
Setups that previously fell back silently now fail at init_state with the underlying error. They were already producing wrong decode tokens.
type(transformer) stays the upstream class under FSDP (trainable_module() wraps transformer.model) and LoRA (inject_adapter_in_model is in-place), so the module lookup is unaffected by training wrappers.
No tests construct HunyuanImage3ARStep with a stub transformer that relied on the None path.

Reviewer Notes

Diffusion's _build_kv_cache is intentionally untouched: predict_noise with state=None is a legitimate uncached mode that recomputes the full sequence every step. AR replay also uses past_key_values=None legitimately (use_cache=False, single full forward) and does not go through _build_kv_cache.
fix(hi3): support batched trainside AR with dynamic KV cache #479 edits _build_kv_cache (adds a dynamic argument), so it will need a rebase: pass dynamic=batch_size == 1 to the inlined constructor in init_state.
fix(hi3): support batched trainside AR with dynamic KV cache #479 has getattr(past_key_values, "dynamic", True) in _cache_attention_mask, whose default currently absorbs this None case. After this lands it can become past_key_values.dynamic.

Checklist

I reviewed the changed code and removed unrelated/generated artifacts.
I updated tests, docs, and configs where needed, or explained why not.

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

…onflict in ar.py)

…lpers Review follow-ups for Tencent-Hunyuan#479 after merging Tencent-Hunyuan#508: - Require conditions.tokenizer_output and fused.prompt_lengths in init_state at every batch size. Upstream sets the decode position_ids from tokenizer_output.real_pos; without it the loop silently re-feeds the shifted prompt, at B=1 too. This makes the real_pos-is-None fallbacks (tail append, logits[:, -1] at prefill, early mask return) unreachable, so they are removed. - Resolve real_pos once into HunyuanImage3ARState from fused.prompt_lengths, which embed_for_ar already normalizes to [B], instead of re-deriving it from tokenizer_output every step. - Inline the single-caller mask and token-placement helpers into step, branching on step_idx instead of on the presence of attention_mask, and read past_key_values.dynamic directly now that the cache is never None. - Drop the pad_id fill: the appended column is always overwritten or never gathered before it is. - README: point the gotchas at the inlined code, note that padded B=1 prompts were also fixed, document the real_pos requirement, and fix the upstream attention class name.

* fix(hi3): support batched trainside AR with dynamic KV cache Cached train-side HunyuanImage3 AR only worked at B=1. Fixes #452. The reported crash is upstream's `assert len(cache_position) == 1` in `HunyuanStaticCache.update`, where `len()` on a 2-D cache position is the batch dim. A dynamic cache trims the returned KV to one scalar end position, which cannot represent per-row right-pad, so `_build_kv_cache` now only requests `dynamic=True` at B=1 and lets B>1 use the full static cache. That makes the assert unreachable rather than patching upstream. Three further defects blocked B=1 parity once the crash was gone: - Upstream deletes the attention mask during `gen_text` decode, which is only safe for a per-row trimmed cache. Over a static cache the decode query also attends right-pad KV and never-written slots, so `_cache_attention_mask` admits row i's keys [0, real_pos_i + step) and right-pads the prefill mask's key axis to the cache width. - Upstream re-reads the next input token with `input_ids.gather(1, position_ids)` at `real_pos`, so a sampled token must be scattered to `real_pos_i + step` (`_place_token`); appending at the padded tail fed short rows their own pad token. - `FusedMultimodalCondition.concat` padded `position_ids` with 0, and those are the cache write indices, so every pad position of a short row index_copy_'d its KV onto slot 0 and silently replaced that row's first real token. `_pad_positions` pads with continuing indices instead. B=1 keeps the dynamic cache and is bit-identical. Batched AR without `real_pos` now raises in `init_state` instead of failing later on an opaque SDPA shape mismatch. * fix(hi3): require real_pos at every B and inline the batched-cache helpers Review follow-ups for #479 after merging #508: - Require conditions.tokenizer_output and fused.prompt_lengths in init_state at every batch size. Upstream sets the decode position_ids from tokenizer_output.real_pos; without it the loop silently re-feeds the shifted prompt, at B=1 too. This makes the real_pos-is-None fallbacks (tail append, logits[:, -1] at prefill, early mask return) unreachable, so they are removed. - Resolve real_pos once into HunyuanImage3ARState from fused.prompt_lengths, which embed_for_ar already normalizes to [B], instead of re-deriving it from tokenizer_output every step. - Inline the single-caller mask and token-placement helpers into step, branching on step_idx instead of on the presence of attention_mask, and read past_key_values.dynamic directly now that the cache is never None. - Drop the pad_id fill: the appended column is always overwritten or never gathered before it is. - README: point the gotchas at the inlined code, note that padded B=1 prompts were also fixed, document the real_pos requirement, and fix the upstream attention class name. * refactor(hi3): drop redundant real_pos state and dead guards from batched AR - Read real_pos from conditions.fused.prompt_lengths in step instead of shadowing it on HunyuanImage3ARState; step already receives the same conditions every call. - Fold the tokenizer_output / prompt_lengths requirement into the existing fused-field check in autoregress, the one validation site. - Drop the int() on max_cache_len (already int) and the clamp on the prefill logits index (real_pos is in [1, L] by construction). - Collapse _pad_positions' early returns; move its rationale pointer to the concat docstring and drop the measured numbers from the README. * docs(hi3): trim README gotchas to the upstream constraints * refactor(hi3): drop redundant helpers and fast paths in fused concat - Remove the leading empty/single-item fast path: the equal-length check already covers both, and Batch.concat raises on empty itself. - Collect sequence lengths as a set from input_ids directly; every producer of this condition sets input_ids. - Replace the three function-local Batch imports and Batch.concat.__func__(cls, ...) with super().concat(...). - Drop _pad_attn: padding the mask's key and query axes is two _pad_seq calls, with identical zero fill. - Build _pad_seq's F.pad spec directly for the negative dims it is called with, and fold its None / already-long early returns. - Rebuild each shard with dataclasses.replace so unpadded fields (scatter indices, prompt_lengths) carry over instead of being re-listed by hand. * refactor(hi3): reuse hydrate and drop the constant pad value in fused concat - Replace the local _materialize duck-type check with the canonical unirl.distributed.tensor.ref.hydrate, which tests isinstance(TensorRef) and maps an empty ref to None instead of torch.empty(0). - Drop _pad_seq's value parameter: every caller passed a zero (0, False, 0.0), which is F.pad's default fill for every dtype here. * refactor(hi3): hoist fused-concat pad helpers to module level _pad_seq and _pad_positions were closures only to capture max_L. Take the target length as an argument instead, matching the module-level _pad_* helpers in types/conditions/text.py and janus_pro/conditions.py, so they are no longer redefined on every concat call. The negative-dim constraint and the README pointer move into their docstrings. * refactor(hi3): drop the default-equal value=False on the prefill mask pad --------- Co-authored-by: CjhHa1 <cjh18671720497@outlook.com> Co-authored-by: Jianghai <72591262+CjhHa1@users.noreply.github.com>

- In modes/t2t.py, call sys.modules[type(transformer).__module__].get_system_prompt directly instead of importing system_prompt dynamically inside a silent try/except fallback. - In diffusion.py, inline _build_kv_cache into predict_noise and construct HunyuanStaticCache directly without the unreachable fused/cache_cls None fallbacks (matching ar.py after Tencent-Hunyuan#508).

) * fix(hi3): fail loudly when the batched chat template drops samples The checkpoint's `apply_general_template(batchify=True)` hands `batch_gen_infer` a one-element `prompt_list` and `zip`s it against the per-sample kwargs, so only the first prompt is encoded and the rest are dropped without an error. `apply_chat_template` with three prompts returns one row (two at `cfg_factor=2`). Today that surfaces later, as a `Part.fill` count mismatch or an it2i row-count check. Raise at the source instead, from the row count the call was asked for. Recipes that pin `rollout.forward_batch_size: 1` are unaffected, since one prompt with `cfg_factor=2` already returns two rows. * fix(hi3): correct the batch-drop gotcha and neutralize the guard message The diffusion path already fails at `split_cfg_fused` with an explicit row-count error, not a `fused.input_ids` shape mismatch, so the guard's new coverage is the AR path. Drop the per-prompt-concat workaround and the padding claim, which #544 contradicts. The error now names the dropped samples instead of declaring the checkpoint unable to batch, so it stays accurate once the tokenizer is repaired. * refactor(hi3): drop dead batch_message_list parameter and refine AR gotcha No caller passes batch_message_list to embed_for_ar or _apply_chat_template, so batch_prompt is always p.texts and the row-count guard reads len(batch_prompt) and tokens.shape[0] directly. Also note in the README gotcha that i2t fails earlier in instantiate_continuous_tokens before reaching Part.fill. * refactor(hi3): consolidate _fused_common into _build_fused Both embed_for_ar and embed_for_gen_image unpacked the 6-tuple from _fused_common, repeated the cond_vae_image_mask and cond_timestep_scatter_index lookups, and built HunyuanImage3FusedMultimodalCondition. Have _build_fused construct the condition container directly. * docs(hi3): align README gotcha and _apply_chat_template docstring with repo guide * refactor(hi3): drop dead parameters and redundant coercions in text_embed - Deduplicate _resolve_build_batch_2d_rope and _resolve_get_system_prompt into _resolve_upstream_symbol and inline the single-caller _resolve_system_prompt_preset. - Drop the unused max_sequence_length attribute on HunyuanImage3TextEmbedStage and the unused max_length parameter on embed_for_ar and _apply_chat_template. - Read gen_config.max_length, output.real_pos.reshape(-1), input_ids.shape, and p.texts directly without redundant getattr fallbacks, int/str coercions, or list copies. * refactor(hi3): resolve upstream helpers via transformer module directly FSDP2 wraps transformer.model rather than the outer HunyuanImage3ForCausalMM wrapper, so sys.modules[type(transformer).__module__] always points to the remote-code module (matching HunyuanImage3ARStep.init_state). * refactor(hi3): resolve get_system_prompt and HunyuanStaticCache directly - In modes/t2t.py, call sys.modules[type(transformer).__module__].get_system_prompt directly instead of importing system_prompt dynamically inside a silent try/except fallback. - In diffusion.py, inline _build_kv_cache into predict_noise and construct HunyuanStaticCache directly without the unreachable fused/cache_cls None fallbacks (matching ar.py after #508). * refactor(hi3): replace _optional_output_tensor with upstream to_device Both Base and Instruct checkpoints define to_device(data, device) on the modeling module and cond_timestep_scatter_index on TokenizerEncodeOutput. Reuse upstream_mod.to_device for optional condition tensors and resolve the vae/vit_image_mask vs cond_vae/vit_image_mask attribute rename inline. * fix(hi3): resolve Instruct stop tokens and drop dead AR/mode scaffolding - Read special tokens from added_tokens_encoder when special_token_map is absent so _stop_tokens_for_bot_task resolves </recaption>, </answer>, </think>, and <img_ratio_0..32> on the Instruct tokenizer as well as Base; fold _cot_stop_tokens into _stop_tokens_for_bot_task. - Drop the 10 unread fields on HunyuanImage3ARParams and the duplicate container constructions in modes/t2t.py, modes/i2t.py, and modes/t2ti.py; inline _pack_text_segment in ar.py. - Deduplicate prepare_seeded_sampling in seed.py and _as_sample_batches in modes/t2t.py across t2i, it2i, and i2t. * refactor(hi3): drop remaining falsy fallbacks, coercions, and ARSamplingParams rebuilds - Read control dicts and fields directly (control.get('ar', {}), bot_task, sys_type, sequence_template, stop_token_ids) without 'or {}' / 'or []' / str(...) coercions. - Forward the frontier ARSamplingParams directly to autoregress instead of rebuilding ARSamplingParams in t2t, i2t, and t2ti. - Read tokenizer special_map and eos_token_id directly in _stop_tokens_for_bot_task, and read out.logits and typed sampling fields directly in ar.py. --------- Co-authored-by: SmartBea_Linlin_r <80329948+GaoRuilin1@users.noreply.github.com> Co-authored-by: CjhHa1 <cjh18671720497@outlook.com>

CjhHa1 requested review from celve and haonan3 as code owners September 24, 2026 16:25

github-actions Bot added the need review Ready and waiting for review label Sep 24, 2026

CjhHa1 added 2 commits September 25, 2026 01:32

refactor(hi3): import sys at module level in AR stage

ebaac30

refactor(hi3): inline AR KV cache construction into init_state

14c63f3

CjhHa1 merged commit 6786296 into Tencent-Hunyuan:main Sep 24, 2026
7 checks passed

CjhHa1 deleted the fix/hi3-ar-kv-cache-fail-fast branch September 24, 2026 17:37

github-actions Bot removed the need review Ready and waiting for review label Sep 24, 2026

CjhHa1 mentioned this pull request Sep 24, 2026

fix(hi3): support batched trainside AR with dynamic KV cache #479

Merged

2 tasks

CjhHa1 added a commit to PushpakAg/UniRL that referenced this pull request Sep 26, 2026

Merge upstream main into fix/issue-452 (resolve Tencent-Hunyuan#508 c…

94108c9

…onflict in ar.py)

Navigation Menu

Repository navigation

fix(hi3): fail fast when the AR KV cache cannot be built - #508

fix(hi3): fail fast when the AR KV cache cannot be built#508
CjhHa1 merged 3 commits into
Tencent-Hunyuan:mainfrom
CjhHa1:fix/hi3-ar-kv-cache-fail-fast

CjhHa1 commented Sep 24, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

CjhHa1 commented Sep 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Related Issue

Test Plan

Compatibility / Risk

Reviewer Notes

Checklist

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

CjhHa1 commented Sep 24, 2026 •

edited

Loading