[bugfix]: complete #1601 Kandinsky5 QAD with kernel/r12 review fixes

docs/training/train_infra.md

@@ -81,7 +81,7 @@ Common model parameters:
 | `disable_custom_init_weights` | `false` | Skip custom weight initialization (use for teacher/critic) |
 | `flow_shift` | `3.0` | Timestep shifting factor |
 | `enable_gradient_checkpointing_type` | `null` | Gradient checkpointing (`"full"` or `null`) |
-| `attention_backend` | `null` | Optional role-local backend for Wan models (for example `ATTN_QAT_TRAIN`); overrides the process default only while this role's transformer is built |
+| `attention_backend` | `null` | Optional role-local backend for supported model plugins (for example `ATTN_QAT_TRAIN`); overrides the process default only while this role's transformer is built |
 
 Which roles are needed depends on the training method:
 

docs/utilities/debugging.md

@@ -77,7 +77,7 @@ If forcing a backend fails, verify optional dependencies are installed:
 - `SAGE_ATTN`: SageAttention package
 - `SAGE_ATTN_THREE`: upstream `sageattn3` package
 - `ATTN_QAT_INFER`: `fastvideo-kernel` checkout/source install that exposes
-  `attn_qat_infer`, AND a consumer-Blackwell (sm_120/sm_121) GPU -- on any
+  `attn_qat_infer`, AND a consumer-Blackwell sm_120 GPU -- on any
   other device the backend reports unavailable (even if a CUDA 13 wheel
   bundles the extension) and selection falls back to FlashAttention
 - `ATTN_QAT_TRAIN`: `fastvideo-kernel`; its runtime-JIT Triton implementation

examples/train/configs/distribution_matching/kandinsky5/distill_dmd_qat.sh

@@ -17,15 +17,9 @@
 #     used as-is with no conversion.
 #
 # The export is passed to all three models.*.init_from overrides: student,
-# teacher, and critic all load the same weights. Teacher/critic are
-# automatically masked back to dense attention and full-precision weights
-# by the _loading_teacher_critic_model gate in
-# fastvideo/models/loader/component_loader.py (family-agnostic, no
-# Kandinsky5-specific handling needed) -- NOT by pointing them at a
-# different export.
-#
-# FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN keeps the student's dense/local
-# attention fake-quantized during distillation.
+# teacher, and critic all load the same weights. The YAML assigns
+# ATTN_QAT_TRAIN to the student and FLASH_ATTN to teacher/critic; the loader's
+# _loading_teacher_critic_model gate also keeps their weights full precision.
 set -euo pipefail
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
@@ -43,10 +37,7 @@ if [[ -f "${CHECKPOINT}/model_index.json" ]]; then
     # Already a diffusers export -- use as-is.
     INIT_FROM="${CHECKPOINT}"
 else
-    # Convert the raw DCP checkpoint. Runs BEFORE ATTN_QAT_TRAIN is
-    # exported below: backend selection for that env var refuses to start
-    # (ImportError) when the training kernel isn't built, and the
-    # conversion itself must not depend on it. --overwrite keeps relaunches
+    # Convert the raw DCP checkpoint before launching training. --overwrite keeps relaunches
     # from silently reusing a stale export for a newer checkpoint;
     # --verify strictly reloads the exported transformer so a key-mapping
     # bug fails here, not deep inside the training launch below.
@@ -60,8 +51,6 @@ else
         --verify
 fi
 
-export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN
-
 bash "${REPO_ROOT}/examples/train/run.sh" \
     "${SCRIPT_DIR}/dmd2_t2v_480p_qat.yaml" \
     --models.student.init_from "${INIT_FROM}" \

examples/train/configs/distribution_matching/kandinsky5/dmd2_t2v_480p_qat.yaml

@@ -9,18 +9,14 @@
 # with the export -- the placeholder paths below only apply when launching
 # the YAML directly.
 #
-# - Teacher: frozen, full-precision Kandinsky5 (loaded from the stage-1 QAT
-#   checkpoint's base weights, or the original pretrained checkpoint --
-#   the _loading_teacher_critic_model gate in component_loader.py masks
-#   quant_config and FASTVIDEO_ATTENTION_BACKEND for teacher/critic
-#   regardless, so they always run full precision / dense attention).
+# - Teacher: frozen, full-precision Kandinsky5 with FLASH_ATTN.
 # - Student: trainable, initialized from the stage-1 QAT finetune checkpoint.
-# - Critic:  trainable, initialized from the same checkpoint as the student.
+# - Critic:  trainable, initialized from the same checkpoint as the student,
+#   with FLASH_ATTN.
 # - Validation: 4-step sampling via Kandinsky5DMDPipeline.
 #
-# Launch with FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN set (student only;
-# teacher/critic are masked to dense automatically) so the student continues
-# learning to absorb fake-quantized attention error during distillation.
+# Role-local backends keep fake-quantized attention on the student and dense
+# attention on teacher/critic without a process-global backend override.
 #
 # Scope: T2V, 480p only, dense/local attention only.
 
