Repository navigation
[Ascend] GLM-5.2 W8A8 ops: fused MC2 MoE, MLAPO decode, on-device NZ conversion - #355
Open
wanfengcxz wants to merge 10 commits into
Open
wanfengcxz wants to merge 10 commits into
wanfengcxz wants to merge 10 commits into
Conversation
mla_preprocess returns q_nope/q_pe separately; routing them through the fused sparse_attention_fwd forced a torch.cat followed by slice + contiguous. Call sparse_attention_fwd_split with the split tensors directly (verified: 12480/12480 decode calls take this path, greedy batch output byte-identical). Co-Authored-By: Claude Code <noreply@anthropic.com>
lmdeploy stock convert_weight_to_nz round-trips every expert tensor through the host (D2H + cpu transpose + H2D); with all TP workers converting concurrently that phase takes minutes after weight loading. Patch LinearWeightsW8A8 to nz_weight.copy_(param.transpose(1, 2)) instead, which performs the ND->NZ format cast as a d2d copy with peak device memory growing by one expert tensor only. lmdeploy main w8a8.py has no NZ support at all, so also wrap FusedMoEW8A8.update_weights to trigger the conversion; the _weight_nz flag keeps it idempotent on branch checkouts that still call convert_weight_to_nz themselves. Verified on GLM-5.2-W8A8 tp16/ep16 against lmdeploy 39907d02 (w8a8.py aligned with main): 16/16 ranks convert on device, post-load conversion ~10s (was ~240s), generation output byte-identical. Co-Authored-By: Claude Code <noreply@anthropic.com>
patch_linear_w8a8_nz_convert probes dispatch_ffn_combine availability before converting: whether the expert weights end up NZ is decided once at weight-load time, and the weight format is that decision. prepare_fused_w8a8_weights shrinks to pack_w8a8_scales (scale packing only) - no more operator probing or NZ re-ensuring in the weight-update path. Verified on 16x 910B tp16/ep16: 16/16 ranks convert at load, greedy outputs match the pre-refactor state. Co-Authored-By: Claude Code <noreply@anthropic.com>
Co-Authored-By: Claude Code <noreply@anthropic.com>
wanfengcxz
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Ascend ops for GLM-5.2 W8A8, paired with the lmdeploy-side PR (engine wiring lives there).
Changes:
vendor/ascend/moe.py,torch_npu_ops.py)writes, split queries fed straight to SFA (
vendor/ascend/mlapo.py,ops/llm.py)weights — replaces the CPU round-trip that stalls 3-4min when 16 TP workers convert
concurrently (
framework/lmdeploy_ext/device/__init__.py)Verified on 16x Ascend 910B (tp=16, ep=16), graph mode: outputs identical to the
branch state before the lmdeploy-side refactor.