feat(glm_moe_dsa): add GLM5.2 context parallel support by HuiyingLi · Pull Request #2695 · NVIDIA-NeMo/Automodel

HuiyingLi · 2026-06-21T17:34:16Z

What does this PR do ?

Adds GLM MoE DSA context-parallel support for TileLang attention, keeping the implementation local to glm_moe_dsa.

Changelog

Adds a GLM-local packed THD CP batch sharder and differentiable K/V activation gather.
Extends GLM DSA TileLang top-k masking for CP query shards and padded THD offsets.
Wires GLM DSA CP capability reporting through the model-owned CP hook.
Adds a 32k Tulu3 CP8 TileLang recipe.

Validation

CP1/CP2 parity

GLM-5.2 TileLang Tulu3 4k CP1 parity: https://wandb.ai/Nemo-automodel/huiyingl_workspace/runs/a7zxyvuf
GLM-5.2 TileLang Tulu3 4k CP2 parity: https://wandb.ai/Nemo-automodel/huiyingl_workspace/runs/zmm8v3di

GLM-5.2 TileLang Tulu3 32k seqlen CP8: https://wandb.ai/Nemo-automodel/huiyingl_workspace/runs/jgb7ci7n

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

Add CPU-safe unit coverage for GLM DSA TileLang dispatch wrappers, validation packing, THD pipeline metas, and layer-level TileLang branches. Exclude copied vendored TileLang kernel definitions from coverage the same way existing vendored kernels are handled. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

Add NVIDIA copyright headers to the vendored GLM TileLang kernel files for consistency with the existing vendored DeepSeek kernels. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

The PP causal-mask cache is an optional mapping on model instances. Plain unittest mocks fabricate missing attributes, so treat non-mapping cache values as absent before checking membership. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

Use the TileLang validation result as an explicit guard in the sparse MLA path instead of discarding it. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

copy-pr-bot · 2026-06-21T17:34:19Z

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.

…ernels Signed-off-by: HuiyingLi <willwin.lee@gmail.com> # Conflicts: # nemo_automodel/components/distributed/pipelining/hf_utils.py

Add CPU unit coverage for GLM DSA context-parallel helpers, TileLang padded/query-index masking, CP gather wiring, capability gates, and the prepacked dataloader path. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

HuiyingLi · 2026-06-22T18:19:14Z

/claude review

HuiyingLi · 2026-06-22T18:25:29Z

/ok to test 874c47d

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

HuiyingLi · 2026-06-22T18:32:57Z

/claude review

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

…o huiyingl/feat/glm52-cp Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

HuiyingLi · 2026-06-23T04:38:44Z

/ok to test 6f70549

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

…o huiyingl/feat/glm52-cp Signed-off-by: HuiyingLi <willwin.lee@gmail.com> # Conflicts: # tests/unit_tests/models/glm_moe_dsa/test_glm_moe_dsa_tilelang.py

HuiyingLi · 2026-06-23T09:34:36Z

/ok to test 1740305

Signed-off-by: HuiyingLi <willwin.lee@gmail.com> # Conflicts: # nemo_automodel/components/models/glm_moe_dsa/layers.py # nemo_automodel/components/models/glm_moe_dsa/model.py # nemo_automodel/components/models/glm_moe_dsa/optimized_kernels.py # tests/unit_tests/models/glm_moe_dsa/test_glm_moe_dsa_tilelang.py

HuiyingLi · 2026-06-23T14:16:30Z

/ok to test 4609c7e

HuiyingLi and others added 9 commits June 21, 2026 00:44

feat(glm_moe_dsa): add TileLang DSA kernels

a8956d7

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

chore(glm_moe_dsa): update TileLang Tulu3 recipe

e2fed03

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

fix(glm_moe_dsa): use public GLM5.2 model id in recipe

cda4554

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

fix(glm_moe_dsa): override GLM5.2 head dim in recipe

4b358b5

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

chore: make tilelang sparse guard explicit

41b83d3

Use the TileLang validation result as an explicit guard in the sparse MLA path instead of discarding it. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

feat(glm_moe_dsa): add TileLang context parallel support

95971b6

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

HuiyingLi requested review from a team as code owners June 21, 2026 17:34

HuiyingLi changed the title ~~feat(glm_moe_dsa): add TileLang context parallel support~~ feat(glm_moe_dsa): add GLM5.2 context parallel support Jun 21, 2026

