fix(awq): stop variable-length calibration collapsing to the last batch

[Leonccaa]

Summary

Fixes #3035.

Variable-length calibration batches cannot be concatenated on dim 0, so _layer_input_features() silently kept only tensors[-1] per module — every earlier batch was discarded from AWQ scale statistics. For mixture-of-experts models this collapses per-expert calibration to a single batch.

This PR adds a model-declared feature-aggregation policy so de-fused pointwise expert projections pack every ragged batch into deterministic, bounded token rows, plus per-module aggregation accounting. Public model classes with a module tree marked :moe inherit the policy; dense-only models remain on the existing path.

What Changed

• New model hook awq_input_feature_aggregation(module_name) returning {"mode": "token_rows", "capture_root": bool} or None, with an optional explicit max_tokens override. BaseQModel enables it for the declared expert root and its de-fused pointwise children; AWQProcessor._feature_aggregation_policy() validates and applies it.
• New AWQProcessor._pack_token_rows(): packs ragged captures into [1, retained_tokens, hidden] under a workload-derived bound. The automatic budget uses one largest observed batch equivalent, raised only when needed to retain at least one row from every contributing batch, and never exceeds raw routed rows. Remaining capacity is distributed proportionally to batch length, with rows sampled evenly inside each batch.
• _record_input_feature() gains dedupe_batch; record_moe_root_input_feature() records the shared pre-router expert input when the policy requests capture_root.
• Expert lifecycle replay makes one guarded call into the root recorder before expert replay.
• Captures are restored to calibration-batch order before collapse so aggregation and kwargs stay aligned when hooks complete out of order.
• Per-module aggregation stats (mode, raw_tokens, retained_tokens, batches) are recorded for both activation and scale features and surfaced in quant log rows. nsamples reports retained tokens when a policy is active and is unchanged otherwise.
• Group fallback compares raw_tokens (corpus coverage routed to the module) rather than the policy's retained sampling bound. For the unchanged latest-batch fallback, the policy reports the actual retained rows.
• Padding rows are excluded before aggregation, and empty captures do not participate in quota allocation.
• Intentionally out of scope: changing the default variable-length path for models without a policy. Sequence-structured replay cannot span ragged batches without aligned masks and positions, so that path keeps its existing contract and now reports what it retained.

Review Follow-ups Included

• Reserve one row per non-empty batch so a short batch cannot disappear during sampling.
• Redistribute leftover quota across multiple passes so saturated short batches cannot strand budget.
• Preserve per-sample quant logs while reporting token-row coverage separately.
• Use raw coverage for token-row fallback decisions and actual retained rows for latest-batch decisions.
• Enable the bounded token-row policy for affected public model classes through their declared expert roots.
• Derive each module's token-row budget from its observed calibration batches instead of fixing a model-level token count; retain an explicit policy override for exceptional model requirements.
• Audit every registered public MoE definition and require one unambiguous top-level :moe root per module-tree variant. The audit corrected the missing mlp marker for GPT-OSS and added Granite MoE Hybrid's Defuser-expanded block_sparse_moe expert paths and marker.
• Keep processor-owned feature accounting state concise (_feature_stats, _scale_feature_by_module, and _feature_task_names); the AWQ-prefixed kwargs key remains namespaced because it crosses into model scaling hooks.
• Centralize successful layer finalization so normal completion and both early-exit paths release child/root capture tasks, aggregation telemetry, replay kwargs, module tracking, and thread-local scale context.

Tests

• Public model-policy coverage for Mixtral, Qwen3 MoE, Qwen3 Next, and MiniMax M2 definitions, including alias-based expert roots and dense-model opt-out.
• Registry-wide expert-root coverage for all 52 detected public MoE model classes/aliases, plus exact pointwise-path expansion checks for the two corrected definitions.
• Variable-length token-row aggregation, padding exclusion, batch ordering, root-capture deduplication, fallback accounting, and bounded-quota property tests.
• Early-completion regression coverage for both an empty scaling configuration and a configuration whose groups are all rejected.
• Combined focused regression selection: 76 passed.
• End-to-end smoke validation confirmed that all de-fused expert projections report token-row aggregation, bounded retention, raw non-padding coverage, batch counts, and root-versus-direct scale-feature routing.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5598cc7e40

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[Qubitium]

