veomni-patchgen-model

name

veomni-patchgen-model

description

Author or refresh a VeOmni model's patchgen-generated modeling under generated/ — GPU and/or NPU config, dense or MoE, text / VLM / Omni. Covers the patchgen decorators, sharing patches across sibling models via name_map, MoE fused-expert weight loading, Ulysses SP in multimodal forwards, __init__.py registration, running codegen, and the test cases. This is the modeling step of adding a new model, not only of refreshing an existing one. Trigger: 'add patchgen for a model', 'write a patch_gen_config', 'regenerate the generated modeling', 'add NPU patchgen', 'port a model to patchgen', 'transformers v5 migration'. Never hand-edit anything under generated/.

VeOmni Patchgen Modeling Protocol

Purpose: add or refresh a model's patchgen-generated modeling under veomni/models/transformers/<model>/generated/. VeOmni pins transformers==5.16.1 and ships patchgen-generated modeling for every supported transformers-family model. The non-transformers architectures (flux, movqgan, wan) have no generated/ directory and are out of scope.

References (read first, load on demand):

• docs/design/patchgen.md — patchgen DSL, CLI, CI drift check
• docs/transformers_v5/transformers_v5_moe_weight_loading.md — MoE fused-expert layout + runtime converter
• docs/transformers_v5/veomni_flash_attention_kernel_adapter.md — FA custom-name adapter
• docs/transformers_v5/testing_new_model.md — test case SOP for a new model

What to read for your model

This file is the spine: it applies to every model. The category-specific material lives in references/ — load only what your model needs.

Your model	Also read
Any model, before Phase 1	references/model-examples.md — pick the closest existing model and mirror it
Has routed experts (MoE)	references/moe.md — Phase 2 expert patches, Phase 3 checkpoint converter, MoE pitfalls
Has a vision / audio / speech tower (VLM or Omni)	references/multimodal.md — SP-aware multimodal forward, metadata precompute, dummy_forward, subtree pruning, VLM/Omni pitfalls
Text-only and dense	neither — the spine plus the examples file is the whole protocol

A text-only dense GPU model therefore reads this file plus the examples, and skips about 380 lines of MoE and multimodal material. A VLM+MoE model reads everything. Read the spine first either way; the reference files add to it and never replace a phase.

---

Phase 0: Environment + Reference Setup

0.1 Verify transformers venv

Patchgen runs against transformers==5.16.1. Before touching code:

source .venv/bin/activate
python -c "import transformers; print(transformers.__version__)"

If not 5.16.1, re-sync the default env:

uv sync --frozen --extra gpu --group dev
source .venv/bin/activate

0.2 (Strongly recommended) Drop HF reference source into .agents_workspace/

.agents_workspace/ is gitignored. Keeping the upstream HF source next to your patchgen config is the single biggest accelerator for catching subtle signature/contract drift while iterating.

Use the pinned version as the directory name so several pins can coexist:

PIN=$(python -c "import transformers; print(transformers.__version__)")
mkdir -p ".agents_workspace/hf_reference/<m>/v${PIN}"

curl -fsSL -o ".agents_workspace/hf_reference/<m>/v${PIN}/modeling_<m>.py" \
  "https://github.com/huggingface/transformers/raw/v${PIN}/src/transformers/models/<m>/modeling_<m>.py"

-f matters: without it a missing tag or renamed module returns 404 and curl writes the error page into modeling_<m>.py with exit status 0, so you would diff against an HTML page and not notice.

For VLMs also grab processing_<m>.py / image_processing_<m>.py / configuration_<m>.py if you expect processor-side or config-shape work.

If you are refreshing an existing patchgen-generated file across a transformers minor bump (the pin the generated file was produced against → the new pin), pull both versions side-by-side and diff to spot contract drift — substitute the <old_ver> / <new_ver> tags with the actual versions you are migrating between:

mkdir -p .agents_workspace/hf_reference/<m>/{old,new}
curl -fsSL -o .agents_workspace/hf_reference/<m>/old/modeling_<m>.py \
  "https://github.com/huggingface/transformers/raw/<old_ver>/src/transformers/models/<m>/modeling_<m>.py"
curl -fsSL -o .agents_workspace/hf_reference/<m>/new/modeling_<m>.py \
  "https://github.com/huggingface/transformers/raw/<new_ver>/src/transformers/models/<m>/modeling_<m>.py"
