NVIDIA / NVIDIA/Megatron-LM

[inference] No inference path for `experimental_attention_variant="dsa"` / AbsorbedMLA (GLM-5.x, DeepSeek-V3.2): asserts in `absorbed_mla.py` and `gpt_layer_specs.py`

Open
#7,106 2 comments 0 reactions 0 assignees View on GitHub
community-request waiting-on-maintainers
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

## Summary
Megatron-Core cannot run inference (text generation / refit-style eval) for models that use the DSA experimental attention variant with absorbed MLA:

- `megatron/core/transformer/experimental_attention_variant/absorbed_mla.py` (`AbsorbedMLASelfAttention.forward`): `assert inference_context is None and inference_params is None, "Inference is not supported for AbsorbedMLA"` (and a second assert "Inference with MLA RoPE fusion is not supported").
- `megatron/core/models/gpt/gpt_layer_specs.py`: the inference-optimized spec (`get_gpt_decoder_layer_specs` / `get_gpt_decoder_block_spec` with `use_inference_optimized...`) asserts `experimental_attention_variant is None`.

Both asserts are present on `main` as of 2026-09-06 (file history since 2026-06 adds no inference support). #6392 (GLM-5.2 tracking) lists only training items; #6998 adds refit CI coverage for DSA models, so there may be partial context already.

## Environment
- `nvcr.io/nvidia/nemo:26.08` (aarch64), Megatron-Core 0.19 (16ad357), Megatron-Bridge v0.6.0; GB200 NVL72
- Model: GLM-5.2 (78 layers, MLA + DSA indexer, MoE 256 experts, MTP 1), Megatron torch_dist checkpoint after SFT

## Reproduction (Megatron-Bridge `examples/inference/gpt/text_generation.py`, 4–8 nodes)
Sequence of blockers we hit while trying to generate from the checkpoint:
1. `--dtype bf16`, `--use-coordinator` (required for EP>1), `--no-cache-mla-latents` (`cache_mla_latents is not supported for AbsorbedMLA`), `--inference-moe-token-dispatcher-type nccl` (default `nvls` OOMs on symmetric memory at EP16), `--use-legacy-generation` (coordinator ZMQ cannot bind hostname inside the container) — then:
2. `AssertionError: Inference is not supported for AbsorbedMLA` (`absorbed_mla.py`).
3. Switching to the inference-optimized spec: `AssertionError: experimental_attention_variant is None` (`gpt_layer_specs.py`).

Workaround: export to HF and serve with vLLM (`GlmMoeDsaForCausalLM`).

## Ask
Is inference for DSA/AbsorbedMLA planned (e.g. a non-absorbed MLA inference path or KV-latent cache support for the absorbed form)? This blocks in-framework eval/RL (NeMo-RL refit + generation) for GLM-5.x and DeepSeek-V3.2-class models. If it is planned, please link the tracking item from #6392; if not, documenting the limitation in the GLM-5 recipe/card would save users several jobs.

Contributor guide

Open the contributing guide

Research direction

Start by reading megatron/core/transformer/experimental_attention_variant/absorbed_mla.py and megatron/core/models/gpt/gpt_layer_specs.py, then reproduce the inference failures using Megatron-Bridge examples/inference/gpt/text_generation.py. Check the context from #6392 and #6998 before defining scope. Done would require either a supported DSA/AbsorbedMLA inference path or clear documentation of the limitation in the GLM-5 recipe/card.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.