[None][feat] Add AutoDeploy custom model for GLM-5 (glm_moe_dsa)

examples/auto_deploy/model_registry/configs/glm_5.yaml

@@ -0,0 +1,36 @@
+# Configuration for GLM-5 (zai-org/GLM-5)
+# Custom model: glm_moe_dsa — prefill-only with MLA+DSA attention and noaux_tc MoE routing
+# 78 layers total (3 dense + 75 MoE), 256 routed experts, 8-way tensor parallelism
+runtime: trtllm
+compile_backend: torch-cudagraph
+max_seq_len: 4096
+max_num_tokens: 4096
+max_batch_size: 64
+world_size: 8
+enable_chunked_prefill: true
+cuda_graph_batch_sizes: [1, 2, 4, 8, 16, 32, 64]
+model_factory: AutoModelForCausalLM
+# Use GLM-4.7-Flash tokenizer since GLM-5's tokenizer_config.json specifies
+# TokenizersBackend which is not a standard transformers class
+tokenizer: zai-org/GLM-4.7-Flash
+model_kwargs:
+  torch_dtype: bfloat16
+kv_cache_config:
+  enable_block_reuse: false
+  free_gpu_memory_fraction: 0.7
+  tokens_per_block: 64
+transforms:
+  match_swiglu_pattern:
+    enabled: true
+  fuse_swiglu:
+    enabled: true
+  # GLM-5 uses torch_dsa (DSA = DeepSeek Sparse Attention) not torch_mla.
+  # Override the MLA cache insertion to use torch_dsa backend so the KV cache
+  # is correctly inserted around the torch_dsa attention nodes.
+  insert_cached_mla_attention:
+    stage: cache_init
+    requires_shape_prop: true
+    backend: torch_dsa
+  multi_stream_mla_attn:
+    stage: compile
+    enabled: true

examples/auto_deploy/model_registry/models.yaml

@@ -224,9 +224,9 @@ models:
   yaml_extra: ['qwen3.5_moe_400b.yaml']
 # --- GLM-5 (Feb 2026) ---
 - name: zai-org/GLM-5
-  yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml']
+  yaml_extra: ['glm_5.yaml']
 - name: zai-org/GLM-5-FP8
-  yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml']
+  yaml_extra: ['glm_5.yaml']
 # --- MiniMax-M2.5 (Feb 2026) ---
 - name: MiniMaxAI/MiniMax-M2.5
   yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml']

tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py

@@ -1392,6 +1392,58 @@ def allocate(self, sequence_info: SequenceInfo) -> torch.Tensor:
         )
 
 