diff -u .agents_workspace/hf_reference/<m>/{old,new}/modeling_<m>.py | less

0.3 For a pin bump: survey signature drift across all configs first

patchgen <config> (without --dry-run) runs the generated file through ruff, so a patch body referencing a symbol upstream no longer defines fails loudly with F821 / F811. That catches removed names. It does not catch a patch whose target still exists but whose signature changed — the patch keeps applying and silently runs against the wrong contract.

Before touching any config, index both upstream versions with ast and compare the parameter lists of every target named in each config's override_method / replace_class / replace_function call. Targets missing from both versions are VeOmni-added methods (patchgen uses override_method to inject them) and should be filtered out, or they drown the real findings.

docs/transformers_v5/upgrade_5_9_to_5_16.md records what that survey turned up for the 5.9 → 5.16 bump and how each class of breakage was resolved — read it before starting a new bump, the categories repeat.

Note that --dry-run returns before the ruff step, so it reports success on files that cannot even import. Never use it as the pass/fail signal.

Things to watch for in upstream contracts:

• @can_return_tuple, @capture_outputs, @merge_with_config_defaults, @auto_docstring decorators → affect behavior of your override_method. When you override_method on a @auto_docstring-decorated method, every parameter you declare in the new signature must also appear in the patched docstring's Args: block — otherwise auto_docstring will emit warnings at import time about "undocumented parameter". For Omni-style overrides that add params like audio_feature_lengths, feature_lens, aftercnn_lens, rope_deltas, image_grid_thw, video_grid_thw, etc., copy the upstream docstring and append minimal one-line entries for every new param.
• attention_mask may be a dict — HF v5 routinely passes attention_mask={"full_attention": <tensor>, ...} keyed by attention type. Any patched forward that forwards attention_mask to compute_3d_position_ids / get_rope_index / other tensor-expecting helpers must defensively unwrap attention_mask.get("full_attention", None) when it's a dict.

VLM and Omni models have four more upstream contracts to check before writing any patch — placeholder masks, get_{image,video}_features return shapes, the packed position-ids layout and mrope shape collapse. See references/multimodal.md, "Phase 0: upstream contracts".

Keep this directory around through commit; delete it after the PR merges (it's already gitignored so it won't leak into the repo).

---

Before You Start: Create a Plan

Track the phases with whatever todo/plan tool the running agent provides. Suggested plan:

Phase 0: Verify venv + drop HF reference files       -> in_progress
Phase 1: Scope & audit upstream surface              -> pending
Phase 2: Draft <model>_gpu_patch_gen_config.py       -> pending
Phase 3: (MoE only) Add checkpoint converter         -> pending
Phase 4: Wire __init__.py to expose generated classes -> pending
Phase 5: Run patchgen + verify diff                   -> pending
Phase 6: Add test cases                               -> pending
Phase 7: Run tests (single-GPU + e2e)                 -> pending
Phase 8: Docs + commit; review before PR or substantive update -> pending

Drop phases that don't apply (e.g. Phase 3 for non-MoE models).

---

Phase 1: Scope & Audit

Input: model name <M> (e.g. qwen3_5, glm_moe_dsa).

Operations:

1. Locate veomni/models/transformers/<M>/. If the directory does not exist yet you are being called as the modeling step of /veomni-new-model: create it, and read that skill's Phase 1 first so the category (text / VLM / Omni, dense / MoE, GPU-only or GPU+NPU) is already decided when you get here.
2. If a patchgen-generated file already exists under veomni/models/transformers/<M>/generated/ you are refreshing an existing config (e.g. picking up upstream changes, adding NPU sibling, fixing a bug). Otherwise you are writing the first config for this model. Either way, the rest of this protocol applies identically.
3. Decide backend coverage:
   • GPU only → one <m>_gpu_patch_gen_config.py + one generated/patched_modeling_<m>_gpu.py.
   • GPU + NPU → add sibling <m>_npu_patch_gen_config.py that writes generated/patched_modeling_<m>_npu.py; mirror the glm_moe_dsa or qwen3_vl layout.
4. Check model category. Each entry below names the closest existing model; references/model-examples.md says what to copy out of it, file by file.
   • Text-only LLM → reference qwen3/ (or llama/ for the minimal example)
   • MoE → reference qwen3_moe/ (plus converter work in Phase 3)
   • VLM (non-MoE) → reference qwen3_vl/
   • VLM + MoE → reference qwen3_vl_moe/ (multimodal forward + SP scatter, ViT dummy forward, Flash-attn kwargs popping, get_position_id_func)
   • Omni (non-MoE thinker + speech subtree to exclude) → reference qwen2_5_omni/ (audio/vision SP + dummy_forward, talker/token2wav/BigVGAN exclusion, log_probs/entropy output dataclass, no parallel_plan/converter)
   • Omni MoE → reference qwen3_omni_moe/
5. Check upstream source (from transformers.models.<m> import modeling_<m>). Confirm class/function names still exist; MoE expert layouts especially diverge between sibling models — see docs/transformers_v5/transformers_v5_moe_weight_loading.md.
6. Note related configs/loaders to preserve: MODELING_REGISTRY, MODEL_CONFIG_REGISTRY in veomni/models/loader.py; any auto-config registrations.
7. Look for a sibling model you can borrow patches from: e.g. qwen3_5_moe reuses GatedDeltaNet/ViT patches from qwen3_5 via direct import + name_map={"Qwen3_5": "Qwen3_5Moe"}. Prefer reuse over copy-paste when the upstream classes are structural duplicates with only a name-prefix difference.
8. Compare upstream and VeOmni parameter keys, including constructor overrides and nested modules. For any mismatch, follow the user-decision rule in veomni-new-model before choosing a model rename or checkpoint conversion. This also applies to refreshes and dependency upgrades. A resolution already authorized in the current task does not require another confirmation.

Validation: you have a concrete list of patches to apply, the reference model directory to mirror, and the backend/category decision pinned down.

---

Phase 2: Draft <M>_gpu_patch_gen_config.py

Create veomni/models/transformers/<M>/<M>_gpu_patch_gen_config.py at the model root.

Skeleton (mirror qwen3_gpu_patch_gen_config.py):

from veomni.patchgen.patch_spec import PatchConfig, create_patch_from_external

config = PatchConfig(
    source_module="transformers.models.<m>.modeling_<m>",
    target_file="patched_modeling_<m>_gpu.py",
    description="<M> with LigerKernel GPU replacements + VeOmni SP/fused-loss patches",
)

Patch primitives:

Effect	patchgen decorator / API
Replace whole class (RMSNorm, MLP, Experts)	@config.replace_class("<Class>") or create_patch_from_external(...) for liger
Replace module-level function (rotary, loss)	@config.replace_function("<name>")
Override a single method (Attention.forward, Model.forward, ForCausalLM.forward)	@config.override_method("<Class>.<method>")
Add attribute / extra super().__init__() wiring	@config.modify_init("<Class>")
Reuse patch from a sibling config (name-prefix difference)	config.override_method("<NewClass>.<m>", replacement=<imported_fn>, name_map={"OldPrefix": "NewPrefix"}) — non-decorator form. Caveat: name_map only rewrites symbol names at the AST level; it does NOT align field sets between sibling output dataclasses (e.g. dense ModelOutputWithPast vs MoE ModelOutputWithPast with extra router_logits). Any <OldClass>Output(...) constructor call in the body gets its name rewritten but keeps the original arg list, silently dropping MoE-only fields. Clone the body when return dataclasses differ.
Supporting import needed in generated file	config.add_import("<module>", names=[...]) (or alias=..., is_from_import=False)
Remove an upstream import the generated file should NOT keep	config.drop_import_names("<symbol>", ...)
Inject raw code (try/except import fallback, helper fn used by patched code) near top of generated file	config.add_post_import_block("""...""")
Remove unused class from output	config.exclude_from_output("<Class>")
Inherit an entire sibling GPU config into an NPU config (reuse helpers / imports / post-import blocks; only override device-specific kernels)	config.helpers.extend(gpu_config.helpers) + config.post_import_blocks.extend(gpu_config.post_import_blocks) + config.additional_imports.extend(gpu_config.additional_imports) + import each <fn>_patched and re-register via config.override_method(...). See qwen3_vl_npu_patch_gen_config.py

Cross-config reuse pattern (qwen3_5_moe reusing qwen3_5):

from veomni.models.transformers.qwen3_5.qwen3_5_gpu_patch_gen_config import (
    qwen3_5_gated_deltanet_forward_patched,
    qwen3_5_vision_model_forward,
    # ...
)

_NAME_MAP = {"Qwen3_5": "Qwen3_5Moe"}
config.override_method(