Skip to content

[Ascend] Add GLM-5.2 W8A8 support in dlinfer - #353

Open
wanfengcxz wants to merge 5 commits into
mainfrom
wq/support_glm52_w8a8
Open

[Ascend] Add GLM-5.2 W8A8 support in dlinfer#353
wanfengcxz wants to merge 5 commits into
mainfrom
wq/support_glm52_w8a8

Conversation

@wanfengcxz

Copy link
Copy Markdown
Collaborator

Summary

This PR adds dlinfer runtime support for GLM-5.2 W8A8 inference on Ascend NPUs. It integrates ModelSlim quantized checkpoints, native Ascend W8A8 kernels, and the DSA sparse MLA attention path, including split KV cache and CUDA graph replay support.

Changes

W8A8 quantization

  • Add linear_w8a8_static for ModelSlim static W8A8 quantization:
    • Per-tensor activation scale and offset.
    • Per-channel weight dequantization scale and quant bias.
    • Support tensor-parallel parameter sharding.
  • Add fused_moe_w8a8 for dynamic W8A8 MoE execution.
    • Supports naive, MC2, and ALLTOALL dispatch paths.
    • Uses public Ascend operators: dynamic quantization, grouped matmul, SwiGLU, and requantization.
  • Update W8A8 linear handling to preserve arbitrary leading dimensions.
  • Extend W8A8 MoE interfaces with EP metadata and expert partitioning.

DSA/MLA attention

  • Add Ascend implementations for:
    • lightning_indexer
    • sparse_flash_attention
    • Interleaved RoPE
  • Add split MLA KV-cache support:
    • K cache stores the RoPE component.
    • V cache stores the latent/noPE component.
    • fill_kv_cache writes both caches with contiguous layouts.
  • Use BF16 for the Ascend DSA indexer cache and enforce a paged-cache block size of 128.
  • Update attention metadata and vendor APIs to use precomputed CPU cumulative sequence lengths instead of q_start_loc and cu_seq_lens_kv.
  • Avoid host-side metadata updates during sparse CUDA graph replay.

GLM-5.2 integration

  • Load ModelSlim metadata from quant_model_description.json.
  • Use separate q_a_proj and kv_a_proj_with_mqa projections because their quantization parameters are independent.
  • Disable the CUDA FlashMLA path on Ascend and route execution through dlinfer Lightning Indexer and sparse attention.
  • Disable unsupported fused/Hadamard DSA indexer preprocessing and fall back to the unfused BF16 path.
  • Ignore the ModelSlim/QuaRot-only rot.weight checkpoint tensor.
  • Replace the generic cache-engine patching with GLM-MoE-DSA-specific split-cache and graph-buffer handling.

Tests

  • Update tests/test_ascend_attention_precision.py for the new attention sequence-length interface.
  • Recommended validation:
    • Ascend native operator tests for Lightning Indexer, sparse attention, split KV cache, and interleaved RoPE.
    • GLM-5.2 W8A8 end-to-end inference on Ascend.
    • CUDA graph capture and replay for DSA decode.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant