feat: support GLM-5.2

docker/prepare_3rdparty_wheel.sh

@@ -15,7 +15,8 @@ fi
 GDRCOPY_VERSION=2.5.1
 DEEP_EP_VERSION=9af0e0d  # v1.2.1
 DEEP_GEMM_VERSION=88965b0
-FLASH_MLA_VERSION=1408756  # no release, pick the latest commit
+FLASH_MLA_VERSION=9241ae3
+FAST_HADAMARD_TRANSFORM_VERSION=v1.1.0.post2
 
 # DeepEP
 if [[ "${CUDA_VERSION_SHORT}" = "cu130" ]]; then
@@ -33,6 +34,10 @@ pip wheel -v --no-build-isolation --no-deps -w /wheels "git+https://github.com/d
 # sm100 compilation for Flash MLA requires NVCC 12.9 or higher
 FLASH_MLA_DISABLE_SM100=1 pip wheel -v --no-build-isolation --no-deps -w /wheels "git+https://github.com/deepseek-ai/FlashMLA.git@${FLASH_MLA_VERSION}"
 
+# fast_hadamard_transform
+pip wheel -v --no-build-isolation --no-deps -w /wheels \
+    "git+https://github.com/Dao-AILab/fast-hadamard-transform.git@${FAST_HADAMARD_TRANSFORM_VERSION}"
+
 # flash_attn_3 (official, LibTorch ABI-stable wheels for torch >= 2.9).
 # Select the FA3 channel matching the PyTorch binary. CUDA 13 FA3 wheels are
 # published in the cu130 channel.

lmdeploy/cli/utils.py

@@ -274,7 +274,8 @@ def _parse(x):
                                    type=_parse,
                                    default=default,
                                    help='KV cache quant policy: none/int4/int8/fp8/fp8_e5m2/'
-                                   'turbo_quant (or 0/4/8/16/17/42). fp8 defaults to fp8_e4m3.')
+                                   'turbo_quant (or 0/4/8/16/17/42). For DSA models, fp8 uses the '
+                                   'fp8_ds_mla layout.')
 
     @staticmethod
     def rope_scaling_factor(parser):

lmdeploy/messages.py

@@ -22,7 +22,7 @@ class QuantPolicy(enum.IntEnum):
     NONE = 0
     INT4 = 4  # 4-bit KV cache
     INT8 = 8  # 8-bit KV cache
-    FP8 = 16  # FP8 KV cache (float8_e4m3fn, per-tensor scale)
+    FP8 = 16  # FP8 KV cache (float8_e4m3fn, per-tensor scale. DSA uses the fp8_ds_mla layout)
     FP8_E5M2 = 17  # FP8 KV cache (float8_e5m2, per-tensor scale)
     TURBO_QUANT = 42  # TurboQuant: K=4bit QJL4 + V=2bit MSE
 
@@ -427,7 +427,7 @@ class PytorchEngineConfig:
             If unspecified, will use the default version.
         quant_policy: default to 0. When k/v is quantized into int4,
             int8, fp8, or fp8_e5m2, set it to 4, 8, 16, or 17,
-            respectively
+            respectively. For DSA models, fp8 selects the fp8_ds_mla layout.
         distributed_executor_backend: backend of distributed backend,
             options: ['uni', 'mp', 'ray']
         empty_init: Whether to load the model weights, you should set

lmdeploy/pytorch/backends/cuda/attention/default.py

@@ -30,6 +30,7 @@ class TritonAttentionMetadata(AttentionMetadata):
         num_splits: Number of splits for Flash MLA.
         cu_seqlens_q: Cumulative query sequence lengths [batch_size + 1].
         cu_seqlens_k: Cumulative KV sequence lengths [batch_size + 1].
+        indexer_kv_seqlens: Per-query causal KV lengths for DSA top-k.
         scheduler_metadata: Scheduler metadata for FA3.
         max_kv_seqlen: Maximum KV sequence length in the batch.
         max_q_seqlen: Maximum query sequence length in the batch.
@@ -47,6 +48,7 @@ class TritonAttentionMetadata(AttentionMetadata):
     num_splits: torch.Tensor = None
     cu_seqlens_q: torch.Tensor = None
     cu_seqlens_k: torch.Tensor = None
+    indexer_kv_seqlens: torch.Tensor = None
     # flash attn
     scheduler_metadata: torch.Tensor = None
     max_kv_seqlen: int = None

lmdeploy/pytorch/backends/cuda/attention/mla.py

@@ -36,26 +36,65 @@ class NSAIndicesUpdater:
 
     def __init__(self):
         self._update_decode_func = None
+        self._update_decode_strided_func = None
         self._update_prefill_func = None
 
-    def _update_decode_impl(self, nsa_indices: torch.Tensor, block_offsets: torch.Tensor,
+    def _update_decode_impl(self, nsa_indices: torch.Tensor, block_offsets: torch.Tensor, max_q_seqlen: int,
                             block_size: int) -> torch.Tensor:
         """Update for decode impl."""
+        batch_size = block_offsets.size(0)
         block_ids = nsa_indices // block_size
         block_ids = block_ids.clamp_min(0)
+        # MTP flattens the query dimension, so repeat each request's block table.
+        if block_ids.size(0) != batch_size:
+            block_offsets = torch.repeat_interleave(block_offsets,
+                                                    max_q_seqlen,
+                                                    dim=0,
+                                                    output_size=block_ids.size(0))
         block_ids = block_offsets.gather(1, block_ids)
         block_remain = nsa_indices % block_size
         ret = block_ids * block_size + block_remain
         ret[nsa_indices < 0] = -1
-        return ret[:, None]
+        return ret.unflatten(0, (batch_size, max_q_seqlen))
 
-    def update_decode(self, nsa_indices: torch.Tensor, block_offsets: torch.Tensor, block_size: int) -> torch.Tensor:
+    def update_decode(self, nsa_indices: torch.Tensor, block_offsets: torch.Tensor, max_q_seqlen: int,
+                      block_size: int) -> torch.Tensor:
         """Update for decode."""
         if self._update_decode_func is None:
             self._update_decode_func = _try_dynamic_compile(self._update_decode_impl, nsa_indices, block_offsets,
-                                                            block_size)
+                                                            max_q_seqlen, block_size)
 
-        return self._update_decode_func(nsa_indices, block_offsets, block_size)
+        return self._update_decode_func(nsa_indices, block_offsets, max_q_seqlen, block_size)
+
+    def _update_decode_strided_impl(self, nsa_indices: torch.Tensor, block_offsets: torch.Tensor, max_q_seqlen: int,
+                                    block_size: int, block_stride: int, token_stride: int,
+                                    index_stride: int) -> torch.Tensor:
+        """Map logical decode indices to an aligned strided cache view."""
+        batch_size = block_offsets.size(0)
+        block_ids = nsa_indices // block_size
+        block_ids = block_ids.clamp_min(0)
+        if block_ids.size(0) != batch_size:
+            block_offsets = torch.repeat_interleave(block_offsets,
+                                                    max_q_seqlen,
+                                                    dim=0,
+                                                    output_size=block_ids.size(0))
+        block_ids = block_offsets.gather(1, block_ids)
+        block_remain = nsa_indices % block_size
+        ret = (block_ids * block_stride + block_remain * token_stride) // index_stride
+        ret[nsa_indices < 0] = -1
+        return ret.unflatten(0, (batch_size, max_q_seqlen))
+
+    def update_decode_strided(self, nsa_indices: torch.Tensor, block_offsets: torch.Tensor, max_q_seqlen: int,
+                              block_size: int, block_stride: int, token_stride: int,
+                              index_stride: int) -> torch.Tensor:
+        """Update decode indices for strided cache storage."""
+        if self._update_decode_strided_func is None:
+            self._update_decode_strided_func = _try_dynamic_compile(self._update_decode_strided_impl, nsa_indices,
+                                                                    block_offsets, max_q_seqlen, block_size,
+                                                                    block_stride, token_stride, index_stride)
+
+        return self._update_decode_strided_func(nsa_indices, block_offsets, max_q_seqlen, block_size, block_stride,
+                                                token_stride, index_stride)
 
     def _update_prefill_impl(self, nsa_indices: torch.Tensor, q_seqlens: torch.Tensor, cu_seqlens_k: torch.Tensor):
         """Update for prefill impl."""
@@ -84,8 +123,9 @@ class FlashMLAImpl(TritonAttentionImpl):
     """Flash MLA (Multi-head Latent Attention) implementation.
 
     This implementation supports multiple execution paths:
-    - Decoding: Uses flash_mla_with_kvcache with paged KV cache
-    - Prefill with NSA: Uses flash_mla_sparse_fwd for sparse attention
+    - Paged FlashMLA decode: Uses flash_mla_with_kvcache with FP8 paged KV cache
+    - Sparse FlashMLA decode: Uses flash_mla_sparse_fwd over a zero-copy BF16 cache view
+    - Sparse FlashMLA prefill: Uses flash_mla_sparse_fwd for NSA attention
     - Prefill with FA3: Uses flash_attn_varlen_func with split q_rope/q_nope
     - Prefill fallback: Uses custom Triton kernel
     """
@@ -94,6 +134,7 @@ class FlashMLAImpl(TritonAttentionImpl):
     _MLA_HEAD_ALIGNMENT = 64  # Query heads must be multiple of 64 for flash_mla
     _MLA_NOPE_SIZE = 512  # Size of non-positional embeddings
     _MLA_SCALE_SIZE = 16  # Size of FP8 quantization scales
+    _BF16_CACHE_INDEX_STRIDE = 64  # Address BF16 cache rows in aligned 128-byte units
 
     def __init__(
         self,
@@ -152,7 +193,7 @@ def _get_flash_mla_sparse_fwd(self):
         except Exception:
             logger.exception('Can not import flash_mla_sparse_fwd from flash_mla.')
 
-    def flash_mla_decoding(
+    def _decode_paged_flash_mla(
         self,
         query: torch.Tensor,
         k_cache: torch.Tensor,
@@ -170,14 +211,22 @@ def flash_mla_decoding(
         max_q_seqlen = query.numel() // (query.size(-1) * query.size(-2))
         max_q_seqlen = max_q_seqlen // batch_size
         query = query.unflatten(0, (batch_size, max_q_seqlen))
+        num_q_heads = query.size(2)
         if kv_seqlens.dtype == torch.int64:
             kv_seqlens = kv_seqlens.to(torch.int32)
 
         # update nsa indice according to flash-mla requirement
         if nsa_indices is not None:
             block_size = k_cache.size(1)
-            nsa_indices = self.nsa_updater.update_decode(nsa_indices, block_offsets, block_size)
+            nsa_indices = self.nsa_updater.update_decode(nsa_indices, block_offsets, max_q_seqlen, block_size)
             causal = False
+            # The new FlashMLASchedMeta API uses a sparse decoder that only
+            # accepts 64 or 128 query heads. The old API stores metadata in a
+            # tensor and supports the unpadded TP head count.
+            if not isinstance(attn_metadata.tile_scheduler_metadata, torch.Tensor):
+                pad_heads = -num_q_heads % self._MLA_HEAD_ALIGNMENT
+                if pad_heads:
+                    query = torch.nn.functional.pad(query, (0, 0, 0, pad_heads))
 
         attn_output, _ = self.flash_mla_with_kvcache(query,
                                                      k_cache=k_cache,
@@ -191,45 +240,58 @@ def flash_mla_decoding(
                                                      is_fp8_kvcache=is_fp8_kvcache,
                                                      indices=nsa_indices)
 
+        attn_output = attn_output[:, :, :num_q_heads]
         attn_output = attn_output.flatten(0, 1)
         return attn_output
 
-    def _prefill_sparse(self, query: torch.Tensor, flatten_k: torch.Tensor, nsa_indices: torch.Tensor,
-                        attn_metadata: TritonAttentionMetadata) -> torch.Tensor:
-        """Sparse prefill using flash_mla_sparse_fwd.
-
-        This path is used when NSA (Non-contiguous Sparse Attention) indices are provided.
-        Requires FP8 KV cache and flash_mla library.
-
-        Args:
-            query: Query tensor.
-            flatten_k: Flattened key cache.
-            nsa_indices: Sparse attention indices.
-            attn_metadata: Attention metadata.
-
-        Returns:
-            Attention output tensor.
-        """
-        q_seqlens = attn_metadata.q_seqlens
+    def _flash_mla_sparse(self, query: torch.Tensor, indexed_k: torch.Tensor,
+                          nsa_indices: torch.Tensor) -> torch.Tensor:
+        """Run sparse FlashMLA over index-addressable BF16 KV storage."""
         flash_mla_sparse_fwd = self._get_flash_mla_sparse_fwd()
-
         num_q_heads = query.size(1)
         # flash_mla_sparse_fwd requires query heads to be multiple of alignment
         if num_q_heads % self._MLA_HEAD_ALIGNMENT != 0:
             padding = self._MLA_HEAD_ALIGNMENT - num_q_heads % self._MLA_HEAD_ALIGNMENT
             query = torch.nn.functional.pad(query, (0, 0, 0, padding))
 
-        nsa_indices = self.nsa_updater.update_prefill(nsa_indices, q_seqlens, attn_metadata.cu_seqlens_k)
         output = flash_mla_sparse_fwd(
             query,
-            flatten_k,
+            indexed_k,
             nsa_indices,
             sm_scale=self.scale,
         )
         attn_output = output[0]
         attn_output = attn_output[:, :num_q_heads]
         return attn_output
 
+    def _prefill_sparse(self, query: torch.Tensor, flatten_k: torch.Tensor, nsa_indices: torch.Tensor,
+                        attn_metadata: TritonAttentionMetadata) -> torch.Tensor:
+        """Run sparse prefill over the flattened BF16 KV cache."""
+        nsa_indices = self.nsa_updater.update_prefill(nsa_indices, attn_metadata.q_seqlens,
+                                                      attn_metadata.cu_seqlens_k)
+        return self._flash_mla_sparse(query, flatten_k, nsa_indices)
+
+    def _decode_bf16_sparse_flash_mla(self, query: torch.Tensor, k_cache: torch.Tensor, nsa_indices: torch.Tensor,
+                                      attn_metadata: TritonAttentionMetadata) -> torch.Tensor:
+        """Run sparse decode over a zero-copy BF16 paged-cache view."""
+        assert query.dtype == torch.bfloat16, 'BF16 sparse MLA requires a bfloat16 query'
+        assert k_cache.dtype == torch.bfloat16, 'BF16 sparse MLA requires a bfloat16 KV cache'
+        block_size = k_cache.size(1)
+        max_q_seqlen = self._get_max_q_seqlen(query, attn_metadata)
+        # The cache shares a block record with indexer state, so flattening it
+        # copies the full cache capacity. Expose the same storage in aligned
+        # address units and let sparse indices point directly at each KV row.
+        index_stride = self._BF16_CACHE_INDEX_STRIDE
+        block_stride, token_stride = k_cache.stride()[:2]
+        last_token_offset = ((k_cache.size(0) - 1) * block_stride + (block_size - 1) * token_stride)
+        storage_rows = last_token_offset // index_stride + 1
+        storage_k = k_cache.as_strided((storage_rows, *k_cache.shape[2:]),
+                                       (index_stride, *k_cache.stride()[2:]))
+        nsa_indices = self.nsa_updater.update_decode_strided(nsa_indices, attn_metadata.block_offsets, max_q_seqlen,
+                                                             block_size, block_stride, token_stride, index_stride)
+        nsa_indices = nsa_indices.flatten(0, 1)[:, None]
+        return self._flash_mla_sparse(query, storage_k, nsa_indices)
+
     def _prefill_triton(
         self,
         query: torch.Tensor,
@@ -392,6 +454,7 @@ def _fill_kv_cache_impl(self,
         """Fill kv cache."""
         is_fp8_kvcache = k_cache.dtype == torch.float8_e4m3fn
         if not is_fp8_kvcache:
+            # The BF16 MLA V cache aliases K, and the base writer skips its duplicate store.
             return super()._fill_kv_cache_impl(
                 key,
                 value,
@@ -456,8 +519,8 @@ def _forward_decoding(
     ) -> torch.Tensor:
         """Forward pass for decoding stage.
 
-        Uses flash_mla_with_kvcache for efficient decoding with paged KV cache.
-        Supports both regular and sparse (NSA) attention patterns.
+        Uses paged FlashMLA decode by default. Sparse NSA with BF16 cache uses
+        sparse FlashMLA decode.
 
         Args:
             query: Query tensor.
@@ -468,7 +531,9 @@ def _forward_decoding(
         Returns:
             Attention output tensor.
         """
-        return self.flash_mla_decoding(query, k_cache, nsa_indices, attn_metadata)
+        if nsa_indices is not None and k_cache.dtype != torch.float8_e4m3fn:
+            return self._decode_bf16_sparse_flash_mla(query, k_cache, nsa_indices, attn_metadata)
+        return self._decode_paged_flash_mla(query, k_cache, nsa_indices, attn_metadata)
 
     def _forward_prefill(
         self,
@@ -483,7 +548,7 @@ def _forward_prefill(
         """Forward pass for prefill stage.
 
         Supports three execution paths:
-        1. Sparse (NSA + FP8): flash_mla_sparse_fwd for sparse attention
+        1. Sparse FlashMLA: flash_mla_sparse_fwd over BF16 flattened KV
         2. FA3 optimized: flash_attn_varlen_func with split q_rope/q_nope
         3. Triton fallback: Custom Triton kernel implementation
 
@@ -534,15 +599,15 @@ def forward(
         """Forward pass for MLA attention computation.
 
         This method handles both prefill and decoding stages by:
-        1. Validating NSA requirements (FP8 KV cache)
-        2. Computing max query sequence length
-        3. Filling KV cache if new key/value are provided
-        4. Dispatching to appropriate stage-specific method
+        1. Computing max query sequence length
+        2. Filling KV cache if new key/value are provided
+        3. Dispatching to the cache-specific stage implementation
 
         Architecture:
-        - Decoding: Uses flash_mla_with_kvcache with paged KV cache
+        - Paged FlashMLA decode: Uses flash_mla_with_kvcache with FP8 paged KV cache
+        - Sparse FlashMLA decode: Uses flash_mla_sparse_fwd over a zero-copy BF16 cache view
         - Prefill: Three paths based on availability and requirements
-          * Sparse (NSA + FP8): flash_mla_sparse_fwd
+          * Sparse FlashMLA: flash_mla_sparse_fwd
           * FA3 optimized: flash_attn_varlen_func with split q_rope/q_nope
           * Triton fallback: Custom triton kernel
 
@@ -560,12 +625,6 @@ def forward(
         Returns:
             Attention output tensor.
         """
-        # Validate NSA requirements
-        is_nsa = nsa_indices is not None
-        if is_nsa:
-            is_fp8_kvcache = k_cache.dtype == torch.float8_e4m3fn
-            assert is_fp8_kvcache, 'NSA sparse attention requires FP8 KV cache'
-
         # Shared preparation
         max_q_seqlen = self._get_max_q_seqlen(query, attn_metadata)
 

lmdeploy/pytorch/backends/cuda/graph_runner.py

@@ -240,7 +240,9 @@ def get_graph_key(self, input_ids: torch.Tensor, position_ids: torch.Tensor, pas
             batch_size = self._get_capture_tokens(batch_size)
         else:
             batch_size = self._get_capture_tokens(meta.padding_batch_size)
-        return (batch_size, is_decoding, enable_microbatch, query_len)
+        graph_key = (batch_size, is_decoding, enable_microbatch, query_len)
+        graph_key += self.model.get_cudagraph_extra_key(**kwargs)
+        return graph_key
 
     def _prepare_inputs(self, **kwargs):
         """Prepare inputs."""