Fix AutoRound gradients for Kimi-VL routed experts

[lesj0610]

Purpose

moonshotai/Kimi-VL-A3B-Thinking-2506 is listed as a supported model, but its routed experts cannot be gradient-tuned by SignRound when iters > 0. The model remains quantizable with RTN, while iterative W4A16_MIXED tuning silently leaves the routed-expert tuning parameters without gradients (or fails at backward when no other differentiable path is present). This matters because the mixed recipe intentionally assigns INT4 only to the routed experts.

Root cause

AutoRound correctly keeps the model in eval mode during block tuning. Kimi-VL's remote DeepseekV3MoE.forward therefore selects moe_infer, which is decorated with @torch.no_grad(). The routed experts execute inside that decorator, so their AutoRound tuning parameters are disconnected from autograd. The shared-expert path is outside moe_infer, which can make a block loss require gradients while still leaving every routed expert untuned.

Switching the block to training mode is not valid: Kimi-VL's noaux_tc routing path asserts inference mode and training mode would change routing semantics. Existing unfused/fused MoE patches do not cover this class because the top-level model type is kimi_vl, and this remote implementation uses a ModuleList rather than the Transformers 5 fused-expert interface.

Design

This change adds an instance-local context manager around SignRoundQuantizer.quantize_block:

• It activates only for model_type == "kimi_vl" and iters > 0.
• It recognizes DeepseekV3MoE through the class MRO and temporarily binds the callable exposed by moe_infer.__wrapped__ to that instance.
• It preserves eval mode and the original routing implementation; it only bypasses the outer torch.no_grad wrapper during optimization.
• It validates all targets before mutation, is reentrant, supports multiple targets, and restores the exact previous instance state on normal and exceptional exits.
• It is a no-op for RTN (iters == 0), other model types, and unrelated MoE classes.
• It explicitly rejects ep_size > 1: the Kimi-VL expert-parallel all_to_all path does not preserve routed-expert gradients. Quantize with ep_size=1 or use iters=0 (RTN).

prepare_run() is validation-only: before block iteration it scans the model to reject unrecognized grad-mode decorators and ep_size > 1, then logs the validated target count. It caches neither target descriptors nor runtime state. Every quantize_block() invocation—including the normal orchestrator path and the public/LLM-Compressor path—lazily scans only the current block and opens exactly one block-scoped context, so the public path behaves identically without prepare_run().

The runtime patch is scoped to block tuning instead of globally modifying remote model source. It also avoids .train() and avoids reimplementing Kimi's routing logic. AdamRound and SignRoundV2 inherit the same quantize_block implementation, so they receive the same scoped behavior.

W4A16_MIXED and MLA preservation

A meta-device inventory of the actual local Kimi-VL model applied the current W4A16_MIXED recipe with group_size=32:

• 5,292 Linear modules in total.
• 4,992 INT4 modules, exactly 26 MoE layers × 64 routed experts × 3 projections; all use group size 32 and integer weight quantization.
• 299 modules remain at 16 bits: 78 shared-expert projections, 108 MLA attention projections, 108 vision linears, 2 multimodal-projector linears, and 3 dense layer-0 MLP projections.
• language_model.lm_head is excluded from layer_config, and the vision patch-embedding convolution is also untouched.

Safetensors metadata was inspected without loading tensor values. All corresponding checkpoint tensors are BF16: 4,992 routed weights, 78 shared weights, 108 MLA projection weights, 329 vision tensors, 6 projector tensors, and the lm_head weight. The patch does not alter this recipe or MLA execution.

Before-fix reproduction

On upstream main at 54b4b84cdfcbd985a00264048a288a224e2a76c0, a CPU reproducer using Kimi's eval-mode MoE path produced loss.requires_grad == False; backward() failed with:

element 0 of tensors does not require grad and does not have a grad_fn

The remote moe_infer method was confirmed to carry the torch.no_grad wrapper, while the top-level config reports model_type="kimi_vl".

Unit and regression tests

Executed locally with CUDA hidden and Hugging Face/Transformers offline mode:

• test/test_cpu/models/test_kimi_vl.py: 22 passed. The tests cover the Kimi-VL moe_infer dataflow, real quantize_block loss/backward/step behavior, active/inactive expert gradient and update behavior, output parity, eval-mode preservation, normal and exceptional restoration, the context manager's explicit reentrancy contract, lazy current-block scoping with sibling-block exclusion, the public/LLM-Compressor path without prepare_run(), prepare_run() rejection of unrecognized decorators and ep_size > 1 before block iteration, multiple targets, exact-MRO inclusion and exclusion, the original scatter-reduce failure mode, inherited quantizers, and the mixed recipe.
• test/test_cpu/models/test_moe_experts_interface.py: 7 passed.
• test/test_cpu/models/test_vlm_ram_reduction.py: 14 passed.
• pre-commit run --all-files passed every configured hook, including Black, isort, Ruff, Bandit, codespell, and license checks.
• Python compilation and git diff --check passed.

Actual Kimi GPU smoke test

