[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`
- 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
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