[Dev] Add configurable attention latent normalization epsilon

[buptzyb]

(Main PR: #6122)

What

• Port Multi-Latent Attention support to HybridModel as the prerequisite already used on main.
• Add MLATransformerConfig.attention_latent_norm_epsilon and --attention-latent-norm-epsilon.
• Default the new option to layernorm_epsilon when it is unset.
• Apply the option to standard MLA/DSA, absorbed MLA, and the primary DSv4 Q/KV latent norms.

This is the dev-targeted counterpart of #6122. The main PR remains open and unchanged.

Scope

The option only selects epsilon; normalization still selects RMSNorm versus LayerNorm. It does not change DSA indexer K normalization, CSA/HCA compressors, DSv4 post-up Q normalization, or the model-wide input/pre-MLP/final norms.

The commits are intentionally split into:

1. the main HybridModel/MLA prerequisite;
2. the generic MLA/DSA option;
3. dev-only DSv4 integration and parity coverage.

Testing

• Black, isort, pylint, and ruff passed through tools/autoformat.sh in check-only mode.
• On one OCI-HSG node (4x GB200), using the CI-dev image:
  • test_mla_qk_norm_config.py: 6 passed per effective rank.
  • test_absorbed_mla.py: 64 passed per effective rank.
  • test_dsv4_hybrid_native_parity.py: 78 passed, 70 skipped per effective rank; fused DSv4 cases retain their existing GB200 skip condition, while the new unfused epsilon parity case passed.
• test_multi_latent_attention.py passed the two new epsilon cases, then hit the existing GB200/TE failure in test_gpu_forward_thd_padded[yarn] (No dot product attention backend is available). Latest origin/dev reproduced the same test and exception on the same node and image, so this is not a regression from this PR.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[buptzyb]

/claude strict-review

[claude]

Strict review summary

CRITICAL: 2 · IMPORTANT: 2 · SUGGESTION: 4

The PR is two logically distinct changes: (1) porting MLA to HybridModel via a new '+' layer symbol and mla_layer spec slot, and (2) introducing attention_latent_norm_epsilon so the Q/KV latent norms can differ from the global layernorm_epsilon. The epsilon half is in good shape — the config field resolves in __post_init__, the CLI flag is wired through the generic core_transformer_config_from_args loop, every new identifier has a real read path, prepare_for_absorption preserves the value via split_te_layernorm_column_parallel_linear(..., eps=fused_layer.eps), and the deliberate retention of layernorm_epsilon at deepseek_v4_hybrid_attention.py:711 matches the stated Scope. QKNormConfigResolver correctly avoids mutating the caller's ModuleSpec (dataclasses.replace), and the test asserts that. No new direct parallel_state.get_*_group() reads. FLOPs accounting picks up '+' automatically because it was added to MLA_ATTENTION.

The findings cluster around the hybrid port, where the new symbol and the new spec slot each have a consumer that was not updated.

Most impactful

1. [CRITICAL Implementation] '+' is advertised for inference but not wired into the dynamic inference layer map. dynamic_context.py:353-364 builds layer_map from ATTENTION/DS_ATTENTION/GDN/MAMBA only. Meanwhile this PR adds an mla_layer to hybrid_inference_stack_spec and makes arguments.py:946-951 infer multi_latent_attention=True from '+'. Net effect for a '+' pattern under dynamic batching: num_attention_layers undercounts (undersized KV cache), then KeyError at self.layer_map[layer_number - 1]. Either add Symbols.MLA to the itemgetter/map/count, or raise NotImplementedError and drop mla_layer from the inference spec.

2. [CRITICAL Correctness] Split K/V up-projections bypass QKNormConfigResolver, and the KV latent norm can vanish silently. In absorbed_mla.py the combined branch now builds from layer_classes["linear_kv_up_proj"], but the split branch still reads submodules.linear_k_up_proj/linear_v_up_proj. With qk_layernorm=True and a trivial kv_layernorm (what the in-tree specs use), the resolver returns the fused norm+linear for a key the split branch never reads, plus kv_layernorm=IdentityOp — so kv_compressed = self.kv_layernorm(kv_compressed) becomes a no-op and the KV latent is never normalized. No exception. Not triggered by in-tree specs today (they use combined K/V) and not covered by the tests (which always pass explicit norms), but reachable from any custom spec, and split K/V does not exist on main — this PR introduces the mismatch.

3. [IMPORTANT Compatibility] ModelOpt hybrid specs silently get mla_layer=IdentityOp. post_training/modelopt/hybrid/model_specs.py:47 and :260 enumerate HybridStackSubmodules fields explicitly and omit the new one. Because IdentityOp.__init__(*args, **kwargs) swallows everything, a '+' pattern there builds an identity layer with no error at all — a model quietly missing attention. Suggest forwarding the field in both, plus an assert in the build branch to close the same latent gap for csa_layer/hca_layer/window_layer.

4. [IMPORTANT Performance] _fuse_mla_down_proj deepcopies the whole submodules tree even with no '+' in the pattern. mla_down_proj_fusion is a global flag, so any pattern with it enabled pays a full recursive copy of every layer spec (including the nested mtp_block_spec) per HybridStack construction — once per PP rank and per virtual-PP chunk — and mutates a spec that is never built. Gate on LayerSymbols.MLA in self.layer_type_list.

Remaining suggestions: missing docstring on _fuse_mla_down_proj; the new eps kwarg undocumented in the Args: blocks of TELayerNormColumnParallelLinear.__init__ and InferenceLayerNormColumnParallelLinear.__init__; test_attention_latent_norm_epsilon_matches_native_reference adds no coverage over the existing parametrizations (the distinct epsilons live in the shared _make_config); and three out-of-scope qk_layernorm = True local-variable hoists in unrelated tests.

Risk assessment

Medium-high. Nothing here breaks an existing configuration: attention_latent_norm_epsilon defaults to layernorm_epsilon, so numerics are bit-identical for every current checkpoint and config, and the '+' symbol is purely additive. The risk is concentrated in the new paths — two of the three consumers of the new symbol/slot (dynamic inference, ModelOpt specs) were not updated, and both fail in the worst way for a training/serving stack: one undersizes a cache then crashes mid-request, the other constructs a structurally wrong model without complaining. The split-K/V resolver gap is the same shape: a missing normalization layer that produces plausible-looking losses. I would want findings 1-3 resolved (or explicitly rejected with a loud error) before this lands; finding 4 and the suggestions are safe follow-ups.

[hxbai]

Why does it need to override in the init func rather than replace the spec directly?

[buptzyb]

I don't know... This is carried over from #4452... This PR needs 4452's mla_qk_norm_config.py.

This is how my agent replies to your question:

The init-time resolution is needed for the config-independent hybrid_stack_spec, since Q/KV norm placement depends on runtime config and backend, and the resolver also validates custom specs. Simply using the original field here would break the static Hybrid MLA/DSA path. I agree that resolving this into a config-aware copied spec before module construction would be cleaner, but that would be a broader #4452 refactor rather than part of the epsilon change.

[hxbai]

OK. I see.

[buptzyb]

/ok to test 3b437a7