hemildesai previously approved these changes Jun 22, 2026

View reviewed changes

HuiyingLi and others added 2 commits June 22, 2026 10:31

Merge remote-tracking branch 'origin/main' into huiyingl/feat/glm52-k…

5d1f71c

…ernels Signed-off-by: HuiyingLi <willwin.lee@gmail.com> # Conflicts: # nemo_automodel/components/distributed/pipelining/hf_utils.py

HuiyingLi dismissed hemildesai’s stale review via 874c47d June 22, 2026 18:18

copy-pr-bot Bot had a problem deploying to nemo-ci June 22, 2026 18:25 Error

copy-pr-bot Bot temporarily deployed to nemo-ci June 22, 2026 18:25 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci June 22, 2026 18:25 Error

copy-pr-bot Bot temporarily deployed to nemo-ci June 22, 2026 18:25 Inactive

copy-pr-bot Bot had a problem deploying to test June 22, 2026 18:26 Error

copy-pr-bot Bot temporarily deployed to public June 22, 2026 18:26 Inactive

copy-pr-bot Bot temporarily deployed to public June 22, 2026 18:29 Inactive

style(glm_moe_dsa): format cp helper

3d312c4

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

HuiyingLi added 2 commits June 22, 2026 19:38

fix: lazy import GLM TileLang kernels

5d3edd6

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

Merge remote-tracking branch 'origin/huiyingl/feat/glm52-kernels' int…

6f70549

…o huiyingl/feat/glm52-cp Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

copy-pr-bot Bot temporarily deployed to nemo-ci June 23, 2026 04:39 Inactive

copy-pr-bot Bot temporarily deployed to test June 23, 2026 04:39 Inactive

copy-pr-bot Bot temporarily deployed to public June 23, 2026 04:39 Inactive

copy-pr-bot Bot temporarily deployed to public June 23, 2026 04:41 Inactive

copy-pr-bot Bot temporarily deployed to public June 23, 2026 04:42 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci June 23, 2026 04:43 Inactive

HuiyingLi added 2 commits June 22, 2026 23:46

test: cover GLM TileLang kernels

3fdd5b4

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

Merge remote-tracking branch 'origin/huiyingl/feat/glm52-kernels' int…

1740305

…o huiyingl/feat/glm52-cp Signed-off-by: HuiyingLi <willwin.lee@gmail.com> # Conflicts: # tests/unit_tests/models/glm_moe_dsa/test_glm_moe_dsa_tilelang.py

copy-pr-bot Bot temporarily deployed to nemo-ci June 23, 2026 09:34 Inactive

copy-pr-bot Bot temporarily deployed to test June 23, 2026 09:35 Inactive

copy-pr-bot Bot temporarily deployed to public June 23, 2026 09:35 Inactive

copy-pr-bot Bot temporarily deployed to public June 23, 2026 09:40 Inactive

hemildesai approved these changes Jun 23, 2026

View reviewed changes

Navigation Menu

Repository navigation

feat(glm_moe_dsa): add GLM5.2 context parallel support - #2695

feat(glm_moe_dsa): add GLM5.2 context parallel support#2695
HuiyingLi merged 22 commits into
mainfrom
huiyingl/feat/glm52-cp

HuiyingLi commented Jun 21, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Jun 21, 2026

Uh oh!

HuiyingLi commented Jun 22, 2026

Uh oh!

HuiyingLi commented Jun 22, 2026

Uh oh!

HuiyingLi commented Jun 22, 2026

Uh oh!

HuiyingLi commented Jun 23, 2026

Uh oh!

HuiyingLi commented Jun 23, 2026

Uh oh!

HuiyingLi commented Jun 23, 2026

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

HuiyingLi commented Jun 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Changelog

Validation

CP1/CP2 parity

Uh oh!

copy-pr-bot Bot commented Jun 21, 2026

Uh oh!

HuiyingLi commented Jun 22, 2026

Uh oh!

HuiyingLi commented Jun 22, 2026

Uh oh!

HuiyingLi commented Jun 22, 2026

Uh oh!

HuiyingLi commented Jun 23, 2026

Uh oh!

HuiyingLi commented Jun 23, 2026

Uh oh!

HuiyingLi commented Jun 23, 2026

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

HuiyingLi commented Jun 21, 2026 •

edited

Loading