+class MLAPagedResourceHandler(ResourceHandler):
+    """Handler for paged resources in MLA that require per-layer contiguous memory.
+
+    While MLA uses paged caching, the underlying FlashMLA kernel uses a uint32_t to track the
+    strides for the cache. The KVCacheManager will allocate a contiguous tensor for the cache
+    across all layers with dim 0 representing the layer index. Hence, the per-layer cache has very
+    large strides to jump between pages which causes overflow in the MLA kernel that uses uint32_t
+    for strides.
+
+    We use a separate handler for this purpose to avoid registering the cache with the
+    KVCacheManager and instead rely on local allocation.
+    """
+
+    @property
+    def is_paged(self) -> bool:
+        """Whether the resource is paged."""
+        return True
+
+    def __init__(self, *token_shape: int, dtype: torch.dtype) -> None:
+        """Initialize the MLAPagedResourceHandler.
+
+        Args:
+            token_shape: The shape of the resource per token.
+            dtype: The dtype of the resource.
+        """
+        self.token_shape = token_shape
+        self.dtype = dtype
+
+    def _get_bytes_per_token(self) -> int:
+        """The size of the resource per token in bytes."""
+        from math import prod
+
+        return prod(self.token_shape) * self.dtype.itemsize
+
+    def allocate(self, sequence_info: SequenceInfo) -> torch.Tensor:
+        """Allocate contiguous paged resource.
+
+        Args:
+            sequence_info: SequenceInfo with device and page information.
+
+        Returns:
+            Contiguous tensor of shape [num_blocks, tokens_per_block, *token_shape].
+        """
+        return torch.empty(
+            sequence_info.num_blocks,
+            sequence_info.tokens_per_block,
+            *self.token_shape,
+            device=sequence_info.device,
+            dtype=self.dtype,
+        )
+
+
 class MHACallable(Protocol):
     def __call__(
         self,

tensorrt_llm/_torch/auto_deploy/custom_ops/mla/__init__.py

@@ -1,21 +1,34 @@
-"""MLA (Multi-head Latent Attention) custom ops.
+"""MLA (Multi-head Latent Attention) and DSA (DeepSeek Sparse Attention) custom ops.
 
 Exports:
 - TorchBackendMLAAttention: Attention descriptor for MLA (registered as "torch_mla")
 - FlashInferMLAAttention: Attention descriptor for FlashInfer MLA (registered as "flashinfer_mla")
+- TorchBackendDSAAttention: Attention descriptor for DSA (registered as "torch_dsa")
+- FlashMLADSAAttention: Attention descriptor for FlashMLA DSA (registered as "flashmla_dsa")
 - torch_mla: Source op for MLA attention
+- torch_dsa: Source op for DSA attention (MLA + Indexer sparse masking)
 - torch_backend_mla_with_cache: Cached backend op with FlashInfer-compatible cache
+- torch_backend_dsa_with_cache: Cached backend op for DSA
 - flashinfer_mla_with_cache: Cached backend op using FlashInfer MLA kernels
+- flash_mla_dsa_with_cache: Cached backend op for DSA using FlashMLA paged kernels
 """
 
 from .flashinfer_mla import FlashInferMLAAttention, flashinfer_mla_with_cache
+from .flashmla_dsa import FlashMLADSAAttention, flash_mla_dsa_with_cache
+from .torch_backend_dsa import TorchBackendDSAAttention, torch_backend_dsa_with_cache
 from .torch_backend_mla import TorchBackendMLAAttention, torch_backend_mla_with_cache
+from .torch_dsa import torch_dsa
 from .torch_mla import torch_mla
 
 __all__ = [
     "TorchBackendMLAAttention",
     "FlashInferMLAAttention",
+    "TorchBackendDSAAttention",
+    "FlashMLADSAAttention",
     "torch_mla",
+    "torch_dsa",
     "torch_backend_mla_with_cache",
+    "torch_backend_dsa_with_cache",
     "flashinfer_mla_with_cache",
+    "flash_mla_dsa_with_cache",
 ]

tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py

@@ -19,7 +19,6 @@
 
 import math
 from dataclasses import dataclass, fields
-from math import prod
 from typing import Dict, List, Literal, Optional, Tuple
 
 import flashinfer
@@ -37,11 +36,10 @@
     AttentionRegistry,
     Constant,
     MHACallable,
+    MLAPagedResourceHandler,
     PrepareMetadataCallable,
     PrepareMetadataHostCallable,
-    ResourceHandler,
     ResourceHandlerDict,
-    SequenceInfo,
 )
 
 
@@ -785,56 +783,6 @@ def flashinfer_mla_with_cache_fake(
     ).contiguous()
 
 
-class MLAPagedResourceHandler(ResourceHandler):
-    """Handler for paged resources in MLA that require per-layer contiguous memory.
-
-    While MLA uses paged caching, the underlying flashinfer MLA kernel uses a uint32_t to track the
-    strides for the cache. The KVCacheManager will allocate a contiguous tensor for the cache
-    across all layers with dim 0 representing the layer index. Hence, the per-layer cache has very
-    large strides to jump between pages which causes overflow in the MLA kernel that uses uint32_t
-    for strides.
-
-    We use a separate handler for this purpose to avoid registering the cache with the
-    KVCacheManager and instead rely on local allocation.
-    """
-
-    @property
-    def is_paged(self) -> bool:
-        """Whether the resource is paged."""
-        return True
-
-    def __init__(self, *token_shape: int, dtype: torch.dtype) -> None:
-        """Initialize the ContiguousPagedResourceHandler.
-
-        Args:
-            token_shape: The shape of the resource per token.
-            dtype: The dtype of the resource.
-        """
-        self.token_shape = token_shape
-        self.dtype = dtype
-
-    def _get_bytes_per_token(self) -> int:
-        """The size of the resource per token in bytes."""
-        return prod(self.token_shape) * self.dtype.itemsize
-
-    def allocate(self, sequence_info: SequenceInfo) -> torch.Tensor:
-        """Allocate contiguous paged resource.
-
-        Args:
-            sequence_info: SequenceInfo with device and page information.
-
-        Returns:
-            Contiguous tensor of shape [num_blocks, tokens_per_block, *token_shape].
-        """
-        return torch.empty(
-            sequence_info.num_blocks,
-            sequence_info.tokens_per_block,
-            *self.token_shape,
-            device=sequence_info.device,
-            dtype=self.dtype,
-        )
-
-
 @AttentionRegistry.register("flashinfer_mla")
 class FlashInferMLAAttention(AttentionDescriptor):
     """Attention descriptor for FlashInfer-based MLA with paged cache.