The first actual MoE decoder block (language_model.model.layers.1) was tested on one RTX 3090 24 GB with PyTorch 2.11.0+cu130, Transformers 5.10.2, and the editable AutoRound checkout.

Settings: W4A16_MIXED, weight INT4, group size 32, activation 16-bit, iters=1, sequence length 32, one sample, batch size 1, seed 42, quant_nontext_module=False. The actual remote DeepseekV3MoE implementation and 197 block tensors were strict-loaded from the local checkpoint (571,080,768 parameters). No model files were modified.

Results:

• 49 routed experts were active and 15 were inactive only because this deliberately tiny single 32-token smoke sample did not cover all 64 experts; the inactive-expert count does not apply to the actual calibration configuration.
• The patched output exactly matched the decorated inference output (max_abs_diff=0) and the model remained in eval mode.
• The lazy quantize_block() path scanned the current block, found one DeepseekV3MoE target, and recorded one active patch at context depth 1.
• Quantization loss required gradients; loss was 0.007146065589040518.
• All 441 active tuning parameters had non-None, finite, nonzero gradients, and all 441 update probes changed.
• All 135 inactive tuning parameters had grad is None, and all 135 full tensors remained unchanged.
• 192 routed-expert linears were quantized; all 3 shared-expert linears remained unquantized.
• Peak CUDA allocated memory: 7,578,559,488 bytes (about 7.06 GiB).
• Peak CUDA reserved memory: 7,604,273,152 bytes (about 7.08 GiB).
• Block quantization time: 2.840883 seconds; internal smoke run: 11.762815 seconds; external wall time: 16.80 seconds.
• External maximum resident set size: 4,491,320 KiB.
• The original config, remote-code, index, and checkpoint files retained their pre-test fingerprints.

Transformers 5.10.2 no longer exposes transformers.utils.import_utils.is_torch_fx_available, which the Transformers-4-era remote code imports. The smoke harness added an in-process compatibility alias returning True; neither AutoRound nor the checkpoint files were modified.

A separate full-model meta instantiation verified all module names and recipe assignments in 13.09 seconds with 2,828,012 KiB maximum RSS. It did not load full weights or run full-model quantization.

Relationship to prior work

• patch for vlm #518 documents the initial Kimi-VL support surface but does not address routed-expert autograd.
• add support for w4a16_mixed  #1326 added the mixed scheme coverage, but its relevant path uses iters=0, so it cannot detect missing routed-expert gradients.
• support qwen3-next/glm4-moe/deepseekv3 monkey patch for transformers 5 #1367 patches the built-in deepseek_v3 model type; Kimi-VL reports kimi_vl, so that patch does not apply.
• Support general MOE replacement for MOE models (Transformers 5.0 compatible) #1374 covers the Transformers 5 fused-expert interface; Kimi-VL's remote MoE uses a ModuleList and has no matching fused-expert interface/decorator.
• fix missing trust_remote_code in unfused_moe #1532 enables trust_remote_code while checking unfused MoE configurations, which is a prerequisite but not a Kimi-VL autograd fix.
• support WOQ model input, such as kimi2.5 #1642 supports pre-quantized WOQ Kimi K2.5-class inputs, a different flow.
• fix bug of qwen and gguf export #1846 adds Kimi/MLA-aware GGUF serialization, not block-tuning gradients.
• [AutoScheme] fix lm_head no_grad issue and support MOE model #1971 improves AutoScheme MoE scoring and inactive-expert score handling with RTN-style scoring; it does not restore SignRound gradients through Kimi's decorated inference method.

Validation not completed

A complete CPU test suite was not practical in this environment:

• A related combined CPU selection reached 26 passing tests and then failed in test/test_cpu/utils/test_alg_ext.py::TestAlgExt::test_alg_ext because the installed gguf package lacks MODEL_ARCH.DFLASH; bundled llama.cpp conversion then reported Qwen3ForCausalLM unsupported. This path does not import or execute the changed Kimi/SignRound code.
• The existing remote-model test_w4a16_mixed_mllm fixture stopped making progress while fetching Qwen2.5-VL data. It was interrupted after 19:06; the preceding selected tests reported 8 passed, 26 deselected. No generated model fixture remains in the worktree.

Full Kimi-VL checkpoint quantization, export/reload, vLLM/Marlin loading, end-to-end multimodal generation with real image inputs, and quality benchmarks were not run. No claim is made here about accuracy, reasoning quality, hallucination rate, or end-to-end compression quality. Maintainers should separately validate this patch with a Transformers combination supported by the repository and run a representative Kimi-VL calibration/evaluation/export cycle before declaring quality parity.

Artifacts and remaining checks

This PR contains only two source files and one CPU regression-test file. It contains no model weights, quantized outputs, caches, logs, or generated artifacts.

The context is designed for AutoRound's serial block-tuning flow and is reentrant. Simultaneous multi-thread tuning of the same module instance was not validated. Maintainers should also confirm the remote-code compatibility matrix for the repository's supported Transformers versions and run full-model quality/export validation on the target checkpoint.