@@ -29,16 +25,19 @@ models:
     _target_: fastvideo.train.models.kandinsky5.Kandinsky5Model
     init_from: outputs/kandinsky5_t2v_qat_finetune/checkpoint-1000-diffusers
     trainable: true
+    attention_backend: ATTN_QAT_TRAIN
   teacher:
     _target_: fastvideo.train.models.kandinsky5.Kandinsky5Model
     init_from: outputs/kandinsky5_t2v_qat_finetune/checkpoint-1000-diffusers
     trainable: false
     disable_custom_init_weights: true
+    attention_backend: FLASH_ATTN
   critic:
     _target_: fastvideo.train.models.kandinsky5.Kandinsky5Model
     init_from: outputs/kandinsky5_t2v_qat_finetune/checkpoint-1000-diffusers
     trainable: true
     disable_custom_init_weights: true
+    attention_backend: FLASH_ATTN
 
 method:
   _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method

examples/train/configs/fine_tuning/kandinsky5/README.md

@@ -52,11 +52,9 @@ bash examples/train/configs/fine_tuning/kandinsky5/finetune_qat.sh \
     --training.data.data_path data/kandinsky5_overfit_preprocessed
 ```
 
-`FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN` (set by the script) routes
-Kandinsky5's dense/local attention through the fake-quantized
-straight-through-estimator Triton kernel, so the DiT learns to absorb
-quantization error instead of fighting it. This is purely env-var driven,
-model-agnostic, and does not touch weights during training -- weight-level
+The YAML assigns `ATTN_QAT_TRAIN` to the student role, routing Kandinsky5's
+dense/local attention through the fake-quantized straight-through-estimator
+Triton kernel so the DiT learns to absorb quantization error. Weight-level
 FP4/FP8 quantization is applied post-hoc at inference time (see below), not
 during either training stage.
 
@@ -88,20 +86,16 @@ then launches `dmd2_t2v_480p_qat.yaml` with
 (Passing an already-exported diffusers directory to the launcher skips the
 conversion and uses it directly.)
 
-`--role student` exports the trained transformer weights (the only role
-that matters here: student/teacher/critic in stage 2 all load from the same
-checkpoint, and `_loading_teacher_critic_model` handles the full-precision
-masking for teacher/critic at load time, not at export time). `--verify`
-strictly reloads the exported transformer immediately, so a key-mapping bug
-fails here instead of deep inside the stage-2 launch.
+`--role student` exports the trained transformer weights; student, teacher,
+and critic in stage 2 all load from the same checkpoint. `--verify` strictly
+reloads the exported transformer immediately, so a key-mapping bug fails here
+instead of deep inside the stage-2 launch.
 
 Only the student is quantized (Attn-QAT); the teacher and critic stay full
-precision / dense attention. This is enforced in the loader
-(`fastvideo/models/loader/component_loader.py`, via the
-`_loading_teacher_critic_model` flag), which masks `quant_config` and clears
-`FASTVIDEO_ATTENTION_BACKEND` for teacher/critic -- the same global env var
-reaches only the student, with no per-model flags or Kandinsky5-specific
-handling. Validation runs the distilled student through
+precision / dense attention. The YAML assigns `ATTN_QAT_TRAIN` to the student
+and `FLASH_ATTN` to teacher/critic. Their `disable_custom_init_weights: true`
+setting also activates the loader's family-agnostic quant-config mask, so
+teacher/critic weights remain full precision. Validation runs the distilled student through
 `Kandinsky5DMDPipeline` at the step counts configured in
 `callbacks.validation.sampling_steps`.
 
@@ -130,35 +124,22 @@ their default layer lists (`fastvideo/layers/quantization/nvfp4_qat_config.py`,
 `fastvideo/layers/quantization/fp8_config.py`), so no `target_layers`
 override is needed:
 
-`dcp_to_diffusers` copies the base T2V checkpoint's `model_index.json`
-unchanged into every export (see "Train stage 2" above), so `_class_name`
-still says the base T2V pipeline and the registry
-(`fastvideo/registry.py`) has no way to auto-detect that a given directory
-is actually a DMD (four-step re-noise sampler) export -- it will resolve to
-`Kandinsky5T2VPipeline` and run the full-length sampler on DMD-distilled
-weights. Pass `override_pipeline_cls_name` and a `Kandinsky5DMDConfig`
-explicitly to select the right pipeline/sampler and get
-`dmd_denoising_steps` set (`Kandinsky5T2VConfig`'s default of `None` makes
-`Kandinsky5DmdDenoisingStage` raise immediately):
+`dcp_to_diffusers` records the configured validation pipeline in the exported
+`model_index.json`. Stage-2 exports therefore resolve automatically to
+`Kandinsky5DMDPipeline` and its four-step `Kandinsky5DMDConfig`:
 
 ```python
 from fastvideo import VideoGenerator
-from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5DMDConfig
 from fastvideo.layers.quantization import get_quantization_config
 
 gen = VideoGenerator.from_pretrained(
     "path/to/kandinsky5_dmd_checkpoint", num_gpus=1,
-    override_pipeline_cls_name="Kandinsky5DMDPipeline",
-    pipeline_config=Kandinsky5DMDConfig(),
     transformer_quant=get_quantization_config("nvfp4_qat")(),
     use_fsdp_inference=False,
 )
 # Kandinsky5DmdDenoisingStage ignores request.sampling.num_inference_steps
 # and guidance_scale -- it's a fixed 4-step, no-CFG sampler driven entirely
-# by pipeline_config.dmd_denoising_steps (set above). Adjust
-# Kandinsky5DMDConfig(dmd_denoising_steps=[...]) instead if the checkpoint
-# was distilled/validated with a different schedule than the default
-# [1000, 750, 500, 250].
+# by the exported pipeline config's default [1000, 750, 500, 250] schedule.
 gen.generate(request={"prompt": "...", "output": {"save_video": True}})
 ```
 
@@ -175,17 +156,16 @@ to a supported dense backend while the FP4 linear layers still run.
   confirms `ATTN_QAT_TRAIN` actually engages (not a silent SDPA fallback)
   and that gradients flow through a forward+backward pass.
 - `fastvideo/tests/api/test_kandinsky5_dmd_pipeline_resolution.py` --
-  confirms an unmodified export resolves to `Kandinsky5T2VPipeline` (the bug)
-  and that `override_pipeline_cls_name="Kandinsky5DMDPipeline"` fixes it (the
-  documented workaround above), plus `Kandinsky5DMDConfig`'s
-  `dmd_denoising_steps` default.
+  confirms a stage-2 export resolves directly to `Kandinsky5DMDPipeline` plus
+  `Kandinsky5DMDConfig`, and that rewriting a hard-linked `model_index.json`
+  does not mutate the base checkpoint.
 - `fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py` -- a few
   steps of both training stages on a single synthetic sample, exercising the
   whole recipe end to end: the `dcp_to_diffusers --verify` conversion
   between the stages AND of the final stage-2 student, then reloading that
-  export through the documented `Kandinsky5DMDPipeline` /
-  `Kandinsky5DMDConfig` override, generating deterministically (fixed
-  seed), and comparing MS-SSIM against a committed reference video. The
+  export through ordinary `VideoGenerator.from_pretrained`, generating
+  deterministically (fixed seed), and comparing MS-SSIM against a committed
+  reference video. The
   test fails (does not skip) if the reference is missing; record it once on
   a sanctioned GPU box with `KANDINSKY5_E2E_WRITE_REFERENCE=1`, review the
   written video, and commit it (see the test's module docstring). Nightly

examples/train/configs/fine_tuning/kandinsky5/finetune_qat.sh

@@ -2,9 +2,8 @@
 # QAD recipe -- quantization-aware finetune of Kandinsky5 T2V (480p) with
 # fake-quant (Attn-QAT) attention.
 #
-# The fake-quantized attention path is selected purely by env var
-# (config-driven, no monkey-patching): FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN
-# routes Kandinsky5's dense/local attention through the fake-quantized
+# The YAML's role-local ATTN_QAT_TRAIN backend routes Kandinsky5's dense/local
+# attention through the fake-quantized
 # straight-through-estimator kernel, so the DiT learns to absorb the
 # quantization error instead of fighting it. No weight quantization happens
 # during this stage; that's applied post-hoc at inference time (see the
@@ -23,8 +22,6 @@
 # the export to all three models.*.init_from overrides automatically.
 set -euo pipefail
 
-export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN   # <-- enables Attn-QAT training
-
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../../.." && pwd)"
 

examples/train/configs/fine_tuning/kandinsky5/t2v_480p_qat.yaml

@@ -1,8 +1,8 @@
 # Kandinsky5 T2V 480p stage-1 QAT finetune config.
 #
-# QAT here means attention-only fake-quant during training: launch with
-# FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN set (see finetune_qat.sh), which
-# routes Kandinsky5's dense/local attention through the fake-quantized
+# QAT here means attention-only fake-quant during training. The student's
+# role-local ATTN_QAT_TRAIN backend below routes Kandinsky5's dense/local
+# attention through the fake-quantized
 # straight-through-estimator kernel. No weight quantization happens during
 # training; that's applied post-hoc at inference time via the nvfp4_qat/FP8
 # quant configs (see the stage-2 config and README for that step).
@@ -19,6 +19,7 @@ models:
     _target_: fastvideo.train.models.kandinsky5.Kandinsky5Model
     init_from: kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers
     trainable: true
+    attention_backend: ATTN_QAT_TRAIN
 
 method:
   _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod

fastvideo/attention/backends/attn_qat_infer.py

@@ -50,9 +50,9 @@ def _get_attn_qat_infer() -> Callable[..., torch.Tensor] | None:
     return _attn_qat_infer
 
 
-# Consumer-Blackwell compute capabilities the modified SageAttention3 FP4
-# kernel is compiled for (sm_120a / sm_121a -- see fastvideo-kernel/README.md).
-_SUPPORTED_DEVICE_CAPABILITIES = frozenset({(12, 0), (12, 1)})
+# Consumer-Blackwell compute capability the modified SageAttention3 FP4
+# kernel is compiled for (sm_120a -- see fastvideo-kernel/CMakeLists.txt).
+_SUPPORTED_DEVICE_CAPABILITIES = frozenset({(12, 0)})
 
 
 def _device_capability_supported() -> bool:
@@ -66,15 +66,15 @@ def _device_capability_supported() -> bool:
 
 def is_attn_qat_infer_available() -> bool:
     """True only when the extension imports AND the active device is a
-    consumer-Blackwell (sm_120/sm_121) GPU the kernel is compiled for.
+    consumer-Blackwell sm_120 GPU the kernel is compiled for.
 
     The import check alone is not sufficient: CUDA 13 wheel builds can
-    carry the sm_120/sm_121 extension on any host (e.g. H100, GB200),
+    carry the sm_120 extension on any host (e.g. H100, GB200, or sm_121),
     where the import succeeds, backend selection picks this backend, and
     the first kernel call then fails with an unsupported-capability error
     instead of ever reaching the documented FlashAttention fallback in
     fastvideo.platforms.cuda. Gating on the active device's capability
-    keeps that fallback working on every non-sm_120/121 GPU.
+    keeps that fallback working on every non-sm_120 GPU.
     """
     return _device_capability_supported() and _get_attn_qat_infer() is not None
 

fastvideo/configs/pipelines/kandinsky5.py

@@ -124,17 +124,6 @@ class Kandinsky5I2VConfig(Kandinsky5T2VConfig):
 
 @dataclass
 class Kandinsky5DMDConfig(Kandinsky5T2VConfig):
-    """Kandinsky-5.0 DMD (few-step distilled) text-to-video pipeline configuration.
-
-    Checkpoints exported by ``fastvideo.train.entrypoint.dcp_to_diffusers``
-    copy their base T2V checkpoint's ``model_index.json`` unchanged, so
-    ``_class_name`` still says the base T2V pipeline and the registry
-    (``fastvideo/registry.py``) cannot auto-detect a DMD export -- pass this
-    config together with ``override_pipeline_cls_name="Kandinsky5DMDPipeline"``
-    explicitly to ``VideoGenerator.from_pretrained``/``from_config`` (see
-    ``examples/train/configs/fine_tuning/kandinsky5/README.md``). Without it,
-    ``Kandinsky5T2VConfig``'s ``dmd_denoising_steps=None`` makes
-    ``Kandinsky5DmdDenoisingStage`` raise immediately.
-    """
+    """Kandinsky-5.0 DMD (few-step) text-to-video configuration."""
 
     dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 750, 500, 250])

fastvideo/registry.py

@@ -28,7 +28,11 @@ from fastvideo.configs.pipelines.hunyuan15 import (Hunyuan15T2V480PConfig, Hunyu
                                                    Hunyuan15T2V720PConfig, Hunyuan15I2V720PConfig,
                                                    Hunyuan15SR1080PConfig)
 from fastvideo.configs.pipelines.hyworld import HYWorldConfig
-from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5I2VConfig, Kandinsky5T2VConfig
+from fastvideo.configs.pipelines.kandinsky5 import (
+    Kandinsky5DMDConfig,
+    Kandinsky5I2VConfig,
+    Kandinsky5T2VConfig,
+)
 from fastvideo.configs.pipelines.lingbot_video import LingBotVideoT2VConfig
 from fastvideo.configs.pipelines.lingbotworld import LingBotWorldI2V480PConfig
 from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V480PConfig
@@ -570,6 +574,17 @@ def _register_configs() -> None:
     _is_kandinsky5_i2v_lite_distilled = _kandinsky5_detector(require=("i2v", "lite", "distilled"), exclude=("t2v", ))
     _is_kandinsky5_i2v_pro_distilled = _kandinsky5_detector(require=("i2v", "pro", "distilled"), exclude=("t2v", ))
 
+    # Register the exported DMD artifact before path-based Kandinsky detectors:
+    # output directories commonly contain "kandinsky5_t2v" too.
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=Kandinsky5DMDConfig,
+        workload_types=(WorkloadType.T2V, ),
+        model_detectors=[lambda path: "kandinsky5dmdpipeline" in path.lower()],
+        model_family="kandinsky5",
+        pipeline_cls_name="Kandinsky5DMDPipeline",
+    )
+
     # Kandinsky5 Lite T2V
     register_configs(
         sampling_param_cls=None,

fastvideo/tests/api/test_attn_qat_infer_capability_gate.py

@@ -3,8 +3,8 @@
 the real platform resolver.
 
 ``is_attn_qat_infer_available()`` used to test only whether the kernel
-extension imports. CUDA 13 wheel builds can carry the sm_120/sm_121
-extension on any host (e.g. H100 sm_90, GB200 sm_100): the import
+extension imports. CUDA 13 wheel builds can carry the sm_120 extension
+on any host (e.g. H100 sm_90, GB200 sm_100, or DGX Spark sm_121): the import
 succeeds, ``CudaPlatformBase.get_attn_backend_cls`` selects the
 consumer-Blackwell backend, and the first kernel call fails with an
 unsupported-capability error -- instead of the FlashAttention fallback
@@ -78,14 +78,20 @@ def test_sm100_host_with_bundled_extension_falls_back(monkeypatch):
     assert _resolve() in FALLBACK_CLASSES
 
 
-@pytest.mark.parametrize("capability", [(12, 0), (12, 1)])
-def test_consumer_blackwell_with_extension_selects_backend(monkeypatch, capability):
-    _fake_gpu(monkeypatch, capability=capability, extension_imports=True)
+def test_sm120_with_extension_selects_backend(monkeypatch):
+    _fake_gpu(monkeypatch, capability=(12, 0), extension_imports=True)
 
     assert is_attn_qat_infer_available()
     assert _resolve() == ATTN_QAT_INFER_CLS
 
 
+def test_sm121_host_with_sm120_extension_falls_back(monkeypatch):
+    _fake_gpu(monkeypatch, capability=(12, 1), extension_imports=True)
+
+    assert not is_attn_qat_infer_available()
+    assert _resolve() in FALLBACK_CLASSES
+
+
 def test_consumer_blackwell_without_extension_falls_back(monkeypatch):
     _fake_gpu(monkeypatch, capability=(12, 0), extension_imports=False)
 

fastvideo/tests/api/test_kandinsky5_dmd_pipeline_resolution.py

@@ -1,41 +1,29 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Regression test for exported Kandinsky5 DMD checkpoints resolving to the
-wrong pipeline.
-
-``fastvideo.train.entrypoint.dcp_to_diffusers`` copies the base T2V
-checkpoint's ``model_index.json`` unchanged into every export (see
-``examples/train/configs/fine_tuning/kandinsky5/README.md``), so
-``_class_name`` still says the base T2V pipeline. Without an explicit
-override, ``fastvideo.registry.get_model_info`` resolves such a directory to
-``Kandinsky5T2VPipeline`` (via the path-based "t2v" fallback detector in
-``fastvideo/registry.py``) instead of ``Kandinsky5DMDPipeline`` -- silently
-running the full-length multi-step sampler on DMD-distilled weights instead
-of the four-step re-noise sampler. ``override_pipeline_cls_name`` bypasses
-that resolution entirely and must be paired with ``Kandinsky5DMDConfig`` (not
-the default ``Kandinsky5T2VConfig``, whose ``dmd_denoising_steps=None`` makes
-``Kandinsky5DmdDenoisingStage`` raise immediately).
-"""
+"""Regression tests for self-identifying Kandinsky5 DMD exports."""
 
 import json
+import os
 
 from fastvideo.configs.pipelines.kandinsky5 import (
     Kandinsky5DMDConfig,
     Kandinsky5T2VConfig,
 )
 from fastvideo.pipelines.basic.kandinsky5.kandinsky5_dmd_pipeline import Kandinsky5DMDPipeline
-from fastvideo.pipelines.basic.kandinsky5.kandinsky5_pipeline import Kandinsky5T2VPipeline
 from fastvideo.pipelines.pipeline_registry import PipelineType
 from fastvideo.registry import get_model_info
+from fastvideo.train.entrypoint.dcp_to_diffusers import (
+    _rewrite_model_index_pipeline_class,
+)
 
 
-def _write_fake_export(tmp_path, name: str):
+def _write_fake_export(tmp_path, name: str, pipeline_cls_name: str):
     """Minimal on-disk stand-in for a dcp_to_diffusers export: just enough
     for verify_model_config_and_directory to accept it."""
     model_dir = tmp_path / name
     (model_dir / "transformer").mkdir(parents=True)
     (model_dir / "model_index.json").write_text(
         json.dumps({
-            "_class_name": "Kandinsky5T2VPipeline",
+            "_class_name": pipeline_cls_name,
             "_diffusers_version": "0.30.0",
         }))
     return model_dir
@@ -46,18 +34,36 @@ def test_kandinsky5_dmd_config_sets_denoising_steps():
     assert Kandinsky5DMDConfig().dmd_denoising_steps == [1000, 750, 500, 250]
 
 
-def test_exported_dmd_checkpoint_resolves_to_t2v_pipeline_without_override(tmp_path):
-    """Documents the bug: an unmodified export resolves to the T2V pipeline."""
-    model_dir = _write_fake_export(tmp_path, "kandinsky5_t2v_dmd2_4steps_qat")
-    info = get_model_info(model_path=str(model_dir), pipeline_type=PipelineType.BASIC)
-    assert info.pipeline_cls is Kandinsky5T2VPipeline
-
-
-def test_override_pipeline_cls_name_resolves_dmd_pipeline(tmp_path):
-    model_dir = _write_fake_export(tmp_path, "kandinsky5_t2v_dmd2_4steps_qat_override")
-    info = get_model_info(
-        model_path=str(model_dir),
-        pipeline_type=PipelineType.BASIC,
-        override_pipeline_cls_name="Kandinsky5DMDPipeline",
+def test_exported_dmd_checkpoint_resolves_without_override(tmp_path):
+    model_dir = _write_fake_export(
+        tmp_path,
+        "kandinsky5_t2v_dmd2_4steps_qat",
+        "Kandinsky5DMDPipeline",
     )
+    info = get_model_info(model_path=str(model_dir), pipeline_type=PipelineType.BASIC)
     assert info.pipeline_cls is Kandinsky5DMDPipeline
+    assert info.pipeline_config_cls is Kandinsky5DMDConfig
+
+
+def test_pipeline_class_rewrite_does_not_mutate_hardlinked_base(tmp_path):
+    base_dir = _write_fake_export(
+        tmp_path,
+        "base",
+        "Kandinsky5T2VPipeline",
+    )
+    export_dir = tmp_path / "export"
+    export_dir.mkdir()
+    os.link(
+        base_dir / "model_index.json",
+        export_dir / "model_index.json",
+    )
+
+    _rewrite_model_index_pipeline_class(
+        str(export_dir),
+        "Kandinsky5DMDPipeline",
+    )
+
+    base_index = json.loads((base_dir / "model_index.json").read_text())
+    export_index = json.loads((export_dir / "model_index.json").read_text())
+    assert base_index["_class_name"] == "Kandinsky5T2VPipeline"
+    assert export_index["_class_name"] == "Kandinsky5DMDPipeline"

fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py

@@ -19,9 +19,7 @@ The full validated chain, each arrow a hard assertion:
     -> dcp_to_diffusers --verify (strict reload)
     -> stage-2 train (student/teacher/critic init_from that export)
     -> stage-2 student DCP -> dcp_to_diffusers --verify (strict reload)
-    -> VideoGenerator.from_pretrained(export,
-           override_pipeline_cls_name="Kandinsky5DMDPipeline",
-           pipeline_config=Kandinsky5DMDConfig())   # the documented recipe
+    -> VideoGenerator.from_pretrained(export)  # ordinary public API
     -> deterministic (fixed-seed) 4-step generation
     -> degeneracy checks + MS-SSIM against the committed reference video.
 
@@ -333,20 +331,14 @@ def _generate_main(export_dir: str, output_dir: str) -> None:
 
     This is the exact inference recipe the QAD README documents for a
     stage-2 export (minus the optional NVFP4/FP8 weight quantization,
-    which needs flashinfer + sm_120 hardware): the export's
-    model_index.json still names the base T2V pipeline, so
-    override_pipeline_cls_name + Kandinsky5DMDConfig are required to get
-    the 4-step re-noise sampler -- see
-    fastvideo/tests/api/test_kandinsky5_dmd_pipeline_resolution.py.
+    which needs flashinfer + sm_120 hardware). The export identifies its
+    DMD pipeline/config through model_index.json.
     """
     from fastvideo import VideoGenerator
-    from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5DMDConfig
 
     generator = VideoGenerator.from_pretrained(
         export_dir,
         num_gpus=1,
-        override_pipeline_cls_name="Kandinsky5DMDPipeline",
-        pipeline_config=Kandinsky5DMDConfig(),
         use_fsdp_inference=False,
     )
     generator.generate_video(
@@ -471,7 +463,7 @@ def test_e2e_kandinsky5_dmd_overfit_single_sample():
             # nondeterminism; structured output does not).
             "--training.optimizer.learning_rate", "1e-6",
         ],
-        env_overrides={"FASTVIDEO_ATTENTION_BACKEND": "ATTN_QAT_TRAIN"},
+        env_overrides={},
     )
     stage1_ckpt = _latest_checkpoint(STAGE1_OUTPUT_DIR)
 
@@ -497,19 +489,24 @@ def test_e2e_kandinsky5_dmd_overfit_single_sample():
             # generator every step instead.
             "--method.generator_update_interval", "1",
         ],
-        env_overrides={"FASTVIDEO_ATTENTION_BACKEND": "ATTN_QAT_TRAIN"},
+        env_overrides={},
     )
     assert any(STAGE2_OUTPUT_DIR.glob("*.mp4")), (
         f"no stage-2 validation video produced under {STAGE2_OUTPUT_DIR}")
 
     # The actual deliverable of this recipe is the exported stage-2
     # student: export it, strict-reload it (--verify), then instantiate it
-    # through the documented Kandinsky5DMDPipeline override and generate.
+    # through the ordinary public API and generate.
     stage2_ckpt = _latest_checkpoint(STAGE2_OUTPUT_DIR)
     stage2_diffusers_dir = STAGE2_DIFFUSERS_DIR / stage2_ckpt.name
     _export_dcp_to_diffusers(stage2_ckpt, stage2_diffusers_dir)
-    assert (stage2_diffusers_dir / "model_index.json").exists(), (
+    model_index_path = stage2_diffusers_dir / "model_index.json"
+    assert model_index_path.exists(), (
         f"dcp_to_diffusers export did not produce a diffusers model dir at {stage2_diffusers_dir}")
+    model_index = json.loads(model_index_path.read_text(encoding="utf-8"))
+    assert model_index.get("_class_name") == "Kandinsky5DMDPipeline", (
+        f"stage-2 export selected {model_index.get('_class_name')!r}, expected "
+        "'Kandinsky5DMDPipeline'")
 
     generated_video = _generate_from_export(stage2_diffusers_dir, GENERATED_VIDEO_DIR)
     _assert_video_not_degenerate(generated_video)

fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py

@@ -9,19 +9,14 @@ Two silent-fallback traps this guards against:
     ``F.scaled_dot_product_attention`` (a real fallback for standalone
     parity tests, not relevant here since we always set a forward context,
     but worth asserting past explicitly).
-  - ``_cached_get_attn_backend`` is process-wide ``@cache``d on
-    ``(head_size, dtype, supported_attention_backends)`` -- NOT on the env
-    var -- so a stale backend selection from an earlier test in the same
-    process can silently linger. Cleared here before selecting, matching
-    the defensive pattern in
-    ``fastvideo/models/loader/component_loader.py``.
+  - ``_cached_get_attn_backend`` is process-wide ``@cache``d, so a stale
+    backend selection from an earlier test in the same process can silently
+    linger. The role-local model option scopes and clears selection while the
+    transformer is built.
 
 This test itself must not become that stale-selection source for whichever
-model test runs next in the same pytest process: ``FASTVIDEO_ATTENTION_BACKEND``
-is set via ``monkeypatch`` (auto-restored at teardown, including on
-failure) and the cache is cleared again in a ``finally`` block, since
-``monkeypatch`` only restores the env var, not the separate process-wide
-cache keyed on it.
+model test runs next in the same pytest process, so the cache is cleared
+again in a ``finally`` block.
 
 Unlike ``fastvideo.platforms.cuda``'s other backends, ATTN_QAT_TRAIN has no
 silent-fallback path at all if the kernel isn't built -- backend selection
@@ -65,7 +60,7 @@ _FIXTURE = str(
 
 
 @pytest.mark.usefixtures("distributed_setup")
-def test_kandinsky5_attn_qat_train_engages_and_backprops(monkeypatch):
+def test_kandinsky5_attn_qat_train_engages_and_backprops():
     if not torch.cuda.is_available():
         pytest.skip("requires CUDA")
 
@@ -74,7 +69,6 @@ def test_kandinsky5_attn_qat_train_engages_and_backprops(monkeypatch):
     if not is_attn_qat_train_available():
         pytest.skip("fastvideo_kernel ATTN_QAT_TRAIN kernel is not built")
 
-    monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "ATTN_QAT_TRAIN")
     _cached_get_attn_backend.cache_clear()
     try:
         from fastvideo.train.models.kandinsky5 import Kandinsky5Model
@@ -85,6 +79,7 @@ def test_kandinsky5_attn_qat_train_engages_and_backprops(monkeypatch):
             init_from=cfg.models["student"]["init_from"],
             training_config=cfg.training,
             trainable=True,
+            attention_backend="ATTN_QAT_TRAIN",
         )
 
         device = torch.device("cuda:0")
@@ -154,9 +149,6 @@ def test_kandinsky5_attn_qat_train_engages_and_backprops(monkeypatch):
                               if hasattr(attn.to_query.weight.grad, "to_local") else
                               attn.to_query.weight.grad).all().item(), "weight grad contains NaN/Inf"
     finally:
-        # monkeypatch restores FASTVIDEO_ATTENTION_BACKEND itself, but the
-        # selector cache is a separate process-wide functools.cache keyed on
-        # (head_size, dtype, supported_backends) -- not on the env var -- so
-        # it must be cleared again here or a later model test in this same
-        # pytest process could reuse this test's ATTN_QAT_TRAIN selection.
+        # Clear the process-wide selector cache so a later model test cannot
+        # reuse this test's ATTN_QAT_TRAIN selection.
         _cached_get_attn_backend.cache_clear()

fastvideo/tests/train/models/test_load_kandinsky5.py

@@ -24,7 +24,9 @@ from pathlib import Path
 import pytest
 import torch
 
+import fastvideo.train.models.kandinsky5.kandinsky5 as kandinsky5_module
 from fastvideo.forward_context import set_forward_context
+from fastvideo.platforms import AttentionBackendEnum
 from fastvideo.train.models.kandinsky5 import Kandinsky5Model
 from fastvideo.train.utils.config import load_run_config
 
@@ -33,6 +35,43 @@ _FIXTURE = str(
     "kandinsky5_t2v_min.yaml")
 
 
+@pytest.mark.parametrize(
+    ("attention_backend", "disable_custom_init_weights"),
+    [
+        ("ATTN_QAT_TRAIN", False),
+        ("FLASH_ATTN", True),
+    ],
+)
+def test_kandinsky5_routes_role_local_attention_backend(
+    monkeypatch,
+    attention_backend,
+    disable_custom_init_weights,
+):
+    cfg = load_run_config(_FIXTURE)
+    captured = {}
+
+    def _fake_load_module_from_path(**kwargs):
+        captured.update(kwargs)
+        return torch.nn.Linear(1, 1)
+
+    monkeypatch.setattr(
+        kandinsky5_module,
+        "load_module_from_path",
+        _fake_load_module_from_path,
+    )
+
+    Kandinsky5Model(
+        init_from=cfg.models["student"]["init_from"],
+        training_config=cfg.training,
+        trainable=not disable_custom_init_weights,
+        disable_custom_init_weights=disable_custom_init_weights,
+        attention_backend=attention_backend,
+    )
+
+    assert captured["attention_backend"] is AttentionBackendEnum[attention_backend]
+    assert captured["disable_custom_init_weights"] is disable_custom_init_weights
+
+
 @pytest.mark.usefixtures("distributed_setup")
 def test_kandinsky5_model_loads_and_forwards():
     if not torch.cuda.is_available():

fastvideo/train/entrypoint/dcp_to_diffusers.py

@@ -38,6 +38,26 @@ from fastvideo.logger import init_logger
 logger = init_logger(__name__)
 
 
+def _rewrite_model_index_pipeline_class(
+    model_dir: str,
+    pipeline_cls_name: str,
+) -> None:
+    """Set the exported pipeline class without mutating a hard-linked base."""
+    import json
+    from pathlib import Path
+
+    model_index_path = Path(model_dir) / "model_index.json"
+    model_index = json.loads(model_index_path.read_text(encoding="utf-8"))
+    model_index["_class_name"] = pipeline_cls_name
+
+    replacement = model_index_path.with_suffix(".json.tmp")
+    replacement.write_text(
+        json.dumps(model_index, indent=2) + "\n",
+        encoding="utf-8",
+    )
+    os.replace(replacement, model_index_path)
+
+
 def _ensure_distributed() -> None:
     """Set up a single-process distributed env if needed.
 
@@ -63,6 +83,7 @@ def _save_role_pretrained(
     module_names: list[str] | None = None,
     overwrite: bool = False,
     model: Any,
+    pipeline_cls_name: str | None = None,
 ) -> str:
     """Export a role's modules into a diffusers-style model dir.
 
@@ -127,6 +148,11 @@ def _save_role_pretrained(
             symlinks=False,
             copy_function=_copy_or_link,
         )
+        if pipeline_cls_name is not None:
+            _rewrite_model_index_pipeline_class(
+                str(dst),
+                pipeline_cls_name,
+            )
 
     _barrier()
 
@@ -316,12 +342,16 @@ def convert(
         output_dir,
         base_model_path,
     )
+    pipeline_target = (cfg.callbacks.get("validation", {}).get("pipeline_target") if role == "student" else None)
+    pipeline_cls_name = (pipeline_target.rsplit(".", 1)[-1]
+                         if isinstance(pipeline_target, str) and pipeline_target.strip() else None)
     result = _save_role_pretrained(
         role=role,
         base_model_path=base_model_path,
         output_dir=output_dir,
         overwrite=overwrite,
         model=model,
+        pipeline_cls_name=pipeline_cls_name,
     )
     logger.info("Export complete: %s", result)
 

fastvideo/train/models/kandinsky5/kandinsky5.py

@@ -43,6 +43,7 @@ from fastvideo.forward_context import set_forward_context
 from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
     FlowMatchEulerDiscreteScheduler, )
 from fastvideo.pipelines import TrainingBatch
+from fastvideo.platforms import AttentionBackendEnum
 from fastvideo.training.activation_checkpoint import (
     apply_activation_checkpointing, )
 from fastvideo.training.training_utils import (
@@ -92,10 +93,12 @@ class Kandinsky5Model(ModelBase):
         transformer_override_safetensor: str
         | None = None,
         lora: LoraConfig | dict[str, Any] | None = None,
+        attention_backend: AttentionBackendEnum | str | None = None,
     ) -> None:
         super().__init__(
             trainable=trainable,
             lora=lora,
+            attention_backend=attention_backend,
         )
         self._init_from = str(init_from)
 
@@ -106,6 +109,7 @@ class Kandinsky5Model(ModelBase):
             enable_gradient_checkpointing_type=(enable_gradient_checkpointing_type),
             training_config=training_config,
             transformer_override_safetensor=(transformer_override_safetensor),
+            attention_backend=self.attention_backend,
         )
 
         self.noise_scheduler = (FlowMatchEulerDiscreteScheduler(shift=float(flow_shift)))
@@ -141,6 +145,7 @@ class Kandinsky5Model(ModelBase):
         enable_gradient_checkpointing_type: str | None,
         training_config: TrainingConfig,
         transformer_override_safetensor: str | None = None,
+        attention_backend: AttentionBackendEnum | str | None = None,
     ) -> torch.nn.Module:
         transformer = load_module_from_path(
             model_path=init_from,
@@ -149,6 +154,7 @@ class Kandinsky5Model(ModelBase):
             disable_custom_init_weights=(disable_custom_init_weights),
             override_transformer_cls_name=(self._transformer_cls_name),
             transformer_override_safetensor=(transformer_override_safetensor),
+            attention_backend=attention_backend,
         )
         ckpt_type = (enable_gradient_checkpointing_type or getattr(
             getattr(training_config, "model", None),