Add GLM-5.2 contrib model (FP8 on trn2.48xlarge)

[qingzwang]

NeuronX Distributed Inference port of GLM-5.2 (zai-org/GLM-5.2-FP8), a ~753B glm_moe_dsa MoE model (MLA attention + 256-expert MoE). Runs the FP8 checkpoint on a single trn2.48xlarge (TP=64, lnc=2, seq_len=2048) via the fused MoE TKG + MLP NKI kernels.

• modeling_glm5.py + mla_attention_nki*.py adapted from the community GLM-5 contribution (github.com/jimburtoft, branch contrib/GLM-5); config-driven, so GLM-5.2's nested rope_theta (8e6) is read automatically.
• DSA indexer disabled (dsa_enabled=False): GLM-5.2 has heterogeneous per-layer indexers and the indexer is a no-op for seq_len <= index_topk (2048), so it runs as standard full MLA.
• FP8 required: BF16 weights (~1506 GB) do not fit the 1536 GB HBM.
• examples/{compile,benchmark}.py; unit tests 11/11 pass.

Measured BS=1: TTFT 6377 ms, ITL 241.5 ms/tok, 2.96 tok/s; output coherent.

Note: The below template includes items meant for model contributions only. For other contributions such as bug fixes, features, etc., only fill out the relevant portions of the form.

Description

Add GLM-5.2 contrib model — NeuronX Distributed Inference port of zai-org/GLM-5.2-FP8, a ~753B-parameter glm_moe_dsa MoE model (MLA attention + 256-expert MoE with sigmoid routing). Runs the FP8 checkpoint on a single trn2.48xlarge (TP=64, lnc=2, seq_len=2048) via the fused MoE TKG + MLP NKI kernels.

Key implementation details:

• Modeling code (modeling_glm5.py + mla_attention_nki*.py) adapted from the community GLM-5 contribution (github.com/jimburtoft, branch contrib/GLM-5); config-driven, so GLM-5.2's nested rope_theta (8e6) is read automatically.
• DSA indexer disabled (dsa_enabled=False): GLM-5.2 has heterogeneous per-layer indexers; for seq_len ≤ index_topk (2048), the indexer is a no-op, so we run standard full MLA.
• FP8 required: BF16 weights (~1506 GB) do not fit the 1536 GB HBM budget.
• MTP layer (speculative decoding draft head) dropped — not needed for standard autoregressive generation.
• Includes examples/{compile,benchmark}.py; unit tests and integration test (11/11 pass).

Model Information

Model Name: GLM-5.2

Model Architecture: glm_moe_dsa — Decoder-only MoE transformer with Multi-head Latent Attention (MLA) + 256-expert sigmoid-routed MoE + DeepSparse Attention (DSA) indexer (disabled in this contribution)

Purpose: Text generation (autoregressive LLM)

Checklist

Required Components

• Accuracy Test (test/integration/test_model.py)
  • Integration test validates model accuracy via logit comparison
  • Test compiles and runs the model on Neuron (trn2.48xlarge)
  • 11/11 tests pass
• README.md with the following sections:
  • Usage Example: Compile + benchmark commands with environment setup
  • Compatibility Matrix: Tested instance types and SDK versions table
  • Example Checkpoints: Link to zai-org/GLM-5.2-FP8 on HuggingFace
  • Testing Instructions: Commands to run unit and integration tests
• Source Code (src/)
  • modeling_glm5.py — full model implementation (~117 KB)
  • mla_attention_nki.py + mla_attention_nki_kernel.py — MLA NKI kernel
  • test_mla_attention_nki.py — kernel validation script
  • __init__.py — module entry point

Optional Components

• Unit Tests (CPU + Neuron-based)
  • test/unit/test_config.py — config parsing tests
  • test/unit/test_weight_conversion.py — weight shard/conversion tests

Folder Structure

/contrib/models/GLM-5.2/
  README.md
  /examples
    compile.py
    benchmark.py
  /src
    __init__.py
    modeling_glm5.py
    mla_attention_nki.py
    mla_attention_nki_kernel.py
    test_mla_attention_nki.py
  /test
    __init__.py
    /unit
      __init__.py
      test_config.py
      test_weight_conversion.py
    /integration
      __init__.py
      test_model.py

Testing

How did you test this change?

Tested on a single trn2.48xlarge instance with:

• Full compile + pre-shard pipeline (examples/compile.py)
• Load + benchmark (examples/benchmark.py)
• Unit tests on CPU (pytest test/unit/ -v)
• Integration test on Neuron (pytest test/integration/test_model.py -v)
• MLA NKI kernel validation (python src/test_mla_attention_nki.py --cpu-only)

Test Results:

Metric	Value
TTFT (2048-token prefill)	6,377 ms
ITL / TPOT (per-token)	241.5 ms/tok
Throughput (BS=1)	2.96 tok/s
Weight load (pre-sharded, 64 ranks)	139 s
Warmup	16 s

Output is coherent and factually correct (e.g. "The capital of France is Paris... the Ile-de-France... where the French language was born").

Compatibility

Tested with:

• Neuron SDK Version(s): 2.24 (neuronx-cc 2.24, NKI 0.3.0)
• Instance Type(s): trn2.48xlarge (TP=64, lnc=2)
• PyTorch Version: 2.9
• NXDI Version: 0.9.0

Additional Information

• FP8 is required — BF16 weights (~1506 GB) exceed the 1536 GB HBM; FP8 checkpoint (~705 GB) fits comfortably.
• nkilib routing fork required — GLM-5.2's sigmoid routing needs selection_bias and routed_scaling_factor support in the NKI fused-MoE router. Install from github.com/jimburtoft/nki-library, branch feature/selection-bias-routing.
• DSA indexer disabled — heterogeneous per-layer indexers cause weight-loading issues; for prompts ≤ 2048 tokens the indexer is mathematically a no-op. Long-context sparse attention is future work.
• MLA NKI kernel included but neutral at BS=1 — decode is bound by MoE expert memory bandwidth, not attention compute. May help at larger batch sizes.
• MTP layer (index 78) dropped — speculative decoding draft head, not needed for standard generation.

Related Issues

N/A — first contribution of this model architecture.

vLLM Integration

• This model/feature is intended for use with vLLM
• Documentation includes vLLM registration instructions

vLLM integration is future work — the current contribution focuses on the core NxDI modeling and compilation path.

---

By submitting this PR, I confirm that:

• I have read and followed the contributing guidelines
• This is a community contribution and may have limited testing compared to officially-supported models
• The code follows best practices and is well-documented
• All required components listed above are included