diff --git a/examples/auto_deploy/model_registry/configs/glm_5.yaml b/examples/auto_deploy/model_registry/configs/glm_5.yaml new file mode 100644 index 000000000000..bc02c8121ff9 --- /dev/null +++ b/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 diff --git a/examples/auto_deploy/model_registry/models.yaml b/examples/auto_deploy/model_registry/models.yaml index 91b2194507ee..c5765743b799 100644 --- a/examples/auto_deploy/model_registry/models.yaml +++ b/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'] diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py index 7272ee591403..89bf47dcb23c 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/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, diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/__init__.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/__init__.py index 261617de1b71..6b2070bcfddb 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/__init__.py +++ b/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", ] diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py index 6985ce99cebd..69e7a81fe445 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py +++ b/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. diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashmla_dsa.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashmla_dsa.py new file mode 100644 index 000000000000..41279ca9b354 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashmla_dsa.py @@ -0,0 +1,673 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""FlashMLA-based DSA (DeepSeek Sparse Attention) backend with paged KV caches. + +Provides: +- flash_mla_dsa_with_cache: cached DSA op using FlashMLA kernels with paged caches +- FlashMLADSAAttention: AttentionDescriptor registered as "flashmla_dsa" + +Cache layout (paged, managed by MLAPagedResourceHandler): + mla_cache: [num_blocks, page_size, 1, kv_lora_rank + qk_rope_head_dim] + index_k_cache: [num_blocks, page_size, index_head_dim] + +Decode: flash_mla_with_kvcache (sparse, causal=False, no Python loops) +Prefill: flash_mla_sparse_fwd (flat ragged, no Python loops) when kv_lora_rank==512; + falls back to flash_mla_with_kvcache(causal=True) + Q-padding otherwise. +""" + +import math +from typing import List, Optional, Tuple + +import torch +from torch._ops import OpOverloadPacket +from torch.fx import Node + +from .....llmapi.llm_args import KvCacheConfig +from ..attention_interface import ( + AttentionDescriptor, + AttentionLayout, + AttentionRegistry, + Constant, + MHACallable, + MLAPagedResourceHandler, + ResourceHandlerDict, +) + +# --------------------------------------------------------------------------- +# Vectorized paged-cache helpers (no Python loops over sequences) +# --------------------------------------------------------------------------- + + +def _write_tokens_to_paged_caches( + compressed_kv: torch.Tensor, # [total_tokens, kv_lora_rank] + kpe: torch.Tensor, # [total_tokens, 1, qk_rope_head_dim] + index_k: torch.Tensor, # [total_tokens, index_head_dim] + mla_cache: torch.Tensor, # [num_blocks, page_size, 1, kv_lora_rank + rope_dim] + index_k_cache: torch.Tensor, # [num_blocks, page_size, index_head_dim] + seq_len: torch.Tensor, # [B] int32 — tokens per seq in this batch + input_pos: torch.Tensor, # [B] int32 — absolute start position in cache per seq + cu_seqlen: torch.Tensor, # [B+1] int32 — cumulative token offsets in flat tensors + cache_loc: torch.Tensor, # [total_pages] int32 — page index array + cu_num_pages: torch.Tensor, # [B+1] int32 — cumulative page counts per seq +) -> None: + """Write all tokens to paged caches using vectorized scatter (no Python loops).""" + B = seq_len.shape[0] + total_tokens = int(cu_seqlen[-1].item()) + page_size = mla_cache.shape[1] + device = mla_cache.device + + # seq_ids[t] = which sequence token t belongs to + seq_ids = torch.repeat_interleave(torch.arange(B, device=device), seq_len) # [total_tokens] + + # local_pos[t] = position within sequence (0-indexed from the batch's q start) + seq_start = cu_seqlen[:-1] # [B] + local_pos = torch.arange(total_tokens, device=device) - seq_start[seq_ids] + + # abs_pos[t] = absolute position in cache for token t + abs_pos = input_pos[seq_ids] + local_pos # [total_tokens] + + # Map to (page_idx, page_offset) via page table + page_k = abs_pos // page_size # [total_tokens] + page_off = abs_pos % page_size # [total_tokens] + gather_base = cu_num_pages[seq_ids] + page_k # [total_tokens] + page_indices = cache_loc[gather_base] # [total_tokens] + + # Combine ckv + kpe into mla_cache entry + kpe_sq = kpe.squeeze(1) # [total_tokens, rope_dim] + flat_kv = torch.cat([compressed_kv, kpe_sq], dim=-1) # [total_tokens, lora+rope] + mla_cache[page_indices, page_off, 0, :] = flat_kv.to(mla_cache.dtype) + index_k_cache[page_indices, page_off, :] = index_k.to(index_k_cache.dtype) + + +def _gather_index_k_dense( + index_k_cache: torch.Tensor, # [num_blocks, page_size, D_idx] + cache_loc: torch.Tensor, # [total_pages] int32 + cu_num_pages: torch.Tensor, # [B+1] int32 + cache_seqlens: torch.Tensor, # [B] int32 — total cached lengths per seq +) -> Tuple[torch.Tensor, torch.Tensor]: + """Gather index_k into dense [B, max_T, D_idx] without Python loops. + + Returns: + dense_index_k: [B, max_T, D_idx] (padded with zeros) + valid_mask: [B, max_T] bool (True = valid KV position) + """ + B = cache_seqlens.shape[0] + max_T = int(cache_seqlens.max().item()) + page_size = index_k_cache.shape[1] + device = index_k_cache.device + + t_range = torch.arange(max_T, device=device) # [max_T] + page_k = t_range.unsqueeze(0) // page_size # [1, max_T] + page_off = t_range.unsqueeze(0) % page_size # [1, max_T] + + # Clamp to avoid out-of-bounds on padded positions + base = cu_num_pages[:-1].unsqueeze(1) # [B, 1] + gather_idx = (base + page_k).clamp(0, cache_loc.shape[0] - 1) # [B, max_T] + page_indices = cache_loc[gather_idx] # [B, max_T] + + dense_index_k = index_k_cache[page_indices, page_off, :] # [B, max_T, D_idx] + + # Mask out positions beyond actual cached length + valid_mask = t_range.unsqueeze(0) < cache_seqlens.unsqueeze(1) # [B, max_T] + return dense_index_k, valid_mask + + +def _gather_mla_kv_dense( + mla_cache: torch.Tensor, # [num_blocks, page_size, 1, D_qk] + cache_loc: torch.Tensor, # [total_pages] int32 + cu_num_pages: torch.Tensor, # [B+1] int32 + cache_seqlens: torch.Tensor, # [B] int32 +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Gather all cached KV tokens into flat dense tensors without Python loops. + + Returns: + dense_kv: [total_kv_tokens, 1, D_qk] for flash_mla_sparse_fwd + kv_cu_seqlens: [B+1] int32 cumulative KV seqlens + """ + B = cache_seqlens.shape[0] + total_kv = int(cache_seqlens.sum().item()) + page_size = mla_cache.shape[1] + device = mla_cache.device + + kv_seq_ids = torch.repeat_interleave(torch.arange(B, device=device), cache_seqlens) + kv_cu_seqlens = torch.cat([cache_seqlens.new_zeros(1), cache_seqlens.cumsum(0)]) + kv_local_pos = torch.arange(total_kv, device=device) - kv_cu_seqlens[kv_seq_ids] + + page_k = kv_local_pos // page_size + page_off = kv_local_pos % page_size + gather_base = cu_num_pages[kv_seq_ids] + page_k + page_indices = cache_loc[gather_base] + + dense_kv = mla_cache[page_indices, page_off, :, :] # [total_kv, 1, D_qk] + return dense_kv, kv_cu_seqlens + + +def _build_block_table( + cache_loc: torch.Tensor, # [total_pages] int32 + cu_num_pages: torch.Tensor, # [B+1] int32 + B: int, + device: torch.device, +) -> torch.Tensor: + """Build block_table [B, max_pages] from cache_loc and cu_num_pages (no Python loop).""" + page_counts = cu_num_pages[1:] - cu_num_pages[:-1] # [B] + max_pages = int(page_counts.max().item()) + block_table = cache_loc.new_zeros(B, max_pages) + col_idx = torch.arange(max_pages, device=device).unsqueeze(0) # [1, max_pages] + valid = col_idx < page_counts.unsqueeze(1) # [B, max_pages] + flat_src_idx = cu_num_pages[:-1].unsqueeze(1) + col_idx # [B, max_pages] + flat_src_idx = flat_src_idx.clamp(0, cache_loc.shape[0] - 1) + block_table[valid] = cache_loc[flat_src_idx[valid]] + return block_table.to(torch.int32) + + +def _compute_batched_index_topk( + index_q: torch.Tensor, # [B, S, H, D_idx] + dense_index_k: torch.Tensor, # [B, max_T, D_idx] + index_weights: torch.Tensor, # [B, S, H] + valid_mask: torch.Tensor, # [B, max_T] bool + index_topk: int, + index_softmax_scale: float, + is_causal: bool, + cache_seqlens: torch.Tensor, # [B] int32 +) -> torch.Tensor: + """Returns sparse_indices [B, S, topk] int32 — no Python loops.""" + S = index_q.shape[1] + max_T = dense_index_k.shape[1] + + # [B, S, H, max_T] + per_head_scores = torch.einsum( + "bshd,btd->bsht", + index_q.float(), + dense_index_k.float(), + ) + # [B, S, max_T] + index_score = torch.einsum("bsht,bsh->bst", per_head_scores, index_weights.float()) + index_score = index_score * index_softmax_scale + + # Mask invalid (padded) KV positions + index_score.masked_fill_(~valid_mask.unsqueeze(1), float("-inf")) + + # Causal mask for prefill + if is_causal and S > 1: + s_range = torch.arange(S, device=index_q.device) # [S] + t_range = torch.arange(max_T, device=index_q.device) # [max_T] + # query s can attend to KV positions t where t < cache_seqlens[b] - S + s + 1 + causal_limit = cache_seqlens.unsqueeze(1).long() - S + s_range.unsqueeze(0) # [B, S] + future_mask = t_range.unsqueeze(0).unsqueeze(0) >= causal_limit.unsqueeze( + 2 + ) # [B, S, max_T] + index_score.masked_fill_(future_mask, float("-inf")) + + effective_topk = min(index_topk, max_T) + topk_indices = index_score.topk(effective_topk, dim=-1).indices # [B, S, topk] + return topk_indices.to(torch.int32) + + +# --------------------------------------------------------------------------- +# Custom op: flash_mla_dsa_with_cache +# --------------------------------------------------------------------------- + + +def _is_sparse_mla_supported() -> bool: + """Check if sparse FlashMLA is supported on the current device (SM100+ / Blackwell).""" + if not torch.cuda.is_available(): + return False + return torch.cuda.get_device_capability()[0] >= 10 + + +@torch.library.custom_op("auto_deploy::flash_mla_dsa_with_cache", mutates_args=()) +def flash_mla_dsa_with_cache( + # 8 QKV tensor args (matches get_num_qkv_args = 8) + q_nope: torch.Tensor, # [B, S, N, qk_nope_head_dim] + q_pe: torch.Tensor, # [B, S, N, qk_rope_head_dim] + compressed_kv: torch.Tensor, # [B, S, kv_lora_rank] + kpe: torch.Tensor, # [B, S, 1, qk_rope_head_dim] + kv_b_proj_weight: torch.Tensor, # [N*(qk_nope+v), kv_lora_rank] + index_q: torch.Tensor, # [B, S, H, D_idx] + index_k: torch.Tensor, # [B, S, D_idx] + index_weights: torch.Tensor, # [B, S, H] + # Standard paged metadata + batch_info_host: torch.Tensor, # [3] int host: [num_prefill, num_prefill_tokens, num_decode] + seq_len: torch.Tensor, # [B] int32 — token counts per seq in this batch + input_pos: torch.Tensor, # [B] int32 — position offset into cache per seq + cu_seqlen: torch.Tensor, # [B+1] int32 — cumulative token offsets in flat tensors + cache_loc: torch.Tensor, # [total_pages] int32 — page index array + cu_num_pages: torch.Tensor, # [B+1] int32 — cumulative page counts per seq + last_page_len: torch.Tensor, # [B] int32 — valid tokens in last page + # Paged caches + mla_cache: torch.Tensor, # [num_blocks, page_size, 1, kv_lora_rank + rope_dim] + index_k_cache: torch.Tensor, # [num_blocks, page_size, index_head_dim] + # Constants + scale: Optional[float] = None, + kv_lora_rank: int = 512, + index_topk: int = 64, +) -> torch.Tensor: + """FlashMLA-based DSA with paged KV caches. + + Decode (S==1): + - SM100+ (Blackwell): flash_mla_with_kvcache with sparse indices + - SM90 (Hopper): flash_mla_with_kvcache dense (no indices); sparsity not supported + Prefill (S>1): + - SM100+ and kv_lora_rank==512: flash_mla_sparse_fwd (flat ragged, no Q-padding) + - Otherwise: flash_mla_with_kvcache(causal=True) + Q-padding + """ + from tensorrt_llm.flash_mla.flash_mla_interface import ( + flash_mla_sparse_fwd, + flash_mla_with_kvcache, + get_mla_metadata, + ) + + b, s = q_nope.shape[:2] + num_heads = q_nope.shape[2] + qk_nope_head_dim = q_nope.shape[3] + qk_rope_head_dim = q_pe.shape[3] + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + + out_features = kv_b_proj_weight.shape[0] + kv_head_dim = out_features // num_heads + v_head_dim = kv_head_dim - qk_nope_head_dim + + index_head_dim = index_q.shape[-1] + index_softmax_scale = 1.0 / math.sqrt(index_head_dim) + + if scale is None: + scale = 1.0 / math.sqrt(qk_head_dim) + + num_prefill, num_prefill_tokens, num_decode = batch_info_host.tolist() + num_seq = num_prefill + num_decode + seq_len_active = seq_len[:num_seq] + input_pos_active = input_pos[:num_seq] + cu_seqlen_active = cu_seqlen[: num_seq + 1] + + # Extract MLA weight matrices + w = kv_b_proj_weight.view(num_heads, kv_head_dim, kv_lora_rank) + w_kn = w[:, :qk_nope_head_dim, :] # [N, nope, lora] + w_v = w[:, qk_nope_head_dim:, :] # [N, v, lora] + + device = q_nope.device + use_sparse = _is_sparse_mla_supported() + + # Flatten inputs to [total_tokens, ...] for cache-write helper + total_tokens = int(cu_seqlen_active[-1].item()) + compressed_kv_flat = compressed_kv.reshape(total_tokens, kv_lora_rank) + kpe_flat = kpe.reshape(total_tokens, 1, qk_rope_head_dim) + index_k_flat = index_k.reshape(total_tokens, index_head_dim) + + # Write new tokens to paged caches + _write_tokens_to_paged_caches( + compressed_kv_flat, + kpe_flat, + index_k_flat, + mla_cache, + index_k_cache, + seq_len_active, + input_pos_active, + cu_seqlen_active, + cache_loc, + cu_num_pages, + ) + + # cache_seqlens = total tokens in cache after write + cache_seqlens = (input_pos_active + seq_len_active).to(torch.int32) + + if s == 1: + # --------------------------------------------------------------- + # DECODE path (all sequences, causal=False) + # --------------------------------------------------------------- + B = num_decode + + # Absorb Q: q_absorbed = einsum("bsnd,ndk->bsnk", q_nope, w_kn) + q_nope_b = q_nope.reshape(B, 1, num_heads, qk_nope_head_dim) + q_absorbed = torch.einsum("bsnd,ndk->bsnk", q_nope_b.float(), w_kn.float()).to( + q_nope.dtype + ) # [B, 1, N, lora] + q_full = torch.cat([q_absorbed, q_pe.reshape(B, 1, num_heads, qk_rope_head_dim)], dim=-1) + # [B, 1, N, lora+rope] + + block_table = _build_block_table(cache_loc, cu_num_pages, B, device) + + if use_sparse: + # Sparse decode: compute Indexer top-k and pass to FlashMLA + dense_index_k, valid_mask = _gather_index_k_dense( + index_k_cache, cache_loc, cu_num_pages, cache_seqlens + ) + index_q_b = index_q.reshape(B, 1, index_q.shape[-2], index_head_dim) + index_weights_b = index_weights.reshape(B, 1, index_weights.shape[-1]) + sparse_indices = _compute_batched_index_topk( + index_q_b, + dense_index_k, + index_weights_b, + valid_mask, + index_topk, + index_softmax_scale, + is_causal=False, + cache_seqlens=cache_seqlens, + ) # [B, 1, topk] + + tile_meta, num_splits = get_mla_metadata( + cache_seqlens, + 1 * num_heads, + num_heads_k=1, + num_heads_q=num_heads, + topk=index_topk, + ) + out_latent, _ = flash_mla_with_kvcache( + q_full, + mla_cache, + block_table, + cache_seqlens, + head_dim_v=kv_lora_rank, + tile_scheduler_metadata=tile_meta, + num_splits=num_splits, + softmax_scale=scale, + causal=False, + indices=sparse_indices, + ) # [B, 1, N, lora] + else: + # Dense decode fallback for SM90 (sparsity not available) + tile_meta, num_splits = get_mla_metadata( + cache_seqlens, + 1 * num_heads, + num_heads_k=1, + num_heads_q=num_heads, + topk=None, + ) + out_latent, _ = flash_mla_with_kvcache( + q_full, + mla_cache, + block_table, + cache_seqlens, + head_dim_v=kv_lora_rank, + tile_scheduler_metadata=tile_meta, + num_splits=num_splits, + softmax_scale=scale, + causal=False, + ) # [B, 1, N, lora] + + out_latent = out_latent.float() + out = torch.einsum("bsnk,nvk->bsnv", out_latent, w_v.float()).to(q_nope.dtype) + return out # [B, 1, N, v_head_dim] + + else: + # --------------------------------------------------------------- + # PREFILL path + # --------------------------------------------------------------- + B = num_prefill + + if use_sparse and kv_lora_rank == 512: + # Use flash_mla_sparse_fwd (flat ragged, no Q-padding) + # Gather dense KV from paged cache + dense_kv, kv_cu_seqlens = _gather_mla_kv_dense( + mla_cache, cache_loc, cu_num_pages, cache_seqlens + ) # dense_kv: [total_kv, 1, lora+rope] + + # Gather index_k for batched Indexer (padded) + dense_index_k, valid_mask = _gather_index_k_dense( + index_k_cache, cache_loc, cu_num_pages, cache_seqlens + ) # [B, max_T, D_idx] + + # Pad query-side tensors to [B, max_S, ...] for batched Indexer + max_S = int(seq_len_active.max().item()) + + # Build padded index_q and index_weights [B, max_S, ...] + index_q_flat = index_q.reshape(total_tokens, index_q.shape[-2], index_head_dim) + index_w_flat = index_weights.reshape(total_tokens, index_weights.shape[-1]) + + index_q_padded = index_q_flat.new_zeros(B, max_S, index_q.shape[-2], index_head_dim) + index_w_padded = index_w_flat.new_zeros(B, max_S, index_weights.shape[-1]) + + for i in range(B): + sl = int(seq_len_active[i].item()) + ss = int(cu_seqlen_active[i].item()) + index_q_padded[i, :sl] = index_q_flat[ss : ss + sl] + index_w_padded[i, :sl] = index_w_flat[ss : ss + sl] + + # Batched Indexer top-k [B, max_S, topk] + local_indices = _compute_batched_index_topk( + index_q_padded, + dense_index_k, + index_w_padded, + valid_mask, + index_topk, + index_softmax_scale, + is_causal=True, + cache_seqlens=cache_seqlens, + ) # [B, max_S, topk] + + # Convert local KV indices to global flat KV indices + total_kv = int(dense_kv.shape[0]) + + # q_seq_ids[t] = batch index, q_local_pos[t] = within-seq query pos + q_seq_ids = torch.repeat_interleave( + torch.arange(B, device=device), seq_len_active + ) # [total_tokens] + q_local_pos = ( + torch.arange(total_tokens, device=device) - cu_seqlen_active[:-1][q_seq_ids] + ) # [total_tokens] + + # global_indices[t, :] = kv_cu_seqlens[b] + local_indices[b, s, :] + global_indices = local_indices[q_seq_ids, q_local_pos, :] # [total_tokens, topk] + global_indices = global_indices.long() + kv_cu_seqlens[q_seq_ids].unsqueeze(1) + # Mark out-of-range as -1 (flash_mla_sparse_fwd treats -1 as invalid) + global_indices[global_indices >= total_kv] = -1 + indices_flat = global_indices.unsqueeze(1).to(torch.int32) # [total_tokens, 1, topk] + + # Absorb Q: q_absorbed = einsum("tnd,ndk->tnk", q_nope_flat, w_kn) + q_nope_flat = q_nope.reshape(total_tokens, num_heads, qk_nope_head_dim) + q_absorbed = torch.einsum("tnd,ndk->tnk", q_nope_flat.float(), w_kn.float()).to( + q_nope.dtype + ) # [total_tokens, N, lora] + q_pe_flat = q_pe.reshape(total_tokens, num_heads, qk_rope_head_dim) + q_flat = torch.cat([q_absorbed, q_pe_flat], dim=-1) # [total_tokens, N, lora+rope] + + # flash_mla_sparse_fwd expects bfloat16 + compute_dtype = torch.bfloat16 + q_fwd = q_flat.to(compute_dtype) + kv_fwd = dense_kv.squeeze(1).to(compute_dtype) # [total_kv, lora+rope] + # flash_mla_sparse_fwd: q=[s_q,h_q,d_qk], kv=[s_kv,h_kv,d_qk] + # h_kv=1 for MLA; indices=[s_q, h_kv, topk] + + out_latent_flat, _, _ = flash_mla_sparse_fwd( + q_fwd, + kv_fwd.unsqueeze(1), # [total_kv, 1, lora+rope] as expected + indices_flat, + scale, + d_v=kv_lora_rank, + ) # [total_tokens, N, lora] + + out_latent_flat = out_latent_flat.float() + out_flat = torch.einsum("tnk,nvk->tnv", out_latent_flat, w_v.float()).to( + q_nope.dtype + ) # [total_tokens, N, v_head_dim] + + return out_flat.reshape(b, s, num_heads, v_head_dim) + + else: + # Fallback: flash_mla_with_kvcache(causal=True) with Q-padding + # Used on SM90 or when kv_lora_rank != 512 + max_S = int(seq_len_active.max().item()) + + # Build padded queries [B, max_S, N, lora+rope] + q_nope_flat = q_nope.reshape(total_tokens, num_heads, qk_nope_head_dim) + q_pe_flat_r = q_pe.reshape(total_tokens, num_heads, qk_rope_head_dim) + + q_absorbed_flat = torch.einsum("tnd,ndk->tnk", q_nope_flat.float(), w_kn.float()).to( + q_nope.dtype + ) + q_full_flat = torch.cat( + [q_absorbed_flat, q_pe_flat_r], dim=-1 + ) # [total_tokens, N, lora+rope] + + # Pad to [B, max_S, N, lora+rope] + q_padded = q_full_flat.new_zeros(B, max_S, num_heads, kv_lora_rank + qk_rope_head_dim) + for i in range(B): + sl = int(seq_len_active[i].item()) + ss = int(cu_seqlen_active[i].item()) + q_padded[i, :sl] = q_full_flat[ss : ss + sl] + + block_table = _build_block_table(cache_loc, cu_num_pages, B, device) + + tile_meta, num_splits = get_mla_metadata( + cache_seqlens, + max_S * num_heads, + num_heads_k=1, + num_heads_q=num_heads, + topk=None, + ) + + out_padded, _ = flash_mla_with_kvcache( + q_padded, + mla_cache, + block_table, + cache_seqlens, + head_dim_v=kv_lora_rank, + tile_scheduler_metadata=tile_meta, + num_splits=num_splits, + softmax_scale=scale, + causal=True, + ) # [B, max_S, N, lora] + + # Unpad and project + out_flat_list = [] + for i in range(B): + sl = int(seq_len_active[i].item()) + out_i = out_padded[i, :sl] # [sl, N, lora] + out_flat_list.append(out_i) + out_latent_flat = torch.cat(out_flat_list, dim=0).float() # [total_tokens, N, lora] + out_flat = torch.einsum("tnk,nvk->tnv", out_latent_flat, w_v.float()).to(q_nope.dtype) + + return out_flat.reshape(b, s, num_heads, v_head_dim) + + +@flash_mla_dsa_with_cache.register_fake +def _( + q_nope: torch.Tensor, + q_pe: torch.Tensor, + compressed_kv: torch.Tensor, + kpe: torch.Tensor, + kv_b_proj_weight: torch.Tensor, + index_q: torch.Tensor, + index_k: torch.Tensor, + index_weights: torch.Tensor, + batch_info_host: torch.Tensor, + seq_len: torch.Tensor, + input_pos: torch.Tensor, + cu_seqlen: torch.Tensor, + cache_loc: torch.Tensor, + cu_num_pages: torch.Tensor, + last_page_len: torch.Tensor, + mla_cache: torch.Tensor, + index_k_cache: torch.Tensor, + scale: Optional[float] = None, + kv_lora_rank: int = 512, + index_topk: int = 64, +) -> torch.Tensor: + """Fake impl for torch.export / graph tracing.""" + B = q_nope.shape[0] + S = q_nope.shape[1] + N = q_nope.shape[2] + qk_nope_head_dim = q_nope.shape[3] + out_features = kv_b_proj_weight.shape[0] + kv_head_dim = out_features // N + v_head_dim = kv_head_dim - qk_nope_head_dim + return q_nope.new_empty(B, S, N, v_head_dim) + + +# --------------------------------------------------------------------------- +# Descriptor +# --------------------------------------------------------------------------- + + +@AttentionRegistry.register("flashmla_dsa") +class FlashMLADSAAttention(AttentionDescriptor): + """Attention descriptor for FlashMLA-based DSA with paged caches. + + Source op: torch_dsa (same as TorchBackendDSAAttention) + Cached op: flash_mla_dsa_with_cache + + Cache layout (paged via MLAPagedResourceHandler): + mla_cache: [num_blocks, page_size, 1, kv_lora_rank + qk_rope_head_dim] + index_k_cache: [num_blocks, page_size, index_head_dim] + """ + + @classmethod + def get_attention_layout(cls) -> AttentionLayout: + return "bsnd" + + @classmethod + def get_num_qkv_args(cls) -> int: + # q_nope, q_pe, compressed_kv, kpe, kv_b_proj_weight, + # index_q, index_k, index_weights + return 8 + + @classmethod + def get_source_attention_op(cls) -> OpOverloadPacket: + return torch.ops.auto_deploy.torch_dsa + + @classmethod + def get_cached_attention_op(cls) -> MHACallable: + return torch.ops.auto_deploy.flash_mla_dsa_with_cache.default + + @classmethod + def get_standard_metadata_args(cls) -> List[str]: + return [ + "batch_info_host", + "seq_len", + "input_pos", + "cu_seqlen", + "cache_loc", + "cu_num_pages", + "last_page_len", + ] + + @classmethod + def get_cache_initializers( + cls, source_attn_node: Node, cache_config: KvCacheConfig + ) -> ResourceHandlerDict: + """Initialize paged mla_cache and index_k_cache.""" + # torch_dsa args: q_nope[0], q_pe[1], compressed_kv[2], kpe[3], + # kv_b_proj_weight[4], index_q[5], index_k[6], index_weights[7] + compressed_kv_fake = source_attn_node.args[2].meta["val"] + kpe_fake = source_attn_node.args[3].meta["val"] + index_q_fake = source_attn_node.args[5].meta["val"] + + kv_lora_rank = compressed_kv_fake.shape[-1] + qk_rope_head_dim = kpe_fake.shape[-1] + index_head_dim = index_q_fake.shape[-1] + + model_dtype = compressed_kv_fake.dtype + cache_dtype = cls.resolve_cache_dtype(cache_config.dtype, model_dtype) + + return { + "mla_cache": MLAPagedResourceHandler( + 1, + kv_lora_rank + qk_rope_head_dim, + dtype=cache_dtype, + ), + "index_k_cache": MLAPagedResourceHandler( + index_head_dim, + dtype=cache_dtype, + ), + } + + @classmethod + def get_constants(cls, source_attn_node: Node) -> List[Constant]: + """Return [scale, kv_lora_rank, index_topk] constants.""" + compressed_kv_fake = source_attn_node.args[2].meta["val"] + kv_lora_rank = compressed_kv_fake.shape[-1] + scale = source_attn_node.kwargs.get("scale", None) + index_topk = source_attn_node.kwargs.get("index_topk", 64) + return [scale, kv_lora_rank, index_topk] diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_backend_dsa.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_backend_dsa.py new file mode 100644 index 000000000000..dc496d0cb283 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_backend_dsa.py @@ -0,0 +1,599 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Torch backend for DeepSeek Sparse Attention (DSA) with KV cache. + +Provides: +- torch_cached_dsa_with_cache: cached DSA op managing two caches: + mla_cache: [max_batch, max_seq, kv_lora_rank + qk_rope_head_dim] (same as MLA) + index_k_cache: [max_batch, max_seq, index_head_dim] (Indexer keys) +- TorchBackendDSAAttention: AttentionDescriptor registered as "torch_dsa" + +Cache layout: + mla_cache: identical to FlashInfer MLA cache layout + index_k_cache: [max_batch, max_seq, index_head_dim] + +Prefill: expand compressed_kv -> full K/V, compute index scores vs cached index_k +Generate: weight absorption for MLA part + index sparse mask from cached index_k +""" + +import math +from typing import List, Optional + +import torch +from torch._ops import OpOverloadPacket +from torch.fx import Node + +from .....llmapi.llm_args import KvCacheConfig +from ..attention_interface import ( + AttentionDescriptor, + AttentionLayout, + AttentionRegistry, + Constant, + MHACallable, + ResourceHandlerDict, + UnpagedResourceHandler, +) +from .torch_backend_mla import _update_mla_cache + + +def _update_index_k_cache( + index_k: torch.Tensor, # [total_tokens, index_head_dim] + index_k_cache: torch.Tensor, # [max_batch, max_seq, index_head_dim] + seq_len: torch.Tensor, + input_pos: torch.Tensor, + slot_idx: torch.Tensor, + seq_start: torch.Tensor, +) -> None: + """Update the Indexer key cache with new token keys.""" + cache_dtype = index_k_cache.dtype + if index_k.dtype != cache_dtype: + index_k = index_k.to(cache_dtype) + + for idx in range(seq_len.shape[0]): + start = seq_start[idx].item() + length = seq_len[idx].item() + cache_idx = slot_idx[idx].item() + pos = input_pos[idx].item() + index_k_cache[cache_idx, pos : pos + length] = index_k[start : start + length] + + +def _compute_seq_index_score( + index_q_seq: torch.Tensor, # [S_q, H, D_idx] + index_k_cached: torch.Tensor, # [T, D_idx] + index_weights_seq: torch.Tensor, # [S_q, H] + index_softmax_scale: float, + index_topk: int, + causal_diagonal: int, +) -> torch.Tensor: + """Compute DSA index mask for a single sequence. + + Returns index_mask [S_q, T] with 0.0 at top-k positions and -inf elsewhere. + causal_diagonal: passed to torch.triu to enforce causality (use kv_seq_len - seq_len + 1). + """ + s_q = index_q_seq.shape[0] + t = index_k_cached.shape[0] + + # per_head_scores: [S_q, H, T] + per_head_scores = torch.einsum( + "shd,td->sht", + index_q_seq.float(), + index_k_cached.float(), + ) + + # weighted sum over heads: [S_q, T] + index_score = torch.einsum("sht,sh->st", per_head_scores, index_weights_seq.float()) + index_score = index_score * index_softmax_scale + + # Causal mask: upper-triangular positions get -inf + if causal_diagonal <= t: + causal_mask = torch.triu( + torch.ones(s_q, t, device=index_q_seq.device, dtype=torch.bool), + diagonal=causal_diagonal, + ) + index_score.masked_fill_(causal_mask, float("-inf")) + + effective_topk = min(index_topk, t) + topk_indices = index_score.topk(effective_topk, dim=-1).indices # [S_q, topk] + + index_mask = index_score.new_full(index_score.shape, float("-inf")) + index_mask.scatter_(-1, topk_indices, 0.0) # [S_q, T] + return index_mask + + +def _torch_dsa_generate_with_absorption( + q_nope: torch.Tensor, # [B, 1, N, qk_nope_head_dim] + q_pe: torch.Tensor, # [B, 1, N, qk_rope_head_dim] + compressed_kv: torch.Tensor, # [B, 1, kv_lora_rank] + kpe: torch.Tensor, # [B, 1, 1, qk_rope_head_dim] + kv_b_proj_weight: torch.Tensor, # [N*(qk_nope+v), kv_lora_rank] + index_q: torch.Tensor, # [B, 1, H, D_idx] + index_k: torch.Tensor, # [B, 1, D_idx] + index_weights: torch.Tensor, # [B, 1, H] + mla_cache: torch.Tensor, # [max_batch, max_seq, kv_lora_rank + qk_rope_head_dim] + index_k_cache: torch.Tensor, # [max_batch, max_seq, index_head_dim] + slot_idx: torch.Tensor, + input_pos: torch.Tensor, + scale: float, + kv_lora_rank: int, + index_topk: int, + num_heads: int, + qk_nope_head_dim: int, + v_head_dim: int, + out: torch.Tensor, +) -> None: + """Generate-phase DSA with MLA weight absorption + sparse index masking. + + Vectorized implementation without .item() calls for CUDA graph compatibility. + Uses full max_seq slices from the cache with a validity mask to avoid + variable-length slicing that would require CPU synchronization. + """ + max_seq = mla_cache.shape[1] + index_head_dim = index_q.shape[-1] + index_softmax_scale = 1.0 / math.sqrt(index_head_dim) + compute_dtype = q_nope.dtype + cache_dtype = mla_cache.dtype + idx_cache_dtype = index_k_cache.dtype + + # Flatten inputs (remove the seq=1 dimension) + compressed_kv_flat = compressed_kv.squeeze(1) # [B, kv_lora_rank] + kpe_flat = kpe.squeeze(1).squeeze(1) # [B, qk_rope_head_dim] + index_k_flat = index_k.squeeze(1) # [B, D_idx] + + # ----------------------------------------------------------------------- + # Cache update — vectorized via advanced indexing (no .item() needed) + # slot_idx: [B], input_pos: [B] are GPU tensors used directly as indices + # ----------------------------------------------------------------------- + mla_write = torch.cat( + [compressed_kv_flat.to(cache_dtype), kpe_flat.to(cache_dtype)], dim=-1 + ) # [B, kv_lora_rank + qk_rope_head_dim] + mla_cache[slot_idx, input_pos] = mla_write + index_k_cache[slot_idx, input_pos] = index_k_flat.to(idx_cache_dtype) + + # ----------------------------------------------------------------------- + # Read full cached slices via batch gather (no variable-length slicing) + # cached_mla: [B, max_seq, kv_lora_rank + qk_rope_head_dim] + # cached_index_k: [B, max_seq, D_idx] + # ----------------------------------------------------------------------- + cached_mla = mla_cache[slot_idx].to(compute_dtype) # [B, max_seq, kv_lora+qk_rope] + cached_index_k = index_k_cache[slot_idx].to(compute_dtype) # [B, max_seq, D_idx] + + compressed_kv_cached = cached_mla[:, :, :kv_lora_rank] # [B, max_seq, kv_lora_rank] + kpe_cached = cached_mla[:, :, kv_lora_rank:] # [B, max_seq, qk_rope_head_dim] + + # ----------------------------------------------------------------------- + # Validity mask: position range 0..input_pos[i] is valid for sequence i + # pos_range: [max_seq], input_pos: [B] -> valid_mask: [B, max_seq] + # ----------------------------------------------------------------------- + pos_range = torch.arange(max_seq, device=input_pos.device, dtype=input_pos.dtype) + valid_mask = pos_range.unsqueeze(0) <= input_pos.unsqueeze(1) # [B, max_seq] + + # ----------------------------------------------------------------------- + # DSA index score — vectorized over batch + # index_q_bh: [B, H, D_idx], cached_index_k: [B, max_seq, D_idx] + # ----------------------------------------------------------------------- + index_q_bh = index_q[:, 0] # [B, H, D_idx] + index_w_bh = index_weights[:, 0] # [B, H] + + # per_head_scores: [B, H, max_seq] + per_head_scores = torch.einsum("bhd,btd->bht", index_q_bh.float(), cached_index_k.float()) + # index_score: [B, max_seq] + index_score = torch.einsum("bht,bh->bt", per_head_scores, index_w_bh.float()) + index_score = index_score * index_softmax_scale + + # Mask out invalid (future) positions before top-k selection + index_score = index_score.masked_fill(~valid_mask, float("-inf")) + + # Top-k selection (fixed shape → CUDA-graph compatible) + effective_topk = min(index_topk, max_seq) + topk_indices = index_score.topk(effective_topk, dim=-1).indices # [B, topk] + index_mask = index_score.new_full(index_score.shape, float("-inf")) + index_mask.scatter_(-1, topk_indices, 0.0) # [B, max_seq] + + # ----------------------------------------------------------------------- + # MLA weight absorption + attention + # ----------------------------------------------------------------------- + weight_reshaped = kv_b_proj_weight.view(num_heads, qk_nope_head_dim + v_head_dim, kv_lora_rank) + w_k_nope = weight_reshaped[:, :qk_nope_head_dim, :] # [N, qk_nope, kv_lora_rank] + w_v = weight_reshaped[:, qk_nope_head_dim:, :] # [N, v_head_dim, kv_lora_rank] + + q_nope_bn = q_nope[:, 0] # [B, N, qk_nope_head_dim] + q_pe_bn = q_pe[:, 0] # [B, N, qk_rope_head_dim] + + # q_absorbed: [B, N, kv_lora_rank] + q_absorbed = torch.einsum("bnd,ndk->bnk", q_nope_bn.float(), w_k_nope.float()) + + # scores_nope: [B, N, max_seq] + scores_nope = torch.matmul(q_absorbed, compressed_kv_cached.float().transpose(1, 2)) + # scores_pe: [B, N, max_seq] + scores_pe = torch.matmul(q_pe_bn.float(), kpe_cached.float().transpose(1, 2)) + + attn_scores = (scores_nope + scores_pe) * scale # [B, N, max_seq] + + # Apply causal validity mask and DSA sparse mask: [B, max_seq] -> [B, 1, max_seq] + attn_scores = attn_scores.masked_fill(~valid_mask.unsqueeze(1), float("-inf")) + attn_scores = attn_scores + index_mask.unsqueeze(1) + + attn_weights = torch.softmax(attn_scores, dim=-1).to(compute_dtype) # [B, N, max_seq] + + # weighted_kv: [B, N, kv_lora_rank] + weighted_kv = torch.matmul(attn_weights, compressed_kv_cached) + # attn_out: [B, N, v_head_dim] + attn_out = torch.einsum("bnk,nvk->bnv", weighted_kv, w_v.to(compute_dtype)) + + out[:] = attn_out + + +def _torch_dsa_context_with_expansion( + q_nope: torch.Tensor, # [total_tokens, N, qk_nope_head_dim] + q_pe: torch.Tensor, # [total_tokens, N, qk_rope_head_dim] + compressed_kv: torch.Tensor, # [total_tokens, kv_lora_rank] + kpe: torch.Tensor, # [total_tokens, 1, qk_rope_head_dim] + kv_b_proj_weight: torch.Tensor, # [N*(qk_nope+v), kv_lora_rank] + index_q: torch.Tensor, # [total_tokens, H, D_idx] + index_k: torch.Tensor, # [total_tokens, D_idx] + index_weights: torch.Tensor, # [total_tokens, H] + mla_cache: torch.Tensor, # [max_batch, max_seq, kv_lora_rank + qk_rope_head_dim] + index_k_cache: torch.Tensor, # [max_batch, max_seq, index_head_dim] + input_pos: torch.Tensor, + slot_idx: torch.Tensor, + seq_len: torch.Tensor, + seq_start: torch.Tensor, + scale: float, + kv_lora_rank: int, + index_topk: int, + num_heads: int, + qk_nope_head_dim: int, + v_head_dim: int, + out: torch.Tensor, +) -> None: + """Context-phase DSA: kv_b_proj expansion + sparse index masking.""" + index_head_dim = index_q.shape[-1] + index_softmax_scale = 1.0 / math.sqrt(index_head_dim) + + kpe_flat = kpe.squeeze(1) # [total_tokens, qk_rope_head_dim] + + # Update MLA cache + _update_mla_cache( + compressed_kv, + kpe_flat, + mla_cache, + seq_len, + input_pos, + slot_idx, + seq_start, + kv_lora_rank, + ) + + # Update index_k cache + _update_index_k_cache( + index_k, + index_k_cache, + seq_len, + input_pos, + slot_idx, + seq_start, + ) + + compute_dtype = q_nope.dtype + attn_outputs = [] + + for idx in range(seq_len.shape[0]): + seq_len_i = seq_len[idx].item() + input_pos_i = input_pos[idx].item() + slot_idx_i = slot_idx[idx].item() + seq_start_i = seq_start[idx].item() + + if seq_len_i == 0: + continue + + kv_seq_len = input_pos_i + seq_len_i + + # Gather query tokens for this sequence + q_nope_seq = q_nope[seq_start_i : seq_start_i + seq_len_i] # [S, N, nope] + q_pe_seq = q_pe[seq_start_i : seq_start_i + seq_len_i] # [S, N, rope] + + # Get cached MLA data + cached_data = mla_cache[slot_idx_i, :kv_seq_len] # [T, kv_lora_rank + qk_rope_head_dim] + compressed_kv_cached = cached_data[:, :kv_lora_rank] + kpe_cached = cached_data[:, kv_lora_rank:] + if compressed_kv_cached.dtype != compute_dtype: + compressed_kv_cached = compressed_kv_cached.to(compute_dtype) + if kpe_cached.dtype != compute_dtype: + kpe_cached = kpe_cached.to(compute_dtype) + + # Get cached index keys + index_k_cached = index_k_cache[slot_idx_i, :kv_seq_len] # [T, D_idx] + if index_k_cached.dtype != compute_dtype: + index_k_cached = index_k_cached.to(compute_dtype) + + # --- DSA index mask --- + index_q_seq = index_q[seq_start_i : seq_start_i + seq_len_i] # [S, H, D_idx] + index_w_seq = index_weights[seq_start_i : seq_start_i + seq_len_i] # [S, H] + index_mask = _compute_seq_index_score( + index_q_seq, + index_k_cached, + index_w_seq, + index_softmax_scale, + index_topk, + causal_diagonal=kv_seq_len - seq_len_i + 1, + ) # [S, T] + + # --- Expand compressed_kv and compute attention --- + kv_expanded = torch.matmul(compressed_kv_cached, kv_b_proj_weight.t()) + kv_expanded = kv_expanded.view(kv_seq_len, num_heads, qk_nope_head_dim + v_head_dim) + k_nope_expanded = kv_expanded[:, :, :qk_nope_head_dim] # [T, N, nope] + v_expanded = kv_expanded[:, :, qk_nope_head_dim:] # [T, N, v] + + kpe_expanded = kpe_cached.unsqueeze(1).expand(-1, num_heads, -1) # [T, N, rope] + + query_full = torch.cat([q_nope_seq, q_pe_seq], dim=-1) # [S, N, qk_head_dim] + key_full = torch.cat([k_nope_expanded, kpe_expanded], dim=-1) # [T, N, qk_head_dim] + + # Transpose to [1, N, S/T, D] for batched matmul + query_t = query_full.transpose(0, 1).unsqueeze(0) # [1, N, S, D] + key_t = key_full.transpose(0, 1).unsqueeze(0) # [1, N, T, D] + + attn_scores = ( + torch.matmul(query_t.float(), key_t.float().transpose(-2, -1)) * scale + ) # [1, N, S, T] in fp32 + + # Causal mask (upper-triangular relative to kv_seq_len) + causal_mask = torch.triu( + torch.ones(seq_len_i, kv_seq_len, device=q_nope.device, dtype=torch.bool), + diagonal=kv_seq_len - seq_len_i + 1, + ) + attn_scores.masked_fill_(causal_mask.unsqueeze(0).unsqueeze(0), float("-inf")) + + # DSA sparse mask: [S, T] -> [1, 1, S, T] + attn_scores = attn_scores + index_mask.unsqueeze(0).unsqueeze(0) + + attn_weights = torch.softmax(attn_scores, dim=-1).to(compute_dtype) # [1, N, S, T] + + v_t = v_expanded.transpose(0, 1).unsqueeze(0) # [1, N, T, v] + attn_out = torch.matmul(attn_weights, v_t) # [1, N, S, v] + attn_out = attn_out[0].transpose(0, 1) # [S, N, v] + + attn_outputs.append(attn_out) + + if len(attn_outputs) == 0: + out.zero_() + elif len(attn_outputs) == 1: + out.copy_(attn_outputs[0]) + else: + out.copy_(torch.cat(attn_outputs, dim=0)) + + +@torch.library.custom_op("auto_deploy::torch_cached_dsa_with_cache", mutates_args=()) +def torch_backend_dsa_with_cache( + # 8 tensor args (get_num_qkv_args = 8) + q_nope: torch.Tensor, # [B, S, N, qk_nope_head_dim] + q_pe: torch.Tensor, # [B, S, N, qk_rope_head_dim] + compressed_kv: torch.Tensor, # [B, S, kv_lora_rank] + kpe: torch.Tensor, # [B, S, 1, qk_rope_head_dim] + kv_b_proj_weight: torch.Tensor, # [N*(qk_nope+v), kv_lora_rank] + index_q: torch.Tensor, # [B, S, H, D_idx] + index_k: torch.Tensor, # [B, S, D_idx] + index_weights: torch.Tensor, # [B, S, H] + # Standard metadata + batch_info_host: torch.Tensor, + seq_len: torch.Tensor, + input_pos: torch.Tensor, + slot_idx: torch.Tensor, + cu_seqlen: torch.Tensor, + # Caches + mla_cache: torch.Tensor, # [max_batch, max_seq, kv_lora_rank + qk_rope_head_dim] + index_k_cache: torch.Tensor, # [max_batch, max_seq, index_head_dim] + # Constants + scale: Optional[float] = None, + kv_lora_rank: int = 512, + index_topk: int = 64, +) -> torch.Tensor: + """Torch backend DSA with KV cache and Indexer key cache. + + Prefill: expand compressed_kv, compute index mask vs cached index_k + Generate: weight absorption for MLA + index sparse mask + """ + b, s = q_nope.shape[:2] + num_heads = q_nope.shape[2] + qk_nope_head_dim = q_nope.shape[3] + qk_rope_head_dim = q_pe.shape[3] + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + + out_features = kv_b_proj_weight.shape[0] + kv_head_dim = out_features // num_heads + v_head_dim = kv_head_dim - qk_nope_head_dim + + num_prefill, num_prefill_tokens, num_decode = batch_info_host.tolist() + num_seq = num_prefill + num_decode + seq_len = seq_len[:num_seq] + input_pos = input_pos[:num_seq] + slot_idx = slot_idx[:num_seq] + seq_start = cu_seqlen[:num_seq] + + if scale is None: + scale = 1.0 / math.sqrt(qk_head_dim) + + output_shape = (b, s, num_heads, v_head_dim) + + if s == 1: + # Generate phase + y = q_nope.new_empty(b, num_heads, v_head_dim).contiguous() + _torch_dsa_generate_with_absorption( + q_nope, + q_pe, + compressed_kv, + kpe, + kv_b_proj_weight, + index_q, + index_k, + index_weights, + mla_cache, + index_k_cache, + slot_idx, + input_pos, + scale, + kv_lora_rank, + index_topk, + num_heads, + qk_nope_head_dim, + v_head_dim, + y, + ) + return y.unsqueeze(1) # [B, 1, N, v_head_dim] + else: + # Prefill / context phase + bs_view = (b * s,) + q_nope_flat = q_nope.contiguous().view(*bs_view, num_heads, qk_nope_head_dim) + q_pe_flat = q_pe.contiguous().view(*bs_view, num_heads, qk_rope_head_dim) + compressed_kv_flat = compressed_kv.contiguous().view(*bs_view, kv_lora_rank) + kpe_flat = kpe.contiguous().view(*bs_view, 1, qk_rope_head_dim) + + index_n_heads = index_q.shape[2] + index_head_dim = index_q.shape[3] + index_q_flat = index_q.contiguous().view(*bs_view, index_n_heads, index_head_dim) + index_k_flat = index_k.contiguous().view(*bs_view, index_head_dim) + index_weights_flat = index_weights.contiguous().view(*bs_view, index_n_heads) + + y = q_nope.new_empty(*bs_view, num_heads, v_head_dim).contiguous() + _torch_dsa_context_with_expansion( + q_nope_flat, + q_pe_flat, + compressed_kv_flat, + kpe_flat, + kv_b_proj_weight, + index_q_flat, + index_k_flat, + index_weights_flat, + mla_cache, + index_k_cache, + input_pos, + slot_idx, + seq_len, + seq_start, + scale, + kv_lora_rank, + index_topk, + num_heads, + qk_nope_head_dim, + v_head_dim, + y, + ) + return y.view(*output_shape) + + +@torch_backend_dsa_with_cache.register_fake +def torch_backend_dsa_with_cache_fake( + q_nope: torch.Tensor, + q_pe: torch.Tensor, + compressed_kv: torch.Tensor, + kpe: torch.Tensor, + kv_b_proj_weight: torch.Tensor, + index_q: torch.Tensor, + index_k: torch.Tensor, + index_weights: torch.Tensor, + batch_info_host: torch.Tensor, + seq_len: torch.Tensor, + input_pos: torch.Tensor, + slot_idx: torch.Tensor, + cu_seqlen: torch.Tensor, + mla_cache: torch.Tensor, + index_k_cache: torch.Tensor, + scale: Optional[float] = None, + kv_lora_rank: int = 512, + index_topk: int = 64, +) -> torch.Tensor: + """Fake implementation for torch_backend_dsa_with_cache.""" + num_heads = q_nope.shape[2] + qk_nope_head_dim = q_nope.shape[-1] + out_features = kv_b_proj_weight.shape[0] + kv_head_dim = out_features // num_heads + v_head_dim = kv_head_dim - qk_nope_head_dim + return q_nope.new_empty( + q_nope.shape[0], q_nope.shape[1], q_nope.shape[2], v_head_dim + ).contiguous() + + +@AttentionRegistry.register("torch_dsa") +class TorchBackendDSAAttention(AttentionDescriptor): + """Attention descriptor for DeepSeek Sparse Attention (DSA). + + Uses torch_dsa as the source op and torch_cached_dsa_with_cache as + the cached op, managing two caches: + - mla_cache: identical layout to MLA cache + - index_k_cache: Indexer key cache [max_batch, max_seq, index_head_dim] + """ + + @classmethod + def get_attention_layout(cls) -> AttentionLayout: + return "bsnd" + + @classmethod + def get_num_qkv_args(cls) -> int: + # q_nope, q_pe, compressed_kv, kpe, kv_b_proj_weight, + # index_q, index_k, index_weights + return 8 + + @classmethod + def get_source_attention_op(cls) -> OpOverloadPacket: + return torch.ops.auto_deploy.torch_dsa + + @classmethod + def get_cached_attention_op(cls) -> MHACallable: + return torch.ops.auto_deploy.torch_cached_dsa_with_cache.default + + @classmethod + def get_standard_metadata_args(cls) -> List[str]: + return ["batch_info_host", "seq_len", "input_pos", "slot_idx", "cu_seqlen"] + + @classmethod + def get_cache_initializers( + cls, source_attn_node: Node, cache_config: KvCacheConfig + ) -> ResourceHandlerDict: + """Initialize mla_cache and index_k_cache.""" + # torch_dsa args: q_nope[0], q_pe[1], compressed_kv[2], kpe[3], kv_b_proj_weight[4], + # index_q[5], index_k[6], index_weights[7] + compressed_kv_fake = source_attn_node.args[2].meta["val"] + kpe_fake = source_attn_node.args[3].meta["val"] + index_q_fake = source_attn_node.args[5].meta["val"] + + kv_lora_rank = compressed_kv_fake.shape[-1] + qk_rope_head_dim = kpe_fake.shape[-1] + index_head_dim = index_q_fake.shape[-1] + + model_dtype = compressed_kv_fake.dtype + cache_dtype = cls.resolve_cache_dtype(cache_config.dtype, model_dtype) + + return { + "mla_cache": UnpagedResourceHandler( + kv_lora_rank + qk_rope_head_dim, + dtype=cache_dtype, + ), + "index_k_cache": UnpagedResourceHandler( + index_head_dim, + dtype=cache_dtype, + ), + } + + @classmethod + def get_constants(cls, source_attn_node: Node) -> List[Constant]: + """Return [scale, kv_lora_rank, index_topk] constants.""" + compressed_kv_fake = source_attn_node.args[2].meta["val"] + kv_lora_rank = compressed_kv_fake.shape[-1] + + scale = source_attn_node.kwargs.get("scale", None) + index_topk = source_attn_node.kwargs.get("index_topk", 64) + + return [scale, kv_lora_rank, index_topk] diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_dsa.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_dsa.py new file mode 100644 index 000000000000..39f1e077c7ff --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_dsa.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Torch reference implementation for DeepSeek Sparse Attention (DSA). + +DSA extends MLA with an Indexer submodule that enforces token sparsity. +The Indexer produces top-k token indices whose sparse mask is added to +the MLA attention scores before softmax, zeroing out non-selected positions. + +This source op accepts pre-computed Indexer tensors (index_q, index_k, +index_weights), all with RoPE already applied by the caller, keeping the +same design philosophy as torch_mla which accepts pre-rotated q_pe/kpe. + +Reference: https://huggingface.co/zai-org/GLM-5 (GlmMoeDsaAttention) + DeepSeek-V3.2 inference/model.py (MLA with Indexer) +""" + +import math +from typing import Optional + +import torch + + +def _compute_dsa_index_mask( + index_q: torch.Tensor, # [B, S_q, H, D_idx] + index_k: torch.Tensor, # [B, S_k, D_idx] + index_weights: torch.Tensor, # [B, S_q, H] + index_topk: int, + is_causal: bool, + softmax_scale: float, +) -> torch.Tensor: + """Compute the DSA sparse mask from Indexer tensors. + + index_score[b,s,t] = softmax_scale * sum_h(index_weights[b,s,h] * dot(index_q[b,s,h], index_k[b,t])) + + Returns index_mask [B, S_q, S_k] with 0.0 at selected top-k positions and -inf elsewhere. + """ + bs, s_q, num_idx_heads, _ = index_q.shape + s_k = index_k.shape[1] + + # scores per head: [B, S_q, H, S_k] + # compute in float32 for numerical stability + per_head_scores = torch.einsum( + "bshd,btd->bsht", + index_q.float(), + index_k.float(), + ) # [B, S_q, H, S_k] + + # weighted sum over heads: [B, S_q, S_k] + index_score = torch.einsum("bsht,bsh->bst", per_head_scores, index_weights.float()) + index_score = index_score * softmax_scale + + # Apply causal mask so future tokens cannot be selected + if is_causal and s_q == s_k: + causal_mask = torch.triu( + torch.ones(s_q, s_k, device=index_q.device, dtype=torch.bool), + diagonal=1, + ) + index_score.masked_fill_(causal_mask.unsqueeze(0), float("-inf")) + + # Select top-k valid positions per query token + effective_topk = min(index_topk, s_k) + topk_indices = index_score.topk(effective_topk, dim=-1).indices # [B, S_q, topk] + + # Build mask: -inf everywhere, 0.0 at selected positions + index_mask = index_score.new_full(index_score.shape, float("-inf")) + index_mask.scatter_(-1, topk_indices, 0.0) # [B, S_q, S_k] + + return index_mask + + +@torch.library.custom_op("auto_deploy::torch_dsa", mutates_args=()) +def torch_dsa( + q_nope: torch.Tensor, # [B, S, N, qk_nope_head_dim] + q_pe: torch.Tensor, # [B, S, N, qk_rope_head_dim] (RoPE applied) + compressed_kv: torch.Tensor, # [B, S, kv_lora_rank] + kpe: torch.Tensor, # [B, S, 1, qk_rope_head_dim] (RoPE applied) + kv_b_proj_weight: torch.Tensor, # [N*(qk_nope_head_dim + v_head_dim), kv_lora_rank] + index_q: torch.Tensor, # [B, S, index_n_heads, index_head_dim] (RoPE applied) + index_k: torch.Tensor, # [B, S, index_head_dim] (shared across heads, RoPE applied) + index_weights: torch.Tensor, # [B, S, index_n_heads] (from weights_proj, pre-scaled) + index_topk: int = 64, + is_causal: bool = True, + scale: Optional[float] = None, + layout: str = "bsnd", +) -> torch.Tensor: + """DeepSeek Sparse Attention (DSA) reference implementation. + + Extends MLA with an Indexer that selects top-k KV positions per query token. + The sparse mask is added to MLA attention scores before softmax. + + Args: + q_nope: Query non-positional component [B, S, N, qk_nope_head_dim] + q_pe: Query positional component (RoPE applied) [B, S, N, qk_rope_head_dim] + compressed_kv: Compressed KV latent [B, S, kv_lora_rank] + kpe: Key positional encoding (RoPE applied) [B, S, 1, qk_rope_head_dim] + kv_b_proj_weight: KV expansion weights [N*(qk_nope+v), kv_lora_rank] + index_q: Indexer query (RoPE applied) [B, S, index_n_heads, index_head_dim] + index_k: Indexer key (RoPE applied) [B, S, index_head_dim] + index_weights: Per-head importance weights [B, S, index_n_heads] + index_topk: Number of KV positions to attend to per query token + is_causal: Whether to apply causal masking + scale: Softmax scale (default: 1/sqrt(qk_nope_head_dim + qk_rope_head_dim)) + layout: Input/output layout, "bsnd" or "bnsd" + + Returns: + Attention output [B, S, N, v_head_dim] (bsnd layout) + """ + if layout not in ("bnsd", "bsnd"): + raise ValueError(f"layout must be 'bnsd' or 'bsnd', got {layout!r}") + + # Infer dimensions + if layout == "bsnd": + bs, s_q, num_heads, qk_nope_head_dim = q_nope.shape + qk_rope_head_dim = q_pe.shape[-1] + else: + bs, num_heads, s_q, qk_nope_head_dim = q_nope.shape + qk_rope_head_dim = q_pe.shape[-1] + + s_k = compressed_kv.shape[1] + out_features = kv_b_proj_weight.shape[0] + kv_head_dim = out_features // num_heads + v_head_dim = kv_head_dim - qk_nope_head_dim + + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + if scale is None: + scale = 1.0 / math.sqrt(qk_head_dim) + + # DSA index softmax scale: 1/sqrt(index_head_dim) + index_head_dim = index_q.shape[-1] + index_softmax_scale = 1.0 / math.sqrt(index_head_dim) + + # ========================================================================= + # Indexer: compute sparse mask + # ========================================================================= + index_mask = _compute_dsa_index_mask( + index_q, index_k, index_weights, index_topk, is_causal, index_softmax_scale + ) # [B, S_q, S_k] + + # ========================================================================= + # MLA: expand compressed_kv and compute attention with sparse mask + # ========================================================================= + # compressed_kv: [B, S, kv_lora_rank] -> [B, S, N, kv_head_dim] + kv = torch.matmul(compressed_kv, kv_b_proj_weight.t()) + kv = kv.view(bs, s_k, num_heads, kv_head_dim) + k_nope, value_states = torch.split(kv, [qk_nope_head_dim, v_head_dim], dim=-1) + + # Convert to [B, N, S, D] for attention computation + k_nope = k_nope.transpose(1, 2).contiguous() + value_states = value_states.transpose(1, 2).contiguous() + + if layout == "bsnd": + q_nope = q_nope.transpose(1, 2).contiguous() + q_pe = q_pe.transpose(1, 2).contiguous() + kpe = kpe.transpose(1, 2).contiguous() + + # kpe: [B, 1, S, rope_head_dim] -> expand to all heads + kpe_expanded = kpe.expand(bs, num_heads, s_k, qk_rope_head_dim) + + # Full query and key: [B, N, S, qk_head_dim] + query_states = torch.cat([q_nope, q_pe], dim=-1) + key_states = torch.cat([k_nope, kpe_expanded], dim=-1) + + # Attention scores: [B, N, S_q, S_k] + attn_scores = torch.matmul(query_states, key_states.transpose(-2, -1)) * scale + + # Apply causal mask (full upper-triangular mask) + if is_causal and s_q == s_k: + causal_mask = torch.triu( + torch.ones(s_q, s_k, device=q_nope.device, dtype=torch.bool), + diagonal=1, + ) + attn_scores.masked_fill_(causal_mask.unsqueeze(0).unsqueeze(0), float("-inf")) + + # Apply DSA sparse mask: [B, S_q, S_k] -> broadcast to [B, N, S_q, S_k] + attn_scores = attn_scores + index_mask.unsqueeze(1) + + # Softmax + output + attn_weights = torch.softmax(attn_scores, dim=-1, dtype=torch.float32).to(q_nope.dtype) + attn_out = torch.matmul(attn_weights, value_states) # [B, N, S_q, v_head_dim] + + if layout == "bsnd": + return attn_out.transpose(1, 2).contiguous() # [B, S, N, v_head_dim] + else: + return attn_out.contiguous() # [B, N, S, v_head_dim] + + +@torch_dsa.register_fake +def torch_dsa_fake( + q_nope: torch.Tensor, + q_pe: torch.Tensor, + compressed_kv: torch.Tensor, + kpe: torch.Tensor, + kv_b_proj_weight: torch.Tensor, + index_q: torch.Tensor, + index_k: torch.Tensor, + index_weights: torch.Tensor, + index_topk: int = 64, + is_causal: bool = True, + scale: Optional[float] = None, + layout: str = "bsnd", +) -> torch.Tensor: + """Fake implementation for torch_dsa.""" + qk_nope_head_dim = q_nope.shape[-1] + num_heads = q_nope.shape[2] if layout == "bsnd" else q_nope.shape[1] + out_features = kv_b_proj_weight.shape[0] + kv_head_dim = out_features // num_heads + v_head_dim = kv_head_dim - qk_nope_head_dim + + if layout == "bsnd": + return q_nope.new_empty( + q_nope.shape[0], q_nope.shape[1], q_nope.shape[2], v_head_dim + ).contiguous() + else: + return q_nope.new_empty( + q_nope.shape[0], q_nope.shape[1], q_nope.shape[2], v_head_dim + ).contiguous() diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py b/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py index eea48c18ad05..5a5808c29f37 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py @@ -3,6 +3,7 @@ from .modeling_deepseek import DeepSeekV3ForCausalLM from .modeling_deepseek_v2 import DeepSeekV2ForCausalLM from .modeling_glm4_moe_lite import Glm4MoeLiteForCausalLM +from .modeling_glm_dsa import GlmDSAForCausalLM from .modeling_granite import GraniteForCausalLM from .modeling_granite_moe_hybrid import GraniteMoeHybridForCausalLM from .modeling_hunyuan_dense import HunYuanDenseForCausalLM @@ -23,9 +24,10 @@ "DeepSeekV2ForCausalLM", "DeepSeekV3ForCausalLM", "Glm4MoeLiteForCausalLM", + "GlmDSAForCausalLM", "GraniteForCausalLM", - "HunYuanDenseForCausalLM", "GraniteMoeHybridForCausalLM", + "HunYuanDenseForCausalLM", "HunYuanMoEForCausalLM", "KimiK2ForCausalLM", "KimiK25ForConditionalGeneration", diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm_dsa.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm_dsa.py new file mode 100644 index 000000000000..da2fe3b7cfb0 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm_dsa.py @@ -0,0 +1,1003 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""AutoDeploy export-ready implementation of GLM MoE DSA (GLM-5). + +Source: https://huggingface.co/zai-org/GLM-5 (model_type = "glm_moe_dsa") + https://github.com/huggingface/transformers/tree/main/src/transformers/models/glm_moe_dsa + +Differences from the HuggingFace reference: +- Prefill-only (no KV caching — caching is handled by AutoDeploy graph transforms) +- RMSNorm uses torch.ops.auto_deploy.torch_rmsnorm canonical op +- Attention uses torch_dsa custom op in BSND layout +- MoE gate uses vanilla PyTorch noaux_tc routing (AD transforms can replace with trtllm kernels) +- Indexer is a separate GlmDSAIndexer submodule matching checkpoint key names +- FP8 quantization and Hadamard rotation are omitted (orthogonal to correctness) +- RoPE weights are de-interleaved at load time via load hooks +- MoE expert weights are expanded from stacked checkpoint format to per-expert ModuleList at load time +- Bundled GlmMoeDsaConfig class (not yet in transformers 4.57) +""" + +import math +from dataclasses import dataclass +from functools import partial +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import AutoConfig, PretrainedConfig +from transformers.activations import ACT2FN +from transformers.generation import GenerationMixin +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import ModelOutput + +from tensorrt_llm._torch.auto_deploy.models.hf import AutoModelForCausalLMFactory +from tensorrt_llm._torch.utils import ActivationType + +# ============================================================================= +# Bundled Config +# ============================================================================= + + +class GlmMoeDsaConfig(PretrainedConfig): + """Configuration class for GLM MoE DSA (GLM-5). + + Bundled here because this model requires transformers 5.0+, but we run on 4.57. + Field names match the checkpoint config.json exactly. + """ + + model_type = "glm_moe_dsa" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size: int = 154880, + hidden_size: int = 6144, + intermediate_size: int = 12288, + moe_intermediate_size: int = 2048, + num_hidden_layers: int = 78, + num_attention_heads: int = 64, + num_key_value_heads: int = 64, + hidden_act: str = "silu", + max_position_embeddings: int = 202752, + initializer_range: float = 0.02, + rms_norm_eps: float = 1e-5, + # MLA parameters + q_lora_rank: int = 2048, + kv_lora_rank: int = 512, + qk_nope_head_dim: int = 192, + qk_rope_head_dim: int = 64, + v_head_dim: int = 256, + # MoE parameters + n_routed_experts: int = 256, + n_shared_experts: int = 1, + num_experts_per_tok: int = 8, + n_group: int = 1, + topk_group: int = 1, + routed_scaling_factor: float = 2.5, + norm_topk_prob: bool = True, + # Layer type control (checkpoint format uses first_k_dense_replace + moe_layer_freq) + first_k_dense_replace: int = 3, + moe_layer_freq: int = 1, + # RoPE — accept both checkpoint format (rope_parameters dict) and direct rope_theta + rope_theta: float = 1000000.0, + rope_scaling: Optional[dict] = None, + rope_parameters: Optional[dict] = None, # checkpoint format + rope_interleave: bool = True, + # Indexer (DSA) parameters + index_topk: int = 2048, + index_head_dim: int = 128, + index_n_heads: int = 32, + indexer_rope_interleave: bool = True, + # Other + attention_bias: bool = False, + attention_dropout: float = 0.0, + tie_word_embeddings: bool = False, + pad_token_id: int = 154820, + # Extra checkpoint fields that we ignore + ep_size: int = 1, + head_dim: int = 64, + qk_head_dim: int = 256, + num_nextn_predict_layers: int = 0, + scoring_func: str = "sigmoid", + topk_method: str = "noaux_tc", + dtype: str = "bfloat16", + pretraining_tp: int = 1, + **kwargs, + ): + # Model dimensions + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + + # MLA + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + + # MoE + self.n_routed_experts = n_routed_experts + self.n_shared_experts = n_shared_experts + self.num_experts_per_tok = num_experts_per_tok + self.n_group = n_group + self.topk_group = topk_group + self.routed_scaling_factor = routed_scaling_factor + self.norm_topk_prob = norm_topk_prob + self.first_k_dense_replace = first_k_dense_replace + self.moe_layer_freq = moe_layer_freq + + # RoPE — extract rope_theta from rope_parameters dict if provided + if rope_parameters is not None and isinstance(rope_parameters, dict): + self.rope_theta = rope_parameters.get("rope_theta", rope_theta) + else: + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.rope_parameters = rope_parameters + self.rope_interleave = rope_interleave + + # Indexer (DSA) + self.index_topk = index_topk + self.index_head_dim = index_head_dim + self.index_n_heads = index_n_heads + self.indexer_rope_interleave = indexer_rope_interleave + + # Other + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + + super().__init__( + tie_word_embeddings=tie_word_embeddings, + pad_token_id=pad_token_id, + **kwargs, + ) + + +AutoConfig.register("glm_moe_dsa", GlmMoeDsaConfig, exist_ok=True) + + +# ============================================================================= +# Building blocks +# ============================================================================= + + +class GlmDSARMSNorm(nn.Module): + """RMS Normalization using the canonical AutoDeploy torch_rmsnorm op.""" + + def __init__(self, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return torch.ops.auto_deploy.torch_rmsnorm( + hidden_states, self.weight, self.variance_epsilon + ) + + +class GlmDSARotaryEmbedding(nn.Module): + """Rotary Position Embedding (non-interleaved, NeoX style after de-interleave at load time).""" + + def __init__(self, dim: int, max_position_embeddings: int = 2048, base: float = 10000.0): + super().__init__() + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + + inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim)) + self.register_buffer("_ad_inv_freq", inv_freq, persistent=False) + self._set_cos_sin_cache(max_position_embeddings) + + def _set_cos_sin_cache(self, seq_len: int): + self.max_seq_len_cached = seq_len + t = torch.arange(seq_len, dtype=self._ad_inv_freq.dtype) + freqs = torch.outer(t, self._ad_inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("_ad_cos_cached", emb.cos(), persistent=False) + self.register_buffer("_ad_sin_cached", emb.sin(), persistent=False) + + def forward( + self, x: torch.Tensor, seq_len: Optional[int] = None + ) -> Tuple[torch.Tensor, torch.Tensor]: + return ( + self._ad_cos_cached.to(dtype=x.dtype, device=x.device), + self._ad_sin_cached.to(dtype=x.dtype, device=x.device), + ) + + +class GlmDSAYarnRotaryEmbedding(GlmDSARotaryEmbedding): + """YaRN-extended rotary embedding.""" + + def __init__( + self, + dim: int, + max_position_embeddings: int = 2048, + base: float = 10000.0, + scaling_factor: float = 1.0, + original_max_position_embeddings: int = 4096, + beta_fast: int = 32, + beta_slow: int = 1, + mscale: float = 1.0, + mscale_all_dim: float = 0.0, + ): + self.scaling_factor = scaling_factor + self.original_max_position_embeddings = original_max_position_embeddings + self.beta_fast = beta_fast + self.beta_slow = beta_slow + self.mscale = mscale + self.mscale_all_dim = mscale_all_dim + super().__init__(dim, max_position_embeddings, base) + + def _set_cos_sin_cache(self, seq_len: int): + self.max_seq_len_cached = seq_len + dim = self.dim + + freq_extra = 1.0 / (self.base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + freq_inter = 1.0 / ( + self.scaling_factor * self.base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim) + ) + + low, high = self._yarn_find_correction_range( + self.beta_fast, self.beta_slow, dim, self.base, self.original_max_position_embeddings + ) + inv_freq_mask = 1.0 - self._yarn_linear_ramp_mask(low, high, dim // 2) + inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask + self.register_buffer("_ad_inv_freq", inv_freq, persistent=False) + + t = torch.arange(seq_len, dtype=torch.float32) + freqs = torch.outer(t, inv_freq) + _mscale = float( + self._yarn_get_mscale(self.scaling_factor, self.mscale) + / self._yarn_get_mscale(self.scaling_factor, self.mscale_all_dim) + ) + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("_ad_cos_cached", (emb.cos() * _mscale), persistent=False) + self.register_buffer("_ad_sin_cached", (emb.sin() * _mscale), persistent=False) + + @staticmethod + def _yarn_find_correction_dim( + num_rotations: float, dim: int, base: float = 10000, max_position_embeddings: int = 2048 + ) -> float: + return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / ( + 2 * math.log(base) + ) + + def _yarn_find_correction_range( + self, low_rot: int, high_rot: int, dim: int, base: float, max_position_embeddings: int + ) -> Tuple[int, int]: + low = math.floor( + self._yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings) + ) + high = math.ceil( + self._yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings) + ) + return max(low, 0), min(high, dim - 1) + + @staticmethod + def _yarn_get_mscale(scale: float = 1.0, mscale: float = 1.0) -> float: + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + @staticmethod + def _yarn_linear_ramp_mask(min_val: float, max_val: float, dim: int) -> torch.Tensor: + if min_val == max_val: + max_val += 0.001 + linear_func = (torch.arange(dim, dtype=torch.float32) - min_val) / (max_val - min_val) + return torch.clamp(linear_func, 0, 1) + + +class GlmDSAMLP(nn.Module): + """MLP with SwiGLU activation.""" + + def __init__( + self, config, hidden_size: Optional[int] = None, intermediate_size: Optional[int] = None + ): + super().__init__() + self.hidden_size = hidden_size or config.hidden_size + self.intermediate_size = intermediate_size or config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +class GlmDSAMoEGate(nn.Module): + """MoE gate with noaux_tc top-k routing — vanilla PyTorch implementation. + + Matches HF GlmMoeDsaTopkRouter + GlmMoeDsaMoE.route_tokens_to_experts logic. + """ + + def __init__(self, config): + super().__init__() + self.top_k = config.num_experts_per_tok + self.n_routed_experts = config.n_routed_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.n_group = config.n_group + self.topk_group = config.topk_group + self.norm_topk_prob = config.norm_topk_prob + + self.weight = nn.Parameter( + torch.empty((self.n_routed_experts, config.hidden_size), dtype=torch.float32) + ) + self.register_buffer( + "e_score_correction_bias", + torch.zeros(self.n_routed_experts, dtype=torch.float32), + ) + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + + def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + bsz, seq_len, hidden_dim = hidden_states.shape + hidden_flat = hidden_states.view(-1, hidden_dim) + + # Router logits and sigmoid scoring (always float32) + router_logits = F.linear(hidden_flat.float(), self.weight.float()) # [T, n_experts] + scores = router_logits.sigmoid() + scores_for_choice = scores + self.e_score_correction_bias # [T, n_experts] + + # Group-level top-k selection + group_scores = ( + scores_for_choice.view(-1, self.n_group, self.n_routed_experts // self.n_group) + .topk(2, dim=-1)[0] + .sum(dim=-1) + ) # [T, n_group] + group_idx = group_scores.topk( + self.topk_group, dim=-1, sorted=False + ).indices # [T, topk_group] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1.0) + score_mask = ( + group_mask.unsqueeze(-1) + .expand(-1, self.n_group, self.n_routed_experts // self.n_group) + .reshape(-1, self.n_routed_experts) + ) # [T, n_experts] + + # Mask out non-selected groups, then take per-token top-k experts + scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), 0.0) + topk_indices = scores_for_choice.topk( + self.top_k, dim=-1, sorted=False + ).indices # [T, top_k] + + # Gather weights from original scores (without correction bias) + topk_weights = scores.gather(1, topk_indices) # [T, top_k] + if self.norm_topk_prob: + topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-20) + topk_weights = topk_weights * self.routed_scaling_factor + + return topk_indices, topk_weights + + +class GlmDSAMoE(nn.Module): + """Mixture of Experts layer.""" + + def __init__(self, config): + super().__init__() + self.num_experts_per_tok = config.num_experts_per_tok + self.experts = nn.ModuleList( + [ + GlmDSAMLP(config, intermediate_size=config.moe_intermediate_size) + for _ in range(config.n_routed_experts) + ] + ) + self.gate = GlmDSAMoEGate(config) + if config.n_shared_experts is not None: + self.shared_experts = GlmDSAMLP( + config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts + ) + else: + self.shared_experts = None + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + identity = hidden_states + orig_shape = hidden_states.shape + topk_indices, topk_weights = self.gate(hidden_states) + + final_hidden_states = torch.ops.auto_deploy.torch_moe( + hidden_states.view(-1, hidden_states.shape[-1]), + topk_indices, + topk_weights, + w1_weight=[e.gate_proj.weight for e in self.experts], + w2_weight=[e.down_proj.weight for e in self.experts], + w3_weight=[e.up_proj.weight for e in self.experts], + is_gated_mlp=True, + act_fn=int(ActivationType.Silu), + ) + + final_hidden_states = final_hidden_states.view(*orig_shape) + if self.shared_experts is not None: + final_hidden_states = final_hidden_states + self.shared_experts(identity) + return final_hidden_states.to(hidden_states.dtype) + + +class GlmDSAIndexer(nn.Module): + """DSA Indexer — computes per-token importance scores and returns top-k index keys. + + Submodule names match the HF GlmMoeDsaIndexer checkpoint keys: + self_attn.indexer.wq_b, .wk, .k_norm, .weights_proj + """ + + def __init__(self, config): + super().__init__() + self.index_n_heads = config.index_n_heads + self.index_head_dim = config.index_head_dim + self.qk_rope_head_dim = config.qk_rope_head_dim + self.q_lora_rank = config.q_lora_rank + + # wq_b: q_lora_rank → index_n_heads * index_head_dim + self.wq_b = nn.Linear( + config.q_lora_rank, config.index_n_heads * config.index_head_dim, bias=False + ) + # wk: hidden_size → index_head_dim + self.wk = nn.Linear(config.hidden_size, config.index_head_dim, bias=False) + # k_norm: LayerNorm on indexer key (eps=1e-6 matches HF) + self.k_norm = nn.LayerNorm(config.index_head_dim, eps=1e-6) + # weights_proj: hidden_size → index_n_heads (per-head importance scalars) + self.weights_proj = nn.Linear(config.hidden_size, config.index_n_heads, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, # [B, S, hidden_size] + qr: torch.Tensor, # [B, S, q_lora_rank] (shared from MLA Q path) + cos: torch.Tensor, # [B, S, qk_rope_head_dim] + sin: torch.Tensor, # [B, S, qk_rope_head_dim] + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Returns (index_q, index_k, index_weights) with RoPE applied.""" + bsz, q_len, _ = hidden_states.shape + + # ---- Indexer Q ---- + # wq_b projects qr → [B, S, index_n_heads, index_head_dim] + # Layout: [rope_dim | nope_dim] within each head (matches HF checkpoint) + index_q_raw = self.wq_b(qr).view(bsz, q_len, self.index_n_heads, self.index_head_dim) + index_q_pe_raw = index_q_raw[:, :, :, : self.qk_rope_head_dim] + index_q_nope = index_q_raw[:, :, :, self.qk_rope_head_dim :] + + # Apply RoPE to indexer Q pe part (use dummy k to satisfy op signature) + index_k_dummy = index_q_raw.new_zeros(bsz, q_len, 1, self.qk_rope_head_dim) + index_q_pe_rotated, _ = torch.ops.auto_deploy.torch_rope_with_explicit_cos_sin( + index_q_pe_raw, index_k_dummy, cos, sin, 2 + ) + # Recombine: [pe_rotated | nope] + index_q = torch.cat([index_q_pe_rotated, index_q_nope], dim=-1) + + # ---- Indexer K ---- + # wk projects hidden_states → [B, S, index_head_dim] + # Layout: [rope_dim | nope_dim] + index_k_raw = self.k_norm(self.wk(hidden_states)) + index_k_pe_raw = index_k_raw[:, :, : self.qk_rope_head_dim] + index_k_nope = index_k_raw[:, :, self.qk_rope_head_dim :] + + # Apply RoPE to indexer K pe part + index_k_pe_4d = index_k_pe_raw.view(bsz, q_len, 1, self.qk_rope_head_dim) + index_q_dummy = index_k_pe_4d.new_zeros(bsz, q_len, 1, self.qk_rope_head_dim) + _, index_k_pe_rotated = torch.ops.auto_deploy.torch_rope_with_explicit_cos_sin( + index_q_dummy, index_k_pe_4d, cos, sin, 2 + ) + index_k_pe_rotated = index_k_pe_rotated.view(bsz, q_len, self.qk_rope_head_dim) + # Recombine: [pe_rotated | nope] + index_k = torch.cat([index_k_pe_rotated, index_k_nope], dim=-1) + + # ---- Indexer importance weights ---- + # weights_proj scaled by n_heads^(-0.5), matching HF + index_weights = self.weights_proj(hidden_states) * (self.index_n_heads**-0.5) + + return index_q, index_k, index_weights + + +class GlmDSAAttention(nn.Module): + """MLA + DSA (DeepSeek Sparse Attention) for GLM-5. + + The Indexer computes per-token importance scores; top-k positions are kept and + the rest are masked to -inf before softmax. RoPE is de-interleaved at load time. + """ + + def __init__(self, config, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.q_lora_rank = config.q_lora_rank + self.kv_lora_rank = config.kv_lora_rank + self.qk_nope_head_dim = config.qk_nope_head_dim + self.qk_rope_head_dim = config.qk_rope_head_dim + self.v_head_dim = config.v_head_dim + self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + + # Softmax scale (with optional YaRN mscale correction) + self.softmax_scale = self.q_head_dim ** (-0.5) + if config.rope_scaling is not None: + mscale_all_dim = config.rope_scaling.get("mscale_all_dim", 0) + scaling_factor = config.rope_scaling["factor"] + if mscale_all_dim: + mscale = GlmDSAYarnRotaryEmbedding._yarn_get_mscale(scaling_factor, mscale_all_dim) + self.softmax_scale = self.softmax_scale * mscale * mscale + + # MLA projections + self.q_a_proj = nn.Linear(self.hidden_size, self.q_lora_rank, bias=False) + self.q_a_layernorm = GlmDSARMSNorm(self.q_lora_rank) + self.q_b_proj = nn.Linear(self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False) + + self.kv_a_proj_with_mqa = nn.Linear( + self.hidden_size, self.kv_lora_rank + self.qk_rope_head_dim, bias=False + ) + self.kv_a_layernorm = GlmDSARMSNorm(self.kv_lora_rank) + self.kv_b_proj = nn.Linear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + ) + self.o_proj = nn.Linear(self.num_heads * self.v_head_dim, self.hidden_size, bias=False) + + # Indexer submodule — names match HF checkpoint: self_attn.indexer.* + self.indexer = GlmDSAIndexer(config) + + self._init_rope() + + def _init_rope(self): + if self.config.rope_scaling is None: + self.rotary_emb = GlmDSARotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + else: + scaling_type = self.config.rope_scaling["type"] + scaling_factor = self.config.rope_scaling["factor"] + if scaling_type == "yarn": + kwargs = { + k: self.config.rope_scaling[k] + for k in [ + "original_max_position_embeddings", + "beta_fast", + "beta_slow", + "mscale", + "mscale_all_dim", + ] + if k in self.config.rope_scaling + } + self.rotary_emb = GlmDSAYarnRotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + scaling_factor=scaling_factor, + base=self.rope_theta, + **kwargs, + ) + else: + self.rotary_emb = GlmDSARotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + bsz, q_len, _ = hidden_states.size() + + # ---- MLA Q path — qr is shared with the Indexer ---- + qr = self.q_a_layernorm(self.q_a_proj(hidden_states)) # [B, S, q_lora_rank] + q = self.q_b_proj(qr).view(bsz, q_len, self.num_heads, self.q_head_dim) + q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + # ---- MLA KV path ---- + kv_a_output = self.kv_a_proj_with_mqa(hidden_states) + compressed_kv, k_pe = torch.split( + kv_a_output, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + compressed_kv = self.kv_a_layernorm(compressed_kv) + k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim) + + # ---- RoPE (weights de-interleaved at load time → NeoX-style) ---- + cos, sin = self.rotary_emb(hidden_states, seq_len=q_len) + cos = cos[position_ids] # [B, S, rope_head_dim] + sin = sin[position_ids] + + q_pe_rotated, kpe = torch.ops.auto_deploy.torch_rope_with_explicit_cos_sin( + q_pe, + k_pe, + cos, + sin, + 2, # unsqueeze_dim=2 for BSND layout + ) + + # ---- Indexer ---- + index_q, index_k, index_weights = self.indexer(hidden_states, qr, cos, sin) + + # ---- DSA attention ---- + attn_output = torch.ops.auto_deploy.torch_dsa( + q_nope, # [B, S, N, qk_nope_head_dim] + q_pe_rotated, # [B, S, N, qk_rope_head_dim] + compressed_kv, # [B, S, kv_lora_rank] + kpe, # [B, S, 1, qk_rope_head_dim] + self.kv_b_proj.weight, # [N*(qk_nope+v), kv_lora_rank] + index_q, # [B, S, index_n_heads, index_head_dim] + index_k, # [B, S, index_head_dim] + index_weights, # [B, S, index_n_heads] + self.config.index_topk, + True, # is_causal + self.softmax_scale, + "bsnd", + ) + + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim) + return self.o_proj(attn_output) + + +class GlmDSADecoderLayer(nn.Module): + """Transformer decoder layer for GLM-5.""" + + def __init__(self, config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.layer_idx = layer_idx + + self.self_attn = GlmDSAAttention(config, layer_idx=layer_idx) + + use_moe = ( + config.n_routed_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ) + self.mlp = GlmDSAMoE(config) if use_moe else GlmDSAMLP(config) + + self.input_layernorm = GlmDSARMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = GlmDSARMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, position_ids) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +# ============================================================================= +# Top-level model classes +# ============================================================================= + + +@dataclass +class GlmDSAOutput(ModelOutput): + last_hidden_state: Optional[torch.FloatTensor] = None + + +@dataclass +class GlmDSACausalLMOutput(ModelOutput): + logits: Optional[torch.FloatTensor] = None + + +class GlmDSAPreTrainedModel(PreTrainedModel): + config_class = GlmMoeDsaConfig + base_model_prefix = "model" + _no_split_modules = ["GlmDSADecoderLayer"] + supports_gradient_checkpointing = False + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +class GlmDSAModel(GlmDSAPreTrainedModel): + """GLM-5 transformer decoder model.""" + + def __init__(self, config): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [GlmDSADecoderLayer(config, layer_idx=idx) for idx in range(config.num_hidden_layers)] + ) + self.norm = GlmDSARMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + **kwargs, + ) -> GlmDSAOutput: + if input_ids is not None and inputs_embeds is not None: + raise ValueError("Cannot specify both input_ids and inputs_embeds") + if input_ids is None and inputs_embeds is None: + raise ValueError("Must specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + assert position_ids is not None, "position_ids must be provided for AD export" + batch_size, seq_length = inputs_embeds.shape[:2] + + hidden_states = inputs_embeds + for decoder_layer in self.layers: + hidden_states = decoder_layer(hidden_states, position_ids) + + hidden_states = self.norm(hidden_states) + return GlmDSAOutput(last_hidden_state=hidden_states) + + +class GlmDSAForCausalLM(GlmDSAPreTrainedModel, GenerationMixin): + """GLM-5 model with language modeling head.""" + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = GlmDSAModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Load hook 1: De-interleave MLA RoPE weights (q_b_proj, kv_a_proj_with_mqa) + self._register_load_state_dict_pre_hook( + partial( + _mla_rope_deinterleave_hook, + qk_rope_head_dim=config.qk_rope_head_dim, + qk_nope_head_dim=config.qk_nope_head_dim, + num_heads=config.num_attention_heads, + kv_lora_rank=config.kv_lora_rank, + num_layers=config.num_hidden_layers, + ) + ) + + # Load hook 2: De-interleave indexer RoPE weights (indexer.wq_b, indexer.wk) + if config.indexer_rope_interleave: + self._register_load_state_dict_pre_hook( + partial( + _indexer_rope_deinterleave_hook, + qk_rope_head_dim=config.qk_rope_head_dim, + index_n_heads=config.index_n_heads, + index_head_dim=config.index_head_dim, + num_layers=config.num_hidden_layers, + ) + ) + + # Load hook 3: Expand stacked MoE expert weights → per-expert ModuleList + self._register_load_state_dict_pre_hook( + partial( + _moe_expert_expand_hook, + n_routed_experts=config.n_routed_experts, + moe_intermediate_size=config.moe_intermediate_size, + first_k_dense_replace=config.first_k_dense_replace, + moe_layer_freq=config.moe_layer_freq, + num_layers=config.num_hidden_layers, + ) + ) + + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def get_decoder(self): + return self.model + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + **kwargs, + ) -> GlmDSACausalLMOutput: + outputs = self.model( + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **kwargs, + ) + logits = self.lm_head(outputs.last_hidden_state).float() + return GlmDSACausalLMOutput(logits=logits) + + +# ============================================================================= +# Load-time weight transformation hooks +# ============================================================================= + + +def _mla_rope_deinterleave_hook( + state_dict, + prefix, + *args, + qk_rope_head_dim: int, + qk_nope_head_dim: int, + num_heads: int, + kv_lora_rank: int, + num_layers: int, +): + """De-interleave MLA RoPE weights from GLM-5 interleaved format to NeoX format. + + For q_b_proj: output shape [num_heads * qk_head_dim, q_lora_rank]. + Each head has [nope_dim | rope_dim]; rope_dim is interleaved and needs reordering. + For kv_a_proj_with_mqa: output shape [kv_lora_rank + qk_rope_head_dim, hidden_size]. + The last qk_rope_head_dim rows are the rope part and need reordering. + """ + d = qk_rope_head_dim + perm = torch.cat([torch.arange(0, d, 2), torch.arange(1, d, 2)]) + qk_head_dim = qk_nope_head_dim + d + + for layer_idx in range(num_layers): + layer_prefix = f"{prefix}model.layers.{layer_idx}.self_attn." + + q_key = layer_prefix + "q_b_proj.weight" + if q_key in state_dict: + w = state_dict[q_key] + w = w.view(num_heads, qk_head_dim, -1) + w_nope = w[:, :qk_nope_head_dim, :] + w_rope = w[:, qk_nope_head_dim:, :] + w_rope = w_rope[:, perm, :] + state_dict[q_key] = torch.cat([w_nope, w_rope], dim=1).view(-1, w.shape[-1]) + + kv_key = layer_prefix + "kv_a_proj_with_mqa.weight" + if kv_key in state_dict: + w = state_dict[kv_key] + w_kv = w[:kv_lora_rank, :] + w_pe = w[kv_lora_rank:, :] + state_dict[kv_key] = torch.cat([w_kv, w_pe[perm, :]], dim=0) + + kv_bias_key = layer_prefix + "kv_a_proj_with_mqa.bias" + if kv_bias_key in state_dict: + b = state_dict[kv_bias_key] + b_kv = b[:kv_lora_rank] + b_pe = b[kv_lora_rank:] + state_dict[kv_bias_key] = torch.cat([b_kv, b_pe[perm]]) + + +def _indexer_rope_deinterleave_hook( + state_dict, + prefix, + *args, + qk_rope_head_dim: int, + index_n_heads: int, + index_head_dim: int, + num_layers: int, +): + """De-interleave indexer RoPE weights from interleaved to non-interleaved format. + + Indexer layout within each head: [rope_dim | nope_dim] (rope is FIRST). + + For indexer.wq_b: shape [index_n_heads * index_head_dim, q_lora_rank]. + Each head has [rope_dim | nope_dim]; rope_dim needs reordering. + For indexer.wk: shape [index_head_dim, hidden_size]. + First rope_dim rows are the rope part and need reordering. + """ + d = qk_rope_head_dim + perm = torch.cat([torch.arange(0, d, 2), torch.arange(1, d, 2)]) + + for layer_idx in range(num_layers): + idx_prefix = f"{prefix}model.layers.{layer_idx}.self_attn.indexer." + + wq_key = idx_prefix + "wq_b.weight" + if wq_key in state_dict: + w = state_dict[wq_key] + # [index_n_heads * index_head_dim, q_lora_rank] → [n_heads, head_dim, q_lora_rank] + w = w.view(index_n_heads, index_head_dim, -1) + w_rope = w[:, :d, :] + w_nope = w[:, d:, :] + w_rope = w_rope[:, perm, :] + state_dict[wq_key] = torch.cat([w_rope, w_nope], dim=1).view(-1, w.shape[-1]) + + wk_key = idx_prefix + "wk.weight" + if wk_key in state_dict: + w = state_dict[wk_key] + # [index_head_dim, hidden_size] + w_rope = w[:d, :] + w_nope = w[d:, :] + state_dict[wk_key] = torch.cat([w_rope[perm, :], w_nope], dim=0) + + +def _moe_expert_expand_hook( + state_dict, + prefix, + *args, + n_routed_experts: int, + moe_intermediate_size: int, + first_k_dense_replace: int, + moe_layer_freq: int, + num_layers: int, +): + """Expand stacked HF expert weights into per-expert format expected by GlmDSAMoE. + + HF checkpoint stores experts as: + mlp.experts.gate_up_proj: [n_experts, 2 * moe_intermediate_size, hidden_size] + mlp.experts.down_proj: [n_experts, hidden_size, moe_intermediate_size] + + Our model expects: + mlp.experts.{i}.gate_proj.weight: [moe_intermediate_size, hidden_size] + mlp.experts.{i}.up_proj.weight: [moe_intermediate_size, hidden_size] + mlp.experts.{i}.down_proj.weight: [hidden_size, moe_intermediate_size] + """ + for layer_idx in range(num_layers): + is_moe = layer_idx >= first_k_dense_replace and layer_idx % moe_layer_freq == 0 + if not is_moe: + continue + + mlp_prefix = f"{prefix}model.layers.{layer_idx}.mlp." + gate_up_key = mlp_prefix + "experts.gate_up_proj" + down_key = mlp_prefix + "experts.down_proj" + + if gate_up_key not in state_dict: + continue + + gate_up = state_dict.pop(gate_up_key) # [n_experts, 2*intermediate, hidden] + down = state_dict.pop(down_key) # [n_experts, hidden, intermediate] + + for i in range(n_routed_experts): + gate_up_i = gate_up[i] # [2*intermediate, hidden] + state_dict[mlp_prefix + f"experts.{i}.gate_proj.weight"] = gate_up_i[ + :moe_intermediate_size + ] + state_dict[mlp_prefix + f"experts.{i}.up_proj.weight"] = gate_up_i[ + moe_intermediate_size: + ] + state_dict[mlp_prefix + f"experts.{i}.down_proj.weight"] = down[i] + + +# ============================================================================= +# Registration +# ============================================================================= + +AutoModelForCausalLMFactory.register_custom_model_cls("GlmMoeDsaConfig", GlmDSAForCausalLM) diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/mla/test_flashmla_dsa_op.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/mla/test_flashmla_dsa_op.py new file mode 100644 index 000000000000..120855808445 --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/mla/test_flashmla_dsa_op.py @@ -0,0 +1,495 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for FlashMLA DSA backend (flash_mla_dsa_with_cache) and FlashMLADSAAttention descriptor. + +FlashMLA requires SM90+ (Hopper/Blackwell). All tests are skipped on older hardware. +""" + +import pytest +import torch + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9, + reason="FlashMLA requires SM90+ (Hopper/Blackwell)", +) + +import tensorrt_llm._torch.auto_deploy # noqa: F401 + +# --------------------------------------------------------------------------- +# Helper: build paged metadata from simple contiguous-page layout +# --------------------------------------------------------------------------- + + +def _make_paged_data( + batch_size: int, + seq_len_per_seq: int, # tokens to write in this batch (1 for decode) + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + kv_lora_rank: int, + v_head_dim: int, + index_n_heads: int, + index_head_dim: int, + page_size: int = 16, + cache_offset: int = 0, + dtype: torch.dtype = torch.bfloat16, + device: str = "cuda", +): + """Build tensors for flash_mla_dsa_with_cache. + + Uses a simple paged layout where each sequence gets its own contiguous set of pages. + """ + B = batch_size + S = seq_len_per_seq + total_tokens = B * S + kv_head_dim = qk_nope_head_dim + v_head_dim + + # --- Input tensors (BSND layout for decode, or flat for prefill) --- + if S == 1: + # Decode: [B, 1, ...] + q_nope = torch.randn(B, 1, num_heads, qk_nope_head_dim, dtype=dtype, device=device) + q_pe = torch.randn(B, 1, num_heads, qk_rope_head_dim, dtype=dtype, device=device) + compressed_kv = torch.randn(B, 1, kv_lora_rank, dtype=dtype, device=device) + kpe = torch.randn(B, 1, 1, qk_rope_head_dim, dtype=dtype, device=device) + index_q = torch.randn(B, 1, index_n_heads, index_head_dim, dtype=dtype, device=device) + index_k = torch.randn(B, 1, index_head_dim, dtype=dtype, device=device) + index_weights = torch.randn(B, 1, index_n_heads, dtype=dtype, device=device) + else: + # Prefill: [1, B*S, ...] (flat batch) + q_nope = torch.randn( + 1, total_tokens, num_heads, qk_nope_head_dim, dtype=dtype, device=device + ) + q_pe = torch.randn(1, total_tokens, num_heads, qk_rope_head_dim, dtype=dtype, device=device) + compressed_kv = torch.randn(1, total_tokens, kv_lora_rank, dtype=dtype, device=device) + kpe = torch.randn(1, total_tokens, 1, qk_rope_head_dim, dtype=dtype, device=device) + index_q = torch.randn( + 1, total_tokens, index_n_heads, index_head_dim, dtype=dtype, device=device + ) + index_k = torch.randn(1, total_tokens, index_head_dim, dtype=dtype, device=device) + index_weights = torch.randn(1, total_tokens, index_n_heads, dtype=dtype, device=device) + + kv_b_proj_weight = torch.randn( + num_heads * kv_head_dim, kv_lora_rank, dtype=dtype, device=device + ) + + # --- Paged cache metadata --- + # Each sequence needs ceil((cache_offset + S) / page_size) pages + max_pos = cache_offset + S + pages_per_seq = (max_pos + page_size - 1) // page_size + total_pages = B * pages_per_seq + + # Simple contiguous page layout: seq i gets pages [i*pages_per_seq, (i+1)*pages_per_seq) + # cache_loc: maps from (seq, page_k) index to actual block index + # We use a simple identity mapping: page i in seq j is stored at block j*pages_per_seq + i + page_table = torch.arange(total_pages, device=device, dtype=torch.int32) + cache_loc = page_table # [total_pages] + + # cu_num_pages[i] = i * pages_per_seq + cu_num_pages = torch.arange(B + 1, device=device, dtype=torch.int32) * pages_per_seq + + # Paged caches (zero-initialized, pre-filled at cache_offset if requested) + mla_cache = torch.zeros( + total_pages, page_size, 1, kv_lora_rank + qk_rope_head_dim, dtype=dtype, device=device + ) + index_k_cache = torch.zeros(total_pages, page_size, index_head_dim, dtype=dtype, device=device) + + # Fill cache at cache_offset positions with random data + if cache_offset > 0: + for b in range(B): + for pos in range(cache_offset): + page_k = pos // page_size + page_off = pos % page_size + blk = int(cache_loc[cu_num_pages[b].item() + page_k].item()) + mla_cache[blk, page_off, 0, :] = torch.randn( + kv_lora_rank + qk_rope_head_dim, dtype=dtype, device=device + ) + index_k_cache[blk, page_off, :] = torch.randn( + index_head_dim, dtype=dtype, device=device + ) + + # last_page_len: valid tokens in last page after writing + last_page_len = torch.full( + (B,), (max_pos - 1) % page_size + 1, device=device, dtype=torch.int32 + ) + + # seq_len and input_pos per sequence + seq_len_tensor = torch.full((B,), S, device=device, dtype=torch.int32) + input_pos_tensor = torch.full((B,), cache_offset, device=device, dtype=torch.int32) + + if S == 1: + # Decode + batch_info_host = torch.tensor([0, 0, B], device=device, dtype=torch.int32) + cu_seqlen = torch.arange(B + 1, device=device, dtype=torch.int32) + else: + # Prefill + batch_info_host = torch.tensor([B, B * S, 0], device=device, dtype=torch.int32) + cu_seqlen = torch.arange(0, B * S + 1, S, device=device, dtype=torch.int32) + + return { + "q_nope": q_nope, + "q_pe": q_pe, + "compressed_kv": compressed_kv, + "kpe": kpe, + "kv_b_proj_weight": kv_b_proj_weight, + "index_q": index_q, + "index_k": index_k, + "index_weights": index_weights, + "batch_info_host": batch_info_host, + "seq_len": seq_len_tensor, + "input_pos": input_pos_tensor, + "cu_seqlen": cu_seqlen, + "cache_loc": cache_loc, + "cu_num_pages": cu_num_pages, + "last_page_len": last_page_len, + "mla_cache": mla_cache, + "index_k_cache": index_k_cache, + "kv_lora_rank": kv_lora_rank, + "v_head_dim": v_head_dim, + "num_heads": num_heads, + } + + +def _run_flash_mla_dsa(data, scale=None, index_topk=64): + """Call flash_mla_dsa_with_cache with the given data dict.""" + return torch.ops.auto_deploy.flash_mla_dsa_with_cache( + data["q_nope"], + data["q_pe"], + data["compressed_kv"], + data["kpe"], + data["kv_b_proj_weight"], + data["index_q"], + data["index_k"], + data["index_weights"], + data["batch_info_host"], + data["seq_len"], + data["input_pos"], + data["cu_seqlen"], + data["cache_loc"], + data["cu_num_pages"], + data["last_page_len"], + data["mla_cache"], + data["index_k_cache"], + scale, + data["kv_lora_rank"], + index_topk, + ) + + +def _run_torch_dsa(data): + """Call the torch reference cached DSA (unpaged) with equivalent inputs. + + Builds an unpaged cache from the paged mla_cache/index_k_cache contents, + then calls torch_cached_dsa_with_cache. + """ + B = data["seq_len"].shape[0] + kv_lora_rank = data["kv_lora_rank"] + index_head_dim = data["index_k_cache"].shape[-1] + qk_rope_head_dim = data["q_pe"].shape[-1] + num_heads = data["num_heads"] + v_head_dim = data["v_head_dim"] + device = data["q_nope"].device + dtype = data["q_nope"].dtype + + seq_len = data["seq_len"][:B] + input_pos = data["input_pos"][:B] + cache_loc_torch = data["cache_loc"] + cu_num_pages = data["cu_num_pages"] + + # Compute cache_seqlens (after write, same as flash_mla_dsa does) + cache_seqlens = (input_pos + seq_len).to(torch.int32) + + # Build unpaged mla_cache and index_k_cache for the torch reference op + max_seq_len = int(cache_seqlens.max().item()) + mla_cache_unpaged = torch.zeros( + B, max_seq_len, kv_lora_rank + qk_rope_head_dim, dtype=dtype, device=device + ) + index_k_cache_unpaged = torch.zeros(B, max_seq_len, index_head_dim, dtype=dtype, device=device) + + page_size = data["mla_cache"].shape[1] + for b in range(B): + for t in range(int(cache_seqlens[b].item())): + page_k = t // page_size + page_off = t % page_size + blk = int(cache_loc_torch[int(cu_num_pages[b].item()) + page_k].item()) + mla_cache_unpaged[b, t] = data["mla_cache"][blk, page_off, 0, :] + index_k_cache_unpaged[b, t] = data["index_k_cache"][blk, page_off, :] + + # slot_idx: contiguous for this simple layout + slot_idx = torch.arange(B, device=device, dtype=torch.int32) + S = data["q_nope"].shape[1] + cu_seqlen = data["cu_seqlen"] + + batch_info_host = data["batch_info_host"] + + return torch.ops.auto_deploy.torch_cached_dsa_with_cache( + data["q_nope"], + data["q_pe"], + data["compressed_kv"], + data["kpe"], + data["kv_b_proj_weight"], + data["index_q"], + data["index_k"], + data["index_weights"], + batch_info_host, + seq_len, + input_pos, + slot_idx, + cu_seqlen, + mla_cache_unpaged, + index_k_cache_unpaged, + None, + kv_lora_rank, + # Use large topk to match full-attend behavior for numerical comparison + int(cache_seqlens.max().item()), + ) + + +# --------------------------------------------------------------------------- +# Class 1: TestFlashMLADSAWithCache +# --------------------------------------------------------------------------- + + +class TestFlashMLADSAWithCache: + """Tests for flash_mla_dsa_with_cache cached op.""" + + @pytest.fixture(autouse=True) + def setup(self): + torch.cuda.empty_cache() + torch.manual_seed(42) + self.dtype = torch.bfloat16 + self.device = "cuda" + self.atol = 5e-2 + + # Standard DSA dims matching DeepSeek/GLM usage + B, S_prefill = 2, 4 + N, nope, rope = 4, 32, 64 + kv_lora_rank = 512 + v_head_dim = 128 + idx_H, idx_D, topk = 1, 16, 4 + page_size = 64 # FlashMLA requires page_block_size == 64 + + def test_decode_shape_and_finite(self): + """Output shape [B, 1, N, v] and all-finite.""" + data = _make_paged_data( + self.B, + 1, + self.N, + self.nope, + self.rope, + self.kv_lora_rank, + self.v_head_dim, + self.idx_H, + self.idx_D, + page_size=self.page_size, + cache_offset=5, + dtype=self.dtype, + device=self.device, + ) + out = _run_flash_mla_dsa(data, index_topk=self.topk) + + assert out.shape == (self.B, 1, self.N, self.v_head_dim), ( + f"Expected ({self.B}, 1, {self.N}, {self.v_head_dim}), got {out.shape}" + ) + assert torch.isfinite(out).all(), "Decode output contains NaN or Inf" + + def test_prefill_shape_and_finite(self): + """Output shape [1, B*S, N, v] and all-finite.""" + data = _make_paged_data( + self.B, + self.S_prefill, + self.N, + self.nope, + self.rope, + self.kv_lora_rank, + self.v_head_dim, + self.idx_H, + self.idx_D, + page_size=self.page_size, + cache_offset=0, + dtype=self.dtype, + device=self.device, + ) + out = _run_flash_mla_dsa(data, index_topk=self.topk) + + expected_shape = (1, self.B * self.S_prefill, self.N, self.v_head_dim) + assert out.shape == expected_shape, f"Expected {expected_shape}, got {out.shape}" + assert torch.isfinite(out).all(), "Prefill output contains NaN or Inf" + + @pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10, + reason="Sparse FlashMLA decode requires SM100+ (Blackwell)", + ) + def test_decode_vs_torch_equivalence(self): + """FlashMLA sparse decode ≈ TorchBackendDSA (atol=5e-2) using large topk.""" + data = _make_paged_data( + self.B, + 1, + self.N, + self.nope, + self.rope, + self.kv_lora_rank, + self.v_head_dim, + self.idx_H, + self.idx_D, + page_size=self.page_size, + cache_offset=5, + dtype=self.dtype, + device=self.device, + ) + out_flash = _run_flash_mla_dsa(data, index_topk=self.topk) + out_torch = _run_torch_dsa(data) + + # Shapes may differ slightly; align to compare + out_flash_cmp = out_flash.reshape(-1, self.N, self.v_head_dim).float() + out_torch_cmp = out_torch.reshape(-1, self.N, self.v_head_dim).float() + + max_diff = (out_flash_cmp - out_torch_cmp).abs().max().item() + assert torch.allclose(out_flash_cmp, out_torch_cmp, atol=self.atol), ( + f"Decode: FlashMLA vs TorchDSA max diff = {max_diff:.4f} > atol={self.atol}" + ) + + @pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10, + reason="Sparse FlashMLA prefill requires SM100+ (Blackwell)", + ) + def test_prefill_vs_torch_equivalence(self): + """FlashMLA sparse prefill ≈ TorchBackendDSA (atol=5e-2) using large topk.""" + data = _make_paged_data( + self.B, + self.S_prefill, + self.N, + self.nope, + self.rope, + self.kv_lora_rank, + self.v_head_dim, + self.idx_H, + self.idx_D, + page_size=self.page_size, + cache_offset=0, + dtype=self.dtype, + device=self.device, + ) + out_flash = _run_flash_mla_dsa(data, index_topk=self.topk) + out_torch = _run_torch_dsa(data) + + out_flash_cmp = out_flash.reshape(-1, self.N, self.v_head_dim).float() + out_torch_cmp = out_torch.reshape(-1, self.N, self.v_head_dim).float() + + max_diff = (out_flash_cmp - out_torch_cmp).abs().max().item() + assert torch.allclose(out_flash_cmp, out_torch_cmp, atol=self.atol), ( + f"Prefill: FlashMLA vs TorchDSA max diff = {max_diff:.4f} > atol={self.atol}" + ) + + def test_both_caches_updated(self): + """mla_cache and index_k_cache must be non-zero after forward.""" + data = _make_paged_data( + self.B, + 1, + self.N, + self.nope, + self.rope, + self.kv_lora_rank, + self.v_head_dim, + self.idx_H, + self.idx_D, + page_size=self.page_size, + cache_offset=0, + dtype=self.dtype, + device=self.device, + ) + mla_before = data["mla_cache"].clone() + idx_before = data["index_k_cache"].clone() + + _run_flash_mla_dsa(data, index_topk=self.topk) + + assert not torch.allclose(data["mla_cache"], mla_before, atol=1e-6), ( + "mla_cache was not updated" + ) + assert not torch.allclose(data["index_k_cache"], idx_before, atol=1e-6), ( + "index_k_cache was not updated" + ) + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_dtype_preservation(self, dtype): + """Output dtype should match input dtype.""" + data = _make_paged_data( + self.B, + 1, + self.N, + self.nope, + self.rope, + self.kv_lora_rank, + self.v_head_dim, + self.idx_H, + self.idx_D, + page_size=self.page_size, + cache_offset=3, + dtype=dtype, + device=self.device, + ) + out = _run_flash_mla_dsa(data, index_topk=self.topk) + assert out.dtype == dtype, f"Expected {dtype}, got {out.dtype}" + + +# --------------------------------------------------------------------------- +# Class 2: TestFlashMLADSADescriptor +# --------------------------------------------------------------------------- + + +class TestFlashMLADSADescriptor: + """Tests for FlashMLADSAAttention descriptor configuration.""" + + def _get_descriptor(self): + from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import AttentionRegistry + + return AttentionRegistry.get("flashmla_dsa") + + def test_descriptor_registration(self): + """FlashMLADSAAttention should be registered under 'flashmla_dsa'.""" + from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import AttentionRegistry + + assert AttentionRegistry.has("flashmla_dsa"), ( + "'flashmla_dsa' not found in AttentionRegistry" + ) + + def test_descriptor_layout(self): + """Descriptor should return 'bsnd' layout.""" + desc = self._get_descriptor() + assert desc.get_attention_layout() == "bsnd" + + def test_descriptor_num_qkv_args(self): + """Descriptor should expect 8 tensor args.""" + desc = self._get_descriptor() + assert desc.get_num_qkv_args() == 8, f"Expected 8, got {desc.get_num_qkv_args()}" + + def test_descriptor_source_op(self): + """Source op should be torch_dsa.""" + desc = self._get_descriptor() + assert desc.get_source_attention_op() == torch.ops.auto_deploy.torch_dsa + + def test_descriptor_cached_op(self): + """Cached op should be flash_mla_dsa_with_cache.default.""" + desc = self._get_descriptor() + assert ( + desc.get_cached_attention_op() == torch.ops.auto_deploy.flash_mla_dsa_with_cache.default + ) + + def test_descriptor_standard_metadata(self): + """Standard metadata should include paged-cache args.""" + desc = self._get_descriptor() + meta = desc.get_standard_metadata_args() + for required in ["cache_loc", "cu_num_pages", "last_page_len"]: + assert required in meta, f"'{required}' missing from standard_metadata_args: {meta}" diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/mla/test_torch_dsa_op.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/mla/test_torch_dsa_op.py new file mode 100644 index 000000000000..ef54cf728aec --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/mla/test_torch_dsa_op.py @@ -0,0 +1,863 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Comprehensive test suite for torch DSA (DeepSeek Sparse Attention) backend operations. + +Tests the torch_dsa source op and torch_cached_dsa_with_cache cached op. +DSA extends MLA with an Indexer that selects top-k KV positions per query token. + +Key features: +- 8 tensor arguments: q_nope, q_pe, compressed_kv, kpe, kv_b_proj_weight, + index_q, index_k, index_weights +- Two caches: mla_cache [max_batch, max_seq, kv_lora_rank + rope_dim] + index_k_cache [max_batch, max_seq, index_head_dim] +- Prefill: Expand compressed_kv, compute sparse index mask +- Generate: Weight absorption for MLA + sparse index mask from cached index_k +""" + +import math + +import numpy as np +import pytest +import torch + +import tensorrt_llm._torch.auto_deploy # noqa: F401 + + +def numpy_dsa_reference( + q_nope: np.ndarray, + q_pe: np.ndarray, + compressed_kv: np.ndarray, + kpe: np.ndarray, + kv_b_proj_weight: np.ndarray, + index_q: np.ndarray, + index_k: np.ndarray, + index_weights: np.ndarray, + mla_cache: np.ndarray, + index_k_cache: np.ndarray, + seq_len: np.ndarray, + input_pos: np.ndarray, + cache_loc: np.ndarray, + seq_start: np.ndarray, + scale: float = None, + kv_lora_rank: int = None, + index_topk: int = 64, + is_generate: bool = False, +): + """Numpy reference implementation of DSA attention with KV cache. + + Mirrors numpy_mla_reference_with_expansion but adds the Indexer step: + per-head scores are computed from index_q / index_k, weighted-summed, + and converted to a sparse mask that is added to MLA attention scores. + """ + if is_generate: + batch_size = q_nope.shape[0] + num_heads = q_nope.shape[2] + qk_nope_head_dim = q_nope.shape[3] + qk_rope_head_dim = q_pe.shape[3] + else: + batch_size = len(seq_len) + num_heads = q_nope.shape[2] + qk_nope_head_dim = q_nope.shape[3] + qk_rope_head_dim = q_pe.shape[3] + + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + + if kv_lora_rank is None: + kv_lora_rank = compressed_kv.shape[-1] + + out_features = kv_b_proj_weight.shape[0] + kv_head_dim = out_features // num_heads + v_head_dim = kv_head_dim - qk_nope_head_dim + index_head_dim = index_q.shape[-1] + index_softmax_scale = 1.0 / math.sqrt(index_head_dim) + + if scale is None: + scale = 1.0 / math.sqrt(qk_head_dim) + + # Update both caches first + if is_generate: + for i in range(batch_size): + cache_idx = cache_loc[i] + pos = input_pos[i] + mla_cache[cache_idx, pos, :kv_lora_rank] = compressed_kv[i, 0] + mla_cache[cache_idx, pos, kv_lora_rank:] = kpe[i, 0, 0] + index_k_cache[cache_idx, pos] = index_k[i, 0] + else: + for i in range(batch_size): + cache_idx = cache_loc[i] + pos = input_pos[i] + seq_len_i = seq_len[i] + seq_start_i = seq_start[i] + for j in range(seq_len_i): + mla_cache[cache_idx, pos + j, :kv_lora_rank] = compressed_kv[seq_start_i + j] + mla_cache[cache_idx, pos + j, kv_lora_rank:] = kpe[seq_start_i + j, 0] + index_k_cache[cache_idx, pos + j] = index_k[seq_start_i + j] + + outputs = [] + + for i in range(batch_size): + cache_idx = cache_loc[i] + pos = input_pos[i] + seq_len_i = seq_len[i] + seq_start_i = seq_start[i] + + if seq_len_i == 0: + continue + + if is_generate: + q_nope_seq = q_nope[i, 0] # [N, qk_nope_head_dim] + q_pe_seq = q_pe[i, 0] # [N, qk_rope_head_dim] + index_q_seq = index_q[i, 0] # [H, index_head_dim] + index_w_seq = index_weights[i, 0] # [H] + else: + q_nope_seq = q_nope[seq_start_i : seq_start_i + seq_len_i] # [S, N, nope] + q_pe_seq = q_pe[seq_start_i : seq_start_i + seq_len_i] # [S, N, rope] + index_q_seq = index_q[seq_start_i : seq_start_i + seq_len_i] # [S, H, D] + index_w_seq = index_weights[seq_start_i : seq_start_i + seq_len_i] # [S, H] + + kv_seq_len = pos + seq_len_i + + # Get cached MLA data + cached_data = mla_cache[cache_idx, :kv_seq_len] + compressed_kv_cached = cached_data[:, :kv_lora_rank] # [T, kv_lora_rank] + kpe_cached = cached_data[:, kv_lora_rank:] # [T, rope_dim] + + # Get cached index_k + index_k_cached = index_k_cache[cache_idx, :kv_seq_len] # [T, D] + + # Expand compressed_kv + kv_expanded = np.matmul(compressed_kv_cached, kv_b_proj_weight.T) + kv_expanded = kv_expanded.reshape(kv_seq_len, num_heads, kv_head_dim) + k_nope = kv_expanded[:, :, :qk_nope_head_dim] # [T, N, nope] + v = kv_expanded[:, :, qk_nope_head_dim:] # [T, N, v] + + kpe_expanded = np.broadcast_to( + kpe_cached[:, None, :], (kv_seq_len, num_heads, qk_rope_head_dim) + ) # [T, N, rope] + + # ====================================================================== + # Indexer: compute sparse mask + # ====================================================================== + # per_head_scores: scores for each (query_token, index_head, kv_token) + if is_generate: + # index_q_seq: [H, D], index_k_cached: [T, D] + per_head_scores = np.einsum("hd,td->ht", index_q_seq, index_k_cached) # [H, T] + # index_score: weighted sum over index heads -> [T] + index_score = np.einsum("ht,h->t", per_head_scores, index_w_seq) * index_softmax_scale + # No causal mask needed for generate (all cached positions are past) + effective_topk = min(index_topk, kv_seq_len) + topk_indices = np.argpartition(index_score, -effective_topk)[-effective_topk:] + index_mask = np.full(kv_seq_len, float("-inf")) + index_mask[topk_indices] = 0.0 # [T] + else: + # index_q_seq: [S, H, D], index_k_cached: [T, D] + per_head_scores = np.einsum("shd,td->sht", index_q_seq, index_k_cached) # [S, H, T] + index_score = ( + np.einsum("sht,sh->st", per_head_scores, index_w_seq) * index_softmax_scale + ) # [S, T] + # Apply causal mask: future positions (relative to current seq) get -inf + causal_mask = np.triu(np.ones((seq_len_i, kv_seq_len)), k=kv_seq_len - seq_len_i + 1) + index_score = np.where(causal_mask, -np.inf, index_score) + # Build sparse mask: -inf everywhere except top-k positions per query token + effective_topk = min(index_topk, kv_seq_len) + index_mask = np.full_like(index_score, float("-inf")) + for s in range(seq_len_i): + valid = np.where(index_score[s] > float("-inf"))[0] + if len(valid) > 0: + scores_valid = index_score[s, valid] + k = min(effective_topk, len(valid)) + topk_local = np.argpartition(scores_valid, -k)[-k:] + index_mask[s, valid[topk_local]] = 0.0 + + # ====================================================================== + # MLA: compute attention scores with sparse mask + # ====================================================================== + if is_generate: + # query_full: [N, qk_head_dim] + query_full = np.concatenate([q_nope_seq, q_pe_seq], axis=-1) + # key_full: [T, N, qk_head_dim] + key_full = np.concatenate([k_nope, kpe_expanded], axis=-1) + attn_scores = np.einsum("nh,tnh->nt", query_full, key_full) * scale # [N, T] + # Add index_mask (broadcast [T] -> [N, T]) + attn_scores = attn_scores + index_mask[None, :] + else: + # query_full: [S, N, qk_head_dim] + query_full = np.concatenate([q_nope_seq, q_pe_seq], axis=-1) + # key_full: [T, N, qk_head_dim] + key_full = np.concatenate([k_nope, kpe_expanded], axis=-1) + attn_scores = np.einsum("snh,tnh->snt", query_full, key_full) * scale # [S, N, T] + # Causal mask for MLA + causal_mask_mla = np.triu( + np.ones((seq_len_i, kv_seq_len)), k=kv_seq_len - seq_len_i + 1 + ) + attn_scores = np.where(causal_mask_mla[:, None, :], -np.inf, attn_scores) + # Add index_mask: [S, T] -> [S, 1, T] -> broadcast over N + attn_scores = attn_scores + index_mask[:, None, :] + + # Softmax + output + attn_scores_max = np.max(attn_scores, axis=-1, keepdims=True) + attn_scores_exp = np.exp(attn_scores - attn_scores_max) + attn_weights = attn_scores_exp / np.sum(attn_scores_exp, axis=-1, keepdims=True) + + if is_generate: + attn_out = np.einsum("nt,tnh->nh", attn_weights, v) # [N, v] + else: + attn_out = np.einsum("snt,tnh->snh", attn_weights, v) # [S, N, v] + + outputs.append(attn_out) + + if len(outputs) == 0: + return np.zeros((1, 0, num_heads, v_head_dim), dtype=np.float32) + elif is_generate: + result = np.stack(outputs, axis=0) + return result[:, None, :, :] # [B, 1, N, v] + else: + result = np.concatenate(outputs, axis=0) + return result[None, :, :, :] # [1, B*S, N, v] + + +class TestTorchDSASourceOp: + """Test torch_dsa source op (without cache).""" + + @pytest.fixture(autouse=True) + def setup_method(self): + """Setup test configuration.""" + self.device = "cuda" + self.dtype = torch.bfloat16 + self.atol = 1e-2 + self.rtol = 1e-2 + + torch.cuda.empty_cache() + torch.manual_seed(42) + np.random.seed(42) + + def _create_dsa_data( + self, + batch_size: int, + seq_len: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + kv_lora_rank: int, + v_head_dim: int, + index_n_heads: int, + index_head_dim: int, + ): + """Create test data for DSA source op (bsnd layout).""" + kv_head_dim = qk_nope_head_dim + v_head_dim + + q_nope = torch.randn( + batch_size, seq_len, num_heads, qk_nope_head_dim, dtype=self.dtype, device=self.device + ) + q_pe = torch.randn( + batch_size, seq_len, num_heads, qk_rope_head_dim, dtype=self.dtype, device=self.device + ) + compressed_kv = torch.randn( + batch_size, seq_len, kv_lora_rank, dtype=self.dtype, device=self.device + ) + kpe = torch.randn( + batch_size, seq_len, 1, qk_rope_head_dim, dtype=self.dtype, device=self.device + ) + kv_b_proj_weight = torch.randn( + num_heads * kv_head_dim, kv_lora_rank, dtype=self.dtype, device=self.device + ) + index_q = torch.randn( + batch_size, seq_len, index_n_heads, index_head_dim, dtype=self.dtype, device=self.device + ) + index_k = torch.randn( + batch_size, seq_len, index_head_dim, dtype=self.dtype, device=self.device + ) + index_weights = torch.randn( + batch_size, seq_len, index_n_heads, dtype=self.dtype, device=self.device + ) + + return { + "q_nope": q_nope, + "q_pe": q_pe, + "compressed_kv": compressed_kv, + "kpe": kpe, + "kv_b_proj_weight": kv_b_proj_weight, + "index_q": index_q, + "index_k": index_k, + "index_weights": index_weights, + } + + def test_basic_functionality(self): + """Test basic DSA source op functionality: shape and finiteness.""" + B, S, N = 1, 4, 4 + nope, rope, kv_lora, v = 32, 16, 128, 32 + idx_H, idx_D, topk = 2, 20, 2 + + data = self._create_dsa_data(B, S, N, nope, rope, kv_lora, v, idx_H, idx_D) + + output = torch.ops.auto_deploy.torch_dsa( + data["q_nope"], + data["q_pe"], + data["compressed_kv"], + data["kpe"], + data["kv_b_proj_weight"], + data["index_q"], + data["index_k"], + data["index_weights"], + topk, + True, + None, + "bsnd", + ) + + expected_shape = (B, S, N, v) + assert output.shape == expected_shape, f"Expected {expected_shape}, got {output.shape}" + assert torch.isfinite(output).all(), "Output contains NaN or Inf values" + + def test_sparse_masking_changes_output(self): + """Test that sparse masking (topk < S) produces different output than MLA (full attend).""" + B, S, N = 1, 4, 4 + nope, rope, kv_lora, v = 32, 16, 128, 32 + idx_H, idx_D = 2, 20 + + data = self._create_dsa_data(B, S, N, nope, rope, kv_lora, v, idx_H, idx_D) + + # DSA with topk < S: mask is applied + output_dsa = torch.ops.auto_deploy.torch_dsa( + data["q_nope"], + data["q_pe"], + data["compressed_kv"], + data["kpe"], + data["kv_b_proj_weight"], + data["index_q"], + data["index_k"], + data["index_weights"], + 2, # topk=2 < S=4 → mask is non-trivial + True, + None, + "bsnd", + ) + + # MLA-only (no sparse mask): use torch_mla + output_mla = torch.ops.auto_deploy.torch_mla( + data["q_nope"], + data["q_pe"], + data["compressed_kv"], + data["kpe"], + data["kv_b_proj_weight"], + True, + None, + "bsnd", + ) + + # DSA output should differ from MLA output when topk < S + assert not torch.allclose(output_dsa, output_mla, atol=1e-3), ( + "DSA with topk < S should produce different output than full MLA" + ) + + def test_topk_full_equals_dense(self): + """Test that DSA with index_topk=S (attend all) ≈ MLA (mask all zeros). + + When every position is selected, the sparse mask is all zeros (no masking), + so DSA should produce the same result as plain MLA. + """ + B, S, N = 1, 4, 4 + nope, rope, kv_lora, v = 32, 16, 128, 32 + idx_H, idx_D = 2, 20 + + # Use float32 for this equality check + self.dtype = torch.float32 + data = self._create_dsa_data(B, S, N, nope, rope, kv_lora, v, idx_H, idx_D) + + output_dsa = torch.ops.auto_deploy.torch_dsa( + data["q_nope"], + data["q_pe"], + data["compressed_kv"], + data["kpe"], + data["kv_b_proj_weight"], + data["index_q"], + data["index_k"], + data["index_weights"], + S, # topk == S → select all → mask is all zeros + True, + None, + "bsnd", + ) + + output_mla = torch.ops.auto_deploy.torch_mla( + data["q_nope"], + data["q_pe"], + data["compressed_kv"], + data["kpe"], + data["kv_b_proj_weight"], + True, + None, + "bsnd", + ) + + max_diff = (output_dsa - output_mla).abs().max().item() + assert torch.allclose(output_dsa, output_mla, atol=1e-5), ( + f"DSA with topk=S should equal MLA. Max diff: {max_diff:.2e}" + ) + + def test_custom_scale(self): + """Test that custom scale changes DSA output.""" + B, S, N = 1, 4, 4 + nope, rope, kv_lora, v = 32, 16, 128, 32 + idx_H, idx_D, topk = 2, 20, 2 + + data = self._create_dsa_data(B, S, N, nope, rope, kv_lora, v, idx_H, idx_D) + + output_default = torch.ops.auto_deploy.torch_dsa( + data["q_nope"], + data["q_pe"], + data["compressed_kv"], + data["kpe"], + data["kv_b_proj_weight"], + data["index_q"], + data["index_k"], + data["index_weights"], + topk, + True, + None, + "bsnd", + ) + + output_custom = torch.ops.auto_deploy.torch_dsa( + data["q_nope"], + data["q_pe"], + data["compressed_kv"], + data["kpe"], + data["kv_b_proj_weight"], + data["index_q"], + data["index_k"], + data["index_weights"], + topk, + True, + 0.5, + "bsnd", + ) + + assert not torch.allclose(output_default, output_custom, atol=1e-3), ( + "Custom scale should affect output" + ) + + +class TestTorchBackendDSAWithCache: + """Test torch_cached_dsa_with_cache cached op.""" + + @pytest.fixture(autouse=True) + def setup_method(self): + """Setup test configuration.""" + self.device = "cuda" + self.dtype = torch.bfloat16 + self.atol = 5e-2 + self.rtol = 5e-2 + + torch.cuda.empty_cache() + torch.manual_seed(42) + np.random.seed(42) + + def _create_cached_dsa_data( + self, + batch_size: int, + seq_len: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + kv_lora_rank: int, + v_head_dim: int, + index_n_heads: int, + index_head_dim: int, + max_seq_len: int, + cache_offset: int = 0, + ): + """Create test data for cached DSA op, mirroring _create_cached_mla_data.""" + kv_head_dim = qk_nope_head_dim + v_head_dim + + # Create input tensors (BSND layout) + q_nope = torch.randn( + batch_size, seq_len, num_heads, qk_nope_head_dim, dtype=self.dtype, device=self.device + ) + q_pe = torch.randn( + batch_size, seq_len, num_heads, qk_rope_head_dim, dtype=self.dtype, device=self.device + ) + compressed_kv = torch.randn( + batch_size, seq_len, kv_lora_rank, dtype=self.dtype, device=self.device + ) + kpe = torch.randn( + batch_size, seq_len, 1, qk_rope_head_dim, dtype=self.dtype, device=self.device + ) + kv_b_proj_weight = torch.randn( + num_heads * kv_head_dim, kv_lora_rank, dtype=self.dtype, device=self.device + ) + index_q = torch.randn( + batch_size, seq_len, index_n_heads, index_head_dim, dtype=self.dtype, device=self.device + ) + index_k = torch.randn( + batch_size, seq_len, index_head_dim, dtype=self.dtype, device=self.device + ) + index_weights = torch.randn( + batch_size, seq_len, index_n_heads, dtype=self.dtype, device=self.device + ) + + # MLA cache: [max_batch, max_seq, kv_lora_rank + qk_rope_head_dim] + mla_cache = torch.zeros( + batch_size, + max_seq_len, + kv_lora_rank + qk_rope_head_dim, + dtype=self.dtype, + device=self.device, + ) + + # Index key cache: [max_batch, max_seq, index_head_dim] + index_k_cache = torch.zeros( + batch_size, + max_seq_len, + index_head_dim, + dtype=self.dtype, + device=self.device, + ) + + if cache_offset > 0: + mla_cache[:, :cache_offset, :] = torch.randn( + batch_size, + cache_offset, + kv_lora_rank + qk_rope_head_dim, + dtype=self.dtype, + device=self.device, + ) + index_k_cache[:, :cache_offset, :] = torch.randn( + batch_size, + cache_offset, + index_head_dim, + dtype=self.dtype, + device=self.device, + ) + + seq_len_tensor = torch.full((batch_size,), seq_len, device=self.device, dtype=torch.int32) + input_pos = torch.full((batch_size,), cache_offset, device=self.device, dtype=torch.int32) + cache_loc = torch.arange(batch_size, device=self.device, dtype=torch.int32) + + if seq_len == 1: + # Generate phase + batch_info_host = torch.tensor( + [0, 0, batch_size], device=self.device, dtype=torch.int32 + ) + cu_seqlen = torch.arange(batch_size, device=self.device, dtype=torch.int32) + else: + # Context phase: flatten inputs + batch_info_host = torch.tensor( + [batch_size, batch_size * seq_len, 0], device=self.device, dtype=torch.int32 + ) + cu_seqlen = torch.arange( + 0, batch_size * seq_len, seq_len, device=self.device, dtype=torch.int32 + ) + q_nope = q_nope.view(1, batch_size * seq_len, num_heads, qk_nope_head_dim) + q_pe = q_pe.view(1, batch_size * seq_len, num_heads, qk_rope_head_dim) + compressed_kv = compressed_kv.view(1, batch_size * seq_len, kv_lora_rank) + kpe = kpe.view(1, batch_size * seq_len, 1, qk_rope_head_dim) + index_q = index_q.view(1, batch_size * seq_len, index_n_heads, index_head_dim) + index_k = index_k.view(1, batch_size * seq_len, index_head_dim) + index_weights = index_weights.view(1, batch_size * seq_len, index_n_heads) + + return { + "q_nope": q_nope, + "q_pe": q_pe, + "compressed_kv": compressed_kv, + "kpe": kpe, + "kv_b_proj_weight": kv_b_proj_weight, + "index_q": index_q, + "index_k": index_k, + "index_weights": index_weights, + "batch_info_host": batch_info_host, + "seq_len": seq_len_tensor, + "input_pos": input_pos, + "cache_loc": cache_loc, + "cu_seqlen": cu_seqlen, + "mla_cache": mla_cache, + "index_k_cache": index_k_cache, + "kv_lora_rank": kv_lora_rank, + "index_head_dim": index_head_dim, + } + + def _run_cached_dsa(self, data, scale=None, index_topk=64): + """Run cached DSA operation.""" + return torch.ops.auto_deploy.torch_cached_dsa_with_cache( + data["q_nope"], + data["q_pe"], + data["compressed_kv"], + data["kpe"], + data["kv_b_proj_weight"], + data["index_q"], + data["index_k"], + data["index_weights"], + data["batch_info_host"], + data["seq_len"], + data["input_pos"], + data["cache_loc"], + data["cu_seqlen"], + data["mla_cache"], + data["index_k_cache"], + scale, + data["kv_lora_rank"], + index_topk, + ) + + def test_context_phase_basic(self): + """Test context (prefill) phase: shape and finiteness.""" + B, S, N = 2, 4, 4 + nope, rope, kv_lora, v = 32, 16, 128, 32 + idx_H, idx_D, topk = 2, 20, 2 + max_seq_len = 64 + + data = self._create_cached_dsa_data( + B, S, N, nope, rope, kv_lora, v, idx_H, idx_D, max_seq_len + ) + output = self._run_cached_dsa(data, index_topk=topk) + + expected_shape = (1, B * S, N, v) + assert output.shape == expected_shape, f"Expected {expected_shape}, got {output.shape}" + assert torch.isfinite(output).all(), "Output contains NaN or Inf values" + + def test_generate_phase_basic(self): + """Test generate phase (single token): shape and finiteness.""" + B, S, N = 2, 1, 4 + nope, rope, kv_lora, v = 32, 16, 128, 32 + idx_H, idx_D, topk = 2, 20, 2 + max_seq_len = 64 + cache_offset = 5 + + data = self._create_cached_dsa_data( + B, S, N, nope, rope, kv_lora, v, idx_H, idx_D, max_seq_len, cache_offset + ) + output = self._run_cached_dsa(data, index_topk=topk) + + expected_shape = (B, S, N, v) + assert output.shape == expected_shape, f"Expected {expected_shape}, got {output.shape}" + assert torch.isfinite(output).all(), "Output contains NaN or Inf values" + + def test_source_vs_cached_equivalence_prefill(self): + """Test that source op == cached prefill op (single batch, no prior cache). + + With cache_offset=0 and batch_size=1, both ops attend to exactly the same + tokens, so their outputs should match numerically in float32. + """ + B, S, N = 1, 4, 4 + nope, rope, kv_lora, v = 32, 16, 128, 32 + idx_H, idx_D, topk = 2, 20, 4 # topk = S → full attend for exact equality + max_seq_len = 64 + + self.dtype = torch.float32 + data = self._create_cached_dsa_data( + B, S, N, nope, rope, kv_lora, v, idx_H, idx_D, max_seq_len, cache_offset=0 + ) + + # Source op (operates on [B, S] layout directly) + # Reconstruct original [B, S] views from the flattened context tensors + q_nope_src = data["q_nope"].view(B, S, N, nope) + q_pe_src = data["q_pe"].view(B, S, N, rope) + compressed_kv_src = data["compressed_kv"].view(B, S, kv_lora) + kpe_src = data["kpe"].view(B, S, 1, rope) + index_q_src = data["index_q"].view(B, S, idx_H, idx_D) + index_k_src = data["index_k"].view(B, S, idx_D) + index_weights_src = data["index_weights"].view(B, S, idx_H) + + output_src = torch.ops.auto_deploy.torch_dsa( + q_nope_src, + q_pe_src, + compressed_kv_src, + kpe_src, + data["kv_b_proj_weight"], + index_q_src, + index_k_src, + index_weights_src, + topk, + True, + None, + "bsnd", + ) # [B, S, N, v] + + output_cached = self._run_cached_dsa(data, index_topk=topk) # [1, B*S, N, v] + + # Reshape for comparison + output_src_flat = output_src.view(1, B * S, N, v) + max_diff = (output_src_flat - output_cached).abs().max().item() + assert torch.allclose(output_src_flat, output_cached, atol=1e-5), ( + f"Source op and cached prefill op differ. Max diff: {max_diff:.2e}" + ) + + def test_both_caches_updated(self): + """Test that after a forward pass, both mla_cache and index_k_cache are updated.""" + B, S, N = 1, 1, 4 + nope, rope, kv_lora, v = 32, 16, 128, 32 + idx_H, idx_D, topk = 2, 20, 2 + max_seq_len = 32 + cache_offset = 5 + + data = self._create_cached_dsa_data( + B, S, N, nope, rope, kv_lora, v, idx_H, idx_D, max_seq_len, cache_offset + ) + + # Record original cache values at target position + orig_mla = data["mla_cache"][0, cache_offset].clone() + orig_idx = data["index_k_cache"][0, cache_offset].clone() + + _ = self._run_cached_dsa(data, index_topk=topk) + + updated_mla = data["mla_cache"][0, cache_offset] + updated_idx = data["index_k_cache"][0, cache_offset] + + assert not torch.allclose(orig_mla, updated_mla, atol=1e-6), ( + "mla_cache should have been updated at the target position" + ) + assert not torch.allclose(orig_idx, updated_idx, atol=1e-6), ( + "index_k_cache should have been updated at the target position" + ) + + def test_generate_with_numpy_reference(self): + """Test generate phase against numpy_dsa_reference within tolerance.""" + # Use small dims + float32 to keep bfloat16 accumulated error below atol + B, S, N = 2, 1, 4 + nope, rope, kv_lora, v = 32, 16, 64, 32 + idx_H, idx_D, topk = 2, 20, 4 + max_seq_len = 64 + cache_offset = 3 + + self.dtype = torch.float32 + + data = self._create_cached_dsa_data( + B, S, N, nope, rope, kv_lora, v, idx_H, idx_D, max_seq_len, cache_offset + ) + + # Disable TF32 so float32 matmuls use full precision and match numpy + prev_tf32 = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = False + try: + output = self._run_cached_dsa(data, index_topk=topk) + finally: + torch.backends.cuda.matmul.allow_tf32 = prev_tf32 + + reference = numpy_dsa_reference( + data["q_nope"].cpu().float().numpy(), + data["q_pe"].cpu().float().numpy(), + data["compressed_kv"].cpu().float().numpy(), + data["kpe"].cpu().float().numpy(), + data["kv_b_proj_weight"].cpu().float().numpy(), + data["index_q"].cpu().float().numpy(), + data["index_k"].cpu().float().numpy(), + data["index_weights"].cpu().float().numpy(), + data["mla_cache"].cpu().float().numpy(), + data["index_k_cache"].cpu().float().numpy(), + data["seq_len"].cpu().numpy(), + data["input_pos"].cpu().numpy(), + data["cache_loc"].cpu().numpy(), + data["cu_seqlen"].cpu().numpy(), + scale=None, + kv_lora_rank=kv_lora, + index_topk=topk, + is_generate=True, + ) + + reference_torch = torch.from_numpy(reference).to(output.device, output.dtype) + max_diff = (output - reference_torch).abs().max().item() + # float32 inputs → tight tolerance; bfloat16 would need ~5e-2 + assert torch.allclose(output, reference_torch, atol=1e-4, rtol=1e-4), ( + f"Generate phase output doesn't match numpy reference. Max diff: {max_diff:.6f}" + ) + + def test_dtype_preservation(self): + """Test that output dtype matches input dtype for float16 and bfloat16.""" + B, S, N = 1, 1, 4 + nope, rope, kv_lora, v = 32, 16, 128, 32 + idx_H, idx_D, topk = 2, 20, 2 + max_seq_len = 32 + + for dtype in [torch.float16, torch.bfloat16]: + self.dtype = dtype + data = self._create_cached_dsa_data( + B, S, N, nope, rope, kv_lora, v, idx_H, idx_D, max_seq_len + ) + output = self._run_cached_dsa(data, index_topk=topk) + assert output.dtype == dtype, f"Expected dtype {dtype}, got {output.dtype}" + + def test_cache_shapes(self): + """Test that both caches have the expected shapes.""" + B, S, N = 2, 1, 4 + nope, rope, kv_lora, v = 32, 16, 128, 32 + idx_H, idx_D = 2, 20 + max_seq_len = 64 + + data = self._create_cached_dsa_data( + B, S, N, nope, rope, kv_lora, v, idx_H, idx_D, max_seq_len + ) + + expected_mla_shape = (B, max_seq_len, kv_lora + rope) + expected_idx_shape = (B, max_seq_len, idx_D) + + assert data["mla_cache"].shape == expected_mla_shape, ( + f"mla_cache shape {data['mla_cache'].shape} != {expected_mla_shape}" + ) + assert data["index_k_cache"].shape == expected_idx_shape, ( + f"index_k_cache shape {data['index_k_cache'].shape} != {expected_idx_shape}" + ) + + +class TestDSADescriptor: + """Test TorchBackendDSAAttention descriptor configuration.""" + + def _get_dsa_descriptor(self): + """Get DSA descriptor from registry.""" + from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import AttentionRegistry + + return AttentionRegistry.get("torch_dsa") + + def test_descriptor_registration(self): + """Test that DSA descriptor is registered under 'torch_dsa'.""" + from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import AttentionRegistry + + assert AttentionRegistry.has("torch_dsa"), "torch_dsa should be registered" + + def test_descriptor_layout(self): + """Test that DSA descriptor returns 'bsnd' layout.""" + dsa_descriptor = self._get_dsa_descriptor() + assert dsa_descriptor.get_attention_layout() == "bsnd", "DSA should use bsnd layout" + + def test_descriptor_num_qkv_args(self): + """Test that DSA descriptor expects 8 tensor args.""" + dsa_descriptor = self._get_dsa_descriptor() + assert dsa_descriptor.get_num_qkv_args() == 8, ( + "DSA should expect 8 tensor args " + "(q_nope, q_pe, compressed_kv, kpe, kv_b_proj_weight, index_q, index_k, index_weights)" + ) + + def test_descriptor_source_op(self): + """Test that DSA descriptor points to torch_dsa source op.""" + dsa_descriptor = self._get_dsa_descriptor() + source_op = dsa_descriptor.get_source_attention_op() + assert source_op == torch.ops.auto_deploy.torch_dsa, "DSA should use torch_dsa as source op" + + def test_descriptor_cached_op(self): + """Test that DSA descriptor points to torch_cached_dsa_with_cache cached op.""" + dsa_descriptor = self._get_dsa_descriptor() + cached_op = dsa_descriptor.get_cached_attention_op() + assert cached_op == torch.ops.auto_deploy.torch_cached_dsa_with_cache.default, ( + "DSA should use torch_cached_dsa_with_cache as cached op" + ) + + def test_descriptor_standard_metadata(self): + """Test that DSA descriptor returns the standard metadata arg names.""" + dsa_descriptor = self._get_dsa_descriptor() + expected_args = ["batch_info_host", "seq_len", "input_pos", "slot_idx", "cu_seqlen"] + actual_args = dsa_descriptor.get_standard_metadata_args() + assert actual_args == expected_args, ( + f"Expected standard metadata {expected_args}, got {actual_args}" + ) diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_glm_dsa_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_glm_dsa_modeling.py new file mode 100644 index 000000000000..cb7e6b7feeeb --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/models/test_glm_dsa_modeling.py @@ -0,0 +1,992 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Hierarchical equivalence tests for the GLM MoE DSA (GLM-5) custom model. + +Since glm_moe_dsa is not in transformers 4.57, this file contains inline +HF reference classes copied from: + https://github.com/huggingface/transformers/tree/main/src/transformers/models/glm_moe_dsa + +Test structure (bottom-up): + 1. Block equivalence — RMSNorm, MLP, MoE gate routing, MoE layer + 2. Layer equivalence — full decoder layer (dense and MoE variants) + 3. Full model — end-to-end logits comparison + 4. Export — torch_export_to_gm, dynamic shapes, finite output +""" + +import pytest +import torch +import torch.nn.functional as F +from _model_test_utils import assert_rmse_close +from torch import nn +from torch.export import Dim + +import tensorrt_llm._torch.auto_deploy # noqa: F401 — registers custom ops +from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm +from tensorrt_llm._torch.auto_deploy.models.custom.modeling_glm_dsa import ( + GlmDSADecoderLayer, + GlmDSAForCausalLM, + GlmDSAIndexer, + GlmDSAMLP, + GlmDSAMoE, + GlmDSARMSNorm, + GlmMoeDsaConfig, +) +from tensorrt_llm._torch.auto_deploy.utils._graph import move_to_device + +# --------------------------------------------------------------------------- +# Small test config +# --------------------------------------------------------------------------- + +_BATCH_AND_SEQ = ((1, 6), (2, 4)) + + +def _small_config(num_hidden_layers: int = 3, first_k_dense_replace: int = 1) -> GlmMoeDsaConfig: + """Return a tiny GLM-5-like config suitable for CPU tests.""" + return GlmMoeDsaConfig( + vocab_size=1000, + hidden_size=64, + intermediate_size=128, + moe_intermediate_size=32, + num_hidden_layers=num_hidden_layers, + num_attention_heads=4, + num_key_value_heads=4, + hidden_act="silu", + max_position_embeddings=512, + rms_norm_eps=1e-5, + # MLA + q_lora_rank=32, + kv_lora_rank=32, + qk_nope_head_dim=8, + qk_rope_head_dim=8, + v_head_dim=16, + # MoE + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=2, + n_group=1, + topk_group=1, + routed_scaling_factor=1.0, + norm_topk_prob=True, + first_k_dense_replace=first_k_dense_replace, + moe_layer_freq=1, + # Indexer + index_topk=4, + index_head_dim=16, + index_n_heads=2, + indexer_rope_interleave=False, # skip de-interleave for random-weight tests + # RoPE + rope_theta=10000.0, + rope_scaling=None, + rope_interleave=False, + # Other + attention_bias=False, + attention_dropout=0.0, + pad_token_id=0, + ) + + +@pytest.fixture(autouse=True) +def seed(): + torch.manual_seed(42) + + +# =========================================================================== +# Inline HF reference classes (copied from transformers glm_moe_dsa) +# These are used as numerical ground truth since the model type is not in +# transformers 4.57. Keep them minimal and identical to the upstream source. +# =========================================================================== + + +class _HFRMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return (self.weight * hidden_states).to(input_dtype) + + +class _HFMoeDsaMLP(nn.Module): + def __init__(self, hidden_size, intermediate_size, act_fn="silu"): + super().__init__() + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) + + def forward(self, x): + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class _HFNaiveMoE(nn.Module): + """Stacked expert storage (HF checkpoint format).""" + + def __init__(self, n_routed_experts, hidden_size, moe_intermediate_size): + super().__init__() + self.gate_up_proj = nn.Parameter( + torch.empty(n_routed_experts, 2 * moe_intermediate_size, hidden_size) + ) + self.down_proj = nn.Parameter( + torch.empty(n_routed_experts, hidden_size, moe_intermediate_size) + ) + + def forward(self, hidden_states, topk_indices, topk_weights): + """Reference MoE dispatch using scatter-reduce.""" + T, top_k = topk_indices.shape + output = torch.zeros_like(hidden_states) + + for t in range(T): + for k in range(top_k): + expert_idx = topk_indices[t, k].item() + w = topk_weights[t, k].item() + gate_up = self.gate_up_proj[expert_idx] # [2*mid, H] + down = self.down_proj[expert_idx] # [H, mid] + mid = gate_up.shape[0] // 2 + x = hidden_states[t : t + 1] # [1, H] + gate_out = F.silu(x @ gate_up[:mid].t()) * (x @ gate_up[mid:].t()) + output[t] += w * (gate_out @ down.t()).squeeze(0) + + return output + + +def _hf_route_tokens( + gate_weight, + e_score_correction_bias, + hidden_flat, + n_group, + topk_group, + top_k, + norm_topk_prob, + routed_scaling_factor, +): + """Vanilla PyTorch noaux_tc routing — mirrors HF route_tokens_to_experts.""" + router_logits = F.linear(hidden_flat.float(), gate_weight.float()) + scores = router_logits.sigmoid() + scores_for_choice = scores + e_score_correction_bias + + group_scores = ( + scores_for_choice.view(-1, n_group, scores_for_choice.shape[-1] // n_group) + .topk(2, dim=-1)[0] + .sum(dim=-1) + ) + group_idx = group_scores.topk(topk_group, dim=-1, sorted=False).indices + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1.0) + score_mask = ( + group_mask.unsqueeze(-1) + .expand(-1, n_group, scores_for_choice.shape[-1] // n_group) + .reshape(-1, scores_for_choice.shape[-1]) + ) + scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), 0.0) + topk_indices = scores_for_choice.topk(top_k, dim=-1, sorted=False).indices + topk_weights = scores.gather(1, topk_indices) + if norm_topk_prob: + topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-20) + topk_weights = topk_weights * routed_scaling_factor + return topk_indices, topk_weights + + +# --------------------------------------------------------------------------- +# Weight conversion helpers (HF stacked → per-expert for loading into custom) +# --------------------------------------------------------------------------- + + +def _stacked_to_per_expert_state_dict(full_sd: dict, config) -> dict: + """Convert HF-style stacked expert weights to our per-expert ModuleList format.""" + out = {} + n = config.n_routed_experts + mid = config.moe_intermediate_size + + for k, v in full_sd.items(): + if ".mlp.experts.gate_up_proj" in k: + prefix = k[: k.index(".experts.gate_up_proj") + len(".experts.")] + for i in range(n): + out[f"{prefix}{i}.gate_proj.weight"] = v[i, :mid] + out[f"{prefix}{i}.up_proj.weight"] = v[i, mid:] + elif ".mlp.experts.down_proj" in k: + prefix = k[: k.index(".experts.down_proj") + len(".experts.")] + for i in range(n): + out[f"{prefix}{i}.down_proj.weight"] = v[i] + else: + out[k] = v + return out + + +# --------------------------------------------------------------------------- +# Inline reference attention, decoder layer, and full model classes. +# These replicate HF GlmMoeDsaAttention math WITHOUT using torch_dsa so +# they serve as independent ground-truth for equivalence tests. +# --------------------------------------------------------------------------- + + +class _HFGlmDsaAttention(nn.Module): + """Reference MLA+DSA attention without torch_dsa custom op. + + Implements the full DSA computation using vanilla PyTorch ops + (matching the math in torch_dsa.py) and is used as ground truth + for test_attention_block_equivalence. + """ + + def __init__(self, config): + super().__init__() + H = config.hidden_size + N = config.num_attention_heads + nope = config.qk_nope_head_dim + rope_dim = config.qk_rope_head_dim + v = config.v_head_dim + kv_lora = config.kv_lora_rank + q_lora = config.q_lora_rank + idx_heads = config.index_n_heads + idx_dim = config.index_head_dim + + self.config = config + + self.q_a_proj = nn.Linear(H, q_lora, bias=False) + self.q_a_layernorm = _HFRMSNorm(q_lora, eps=config.rms_norm_eps) + self.q_b_proj = nn.Linear(q_lora, N * (nope + rope_dim), bias=False) + self.kv_a_proj_with_mqa = nn.Linear(H, kv_lora + rope_dim, bias=False) + self.kv_a_layernorm = _HFRMSNorm(kv_lora, eps=config.rms_norm_eps) + self.kv_b_proj = nn.Linear(kv_lora, N * (nope + v), bias=False) + self.o_proj = nn.Linear(N * v, H, bias=False) + + # Indexer (flat, not nested under .indexer) + self.wq_b = nn.Linear(q_lora, idx_heads * idx_dim, bias=False) + self.wk = nn.Linear(H, idx_dim, bias=False) + self.k_norm = nn.LayerNorm(idx_dim, eps=1e-6) + self.weights_proj = nn.Linear(H, idx_heads, bias=False) + + self.softmax_scale = (nope + rope_dim) ** (-0.5) + + @staticmethod + def _rope_bsnd(x, cos_bsd, sin_bsd): + """NeoX RoPE on [B, S, N, D]; cos/sin are [B, S, D] (full rope_dim). + + Uses the rotate-half formula: out = x * cos + rotate_half(x) * sin + where rotate_half(x) = cat(-x[..., D//2:], x[..., :D//2]). + """ + cos = cos_bsd.unsqueeze(2) # [B, S, 1, D] + sin = sin_bsd.unsqueeze(2) + half = x.shape[-1] // 2 + x_rot = torch.cat([-x[..., half:], x[..., :half]], dim=-1) + return x * cos + x_rot * sin + + @staticmethod + def _rope_bsd(x, cos_bsd, sin_bsd): + """NeoX RoPE on [B, S, D]; cos/sin are [B, S, D] (full rope_dim).""" + half = x.shape[-1] // 2 + x_rot = torch.cat([-x[..., half:], x[..., :half]], dim=-1) + return x * cos_bsd + x_rot * sin_bsd + + def _cos_sin(self, seq_len, device, dtype): + """Vanilla RoPE table (no YaRN).""" + rope_dim = self.config.qk_rope_head_dim + inv_freq = 1.0 / ( + self.config.rope_theta + ** (torch.arange(0, rope_dim, 2, dtype=torch.float32, device=device) / rope_dim) + ) + t = torch.arange(seq_len, dtype=torch.float32, device=device) + freqs = torch.outer(t, inv_freq) + emb = torch.cat([freqs, freqs], dim=-1) + return emb.cos().to(dtype), emb.sin().to(dtype) # [S, rope_dim] + + def forward(self, hidden_states, position_ids): + bsz, q_len = hidden_states.shape[:2] + device = hidden_states.device + dtype = hidden_states.dtype + N = self.config.num_attention_heads + nope = self.config.qk_nope_head_dim + rope_dim = self.config.qk_rope_head_dim + v = self.config.v_head_dim + kv_lora = self.config.kv_lora_rank + idx_heads = self.config.index_n_heads + idx_dim = self.config.index_head_dim + + # Q path + qr = self.q_a_layernorm(self.q_a_proj(hidden_states)) + q = self.q_b_proj(qr).view(bsz, q_len, N, nope + rope_dim) + q_nope = q[..., :nope] + q_pe = q[..., nope:] + + # KV path + kv_a = self.kv_a_proj_with_mqa(hidden_states) + compressed_kv = self.kv_a_layernorm(kv_a[..., :kv_lora]) + k_pe_raw = kv_a[..., kv_lora:] # [B, S, rope_dim] + + # RoPE + cos_full, sin_full = self._cos_sin(q_len, device, dtype) + cos = cos_full[position_ids] # [B, S, rope_dim] + sin = sin_full[position_ids] + q_pe_rot = self._rope_bsnd(q_pe, cos, sin) + k_pe_rot = self._rope_bsd(k_pe_raw, cos, sin) # [B, S, rope_dim] + + # Indexer Q — layout: [rope | nope] within each head + idx_q = self.wq_b(qr).view(bsz, q_len, idx_heads, idx_dim) + idx_q_rope, idx_q_nope = idx_q[..., :rope_dim], idx_q[..., rope_dim:] + idx_q = torch.cat([self._rope_bsnd(idx_q_rope, cos, sin), idx_q_nope], dim=-1) + + # Indexer K + idx_k_raw = self.k_norm(self.wk(hidden_states)) # [B, S, idx_dim] + idx_k_rope, idx_k_nope = idx_k_raw[..., :rope_dim], idx_k_raw[..., rope_dim:] + idx_k = torch.cat([self._rope_bsd(idx_k_rope, cos, sin), idx_k_nope], dim=-1) + + # Indexer weights (pre-scaled) + idx_weights = self.weights_proj(hidden_states) * (idx_heads**-0.5) # [B, S, idx_heads] + + # Expand KV via kv_b_proj + kv = torch.matmul(compressed_kv, self.kv_b_proj.weight.t()) # [B, S, N*(nope+v)] + kv = kv.view(bsz, q_len, N, nope + v) + k_nope_t = kv[..., :nope].transpose(1, 2) # [B, N, S, nope] + value_states = kv[..., nope:].transpose(1, 2) # [B, N, S, v] + + # Full Q, K [B, N, S, qk_head_dim] + q_full = torch.cat([q_nope.transpose(1, 2), q_pe_rot.transpose(1, 2)], dim=-1) + k_pe_expand = k_pe_rot.unsqueeze(1).expand(bsz, N, q_len, rope_dim) + k_full = torch.cat([k_nope_t, k_pe_expand], dim=-1) + + # Attention scores [B, N, S_q, S_k] + attn = torch.matmul(q_full, k_full.transpose(-2, -1)) * self.softmax_scale + + # Causal mask + causal = torch.triu(torch.ones(q_len, q_len, device=device, dtype=torch.bool), diagonal=1) + attn.masked_fill_(causal[None, None], float("-inf")) + + # DSA index mask (same math as _compute_dsa_index_mask in torch_dsa.py) + per_head = torch.einsum("bshd,btd->bsht", idx_q.float(), idx_k.float()) + idx_score = torch.einsum("bsht,bsh->bst", per_head, idx_weights.float()) + idx_score = idx_score * (idx_dim**-0.5) + idx_score.masked_fill_(causal[None], float("-inf")) + eff_topk = min(self.config.index_topk, q_len) + topk_idx = idx_score.topk(eff_topk, dim=-1).indices + idx_mask = idx_score.new_full(idx_score.shape, float("-inf")) + idx_mask.scatter_(-1, topk_idx, 0.0) + + attn = attn + idx_mask[:, None, :, :] + w = torch.softmax(attn, dim=-1, dtype=torch.float32).to(dtype) + out = torch.matmul(w, value_states) # [B, N, S, v] + out = out.transpose(1, 2).reshape(bsz, q_len, N * v) + return self.o_proj(out) + + +def _copy_attn_weights(custom_attn, ref_attn): + """Copy weights from GlmDSAAttention → _HFGlmDsaAttention.""" + for name in ["q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj", "o_proj"]: + getattr(ref_attn, name).weight = getattr(custom_attn, name).weight + ref_attn.q_a_layernorm.weight = custom_attn.q_a_layernorm.weight + ref_attn.kv_a_layernorm.weight = custom_attn.kv_a_layernorm.weight + ref_attn.wq_b.weight = custom_attn.indexer.wq_b.weight + ref_attn.wk.weight = custom_attn.indexer.wk.weight + ref_attn.k_norm.weight = custom_attn.indexer.k_norm.weight + ref_attn.k_norm.bias = custom_attn.indexer.k_norm.bias + ref_attn.weights_proj.weight = custom_attn.indexer.weights_proj.weight + + +class _HFDecoderLayer(nn.Module): + """Reference decoder layer using _HFGlmDsaAttention + HF MLP references.""" + + def __init__(self, config, is_moe=False): + super().__init__() + H = config.hidden_size + self.is_moe = is_moe + self.config = config + + self.input_layernorm = _HFRMSNorm(H, eps=config.rms_norm_eps) + self.self_attn = _HFGlmDsaAttention(config) + self.post_attention_layernorm = _HFRMSNorm(H, eps=config.rms_norm_eps) + + if is_moe: + self.gate_weight = nn.Parameter(torch.empty(config.n_routed_experts, H)) + self.e_score_correction_bias = nn.Parameter(torch.zeros(config.n_routed_experts)) + self.routed_experts = _HFNaiveMoE( + config.n_routed_experts, H, config.moe_intermediate_size + ) + self.shared_expert = _HFMoeDsaMLP( + H, config.moe_intermediate_size * config.n_shared_experts + ) + else: + self.mlp = _HFMoeDsaMLP(H, config.intermediate_size) + + def forward(self, hidden_states, position_ids): + residual = hidden_states + hidden_states = self.self_attn(self.input_layernorm(hidden_states), position_ids) + hidden_states = residual + hidden_states + + residual = hidden_states + normed = self.post_attention_layernorm(hidden_states) + + if self.is_moe: + cfg = self.config + T = normed.shape[0] * normed.shape[1] + hidden_flat = normed.view(T, -1) + topk_indices, topk_weights = _hf_route_tokens( + self.gate_weight, + self.e_score_correction_bias, + hidden_flat, + cfg.n_group, + cfg.topk_group, + cfg.num_experts_per_tok, + cfg.norm_topk_prob, + cfg.routed_scaling_factor, + ) + routed = self.routed_experts(hidden_flat.float(), topk_indices, topk_weights.float()) + routed = routed.view(*normed.shape) + shared = self.shared_expert(normed.float()).to(hidden_states.dtype) + mlp_out = (routed + shared).to(hidden_states.dtype) + else: + mlp_out = self.mlp(normed) + + return residual + mlp_out + + +def _build_hf_decoder_layer(custom_layer, config, is_moe): + """Build _HFDecoderLayer with weights copied from GlmDSADecoderLayer.""" + ref = _HFDecoderLayer(config, is_moe=is_moe) + ref.input_layernorm.weight = custom_layer.input_layernorm.weight + ref.post_attention_layernorm.weight = custom_layer.post_attention_layernorm.weight + _copy_attn_weights(custom_layer.self_attn, ref.self_attn) + + if is_moe: + cm = custom_layer.mlp + ref.gate_weight = nn.Parameter(cm.gate.weight.data.clone()) + ref.e_score_correction_bias = nn.Parameter(cm.gate.e_score_correction_bias.data.clone()) + n = config.n_routed_experts + mid = config.moe_intermediate_size + for j in range(n): + ref.routed_experts.gate_up_proj.data[j, :mid] = cm.experts[j].gate_proj.weight.data + ref.routed_experts.gate_up_proj.data[j, mid:] = cm.experts[j].up_proj.weight.data + ref.routed_experts.down_proj.data[j] = cm.experts[j].down_proj.weight.data + ref.shared_expert.gate_proj.weight = cm.shared_experts.gate_proj.weight + ref.shared_expert.up_proj.weight = cm.shared_experts.up_proj.weight + ref.shared_expert.down_proj.weight = cm.shared_experts.down_proj.weight + else: + cm = custom_layer.mlp + ref.mlp.gate_proj.weight = cm.gate_proj.weight + ref.mlp.up_proj.weight = cm.up_proj.weight + ref.mlp.down_proj.weight = cm.down_proj.weight + + return ref + + +class _HFGlmDsaForCausalLM(nn.Module): + """Reference full model using all HF reference components.""" + + def __init__(self, config): + super().__init__() + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [ + _HFDecoderLayer( + config, + is_moe=( + config.n_routed_experts is not None + and i >= config.first_k_dense_replace + and i % config.moe_layer_freq == 0 + ), + ) + for i in range(config.num_hidden_layers) + ] + ) + self.norm = _HFRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + def forward(self, input_ids, position_ids): + hidden_states = self.embed_tokens(input_ids) + for layer in self.layers: + hidden_states = layer(hidden_states, position_ids) + return self.lm_head(self.norm(hidden_states)) + + +def _build_hf_model(custom_model, config): + """Build _HFGlmDsaForCausalLM with weights copied from GlmDSAForCausalLM.""" + ref = _HFGlmDsaForCausalLM(config) + ref.embed_tokens.weight = custom_model.model.embed_tokens.weight + ref.norm.weight = custom_model.model.norm.weight + ref.lm_head.weight = custom_model.lm_head.weight + + for i, (cust_layer, ref_layer) in enumerate(zip(custom_model.model.layers, ref.layers)): + ref_layer.input_layernorm.weight = cust_layer.input_layernorm.weight + ref_layer.post_attention_layernorm.weight = cust_layer.post_attention_layernorm.weight + _copy_attn_weights(cust_layer.self_attn, ref_layer.self_attn) + + if ref_layer.is_moe: + cm = cust_layer.mlp + ref_layer.gate_weight = nn.Parameter(cm.gate.weight.data.clone()) + ref_layer.e_score_correction_bias = nn.Parameter( + cm.gate.e_score_correction_bias.data.clone() + ) + n = config.n_routed_experts + mid = config.moe_intermediate_size + for j in range(n): + ref_layer.routed_experts.gate_up_proj.data[j, :mid] = cm.experts[ + j + ].gate_proj.weight.data + ref_layer.routed_experts.gate_up_proj.data[j, mid:] = cm.experts[ + j + ].up_proj.weight.data + ref_layer.routed_experts.down_proj.data[j] = cm.experts[j].down_proj.weight.data + ref_layer.shared_expert.gate_proj.weight = cm.shared_experts.gate_proj.weight + ref_layer.shared_expert.up_proj.weight = cm.shared_experts.up_proj.weight + ref_layer.shared_expert.down_proj.weight = cm.shared_experts.down_proj.weight + else: + cm = cust_layer.mlp + ref_layer.mlp.gate_proj.weight = cm.gate_proj.weight + ref_layer.mlp.up_proj.weight = cm.up_proj.weight + ref_layer.mlp.down_proj.weight = cm.down_proj.weight + + return ref + + +# =========================================================================== +# 1. Block-level equivalence tests +# =========================================================================== + + +@torch.no_grad() +def test_rmsnorm_equivalence(): + """Custom RMSNorm matches HF reference (identical math).""" + config = _small_config() + custom = GlmDSARMSNorm(config.hidden_size, eps=config.rms_norm_eps) + ref = _HFRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + ref.weight = custom.weight # share weights + + x = torch.randn(2, 6, config.hidden_size) + torch.testing.assert_close(custom(x), ref(x), rtol=1e-3, atol=1e-3) + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQ) +@torch.no_grad() +def test_mlp_equivalence(B, S): + """Custom MLP matches HF MLP (identical math).""" + config = _small_config() + custom = GlmDSAMLP(config) + ref = _HFMoeDsaMLP(config.hidden_size, config.intermediate_size) + ref.gate_proj.weight = custom.gate_proj.weight + ref.up_proj.weight = custom.up_proj.weight + ref.down_proj.weight = custom.down_proj.weight + + x = torch.randn(B, S, config.hidden_size) + torch.testing.assert_close(custom(x), ref(x), rtol=1e-3, atol=1e-3) + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQ) +@torch.no_grad() +def test_moe_gate_routing_equivalence(B, S): + """Custom MoE gate routing matches HF reference routing (noaux_tc).""" + config = _small_config() + custom_moe = GlmDSAMoE(config) + custom_moe.gate.weight = nn.Parameter(torch.randn_like(custom_moe.gate.weight)) + + x = torch.randn(B, S, config.hidden_size) + T = B * S + hidden_flat = x.view(T, -1) + + # Custom model routing + custom_indices, custom_weights = custom_moe.gate(x) + + # HF reference routing + ref_indices, ref_weights = _hf_route_tokens( + custom_moe.gate.weight, + custom_moe.gate.e_score_correction_bias, + hidden_flat, + config.n_group, + config.topk_group, + config.num_experts_per_tok, + config.norm_topk_prob, + config.routed_scaling_factor, + ) + + # Indices might differ in order (topk sorted=False), but selected expert sets should match per token + custom_sorted = custom_indices.sort(dim=-1).values + ref_sorted = ref_indices.sort(dim=-1).values + torch.testing.assert_close(custom_sorted, ref_sorted) + # Weights (after sorting by index) should match + custom_w_sorted = custom_weights.gather(1, custom_indices.argsort(dim=-1)) + ref_w_sorted = ref_weights.gather(1, ref_indices.argsort(dim=-1)) + torch.testing.assert_close(custom_w_sorted.float(), ref_w_sorted.float(), rtol=1e-4, atol=1e-4) + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQ) +@torch.no_grad() +def test_moe_layer_equivalence(B, S): + """Custom MoE output matches reference MoE using same weights and routing.""" + config = _small_config() + custom_moe = GlmDSAMoE(config) + custom_moe.gate.weight = nn.Parameter(torch.randn_like(custom_moe.gate.weight)) + + # Build HF-style reference: use same weights, route identically, compute via loop + hf_naive = _HFNaiveMoE( + config.n_routed_experts, config.hidden_size, config.moe_intermediate_size + ) + for i in range(config.n_routed_experts): + hf_naive.gate_up_proj.data[i, : config.moe_intermediate_size] = custom_moe.experts[ + i + ].gate_proj.weight.data + hf_naive.gate_up_proj.data[i, config.moe_intermediate_size :] = custom_moe.experts[ + i + ].up_proj.weight.data + hf_naive.down_proj.data[i] = custom_moe.experts[i].down_proj.weight.data + + hf_shared_gate = custom_moe.shared_experts.gate_proj.weight.data.clone() + hf_shared_up = custom_moe.shared_experts.up_proj.weight.data.clone() + hf_shared_down = custom_moe.shared_experts.down_proj.weight.data.clone() + + x = torch.randn(B, S, config.hidden_size) + T = B * S + + # Custom forward + custom_out = custom_moe(x) + + # Reference forward + topk_indices, topk_weights = _hf_route_tokens( + custom_moe.gate.weight, + custom_moe.gate.e_score_correction_bias, + x.view(T, -1), + config.n_group, + config.topk_group, + config.num_experts_per_tok, + config.norm_topk_prob, + config.routed_scaling_factor, + ) + ref_routed = hf_naive(x.view(T, -1).float(), topk_indices, topk_weights.float()) + ref_routed = ref_routed.view(B, S, config.hidden_size) + + # Shared expert + x_shared = x.float() + shared_out = F.silu(x_shared @ hf_shared_gate.float().t()) * ( + x_shared @ hf_shared_up.float().t() + ) + shared_out = shared_out @ hf_shared_down.float().t() + + ref_out = (ref_routed + shared_out).to(x.dtype) + + assert_rmse_close(custom_out.float(), ref_out.float(), rmse_ratio_tol=0.02, msg="MoE layer: ") + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQ) +@torch.no_grad() +def test_attention_block_equivalence(B, S): + """GlmDSAAttention (uses torch_dsa) matches _HFGlmDsaAttention (vanilla torch). + + This test verifies that the projection/RoPE/indexer wiring in GlmDSAAttention + correctly feeds into the torch_dsa custom op by comparing against an independent + reference that implements DSA using raw PyTorch ops. + """ + config = _small_config() + custom_attn = GlmDSADecoderLayer(config, layer_idx=0).self_attn + custom_attn.eval() + + ref_attn = _HFGlmDsaAttention(config) + _copy_attn_weights(custom_attn, ref_attn) + ref_attn.eval() + + x = torch.randn(B, S, config.hidden_size) + pos = torch.arange(S).unsqueeze(0).expand(B, -1) + + custom_out = custom_attn(x, pos) + ref_out = ref_attn(x, pos) + + assert custom_out.shape == ref_out.shape + assert_rmse_close( + custom_out.float(), ref_out.float(), rmse_ratio_tol=0.10, msg="Attention block: " + ) + + +# =========================================================================== +# 2. Layer-level equivalence tests +# =========================================================================== + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQ) +@torch.no_grad() +def test_dense_decoder_layer_equivalence(B, S): + """Dense decoder layer matches HF reference decoder layer with same weights (CPU).""" + config = _small_config(num_hidden_layers=3, first_k_dense_replace=1) + + layer = GlmDSADecoderLayer(config, layer_idx=0) # dense layer + assert isinstance(layer.mlp, GlmDSAMLP), "Layer 0 should be dense" + layer.eval() + + ref = _build_hf_decoder_layer(layer, config, is_moe=False) + ref.eval() + + x = torch.randn(B, S, config.hidden_size) + pos = torch.arange(S).unsqueeze(0).expand(B, -1) + + custom_out = layer(x, pos) + ref_out = ref(x, pos) + + assert custom_out.shape == x.shape + assert torch.isfinite(custom_out).all() + assert_rmse_close( + custom_out.float(), ref_out.float(), rmse_ratio_tol=0.05, msg="Dense decoder layer: " + ) + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQ) +@torch.no_grad() +def test_moe_decoder_layer_equivalence(B, S): + """MoE decoder layer matches HF reference decoder layer with same weights (CPU).""" + config = _small_config(num_hidden_layers=3, first_k_dense_replace=1) + + layer = GlmDSADecoderLayer(config, layer_idx=1) # MoE layer + assert isinstance(layer.mlp, GlmDSAMoE), "Layer 1 should be MoE" + layer.eval() + + ref = _build_hf_decoder_layer(layer, config, is_moe=True) + ref.eval() + + x = torch.randn(B, S, config.hidden_size) + pos = torch.arange(S).unsqueeze(0).expand(B, -1) + + custom_out = layer(x, pos) + ref_out = ref(x, pos) + + assert custom_out.shape == x.shape + assert torch.isfinite(custom_out).all() + assert_rmse_close( + custom_out.float(), ref_out.float(), rmse_ratio_tol=0.05, msg="MoE decoder layer: " + ) + + +# =========================================================================== +# 3. Full-model equivalence test +# =========================================================================== + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQ) +@torch.no_grad() +def test_full_model_equivalence(B, S): + """Full model logits match HF reference model with same weights (CPU).""" + config = _small_config(num_hidden_layers=3, first_k_dense_replace=1) + + model = GlmDSAForCausalLM(config) + model.eval() + + ref = _build_hf_model(model, config) + ref.eval() + + input_ids = torch.randint(0, config.vocab_size, (B, S)) + pos = torch.arange(S).unsqueeze(0).expand(B, -1) + + custom_logits = model(input_ids=input_ids, position_ids=pos).logits + ref_logits = ref(input_ids, pos) + + assert custom_logits.shape == (B, S, config.vocab_size) + assert torch.isfinite(custom_logits).all(), "Logits contain NaN or Inf" + assert_rmse_close( + custom_logits.float(), ref_logits.float(), rmse_ratio_tol=0.05, msg="Full model: " + ) + + +@torch.no_grad() +def test_full_model_self_consistency(): + """Two forward passes with the same input produce identical outputs (CPU).""" + config = _small_config(num_hidden_layers=3, first_k_dense_replace=1) + + model = GlmDSAForCausalLM(config) + model.eval() + + B, S = 2, 6 + input_ids = torch.randint(0, config.vocab_size, (B, S)) + pos = torch.arange(S).unsqueeze(0).expand(B, -1) + + out1 = model(input_ids=input_ids, position_ids=pos) + out2 = model(input_ids=input_ids, position_ids=pos) + torch.testing.assert_close(out1.logits, out2.logits) + + +# =========================================================================== +# 4. Structural / config tests +# =========================================================================== + + +def test_config_registration(): + """Config model_type is correct and has expected attributes.""" + config = _small_config() + assert config.model_type == "glm_moe_dsa" + assert hasattr(config, "index_topk") + assert hasattr(config, "index_n_heads") + assert hasattr(config, "index_head_dim") + assert hasattr(config, "kv_lora_rank") + assert hasattr(config, "qk_rope_head_dim") + + +def test_config_rope_theta_from_rope_parameters(): + """rope_theta is correctly extracted from rope_parameters dict.""" + config = GlmMoeDsaConfig(rope_parameters={"rope_theta": 500000.0, "rope_type": "default"}) + assert config.rope_theta == 500000.0 + + +def test_layer_types(): + """Layer 0..first_k_dense_replace-1 are dense, rest are MoE.""" + config = _small_config(num_hidden_layers=4, first_k_dense_replace=2) + model = GlmDSAForCausalLM(config) + + for i in range(2): + assert isinstance(model.model.layers[i].mlp, GlmDSAMLP), f"Layer {i} should be dense" + for i in range(2, 4): + assert isinstance(model.model.layers[i].mlp, GlmDSAMoE), f"Layer {i} should be MoE" + + +def test_expert_structure(): + """MoE expert list has correct structure for checkpoint loading.""" + config = _small_config() + moe = GlmDSAMoE(config) + + assert isinstance(moe.experts, nn.ModuleList) + assert len(moe.experts) == config.n_routed_experts + + sd = moe.state_dict() + for i in range(config.n_routed_experts): + assert f"experts.{i}.gate_proj.weight" in sd + assert f"experts.{i}.up_proj.weight" in sd + assert f"experts.{i}.down_proj.weight" in sd + + +def test_indexer_structure(): + """GlmDSAIndexer has correct submodule names for checkpoint key matching.""" + config = _small_config() + indexer = GlmDSAIndexer(config) + + assert hasattr(indexer, "wq_b") + assert hasattr(indexer, "wk") + assert hasattr(indexer, "k_norm") + assert hasattr(indexer, "weights_proj") + assert isinstance(indexer.k_norm, nn.LayerNorm) + + +def test_attention_indexer_submodule(): + """GlmDSAAttention exposes indexer as a submodule (for checkpoint key: self_attn.indexer.*).""" + config = _small_config() + model = GlmDSAForCausalLM(config) + attn = model.model.layers[0].self_attn + + assert hasattr(attn, "indexer"), "indexer must be a submodule of self_attn" + assert isinstance(attn.indexer, GlmDSAIndexer) + + +def test_moe_weight_expand_hook(): + """_moe_expert_expand_hook correctly expands stacked expert weights at load time.""" + config = _small_config(num_hidden_layers=2, first_k_dense_replace=1) + model = GlmDSAForCausalLM(config) + + n = config.n_routed_experts + mid = config.moe_intermediate_size + H = config.hidden_size + + # Build a fake stacked-format state_dict (layer 1 is MoE) + original_sd = model.state_dict() + stacked_sd = {} + for k, v in original_sd.items(): + if ".mlp.experts." in k and "gate_proj" in k: + # We'll replace per-expert keys with stacked ones + pass + else: + stacked_sd[k] = v + + # Stack the expert weights back (simulate HF checkpoint format) + gate_up = torch.zeros(n, 2 * mid, H) + down = torch.zeros(n, H, mid) + for i in range(n): + gate_up[i, :mid] = original_sd[f"model.layers.1.mlp.experts.{i}.gate_proj.weight"] + gate_up[i, mid:] = original_sd[f"model.layers.1.mlp.experts.{i}.up_proj.weight"] + down[i] = original_sd[f"model.layers.1.mlp.experts.{i}.down_proj.weight"] + + stacked_sd["model.layers.1.mlp.experts.gate_up_proj"] = gate_up + stacked_sd["model.layers.1.mlp.experts.down_proj"] = down + + # Load — the hook should expand stacked → per-expert + model2 = GlmDSAForCausalLM(config) + model2.load_state_dict(stacked_sd) + + # Verify weights match + for i in range(n): + torch.testing.assert_close( + model2.model.layers[1].mlp.experts[i].gate_proj.weight, + original_sd[f"model.layers.1.mlp.experts.{i}.gate_proj.weight"], + ) + + +# =========================================================================== +# 5. Export test +# =========================================================================== + + +def test_model_export(): + """Model exports with torch_export_to_gm, produces finite output on two shapes.""" + device = "cuda" + dtype = torch.bfloat16 + config = _small_config(num_hidden_layers=2, first_k_dense_replace=1) + + model = GlmDSAForCausalLM(config) + model.to(device=device, dtype=dtype) + model.eval() + + B, S = 2, 8 + input_ids = torch.randint(0, config.vocab_size, (B, S), device=device) + pos = torch.arange(S, device=device).unsqueeze(0).expand(B, -1) + + dynamic_shapes = ( + {0: Dim.DYNAMIC, 1: Dim.DYNAMIC}, + {0: Dim.DYNAMIC, 1: Dim.DYNAMIC}, + ) + gm = torch_export_to_gm( + model, + args=(), + kwargs={"input_ids": input_ids, "position_ids": pos}, + dynamic_shapes=dynamic_shapes, + ) + move_to_device(gm, device) + + with torch.inference_mode(): + out = gm(input_ids=input_ids, position_ids=pos) + + assert "logits" in out + assert out["logits"].shape == (B, S, config.vocab_size) + assert torch.isfinite(out["logits"]).all() + + # Numerical equivalence between eager and exported graph + with torch.no_grad(): + eager_out = model(input_ids=input_ids, position_ids=pos) + assert_rmse_close( + out["logits"].float(), + eager_out.logits.float(), + rmse_ratio_tol=0.05, + msg="Export vs eager (shape 1): ", + ) + + # Test second shape + B2, S2 = 1, 4 + input_ids2 = torch.randint(0, config.vocab_size, (B2, S2), device=device) + pos2 = torch.arange(S2, device=device).unsqueeze(0).expand(B2, -1) + + with torch.inference_mode(): + out2 = gm(input_ids=input_ids2, position_ids=pos2) + + assert out2["logits"].shape == (B2, S2, config.vocab_size) + assert torch.isfinite(out2["logits"]).all() + + with torch.no_grad(): + eager_out2 = model(input_ids=input_ids2, position_ids=pos2) + assert_rmse_close( + out2["logits"].float(), + eager_out2.logits.float(), + rmse_ratio_tol=0.05, + msg="Export vs eager (shape 2): ", + )