@Leonccaa Codex review:

• P1 — Fallback behavior regressed for legacy ragged inputs.
  _should_fallback_group always uses raw_tokens. For latest_batch, raw tokens include discarded batches while AWQ uses only the final batch, so the default fallback threshold can be incorrectly bypassed. Use raw counts only for token_rows; otherwise use retained rows.

• P1 — The advertised fix is inactive for all in-tree models.
  The patch adds an optional hook lookup, but no public model defines awq_input_feature_aggregation. Therefore every current model still follows latest_batch. Add model implementations and integration tests, or narrow the PR scope.

• P1 — Padding tokens can contaminate the new statistics.
  record_moe_root_input_feature records hidden_states without applying the active keep mask. Padded rows can enter token-row sampling and inflate coverage. Apply masking consistently before root capture.

• P2 — nsamples reporting changed for models without a policy.
  apply_quant uses retained rows unconditionally, so legacy latest_batch modules now report only the final batch instead of total calibration samples.

• P2 — Empty captures underfill the token budget.
  _pack_token_rows fails for zero-row tensors; [0, 10] with max_tokens=1 returns zero rows. Filter empty batches or distribute quotas only across non-empty batches.

Validation: six targeted new unit tests passed; compilation and whitespace checks passed. Full integration testing was blocked by an existing CUDA/CPU device mismatch, and GitHub reports no PR checks.

[Qubitium]

@Leonccaa Also I don't see a model or model def that actually enables the new aggregation policy so I cannot run ci tests on this PR with a real (small) model. Can you push a model def change that enables this?

[Qubitium]

@Leonccaa Currently doing PR cleanup and refractor. Will let you know when I am done and you can re-test on your end.

[Qubitium]

@Leonccaa I have pushed all the changes.

1. The fix is now automatically applied to all (well-defined) models that has the correct :moe tagging in the model def.
2. User do not need to manually set any max_tokens value. It is auto derived now from calirbation data pass. You can override this.

The rest of small pattern fixes.

Please test and verify on your end and push any regression you find.

[Qubitium]

Re-reviewed the current head. The earlier requested changes are resolved, and successful AWQ layer finalization now also releases child/root capture tasks, aggregation telemetry, replay kwargs, module tracking, and scale context on both early-completion paths. Focused regression coverage passes with no remaining code findings.

[Leonccaa]

Re-reviewed the current head. The earlier requested changes are resolved, and successful AWQ layer finalization now also releases child/root capture tasks, aggregation telemetry, replay kwargs, module tracking, and scale context on both early-completion paths. Focused regression coverage passes with no remaining code findings.

Thanks for pushing these changes and the final cleanup.

We reviewed the latest head. The automatic :moe enablement and calibration-derived token-row budget address our remaining concerns, including removal of the fixed 512-token cap.

Our focused validation has passed, including a tiny Qwen3-MoE AWQ ExpertsRoutingBypass quantize → save → reload smoke test. We are currently running a full Qwen 3.8-Flash-Next AWQ quantization with this fix. Qwen 3.8-Flash-Next uses ExpertsRoutingBypass, so this run directly exercises the new root-capture path.

We’ll report back after the full quantization and post-quantization checks finish, and will push any regression we find. Thanks again.

[Qubitium]

@Leonccaa e2e testing on real model found more lifecycle bugs in the AWQ processor when this new pr fix/feature is activated. Pending fix incoming.

[Qubitium]

@Leonccaa Bug fixed and more refractor passed tests. You can run your own tests now.

[Qubitium]

Validation update for commit 2c1feed4:

Check	Result
H100 ARC accuracy (1,172 samples)	0.5273037542662116
H100 ARC normalized accuracy	0.552901023890785
Focused AWQ mask tests	2 passed
Broader AWQ/lifecycle/model-policy suite	127 passed, 1 skipped

The review findings are covered in the current head: active keep-mask application during root capture, raw-vs-retained fallback accounting, legacy nsamples preservation, empty-batch filtering and quota packing, and automatic public :moe model policy enablement. The full-model quality gates remain 0.5094 / 0.5486; the ARC values above are the measured q-projection snapshot results.

[Qubitium]

@Leonccaa Please create new PR if you find any bugs.