diff --git a/3rdparty/fetch_content.json b/3rdparty/fetch_content.json index 6a4679db5262..e535a2797eeb 100644 --- a/3rdparty/fetch_content.json +++ b/3rdparty/fetch_content.json @@ -44,8 +44,8 @@ }, { "name": "flashmla", - "git_repository": "https://github.com/deepseek-ai/FlashMLA.git", - "git_tag": "1408756a88e52a25196b759eaf8db89d2b51b5a1", + "git_repository": "https://github.com/vllm-project/FlashMLA.git", + "git_tag": "a6ec2ba7bd0a7dff98b3f4d3e6b52b159c48d78b", "git_submodules_recurse": true, "source_subdir": "dont-add-this-project-with-add-subdirectory" }, diff --git a/cpp/tensorrt_llm/flash_mla/CMakeLists.txt b/cpp/tensorrt_llm/flash_mla/CMakeLists.txt index e87f12275f60..73c0dbf86e07 100644 --- a/cpp/tensorrt_llm/flash_mla/CMakeLists.txt +++ b/cpp/tensorrt_llm/flash_mla/CMakeLists.txt @@ -106,12 +106,17 @@ foreach(SOURCE_FILE ${FLASH_MLA_ALL_FILES}) "${_content}") # Replace absolute imports with relative imports for internal modules string(REPLACE "from flash_mla." "from ." _content "${_content}") + # vLLM's FlashMLA fork imports this helper from __init__.py, but it is + # provided by vLLM's separate _flashmla_extension_C wrapper, not by the + # FlashMLA pybind interface built here. + string(REPLACE " flash_mla_with_kvcache_fp8,\n" "" _content + "${_content}") # Add adaptation header string( PREPEND _content - "# Adapted from https://github.com/deepseek-ai/FlashMLA/blob/main/flash_mla/${REL_PATH}\n" + "# Adapted from https://github.com/vllm-project/FlashMLA/blob/main/flash_mla/${REL_PATH}\n" ) # Write modified content @@ -137,22 +142,31 @@ endforeach() find_library(TORCH_PYTHON_LIB torch_python REQUIRED HINTS ${TORCH_INSTALL_PREFIX}/lib) -# Define source files matching FlashMLA's setup.py Note: pybind.cpp has runtime -# checks (arch.is_sm90()), but still references all kernel symbols We compile -# all sources to avoid undefined symbols at module load time +# Define source files matching FlashMLA's setup.py. The API layer references all +# registered kernels, so keep this list synchronized with the pinned FlashMLA +# dependency. set(FLASH_MLA_SOURCES - ${FLASH_MLA_SOURCE_DIR}/csrc/pybind.cpp - ${FLASH_MLA_SOURCE_DIR}/csrc/smxx/get_mla_metadata.cu - ${FLASH_MLA_SOURCE_DIR}/csrc/smxx/mla_combine.cu) + ${FLASH_MLA_SOURCE_DIR}/csrc/api/api.cpp + ${FLASH_MLA_SOURCE_DIR}/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu + ${FLASH_MLA_SOURCE_DIR}/csrc/smxx/decode/combine/combine.cu) # Add SM90 sources (always include if CUDA >= 12.3) if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3) list( APPEND FLASH_MLA_SOURCES - ${FLASH_MLA_SOURCE_DIR}/csrc/sm90/decode/dense/splitkv_mla.cu - ${FLASH_MLA_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/splitkv_mla.cu - ${FLASH_MLA_SOURCE_DIR}/csrc/sm90/prefill/sparse/fwd.cu) + ${FLASH_MLA_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/fp16.cu + ${FLASH_MLA_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/bf16.cu + ${FLASH_MLA_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h64.cu + ${FLASH_MLA_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h128.cu + ${FLASH_MLA_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h64.cu + ${FLASH_MLA_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h128.cu + ${FLASH_MLA_SOURCE_DIR}/csrc/sm90/prefill/sparse/fwd.cu + ${FLASH_MLA_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512.cu + ${FLASH_MLA_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu + ${FLASH_MLA_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576.cu + ${FLASH_MLA_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu + ) endif() # Add SM100 sources (always include if CUDA >= 12.9) @@ -160,10 +174,17 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) list( APPEND FLASH_MLA_SOURCES - ${FLASH_MLA_SOURCE_DIR}/csrc/sm100/decode/sparse_fp8/splitkv_mla.cu ${FLASH_MLA_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu ${FLASH_MLA_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu - ${FLASH_MLA_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd.cu) + ${FLASH_MLA_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu + ${FLASH_MLA_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu + ${FLASH_MLA_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu + ${FLASH_MLA_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu + ${FLASH_MLA_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu + ${FLASH_MLA_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32.cu + ${FLASH_MLA_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/model1.cu + ${FLASH_MLA_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu + ) endif() # Disable LTO before creating target (similar to DeepEP) Let CMake generate @@ -175,9 +196,9 @@ pybind11_add_module(flash_mla_cpp_tllm ${FLASH_MLA_SOURCES}) set_target_properties( flash_mla_cpp_tllm PROPERTIES CXX_STANDARD_REQUIRED ON - CXX_STANDARD 17 + CXX_STANDARD 20 CXX_SCAN_FOR_MODULES OFF - CUDA_STANDARD 17 + CUDA_STANDARD 20 CUDA_STANDARD_REQUIRED ON CUDA_SEPARABLE_COMPILATION ON CUDA_RESOLVE_DEVICE_SYMBOLS ON @@ -201,15 +222,20 @@ target_compile_options( flash_mla_cpp_tllm PRIVATE ${TORCH_CXX_FLAGS} - $<$:-std=c++17> + $<$:-std=c++20> $<$:-O3> $<$:-fPIC> $<$:-DNDEBUG> $<$:-Wno-deprecated-declarations> $<$:-Wno-c++11-narrowing> $<$:-fno-lto> + # The FlashMLA API source only includes pybind11/pybind11.h, but the sparse + # interface exposes std::optional. Force Torch's pybind casters + # into that translation unit without patching the fetched source. + $<$:-include> + $<$:torch/csrc/utils/pybind.h> $<$:-O3> - $<$:-std=c++17> + $<$:-std=c++20> $<$:-DNDEBUG> $<$:-D_USE_MATH_DEFINES> $<$:-U__CUDA_NO_HALF_OPERATORS__> @@ -230,6 +256,7 @@ target_include_directories( flash_mla_cpp_tllm PRIVATE ${CUDA_INCLUDE_DIRS} ${FLASH_MLA_SOURCE_DIR}/csrc + ${FLASH_MLA_SOURCE_DIR}/csrc/kerutils/include ${FLASH_MLA_SOURCE_DIR}/csrc/sm90 ${FLASH_MLA_SOURCE_DIR}/csrc/cutlass/include ${FLASH_MLA_SOURCE_DIR}/csrc/cutlass/tools/util/include) diff --git a/examples/auto_deploy/model_registry/configs/deepseek_v4_flash.yaml b/examples/auto_deploy/model_registry/configs/deepseek_v4_flash.yaml new file mode 100644 index 000000000000..778112da979d --- /dev/null +++ b/examples/auto_deploy/model_registry/configs/deepseek_v4_flash.yaml @@ -0,0 +1,54 @@ +# 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. + +# DeepSeek V4 Flash full-layer config for generation quality validation with monolithic decode CUDA graph. +compile_backend: torch-cudagraph +model_factory: DeepseekV4AutoModelForCausalLM +max_batch_size: 2 +max_num_tokens: 512 +enable_chunked_prefill: true +cuda_graph_config: + max_batch_size: 2 + batch_sizes: [1, 2] +transforms: + load_weights: + disable_preload: true + insert_cached_deepseek_v4_sparse_attention: + stage: cache_init + backend: flashmla_deepseek_v4_sparse + enabled: true + fuse_rmsnorm: + rmsnorm_backend: torch + detect_sharding: + enabled: false + sharding_transform_executor: + enabled: false + apply_sharding_hints: + enabled: true + dist_mapping: + tp: 8 + moe_ep: 8 + moe_tp: 1 + moe_cluster: 1 + enable_attention_dp: false + shard_layers: ["mla", "moe"] + fuse_finegrained_fp8_linear: + enabled: false + quantize_mxfp4_moe: + enabled: true + compile_model: + piecewise_enabled: false + multi_stream_moe: + enabled: false diff --git a/examples/auto_deploy/model_registry/models.yaml b/examples/auto_deploy/model_registry/models.yaml index c301d17171a7..e01571065d1a 100644 --- a/examples/auto_deploy/model_registry/models.yaml +++ b/examples/auto_deploy/model_registry/models.yaml @@ -312,6 +312,9 @@ models: - name: deepseek-ai/DeepSeek-R1-0528 config_id: deepseek_r1 yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml', 'deepseek-r1.yaml', 'enable_sharder_ir.yaml'] +- name: deepseek-ai/DeepSeek-V4-Flash + config_id: deepseek_v4_flash + yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml', 'deepseek_v4_flash.yaml'] # OOM during AutoDeploy run. # - name: deepseek-ai/DeepSeek-Coder-V2-Instruct # config_id: deepseek_v2_ep diff --git a/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py b/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py index f9f7981b1e1b..9b94e9982c37 100644 --- a/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py +++ b/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py @@ -558,9 +558,8 @@ def prepare(self) -> None: ): trailing_idx = static_indices.pop() ad_logger.info( - "PiecewiseCapturedGraph: excluding trailing static submod_%d " - "(lm_head) from capture — will run eagerly", - trailing_idx, + f"PiecewiseCapturedGraph: excluding trailing static submod_{trailing_idx} " + "(lm_head) from capture — will run eagerly" ) runner_by_idx: Dict[int, ADPiecewiseRunner] = {} diff --git a/tensorrt_llm/_torch/auto_deploy/compile/piecewise_utils.py b/tensorrt_llm/_torch/auto_deploy/compile/piecewise_utils.py index 4d875e58ab83..d7f1116e3890 100644 --- a/tensorrt_llm/_torch/auto_deploy/compile/piecewise_utils.py +++ b/tensorrt_llm/_torch/auto_deploy/compile/piecewise_utils.py @@ -99,6 +99,7 @@ class DynamicOpPolicy(Enum): "auto_deploy::flashinfer_trtllm_mla_with_cache": DynamicOpPolicy.OUT_BUFFER, "auto_deploy::torch_cached_mla_with_cache": DynamicOpPolicy.OUT_BUFFER, "auto_deploy::trtllm_mla_with_cache": DynamicOpPolicy.OUT_BUFFER, + "auto_deploy::torch_deepseek_v4_sparse_attention_with_cache": DynamicOpPolicy.OUT_BUFFER, "auto_deploy::triton_cached_ssm": DynamicOpPolicy.OUT_BUFFER, "auto_deploy::torch_cached_ssm": DynamicOpPolicy.OUT_BUFFER, "auto_deploy::flashinfer_cached_ssm": DynamicOpPolicy.OUT_BUFFER, diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml index 6952ff96c64c..04849ac85efd 100644 --- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml +++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml @@ -283,6 +283,10 @@ transforms: stage: cache_init requires_shape_prop: true backend: flashinfer_mla + insert_cached_deepseek_v4_sparse_attention: + stage: cache_init + backend: deepseek_v4_sparse + enabled: false insert_cached_ssm_attention: stage: cache_init backend: triton_ssm diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/README.md b/tensorrt_llm/_torch/auto_deploy/custom_ops/README.md index 2e0ad3e5a2ce..f411a800b5b7 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/README.md +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/README.md @@ -59,6 +59,10 @@ The table below lists the operators grouped by category. | `torch.ops.auto_deploy.torch_moe_dense_mlp` | Dense MLP implementation for MoE (PyTorch backend) | | `torch.ops.auto_deploy.torch_quant_fp8_moe` | FP8 quantized MoE (PyTorch backend) | | `torch.ops.auto_deploy.torch_quant_nvfp4_moe` | NVFP4 quantized MoE (PyTorch backend) | +| `torch.ops.auto_deploy.torch_mxfp4_moe` | MXFP4 MoE reference implementation (PyTorch backend) | +| `torch.ops.auto_deploy.torch_mxfp4_moe_from_routing` | Routing-driven MXFP4 MoE reference implementation (PyTorch backend) | +| `torch.ops.auto_deploy.torch_mxfp4_moe_from_routing_ep` | Routing-driven MXFP4 MoE reference implementation with Expert Parallelism (PyTorch backend) | +| `torch.ops.auto_deploy.torch_mxfp4_moe_ep` | MXFP4 MoE reference implementation with Expert Parallelism (PyTorch backend) | | `torch.ops.auto_deploy.triton_moe_fused` | Fused MoE (Triton backend) | | `torch.ops.auto_deploy.triton_quant_fp8_moe` | FP8 quantized MoE (Triton backend) | | `torch.ops.auto_deploy.triton_mxfp4_moe` | MXFP4 MoE with triton-kernels matmul_ogs | diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/__init__.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/__init__.py index b76beb693b98..52e4e23f7a12 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/__init__.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/__init__.py @@ -24,6 +24,7 @@ """ __all__ = [ + "deepseek_v4_sparse_attention", "torch_attention", "torch_backend_attention", "flashinfer_attention", diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/deepseek_v4_sparse_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/deepseek_v4_sparse_attention.py new file mode 100644 index 000000000000..a4beb46ce66e --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/deepseek_v4_sparse_attention.py @@ -0,0 +1,3801 @@ +# 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. + +"""DeepSeek V4 sparse attention source and cached reference ops.""" + +import inspect +from functools import lru_cache +from typing import Callable, NamedTuple, Optional + +import torch +from torch._ops import OpOverloadPacket +from torch._subclasses import FakeTensor +from torch.fx import Node + +from ..._compat import KvCacheConfig +from ...distributed import common as dist_common +from ...utils.node_utils import extract_op_args +from ...utils.quantization_utils import fake_fp4_act_quant as _fake_fp4_act_quant +from ...utils.quantization_utils import fake_fp8_act_quant as _fake_fp8_act_quant +from ...utils.quantization_utils import hadamard_rotate as _hadamard_rotate +from ..attention_interface import ( + AttentionDescriptor, + AttentionLayout, + AttentionRegistry, + BatchInfo, + Constant, + MHACallable, + PagedResourceHandler, + ResourceHandlerDict, +) + +__all__ = [ + "DeepSeekV4SparseAttention", + "FlashMLADeepSeekV4SparseAttention", + "flashmla_deepseek_v4_sparse_attention_with_cache", + "torch_deepseek_v4_sparse_attention", + "torch_deepseek_v4_sparse_attention_with_cache", +] + +_SPARSE_ATTENTION_CHUNK_TARGET_BYTES = 512 * 1024 * 1024 +_SPARSE_ATTENTION_MAX_CHUNK_TOKENS = 64 +_FLASH_MLA_PREFILL_HEAD_DIM = 512 +_FLASH_MLA_PREFILL_TOPK_MULTIPLE = 128 +_FLASH_MLA_DECODE_TOPK_MULTIPLE = 128 +_FLASH_MLA_MODEL1_HEAD_DIM = 512 +_FLASH_MLA_MODEL1_NOPE_DIM = 448 +_FLASH_MLA_MODEL1_ROPE_DIM = 64 +_FLASH_MLA_MODEL1_QUANT_TILE = 64 +_FLASH_MLA_MODEL1_NUM_TILES = 7 +_FLASH_MLA_MODEL1_TOKEN_DATA_BYTES = _FLASH_MLA_MODEL1_NOPE_DIM + 2 * _FLASH_MLA_MODEL1_ROPE_DIM +_FLASH_MLA_MODEL1_SCALE_BYTES = 8 +_FLASH_MLA_MODEL1_BYTES_PER_TOKEN = ( + _FLASH_MLA_MODEL1_TOKEN_DATA_BYTES + _FLASH_MLA_MODEL1_SCALE_BYTES +) +_FLASH_MLA_FP8_MAX = 448.0 +_COMPRESS_RATIO_DISABLED = 0 +_COMPRESS_RATIO_OVERLAP_INDEXER = 4 +_COMPRESS_RATIO_DENSE = 128 + + +class _CompressionMode(NamedTuple): + ratio: int + enabled: bool + overlap: bool + uses_indexer: bool + channels: int + + +_COMPRESSION_MODES = { + _COMPRESS_RATIO_DISABLED: _CompressionMode( + ratio=_COMPRESS_RATIO_DISABLED, + enabled=False, + overlap=False, + uses_indexer=False, + channels=1, + ), + _COMPRESS_RATIO_OVERLAP_INDEXER: _CompressionMode( + ratio=_COMPRESS_RATIO_OVERLAP_INDEXER, + enabled=True, + overlap=True, + uses_indexer=True, + channels=2, + ), + _COMPRESS_RATIO_DENSE: _CompressionMode( + ratio=_COMPRESS_RATIO_DENSE, + enabled=True, + overlap=False, + uses_indexer=False, + channels=1, + ), +} +_SUPPORTED_COMPRESS_RATIOS = tuple(_COMPRESSION_MODES) +_SOURCE_TENSOR_ARG_NAMES = ( + "q", + "kv", + "attn_sink", + "topk_idxs", + "compressor_kv", + "compressor_gate", + "compressor_ape", + "compressor_norm_weight", + "cos_table", + "sin_table", + "position_ids", + "indexer_q", + "indexer_weights", + "indexer_compressor_kv", + "indexer_compressor_gate", + "indexer_compressor_ape", + "indexer_compressor_norm_weight", +) + + +def _validate_rank(name: str, tensor: torch.Tensor, rank: int) -> None: + if tensor.dim() != rank: + raise ValueError(f"{name} must have rank {rank}, got rank {tensor.dim()}") + + +def _compression_mode(compress_ratio: int) -> _CompressionMode: + mode = _COMPRESSION_MODES.get(compress_ratio) + if mode is None: + raise ValueError( + "DeepSeek V4 cached sparse attention supports " + f"compress_ratio in {_SUPPORTED_COMPRESS_RATIOS}, got {compress_ratio}" + ) + return mode + + +def _validate_compress_ratio(compress_ratio: int) -> None: + _compression_mode(compress_ratio) + + +def _validate_deepseek_v4_sparse_attention_inputs( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, +) -> None: + _validate_rank("q", q, 4) + _validate_rank("kv", kv, 3) + _validate_rank("attn_sink", attn_sink, 1) + _validate_rank("topk_idxs", topk_idxs, 3) + + if not q.is_floating_point(): + raise TypeError(f"q must be floating point, got {q.dtype}") + if not kv.is_floating_point(): + raise TypeError(f"kv must be floating point, got {kv.dtype}") + if not attn_sink.is_floating_point(): + raise TypeError(f"attn_sink must be floating point, got {attn_sink.dtype}") + if q.dtype != kv.dtype: + raise TypeError(f"q and kv must have the same dtype, got {q.dtype} and {kv.dtype}") + if topk_idxs.dtype not in (torch.int32, torch.int64): + raise TypeError(f"topk_idxs must be int32 or int64, got {topk_idxs.dtype}") + + if kv.device != q.device: + raise ValueError(f"kv must be on {q.device}, got {kv.device}") + if attn_sink.device != q.device: + raise ValueError(f"attn_sink must be on {q.device}, got {attn_sink.device}") + if topk_idxs.device != q.device: + raise ValueError(f"topk_idxs must be on {q.device}, got {topk_idxs.device}") + + batch_size, seq_len, num_heads, head_dim = q.shape + kv_batch_size, _, kv_head_dim = kv.shape + topk_batch_size, topk_seq_len, _ = topk_idxs.shape + + if kv_batch_size != batch_size: + raise ValueError(f"kv batch dimension must be {batch_size}, got {kv_batch_size}") + if topk_batch_size != batch_size: + raise ValueError(f"topk_idxs batch dimension must be {batch_size}, got {topk_batch_size}") + if topk_seq_len != seq_len: + raise ValueError(f"topk_idxs sequence dimension must be {seq_len}, got {topk_seq_len}") + if kv_head_dim != head_dim: + raise ValueError(f"kv head dimension must be {head_dim}, got {kv_head_dim}") + if attn_sink.shape[0] != num_heads: + raise ValueError(f"attn_sink length must be {num_heads}, got {attn_sink.shape[0]}") + + +def _validate_swa_cache_inputs(q: torch.Tensor, kv: torch.Tensor, swa_cache: torch.Tensor) -> None: + _validate_rank("swa_cache", swa_cache, 3) + if not swa_cache.is_floating_point(): + raise TypeError(f"swa_cache must be floating point, got {swa_cache.dtype}") + if swa_cache.device != q.device: + raise ValueError(f"swa_cache must be on {q.device}, got {swa_cache.device}") + if swa_cache.shape[-1] != kv.shape[-1]: + raise ValueError( + f"swa_cache head dimension must be {kv.shape[-1]}, got {swa_cache.shape[-1]}" + ) + + +def _is_flash_mla_model1_cache(cache: torch.Tensor) -> bool: + return ( + cache.dim() == 3 + and cache.dtype == torch.uint8 + and cache.shape[-1] == _FLASH_MLA_MODEL1_BYTES_PER_TOKEN + ) + + +def _validate_flash_mla_model1_cache_inputs( + q: torch.Tensor, + kv: torch.Tensor, + name: str, + cache: torch.Tensor, +) -> None: + _validate_rank(name, cache, 3) + if cache.device != q.device: + raise ValueError(f"{name} must be on {q.device}, got {cache.device}") + if kv.shape[-1] != _FLASH_MLA_MODEL1_HEAD_DIM: + raise ValueError( + "FlashMLA DeepSeek V4 sparse backend requires kv head dimension " + f"{_FLASH_MLA_MODEL1_HEAD_DIM}, got {kv.shape[-1]}" + ) + if not _is_flash_mla_model1_cache(cache): + raise TypeError( + f"{name} must be a uint8 FlashMLA MODEL1 cache with token shape " + f"[{_FLASH_MLA_MODEL1_BYTES_PER_TOKEN}], got dtype={cache.dtype} " + f"and shape={tuple(cache.shape)}" + ) + + +def _pack_flash_mla_model1_rows( + rows: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if rows.dim() != 2 or rows.shape[-1] != _FLASH_MLA_MODEL1_HEAD_DIM: + raise ValueError( + "FlashMLA MODEL1 packing requires rows with shape [num_rows, " + f"{_FLASH_MLA_MODEL1_HEAD_DIM}], got {tuple(rows.shape)}" + ) + + rows_bf16 = rows.to(torch.bfloat16).contiguous() + nope = rows_bf16[:, :_FLASH_MLA_MODEL1_NOPE_DIM].float() + nope_tiles = nope.view(-1, _FLASH_MLA_MODEL1_NUM_TILES, _FLASH_MLA_MODEL1_QUANT_TILE) + absmax = nope_tiles.abs().amax(dim=-1) + exponent = torch.ceil(torch.log2(torch.clamp_min(absmax / _FLASH_MLA_FP8_MAX, 1.0e-4))) + scales = torch.pow(torch.full_like(exponent, 2.0), exponent) + quantized = torch.clamp( + nope_tiles / scales.unsqueeze(-1), -_FLASH_MLA_FP8_MAX, _FLASH_MLA_FP8_MAX + ) + nope_bytes = quantized.to(torch.float8_e4m3fn).view(torch.uint8) + nope_bytes = nope_bytes.reshape(rows.shape[0], _FLASH_MLA_MODEL1_NOPE_DIM).contiguous() + + rope = rows_bf16[:, _FLASH_MLA_MODEL1_NOPE_DIM:].contiguous() + rope_bytes = rope.view(torch.uint8).reshape(rows.shape[0], 2 * _FLASH_MLA_MODEL1_ROPE_DIM) + scale_bytes = torch.clamp(exponent + 127.0, 0.0, 255.0).to(torch.uint8).contiguous() + return nope_bytes, rope_bytes, scale_bytes + + +def _write_flash_mla_model1_rows_to_cache( + cache: torch.Tensor, + values: torch.Tensor, + page_ids: torch.Tensor, + page_offsets: torch.Tensor, + valid: Optional[torch.Tensor] = None, +) -> None: + if values.numel() == 0: + return + if not _is_flash_mla_model1_cache(cache): + raise TypeError("cache must be a FlashMLA MODEL1 uint8 cache") + + flat_values = values.reshape(-1, values.shape[-1]) + flat_page_ids = page_ids.reshape(-1).to(torch.long) + flat_page_offsets = page_offsets.reshape(-1).to(torch.long) + if valid is None: + flat_valid = None + else: + flat_valid = valid.reshape(-1).to(torch.bool) + + nope_bytes, rope_bytes, scale_bytes = _pack_flash_mla_model1_rows(flat_values) + zero_scale = torch.zeros( + scale_bytes.shape[0], + 1, + dtype=torch.uint8, + device=scale_bytes.device, + ) + scale_bytes = torch.cat((scale_bytes, zero_scale), dim=-1) + flat_cache = cache.view(-1) + block_size = int(cache.shape[1]) + block_offsets = flat_page_ids * int(cache.stride(0)) + token_data_start = block_offsets + flat_page_offsets * _FLASH_MLA_MODEL1_TOKEN_DATA_BYTES + scale_start = block_offsets + block_size * _FLASH_MLA_MODEL1_TOKEN_DATA_BYTES + scale_start = scale_start + flat_page_offsets * _FLASH_MLA_MODEL1_SCALE_BYTES + + def scatter_bytes(start_offsets: torch.Tensor, packed: torch.Tensor) -> None: + byte_offsets = torch.arange(packed.shape[-1], dtype=torch.long, device=packed.device) + offsets = start_offsets.unsqueeze(-1) + byte_offsets + flat_offsets = offsets.reshape(-1) + flat_packed = packed.reshape(-1) + if flat_valid is not None: + valid_bytes = flat_valid.unsqueeze(-1).expand_as(packed).reshape(-1) + existing = flat_cache.gather(0, flat_offsets) + flat_packed = torch.where(valid_bytes, flat_packed, existing) + flat_cache.scatter_(0, flat_offsets, flat_packed) + + scatter_bytes(token_data_start, nope_bytes) + scatter_bytes(token_data_start + _FLASH_MLA_MODEL1_NOPE_DIM, rope_bytes) + scatter_bytes(scale_start, scale_bytes) + + +def _rms_norm_ref(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + compute = x.to(torch.float32) + output = compute * torch.rsqrt(compute.square().mean(dim=-1, keepdim=True) + eps) + if weight.numel() != 0: + output = output * weight.to(device=x.device, dtype=torch.float32) + return output.to(x.dtype) + + +def _apply_interleaved_rope_ref( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> torch.Tensor: + if x.shape[-1] == 0: + return x.contiguous() + x_even = x[..., 0::2] + x_odd = x[..., 1::2] + out_even = x_even * cos - x_odd * sin + out_odd = x_even * sin + x_odd * cos + return torch.stack((out_even, out_odd), dim=-1).flatten(-2).to(x.dtype) + + +def _apply_compressed_rope_and_quantize( + compressed: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + rope_dim: int, + rotate: bool = False, +) -> torch.Tensor: + if rope_dim < 0 or rope_dim > compressed.shape[-1]: + raise ValueError(f"rope_dim must be in [0, {compressed.shape[-1]}], got {rope_dim}") + nope_dim = compressed.shape[-1] - rope_dim + nope, pe = torch.split(compressed, [nope_dim, rope_dim], dim=-1) + pe = _apply_interleaved_rope_ref(pe, cos, sin) + compressed = torch.cat((nope, pe), dim=-1) + if rotate: + return _fake_fp4_act_quant(_hadamard_rotate(compressed), block_size=32) + nope, pe = torch.split(compressed, [nope_dim, rope_dim], dim=-1) + nope = _fake_fp8_act_quant(nope, block_size=64) + return torch.cat((nope, pe), dim=-1) + + +def _overlap_transform_projected( + tensor: torch.Tensor, + head_dim: int, + value: float, +) -> torch.Tensor: + batch_size, compressed_len, ratio, _ = tensor.shape + previous = tensor[:, :, :, :head_dim] + current = tensor[:, :, :, head_dim:] + prefix = tensor.new_full((batch_size, 1, ratio, head_dim), value) + previous = torch.cat((prefix, previous[:, :-1]), dim=1) + return torch.cat((previous, current), dim=2) + + +def _build_full_compressed_kv( + compressor_kv: torch.Tensor, + compressor_gate: torch.Tensor, + compressor_ape: torch.Tensor, + compressor_norm_weight: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + position_ids: torch.Tensor, + rms_norm_eps: float, + rope_dim: int, + compress_ratio: int, + max_compressed_len: int, + rotate: bool = False, +) -> torch.Tensor: + mode = _compression_mode(compress_ratio) + if not mode.enabled: + return compressor_kv.new_empty(compressor_kv.shape[0], 0, compressor_kv.shape[-1]) + + _validate_rank("compressor_kv", compressor_kv, 3) + _validate_rank("compressor_gate", compressor_gate, 3) + _validate_rank("compressor_ape", compressor_ape, 2) + _validate_rank("compressor_norm_weight", compressor_norm_weight, 1) + _validate_rank("cos_table", cos_table, 2) + _validate_rank("sin_table", sin_table, 2) + _validate_rank("position_ids", position_ids, 2) + if compressor_kv.shape != compressor_gate.shape: + raise ValueError( + "compressor_kv and compressor_gate must have matching shapes, " + f"got {tuple(compressor_kv.shape)} and {tuple(compressor_gate.shape)}" + ) + if max_compressed_len <= 0: + raise ValueError(f"max_compressed_len must be positive, got {max_compressed_len}") + + batch_size, seq_len, state_dim = compressor_kv.shape + if state_dim % mode.channels != 0: + raise ValueError(f"compressor state dim {state_dim} is not divisible by {mode.channels}") + head_dim = state_dim // mode.channels + max_compressed_tokens = max_compressed_len * compress_ratio + if seq_len > max_compressed_tokens: + raise ValueError(f"seq_len {seq_len} exceeds compressed capacity {max_compressed_tokens}") + if seq_len == 0: + return compressor_kv.new_empty(batch_size, max_compressed_len, head_dim) + + row_offsets = torch.arange(max_compressed_len, device=compressor_kv.device) + token_offsets = torch.arange(compress_ratio, device=compressor_kv.device) + gather_idxs = row_offsets.unsqueeze(1) * compress_ratio + token_offsets + valid = gather_idxs < seq_len + gather_idxs = torch.where(valid, gather_idxs, torch.zeros_like(gather_idxs)) + flat_idxs = gather_idxs.reshape(-1) + + kv = compressor_kv[:, flat_idxs].view(batch_size, max_compressed_len, compress_ratio, state_dim) + gate = compressor_gate[:, flat_idxs].view( + batch_size, max_compressed_len, compress_ratio, state_dim + ) + gate = gate + compressor_ape.to(device=gate.device, dtype=gate.dtype) + gate = torch.where( + valid.view(1, max_compressed_len, compress_ratio, 1), + gate, + gate.new_full((), -1.0e20), + ) + if mode.overlap: + kv = _overlap_transform_projected(kv, head_dim, 0.0) + gate = _overlap_transform_projected(gate, head_dim, -1.0e20) + + compressed = (kv * gate.softmax(dim=2)).sum(dim=2) + compressed = _rms_norm_ref(compressed, compressor_norm_weight, rms_norm_eps) + + row_start = row_offsets * compress_ratio + row_start = torch.minimum(row_start, torch.full_like(row_start, seq_len - 1)) + compressed_position_ids = position_ids[:, row_start] + cos = cos_table[compressed_position_ids] + sin = sin_table[compressed_position_ids] + return _apply_compressed_rope_and_quantize(compressed, cos, sin, rope_dim, rotate=rotate) + + +def _gather_selected_kv( + kv: torch.Tensor, + topk_idxs: torch.Tensor, + batch_idxs: Optional[torch.Tensor] = None, +) -> torch.Tensor: + kv_rows = kv.shape[1] + if kv_rows == 0: + return kv.new_zeros(*topk_idxs.shape, kv.shape[-1]) + + gather_topk_idxs = topk_idxs.to(torch.long).clamp(min=0, max=kv_rows - 1) + if batch_idxs is not None: + return kv[batch_idxs.to(torch.long).unsqueeze(1), gather_topk_idxs] + + batch_size, seq_len, k_select = topk_idxs.shape + head_dim = kv.shape[-1] + gather_idx = gather_topk_idxs.unsqueeze(-1).expand(batch_size, seq_len, k_select, head_dim) + expanded_kv = kv.unsqueeze(1).expand(batch_size, seq_len, kv.shape[1], head_dim) + return torch.gather(expanded_kv, dim=2, index=gather_idx) + + +def _to_host_long(name: str, tensor: torch.Tensor, length: int) -> torch.Tensor: + flat = tensor.detach().cpu().to(torch.long).flatten() + if flat.numel() < length: + raise ValueError(f"{name} must have at least {length} elements, got {flat.numel()}") + return flat[:length] + + +def _host_page_id_and_offset( + cache: torch.Tensor, + seq_idx: int, + logical_pos: int, + cu_num_pages_host: torch.Tensor, + cache_loc_host: torch.Tensor, +) -> tuple[int, int]: + if logical_pos < 0: + raise ValueError(f"logical_pos must be non-negative, got {logical_pos}") + tokens_per_block = int(cache.shape[1]) + page_ordinal = logical_pos // tokens_per_block + page_offset = logical_pos % tokens_per_block + page_start = int(cu_num_pages_host[seq_idx].item()) + page_end = int(cu_num_pages_host[seq_idx + 1].item()) + page_table_idx = page_start + page_ordinal + if page_table_idx >= page_end: + raise ValueError( + f"Sequence {seq_idx} logical position {logical_pos} needs page ordinal " + f"{page_ordinal}, but only {page_end - page_start} page(s) are active" + ) + return int(cache_loc_host[page_table_idx].item()), page_offset + + +def _host_position_is_valid( + cache: torch.Tensor, + seq_idx: int, + logical_pos: int, + cu_num_pages_host: torch.Tensor, +) -> bool: + if logical_pos < 0: + return False + tokens_per_block = int(cache.shape[1]) + page_ordinal = logical_pos // tokens_per_block + page_start = int(cu_num_pages_host[seq_idx].item()) + page_end = int(cu_num_pages_host[seq_idx + 1].item()) + return page_start + page_ordinal < page_end + + +def _host_page_ids_and_offsets( + cache: torch.Tensor, + seq_idx: int, + logical_positions_host: torch.Tensor, + cu_num_pages_host: torch.Tensor, + cache_loc_host: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + logical_positions_host = logical_positions_host.to(torch.long).flatten() + if logical_positions_host.numel() == 0: + return logical_positions_host, logical_positions_host + + negative_mask = logical_positions_host < 0 + if negative_mask.any(): + first_bad = int(logical_positions_host[negative_mask.nonzero()[0].item()].item()) + raise ValueError(f"logical_pos must be non-negative, got {first_bad}") + + tokens_per_block = int(cache.shape[1]) + page_ordinals = torch.div(logical_positions_host, tokens_per_block, rounding_mode="floor") + page_offsets = logical_positions_host.remainder(tokens_per_block) + page_start = int(cu_num_pages_host[seq_idx].item()) + page_end = int(cu_num_pages_host[seq_idx + 1].item()) + page_table_idxs = page_start + page_ordinals + invalid_mask = page_table_idxs >= page_end + if invalid_mask.any(): + first_bad_idx = int(invalid_mask.nonzero()[0].item()) + logical_pos = int(logical_positions_host[first_bad_idx].item()) + page_ordinal = int(page_ordinals[first_bad_idx].item()) + raise ValueError( + f"Sequence {seq_idx} logical position {logical_pos} needs page ordinal " + f"{page_ordinal}, but only {page_end - page_start} page(s) are active" + ) + + page_ids = cache_loc_host.index_select(0, page_table_idxs.to(torch.long)).to(torch.long) + return page_ids, page_offsets + + +def _write_paged_cache_rows( + values: torch.Tensor, + cache: torch.Tensor, + seq_idx: int, + input_pos: int, + cu_num_pages_host: torch.Tensor, + cache_loc_host: torch.Tensor, +) -> None: + if input_pos < 0: + raise ValueError(f"input_pos must be non-negative, got {input_pos}") + if values.numel() == 0: + return + logical_positions_host = torch.arange( + input_pos, + input_pos + values.shape[0], + dtype=torch.long, + ) + page_ids, page_offsets = _host_page_ids_and_offsets( + cache, + seq_idx, + logical_positions_host, + cu_num_pages_host, + cache_loc_host, + ) + if _is_flash_mla_model1_cache(cache): + _write_flash_mla_model1_rows_to_cache( + cache, + values, + page_ids.to(device=values.device), + page_offsets.to(device=values.device), + ) + return + + cache[ + page_ids.to(device=cache.device), + page_offsets.to(device=cache.device), + ] = values.to(cache.dtype) + + +def _write_paged_cache_rows_at_positions( + values: torch.Tensor, + cache: torch.Tensor, + seq_idx: int, + logical_positions: torch.Tensor, + cu_num_pages_host: torch.Tensor, + cache_loc_host: torch.Tensor, +) -> None: + if values.numel() == 0: + return + positions_host = logical_positions.detach().cpu().to(torch.long).flatten() + page_ids, page_offsets = _host_page_ids_and_offsets( + cache, + seq_idx, + positions_host, + cu_num_pages_host, + cache_loc_host, + ) + if _is_flash_mla_model1_cache(cache): + _write_flash_mla_model1_rows_to_cache( + cache, + values, + page_ids.to(device=values.device), + page_offsets.to(device=values.device), + ) + return + + cache[ + page_ids.to(device=cache.device), + page_offsets.to(device=cache.device), + ] = values.to(cache.dtype) + + +def _slice_sequence_tokens( + tensor: torch.Tensor, + seq_idx: int, + flat_start: int, + seq_len: int, +) -> torch.Tensor: + if tensor.numel() == 0: + return tensor.new_empty(seq_len, *tensor.shape[2:]) + if tensor.shape[0] > seq_idx and tensor.shape[0] != 1: + return tensor[seq_idx, :seq_len] + return tensor.reshape(-1, *tensor.shape[2:])[flat_start : flat_start + seq_len] + + +def _prefill_kv_source( + kv: torch.Tensor, + kv_seq: torch.Tensor, + seq_idx: int, + num_seq: int, +) -> torch.Tensor: + if kv.shape[0] > seq_idx and kv.shape[0] != 1: + return kv[seq_idx : seq_idx + 1] + if num_seq == 1: + return kv + return kv_seq.unsqueeze(0) + + +def _slice_sequence_positions( + position_ids: torch.Tensor, + seq_idx: int, + flat_start: int, + seq_len: int, +) -> torch.Tensor: + if position_ids.shape[0] > seq_idx and position_ids.shape[0] != 1: + return position_ids[seq_idx : seq_idx + 1, :seq_len] + return position_ids.reshape(1, -1)[:, flat_start : flat_start + seq_len] + + +def _slice_sequence_kv_rows( + kv: torch.Tensor, + seq_idx: int, + flat_start: int, + seq_len: int, + num_seq: int, + compress_ratio: int, +) -> torch.Tensor: + if compress_ratio == 0: + return _slice_sequence_tokens(kv, seq_idx, flat_start, seq_len) + if kv.shape[0] > seq_idx and kv.shape[0] != 1: + return kv[seq_idx] + if num_seq == 1: + return kv.reshape(-1, kv.shape[-1]) + raise ValueError( + "Flattened compressed DeepSeek V4 sparse attention KV rows are not supported; " + f"pass batched kv for compress_ratio={compress_ratio}." + ) + + +def _cached_sparse_attention_from_positions( + q_token: torch.Tensor, + attn_sink: torch.Tensor, + swa_cache: torch.Tensor, + seq_idx: int, + positions: torch.Tensor, + cu_num_pages_host: torch.Tensor, + cache_loc_host: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + selected_kv, valid_rows = _gather_paged_rows_from_positions( + swa_cache, + seq_idx, + positions, + cu_num_pages_host, + cache_loc_host, + q_token.dtype, + ) + selected_kv = selected_kv.unsqueeze(0) + local_topk = torch.arange(positions.numel(), dtype=torch.long, device=q_token.device).view( + 1, 1, -1 + ) + local_topk = torch.where(valid_rows.view(1, 1, -1), local_topk, torch.full_like(local_topk, -1)) + return _deepseek_v4_sparse_attention( + q_token.view(1, 1, *q_token.shape), + selected_kv, + attn_sink, + local_topk, + softmax_scale, + ).view(*q_token.shape) + + +def _cached_local_window_attention( + q_seq: torch.Tensor, + attn_sink: torch.Tensor, + swa_cache: torch.Tensor, + seq_idx: int, + input_pos: int, + cu_num_pages_host: torch.Tensor, + cache_loc_host: torch.Tensor, + window_size: int, + softmax_scale: float, +) -> torch.Tensor: + outputs = [] + for token_offset in range(q_seq.shape[0]): + query_pos = input_pos + token_offset + start_pos = max(0, query_pos - window_size + 1) + positions = torch.arange(start_pos, query_pos + 1, device=q_seq.device) + outputs.append( + _cached_sparse_attention_from_positions( + q_seq[token_offset], + attn_sink, + swa_cache, + seq_idx, + positions, + cu_num_pages_host, + cache_loc_host, + softmax_scale, + ) + ) + if not outputs: + return q_seq.new_empty(q_seq.shape) + return torch.stack(outputs, dim=0) + + +def _gather_paged_rows_from_positions( + cache: torch.Tensor, + seq_idx: int, + positions: torch.Tensor, + cu_num_pages_host: torch.Tensor, + cache_loc_host: torch.Tensor, + dtype: torch.dtype, + width: Optional[int] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + positions_host = positions.detach().cpu().to(torch.long).flatten() + rows = [] + valid_rows = [] + row_width = cache.shape[-1] if width is None else width + zero = cache.new_zeros(row_width) + for logical_pos_tensor in positions_host: + logical_pos = int(logical_pos_tensor.item()) + is_valid = _host_position_is_valid(cache, seq_idx, logical_pos, cu_num_pages_host) + valid_rows.append(is_valid) + if is_valid: + page_id, page_offset = _host_page_id_and_offset( + cache, seq_idx, logical_pos, cu_num_pages_host, cache_loc_host + ) + row = cache[page_id, page_offset] + if width is not None: + row = row[..., :width] + rows.append(row.to(dtype)) + else: + rows.append(zero.to(dtype)) + + if rows: + gathered = torch.stack(rows, dim=0) + else: + gathered = cache.new_empty(0, row_width, dtype=dtype) + valid = torch.tensor(valid_rows, dtype=torch.bool, device=positions.device) + return gathered.view(*positions.shape, row_width), valid.view(positions.shape) + + +def _gather_paged_rows( + cache: torch.Tensor, + seq_idx: int, + start_pos: int, + end_pos: int, + cu_num_pages_host: torch.Tensor, + cache_loc_host: torch.Tensor, + dtype: torch.dtype, + width: Optional[int] = None, +) -> torch.Tensor: + if start_pos < 0 or end_pos < start_pos: + raise ValueError(f"Invalid cache slice [{start_pos}, {end_pos})") + positions = torch.arange(start_pos, end_pos, dtype=torch.long, device=cache.device) + rows, _ = _gather_paged_rows_from_positions( + cache, seq_idx, positions, cu_num_pages_host, cache_loc_host, dtype, width=width + ) + return rows + + +def _compressed_row_from_paged_state( + compressor_kv_cache: torch.Tensor, + compressor_gate_cache: torch.Tensor, + seq_idx: int, + row_idx: int, + row_position_id: int, + cu_num_pages_host: torch.Tensor, + cache_loc_host: torch.Tensor, + compressor_ape: torch.Tensor, + compressor_norm_weight: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + rms_norm_eps: float, + rope_dim: int, + compress_ratio: int, + head_dim: int, + state_dim: int, + dtype: torch.dtype, + rotate: bool = False, +) -> torch.Tensor: + anchor = row_idx * compress_ratio + kv_rows = [] + gate_rows = [] + mode = _compression_mode(compress_ratio) + if mode.overlap: + for offset in range(compress_ratio): + position = anchor - compress_ratio + offset + if position < 0: + kv_rows.append( + torch.zeros(head_dim, dtype=dtype, device=compressor_kv_cache.device) + ) + gate_rows.append( + torch.full( + (head_dim,), + -1.0e20, + dtype=dtype, + device=compressor_gate_cache.device, + ) + ) + continue + kv_state = _gather_paged_rows( + compressor_kv_cache, + seq_idx, + position, + position + 1, + cu_num_pages_host, + cache_loc_host, + dtype, + ).squeeze(0) + gate_state = _gather_paged_rows( + compressor_gate_cache, + seq_idx, + position, + position + 1, + cu_num_pages_host, + cache_loc_host, + dtype, + ).squeeze(0) + kv_rows.append(kv_state[:head_dim]) + gate_rows.append(gate_state[:head_dim] + compressor_ape[offset, :head_dim].to(dtype)) + + for offset in range(compress_ratio): + position = anchor + offset + kv_state = _gather_paged_rows( + compressor_kv_cache, + seq_idx, + position, + position + 1, + cu_num_pages_host, + cache_loc_host, + dtype, + ).squeeze(0) + gate_state = _gather_paged_rows( + compressor_gate_cache, + seq_idx, + position, + position + 1, + cu_num_pages_host, + cache_loc_host, + dtype, + ).squeeze(0) + kv_rows.append(kv_state[head_dim : 2 * head_dim]) + gate_rows.append( + gate_state[head_dim : 2 * head_dim] + + compressor_ape[offset, head_dim : 2 * head_dim].to(dtype) + ) + else: + for offset in range(compress_ratio): + position = anchor + offset + kv_state = _gather_paged_rows( + compressor_kv_cache, + seq_idx, + position, + position + 1, + cu_num_pages_host, + cache_loc_host, + dtype, + ).squeeze(0) + gate_state = _gather_paged_rows( + compressor_gate_cache, + seq_idx, + position, + position + 1, + cu_num_pages_host, + cache_loc_host, + dtype, + ).squeeze(0) + kv_rows.append(kv_state[:head_dim]) + gate_rows.append(gate_state[:head_dim] + compressor_ape[offset, :head_dim].to(dtype)) + + kv = torch.stack(kv_rows, dim=0) + gate = torch.stack(gate_rows, dim=0) + pooled = (kv * gate.softmax(dim=0)).sum(dim=0) + pooled = _rms_norm_ref(pooled.unsqueeze(0), compressor_norm_weight, rms_norm_eps).squeeze(0) + del state_dim + row_position_id = max(0, min(row_position_id, cos_table.shape[0] - 1)) + cos = cos_table[row_position_id].unsqueeze(0) + sin = sin_table[row_position_id].unsqueeze(0) + return _apply_compressed_rope_and_quantize( + pooled.unsqueeze(0), + cos, + sin, + rope_dim, + rotate=rotate, + ).squeeze(0) + + +def _update_compressed_paged_caches( + compressor_kv_seq: torch.Tensor, + compressor_gate_seq: torch.Tensor, + position_ids_seq: torch.Tensor, + compressor_ape: torch.Tensor, + compressor_norm_weight: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + seq_idx: int, + input_pos: int, + cu_num_pages_host: torch.Tensor, + cache_loc_host: torch.Tensor, + mhc_cache: torch.Tensor, + compressor_kv_cache: torch.Tensor, + compressor_gate_cache: torch.Tensor, + rms_norm_eps: float, + rope_dim: int, + compress_ratio: int, + max_compressed_len: int, +) -> None: + mode = _compression_mode(compress_ratio) + if not mode.enabled or compressor_kv_seq.numel() == 0: + return + + if compressor_kv_seq.shape != compressor_gate_seq.shape: + raise ValueError( + "compressor_kv and compressor_gate sequence slices must have matching shapes, " + f"got {tuple(compressor_kv_seq.shape)} and {tuple(compressor_gate_seq.shape)}" + ) + state_dim = int(compressor_kv_seq.shape[-1]) + head_dim = state_dim // mode.channels + _write_paged_cache_rows( + compressor_kv_seq, + compressor_kv_cache, + seq_idx, + input_pos, + cu_num_pages_host, + cache_loc_host, + ) + _write_paged_cache_rows( + compressor_gate_seq, + compressor_gate_cache, + seq_idx, + input_pos, + cu_num_pages_host, + cache_loc_host, + ) + + old_completed = min(input_pos // compress_ratio, max_compressed_len) + new_completed = min( + (input_pos + compressor_kv_seq.shape[0]) // compress_ratio, max_compressed_len + ) + compressed_rows = [] + flat_position_ids = position_ids_seq.reshape(-1) + first_position_id = int(flat_position_ids[0].item()) + for row_idx in range(old_completed, new_completed): + row_token_offset = row_idx * compress_ratio - input_pos + if 0 <= row_token_offset < flat_position_ids.numel(): + row_position_id = int(flat_position_ids[row_token_offset].item()) + else: + row_position_id = first_position_id + row_token_offset + compressed_rows.append( + _compressed_row_from_paged_state( + compressor_kv_cache, + compressor_gate_cache, + seq_idx, + row_idx, + row_position_id, + cu_num_pages_host, + cache_loc_host, + compressor_ape, + compressor_norm_weight, + cos_table, + sin_table, + rms_norm_eps, + rope_dim, + compress_ratio, + head_dim, + state_dim, + compressor_kv_seq.dtype, + ) + ) + if compressed_rows: + logical_positions = torch.arange( + old_completed, + old_completed + len(compressed_rows), + dtype=torch.long, + device=mhc_cache.device, + ) + logical_positions = logical_positions * compress_ratio + _write_paged_cache_rows_at_positions( + torch.stack(compressed_rows, dim=0), + mhc_cache, + seq_idx, + logical_positions, + cu_num_pages_host, + cache_loc_host, + ) + + +def _update_raw_paged_caches( + compressor_kv_seq: torch.Tensor, + compressor_gate_seq: torch.Tensor, + compressor_kv_cache: torch.Tensor, + compressor_gate_cache: torch.Tensor, + seq_idx: int, + input_pos: int, + cu_num_pages_host: torch.Tensor, + cache_loc_host: torch.Tensor, +) -> None: + if compressor_kv_seq.numel() == 0: + return + if compressor_kv_seq.shape != compressor_gate_seq.shape: + raise ValueError( + "compressor_kv and compressor_gate sequence slices must have matching shapes, " + f"got {tuple(compressor_kv_seq.shape)} and {tuple(compressor_gate_seq.shape)}" + ) + _write_paged_cache_rows( + compressor_kv_seq, + compressor_kv_cache, + seq_idx, + input_pos, + cu_num_pages_host, + cache_loc_host, + ) + _write_paged_cache_rows( + compressor_gate_seq, + compressor_gate_cache, + seq_idx, + input_pos, + cu_num_pages_host, + cache_loc_host, + ) + + +def _select_ratio4_indexer_rows( + q_index: torch.Tensor, + indexer_weights: torch.Tensor, + indexer_compressor_kv_cache: torch.Tensor, + indexer_compressor_gate_cache: torch.Tensor, + seq_idx: int, + query_pos: int, + query_position_id: int, + index_topk: int, + cu_num_pages_host: torch.Tensor, + cache_loc_host: torch.Tensor, + indexer_compressor_ape: torch.Tensor, + indexer_compressor_norm_weight: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + rms_norm_eps: float, + rope_dim: int, + max_compressed_len: int, +) -> torch.Tensor: + if index_topk <= 0: + return torch.empty(0, dtype=torch.int64, device=q_index.device) + + visible_len = min((query_pos + 1) // 4, max_compressed_len) + if visible_len <= 0: + return torch.full((index_topk,), -1, dtype=torch.int64, device=q_index.device) + + index_head_dim = int(q_index.shape[-1]) + state_dim = int(indexer_compressor_kv_cache.shape[-1]) + index_k = torch.stack( + [ + _compressed_row_from_paged_state( + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + seq_idx, + row_idx, + query_position_id - (query_pos - row_idx * 4), + cu_num_pages_host, + cache_loc_host, + indexer_compressor_ape, + indexer_compressor_norm_weight, + cos_table, + sin_table, + rms_norm_eps, + rope_dim, + 4, + index_head_dim, + state_dim, + q_index.dtype, + rotate=True, + ) + for row_idx in range(visible_len) + ], + dim=0, + ) + index_score = torch.matmul(q_index, index_k.transpose(-1, -2)).float() + index_score = (index_score.relu() * indexer_weights.float().unsqueeze(-1)).sum(dim=0) + if dist_common.is_initialized() and dist_common.get_world_size() > 1: + dist_common.all_reduce(index_score, op=dist_common.ReduceOp.SUM) + + topk_count = min(index_topk, visible_len) + selected = index_score.topk(topk_count, dim=-1).indices.to(torch.int64) + if topk_count < index_topk: + pad = torch.full( + (index_topk - topk_count,), + -1, + dtype=selected.dtype, + device=selected.device, + ) + selected = torch.cat((selected, pad), dim=0) + return selected + + +def _cached_compressed_attention( + q_seq: torch.Tensor, + attn_sink: torch.Tensor, + swa_cache: torch.Tensor, + mhc_cache: torch.Tensor, + seq_idx: int, + input_pos: int, + cu_num_pages_host: torch.Tensor, + cache_loc_host: torch.Tensor, + position_ids_seq: torch.Tensor, + window_size: int, + compress_ratio: int, + max_compressed_len: int, + softmax_scale: float, + topk_seq: Optional[torch.Tensor] = None, + indexer_q_seq: Optional[torch.Tensor] = None, + indexer_weights_seq: Optional[torch.Tensor] = None, + indexer_compressor_kv_cache: Optional[torch.Tensor] = None, + indexer_compressor_gate_cache: Optional[torch.Tensor] = None, + indexer_compressor_ape: Optional[torch.Tensor] = None, + indexer_compressor_norm_weight: Optional[torch.Tensor] = None, + cos_table: Optional[torch.Tensor] = None, + sin_table: Optional[torch.Tensor] = None, + rms_norm_eps: float = 1e-6, + rope_dim: Optional[int] = None, +) -> torch.Tensor: + outputs = [] + flat_position_ids = position_ids_seq.reshape(-1) + for token_offset in range(q_seq.shape[0]): + query_pos = input_pos + token_offset + query_position_id = int(flat_position_ids[token_offset].item()) + local_start = max(0, query_pos - window_size + 1) + local_kv = _gather_paged_rows( + swa_cache, + seq_idx, + local_start, + query_pos + 1, + cu_num_pages_host, + cache_loc_host, + q_seq.dtype, + ) + local_idxs = torch.arange(local_kv.shape[0], dtype=torch.int64, device=q_seq.device) + mode = _compression_mode(compress_ratio) + if mode.uses_indexer: + if ( + topk_seq is None + or indexer_q_seq is None + or indexer_weights_seq is None + or indexer_compressor_kv_cache is None + or indexer_compressor_gate_cache is None + or indexer_compressor_ape is None + or indexer_compressor_norm_weight is None + or cos_table is None + or sin_table is None + or rope_dim is None + ): + raise ValueError( + "Overlap/indexer cached decode requires indexer tensors and caches." + ) + index_topk = max(int(topk_seq.shape[-1]) - int(window_size), 0) + selected_rows = _select_ratio4_indexer_rows( + indexer_q_seq[token_offset], + indexer_weights_seq[token_offset], + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + seq_idx, + query_pos, + query_position_id, + index_topk, + cu_num_pages_host, + cache_loc_host, + indexer_compressor_ape, + indexer_compressor_norm_weight, + cos_table, + sin_table, + rms_norm_eps, + rope_dim, + max_compressed_len, + ) + compressed_positions = selected_rows.clamp(min=0) * compress_ratio + compressed_kv, compressed_valid = _gather_paged_rows_from_positions( + mhc_cache, + seq_idx, + compressed_positions, + cu_num_pages_host, + cache_loc_host, + q_seq.dtype, + ) + compressed_valid = compressed_valid & (selected_rows >= 0) + compressed_idxs = torch.where( + compressed_valid, + torch.arange(selected_rows.numel(), dtype=torch.int64, device=q_seq.device) + + local_kv.shape[0], + torch.full_like(selected_rows, -1), + ) + else: + compressed_len = min((query_pos + 1) // compress_ratio, max_compressed_len) + compressed_positions = ( + torch.arange(compressed_len, dtype=torch.long, device=q_seq.device) * compress_ratio + ) + compressed_kv, compressed_valid = _gather_paged_rows_from_positions( + mhc_cache, + seq_idx, + compressed_positions, + cu_num_pages_host, + cache_loc_host, + q_seq.dtype, + ) + compressed_idxs = torch.arange( + compressed_kv.shape[0], dtype=torch.int64, device=q_seq.device + ) + compressed_idxs = compressed_idxs + local_kv.shape[0] + compressed_idxs = torch.where( + compressed_valid, + compressed_idxs, + torch.full_like(compressed_idxs, -1), + ) + topk = torch.cat((local_idxs, compressed_idxs), dim=0).view(1, 1, -1) + kv = torch.cat((local_kv, compressed_kv), dim=0) + out = _deepseek_v4_sparse_attention( + q_seq[token_offset : token_offset + 1].unsqueeze(0), + kv.unsqueeze(0), + attn_sink, + topk, + softmax_scale, + ) + outputs.append(out.squeeze(0).squeeze(0)) + if not outputs: + return q_seq.new_empty(q_seq.shape) + return torch.stack(outputs, dim=0) + + +def _cached_topk_attention( + q_seq: torch.Tensor, + attn_sink: torch.Tensor, + topk_seq: torch.Tensor, + swa_cache: torch.Tensor, + seq_idx: int, + cu_num_pages_host: torch.Tensor, + cache_loc_host: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + outputs = [] + for token_offset in range(q_seq.shape[0]): + outputs.append( + _cached_sparse_attention_from_positions( + q_seq[token_offset], + attn_sink, + swa_cache, + seq_idx, + topk_seq[token_offset], + cu_num_pages_host, + cache_loc_host, + softmax_scale, + ) + ) + if not outputs: + return q_seq.new_empty(q_seq.shape) + return torch.stack(outputs, dim=0) + + +def _flatten_decode_tokens( + tensor: torch.Tensor, + num_decode: int, + token_start: int = 0, + seq_start: int = 0, +) -> torch.Tensor: + if tensor.numel() == 0: + return tensor.new_empty(num_decode, *tensor.shape[2:]) + if tensor.shape[0] > seq_start and tensor.shape[0] != 1: + return tensor[seq_start : seq_start + num_decode, :1].reshape(num_decode, *tensor.shape[2:]) + return tensor.reshape(-1, *tensor.shape[2:])[token_start : token_start + num_decode] + + +def _flatten_decode_position_ids( + position_ids: torch.Tensor, + num_decode: int, + token_start: int = 0, + seq_start: int = 0, +) -> torch.Tensor: + if position_ids.shape[0] > seq_start and position_ids.shape[0] != 1: + return position_ids[seq_start : seq_start + num_decode, :1].reshape(-1).to(torch.long) + return position_ids.reshape(-1)[token_start : token_start + num_decode].to(torch.long) + + +def _decode_page_ids_and_offsets( + cache: torch.Tensor, + seq_idx: torch.Tensor, + positions: torch.Tensor, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if cache_loc.numel() == 0: + raise ValueError("cache_loc must contain at least one page id") + + positions_long = positions.to(torch.long) + safe_positions = positions_long.clamp(min=0) + tokens_per_block = int(cache.shape[1]) + page_ordinals = safe_positions // tokens_per_block + page_offsets = safe_positions % tokens_per_block + + seq_idx_long = seq_idx.to(torch.long) + while seq_idx_long.dim() < positions_long.dim(): + seq_idx_long = seq_idx_long.unsqueeze(-1) + page_start = cu_num_pages[seq_idx_long].to(torch.long) + page_end = cu_num_pages[seq_idx_long + 1].to(torch.long) + page_table_idx = page_start + page_ordinals + valid = (positions_long >= 0) & (page_table_idx < page_end) + safe_page_table_idx = torch.where(valid, page_table_idx, page_start) + safe_page_table_idx = safe_page_table_idx.clamp(min=0, max=cache_loc.numel() - 1) + page_ids = cache_loc[safe_page_table_idx].to(torch.long) + return page_ids, page_offsets, valid + + +def _write_decode_cache_rows( + cache: torch.Tensor, + values: torch.Tensor, + seq_idx: torch.Tensor, + input_pos: torch.Tensor, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + valid: Optional[torch.Tensor] = None, +) -> None: + if values.numel() == 0: + return + page_ids, page_offsets, page_valid = _decode_page_ids_and_offsets( + cache, seq_idx, input_pos, cu_num_pages, cache_loc + ) + write_valid = (input_pos.to(torch.long) >= 0) & page_valid + if valid is not None: + write_valid = write_valid & valid.reshape(write_valid.shape).to(torch.bool) + if _is_flash_mla_model1_cache(cache): + _write_flash_mla_model1_rows_to_cache(cache, values, page_ids, page_offsets, write_valid) + return + rows = values.to(cache.dtype) + if valid is not None: + previous_rows = cache[page_ids, page_offsets] + row_valid = write_valid + while row_valid.dim() < rows.dim(): + row_valid = row_valid.unsqueeze(-1) + rows = torch.where(row_valid, rows, previous_rows) + cache[page_ids, page_offsets] = rows + + +def _decode_cache_rows_from_positions( + cache: torch.Tensor, + seq_idx: torch.Tensor, + positions: torch.Tensor, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + page_ids, page_offsets, valid = _decode_page_ids_and_offsets( + cache, seq_idx, positions, cu_num_pages, cache_loc + ) + return cache[page_ids, page_offsets].to(dtype), valid + + +def _decode_attention_from_rows( + q_decode: torch.Tensor, + selected_kv: torch.Tensor, + valid_rows: torch.Tensor, + attn_sink: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + rel_topk = torch.arange(selected_kv.shape[1], dtype=torch.int64, device=q_decode.device) + rel_topk = rel_topk.view(1, -1).expand(q_decode.shape[0], -1) + rel_topk = torch.where(valid_rows, rel_topk, torch.full_like(rel_topk, -1)) + output = _deepseek_v4_sparse_attention( + q_decode.unsqueeze(1), + selected_kv, + attn_sink, + rel_topk.unsqueeze(1), + softmax_scale, + ) + return output.squeeze(1) + + +def _decode_local_cache_rows( + swa_cache: torch.Tensor, + seq_idx: torch.Tensor, + input_pos: torch.Tensor, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + window_size: int, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + offsets = torch.arange(window_size, dtype=torch.long, device=input_pos.device) + positions = input_pos.unsqueeze(1) - window_size + 1 + offsets.view(1, -1) + valid = (positions >= 0) & (positions <= input_pos.unsqueeze(1)) + rows, page_valid = _decode_cache_rows_from_positions( + swa_cache, seq_idx, positions, cu_num_pages, cache_loc, dtype + ) + valid = valid & page_valid + return rows, valid + + +def _decode_topk_cache_attention( + q_decode: torch.Tensor, + attn_sink: torch.Tensor, + topk_decode: torch.Tensor, + swa_cache: torch.Tensor, + seq_idx: torch.Tensor, + input_pos: torch.Tensor, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + softmax_scale: float, + window_size: Optional[int], +) -> torch.Tensor: + if window_size is not None: + selected_kv, valid_rows = _decode_local_cache_rows( + swa_cache, + seq_idx, + input_pos, + cu_num_pages, + cache_loc, + window_size, + q_decode.dtype, + ) + return _decode_attention_from_rows( + q_decode, + selected_kv, + valid_rows, + attn_sink, + softmax_scale, + ) + + positions = topk_decode.to(torch.long) + selected_kv, page_valid = _decode_cache_rows_from_positions( + swa_cache, seq_idx, positions, cu_num_pages, cache_loc, q_decode.dtype + ) + valid_rows = (positions >= 0) & page_valid + return _decode_attention_from_rows( + q_decode, + selected_kv, + valid_rows, + attn_sink, + softmax_scale, + ) + + +def _batched_compressed_rows_from_paged_state( + compressor_kv_cache: torch.Tensor, + compressor_gate_cache: torch.Tensor, + seq_idx: torch.Tensor, + row_idx: torch.Tensor, + row_position_id: torch.Tensor, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + compressor_ape: torch.Tensor, + compressor_norm_weight: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + rms_norm_eps: float, + rope_dim: int, + compress_ratio: int, + head_dim: int, + dtype: torch.dtype, + rotate: bool = False, +) -> torch.Tensor: + offsets = torch.arange(compress_ratio, dtype=torch.long, device=row_idx.device) + anchor = row_idx.to(torch.long) * compress_ratio + + mode = _compression_mode(compress_ratio) + if mode.overlap: + previous_positions = anchor.unsqueeze(1) - compress_ratio + offsets.view(1, -1) + current_positions = anchor.unsqueeze(1) + offsets.view(1, -1) + previous_valid = previous_positions >= 0 + + previous_kv_state, previous_page_valid = _decode_cache_rows_from_positions( + compressor_kv_cache, + seq_idx, + previous_positions, + cu_num_pages, + cache_loc, + dtype, + ) + previous_gate_state, _ = _decode_cache_rows_from_positions( + compressor_gate_cache, + seq_idx, + previous_positions, + cu_num_pages, + cache_loc, + dtype, + ) + current_kv_state, _ = _decode_cache_rows_from_positions( + compressor_kv_cache, + seq_idx, + current_positions, + cu_num_pages, + cache_loc, + dtype, + ) + current_gate_state, _ = _decode_cache_rows_from_positions( + compressor_gate_cache, + seq_idx, + current_positions, + cu_num_pages, + cache_loc, + dtype, + ) + previous_valid = previous_valid & previous_page_valid + + previous_kv = previous_kv_state[..., :head_dim] + previous_gate = previous_gate_state[..., :head_dim] + previous_gate = previous_gate + compressor_ape[:, :head_dim].to( + device=previous_gate.device, dtype=dtype + ) + previous_kv = torch.where( + previous_valid.unsqueeze(-1), previous_kv, previous_kv.new_zeros(()) + ) + previous_gate = torch.where( + previous_valid.unsqueeze(-1), + previous_gate, + previous_gate.new_full((), -1.0e20), + ) + + current_kv = current_kv_state[..., head_dim : 2 * head_dim] + current_gate = current_gate_state[..., head_dim : 2 * head_dim] + current_gate = current_gate + compressor_ape[:, head_dim : 2 * head_dim].to( + device=current_gate.device, dtype=dtype + ) + kv = torch.cat((previous_kv, current_kv), dim=1) + gate = torch.cat((previous_gate, current_gate), dim=1) + else: + positions = anchor.unsqueeze(1) + offsets.view(1, -1) + kv_state, _ = _decode_cache_rows_from_positions( + compressor_kv_cache, + seq_idx, + positions, + cu_num_pages, + cache_loc, + dtype, + ) + gate_state, _ = _decode_cache_rows_from_positions( + compressor_gate_cache, + seq_idx, + positions, + cu_num_pages, + cache_loc, + dtype, + ) + kv = kv_state[..., :head_dim] + gate = gate_state[..., :head_dim] + gate = gate + compressor_ape[:, :head_dim].to(device=gate.device, dtype=dtype) + + pooled = (kv * gate.softmax(dim=1)).sum(dim=1) + pooled = _rms_norm_ref(pooled, compressor_norm_weight, rms_norm_eps) + row_position_id = row_position_id.to(torch.long).clamp(min=0, max=cos_table.shape[0] - 1) + cos = cos_table[row_position_id] + sin = sin_table[row_position_id] + return _apply_compressed_rope_and_quantize( + pooled, + cos, + sin, + rope_dim, + rotate=rotate, + ) + + +def _update_decode_compressed_caches( + compressor_kv_decode: torch.Tensor, + compressor_gate_decode: torch.Tensor, + position_ids_decode: torch.Tensor, + compressor_ape: torch.Tensor, + compressor_norm_weight: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + seq_idx: torch.Tensor, + input_pos: torch.Tensor, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + mhc_cache: torch.Tensor, + compressor_kv_cache: torch.Tensor, + compressor_gate_cache: torch.Tensor, + rms_norm_eps: float, + rope_dim: int, + compress_ratio: int, + max_compressed_len: int, +) -> None: + mode = _compression_mode(compress_ratio) + if not mode.enabled or compressor_kv_decode.numel() == 0: + return + + state_dim = int(compressor_kv_decode.shape[-1]) + head_dim = state_dim // mode.channels + _write_decode_cache_rows( + compressor_kv_cache, compressor_kv_decode, seq_idx, input_pos, cu_num_pages, cache_loc + ) + _write_decode_cache_rows( + compressor_gate_cache, compressor_gate_decode, seq_idx, input_pos, cu_num_pages, cache_loc + ) + + old_completed = input_pos // compress_ratio + new_completed = (input_pos + 1) // compress_ratio + row_valid = (new_completed > old_completed) & (old_completed < max_compressed_len) + row_idx = old_completed.clamp(min=0, max=max_compressed_len - 1) + row_position_id = position_ids_decode.to(torch.long) - (input_pos - row_idx * compress_ratio) + compressed_rows = _batched_compressed_rows_from_paged_state( + compressor_kv_cache, + compressor_gate_cache, + seq_idx, + row_idx, + row_position_id, + cu_num_pages, + cache_loc, + compressor_ape, + compressor_norm_weight, + cos_table, + sin_table, + rms_norm_eps, + rope_dim, + compress_ratio, + head_dim, + compressor_kv_decode.dtype, + ) + row_logical_pos = row_idx * compress_ratio + if _is_flash_mla_model1_cache(mhc_cache): + _write_decode_cache_rows( + mhc_cache, + compressed_rows, + seq_idx, + row_logical_pos, + cu_num_pages, + cache_loc, + valid=row_valid, + ) + return + + previous_rows, _ = _decode_cache_rows_from_positions( + mhc_cache, seq_idx, row_logical_pos, cu_num_pages, cache_loc, mhc_cache.dtype + ) + rows_to_write = torch.where( + row_valid.unsqueeze(-1), + compressed_rows.to(mhc_cache.dtype), + previous_rows, + ) + _write_decode_cache_rows( + mhc_cache, rows_to_write, seq_idx, row_logical_pos, cu_num_pages, cache_loc + ) + + +def _select_decode_ratio4_indexer_rows( + q_index: torch.Tensor, + indexer_weights: torch.Tensor, + indexer_compressor_kv_cache: torch.Tensor, + indexer_compressor_gate_cache: torch.Tensor, + seq_idx: torch.Tensor, + input_pos: torch.Tensor, + position_ids_decode: torch.Tensor, + index_topk: int, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + indexer_compressor_ape: torch.Tensor, + indexer_compressor_norm_weight: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + rms_norm_eps: float, + rope_dim: int, + max_compressed_len: int, +) -> tuple[torch.Tensor, torch.Tensor]: + if index_topk <= 0: + empty_rows = torch.empty(q_index.shape[0], 0, dtype=torch.int64, device=q_index.device) + empty_valid = torch.empty(q_index.shape[0], 0, dtype=torch.bool, device=q_index.device) + return empty_rows, empty_valid + + candidate_rows = torch.arange(max_compressed_len, dtype=torch.long, device=q_index.device) + candidate_rows = candidate_rows.view(1, -1).expand(q_index.shape[0], -1) + flat_seq_idx = seq_idx.unsqueeze(1).expand_as(candidate_rows).reshape(-1) + flat_rows = candidate_rows.reshape(-1) + row_position_id = position_ids_decode.unsqueeze(1) - ( + input_pos.unsqueeze(1) - candidate_rows * 4 + ) + flat_row_position_id = row_position_id.reshape(-1) + index_head_dim = int(q_index.shape[-1]) + index_k = _batched_compressed_rows_from_paged_state( + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + flat_seq_idx, + flat_rows, + flat_row_position_id, + cu_num_pages, + cache_loc, + indexer_compressor_ape, + indexer_compressor_norm_weight, + cos_table, + sin_table, + rms_norm_eps, + rope_dim, + 4, + index_head_dim, + q_index.dtype, + rotate=True, + ) + index_k = index_k.view(q_index.shape[0], max_compressed_len, index_head_dim) + index_score = torch.matmul(q_index, index_k.transpose(-1, -2)).float() + index_score = (index_score.relu() * indexer_weights.float().unsqueeze(-1)).sum(dim=1) + if dist_common.is_initialized() and dist_common.get_world_size() > 1: + dist_common.all_reduce(index_score, op=dist_common.ReduceOp.SUM) + + visible_len = ((input_pos + 1) // 4).clamp(max=max_compressed_len) + visible = candidate_rows < visible_len.unsqueeze(1) + index_score = index_score.masked_fill(~visible, float("-inf")) + topk_count = min(index_topk, max_compressed_len) + topk_values, topk_rows = index_score.topk(topk_count, dim=-1) + topk_valid = torch.isfinite(topk_values) + topk_rows = torch.where(topk_valid, topk_rows.to(torch.int64), torch.full_like(topk_rows, -1)) + if topk_count < index_topk: + pad_shape = (q_index.shape[0], index_topk - topk_count) + row_pad = torch.full(pad_shape, -1, dtype=torch.int64, device=q_index.device) + valid_pad = torch.zeros(pad_shape, dtype=torch.bool, device=q_index.device) + topk_rows = torch.cat((topk_rows, row_pad), dim=-1) + topk_valid = torch.cat((topk_valid, valid_pad), dim=-1) + return topk_rows, topk_valid + + +def _decode_compressed_cache_attention( + q_decode: torch.Tensor, + attn_sink: torch.Tensor, + topk_decode: torch.Tensor, + indexer_q_decode: torch.Tensor, + indexer_weights_decode: torch.Tensor, + swa_cache: torch.Tensor, + mhc_cache: torch.Tensor, + indexer_compressor_kv_cache: torch.Tensor, + indexer_compressor_gate_cache: torch.Tensor, + indexer_compressor_ape: torch.Tensor, + indexer_compressor_norm_weight: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + seq_idx: torch.Tensor, + input_pos: torch.Tensor, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + position_ids_decode: torch.Tensor, + window_size: int, + compress_ratio: int, + max_compressed_len: int, + softmax_scale: float, + rms_norm_eps: float, + rope_dim: int, +) -> torch.Tensor: + local_kv, local_valid = _decode_local_cache_rows( + swa_cache, + seq_idx, + input_pos, + cu_num_pages, + cache_loc, + window_size, + q_decode.dtype, + ) + mode = _compression_mode(compress_ratio) + if mode.uses_indexer: + index_topk = max(int(topk_decode.shape[-1]) - int(window_size), 0) + selected_rows, compressed_valid = _select_decode_ratio4_indexer_rows( + indexer_q_decode, + indexer_weights_decode, + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + seq_idx, + input_pos, + position_ids_decode, + index_topk, + cu_num_pages, + cache_loc, + indexer_compressor_ape, + indexer_compressor_norm_weight, + cos_table, + sin_table, + rms_norm_eps, + rope_dim, + max_compressed_len, + ) + compressed_positions = selected_rows.clamp(min=0) * compress_ratio + compressed_kv, page_valid = _decode_cache_rows_from_positions( + mhc_cache, + seq_idx, + compressed_positions, + cu_num_pages, + cache_loc, + q_decode.dtype, + ) + compressed_valid = compressed_valid & page_valid & (selected_rows >= 0) + else: + candidate_rows = torch.arange( + max_compressed_len, + dtype=torch.long, + device=q_decode.device, + ) + selected_rows = candidate_rows.view(1, -1).expand(q_decode.shape[0], -1) + compressed_len = ((input_pos + 1) // compress_ratio).clamp(max=max_compressed_len) + compressed_valid = selected_rows < compressed_len.unsqueeze(1) + compressed_positions = selected_rows * compress_ratio + compressed_kv, page_valid = _decode_cache_rows_from_positions( + mhc_cache, + seq_idx, + compressed_positions, + cu_num_pages, + cache_loc, + q_decode.dtype, + ) + compressed_valid = compressed_valid & page_valid + + selected_kv = torch.cat((local_kv, compressed_kv), dim=1) + valid_rows = torch.cat((local_valid, compressed_valid), dim=1) + return _decode_attention_from_rows( + q_decode, + selected_kv, + valid_rows, + attn_sink, + softmax_scale, + ) + + +def _deepseek_v4_sparse_attention_decode_with_cache( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + compressor_kv: torch.Tensor, + compressor_gate: torch.Tensor, + compressor_ape: torch.Tensor, + compressor_norm_weight: torch.Tensor, + indexer_q: torch.Tensor, + indexer_weights: torch.Tensor, + indexer_compressor_kv: torch.Tensor, + indexer_compressor_gate: torch.Tensor, + indexer_compressor_ape: torch.Tensor, + indexer_compressor_norm_weight: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + position_ids: torch.Tensor, + input_pos: torch.Tensor, + slot_idx: torch.Tensor, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + swa_cache: torch.Tensor, + mhc_cache: torch.Tensor, + compressor_kv_cache: torch.Tensor, + compressor_gate_cache: torch.Tensor, + indexer_compressor_kv_cache: torch.Tensor, + indexer_compressor_gate_cache: torch.Tensor, + num_decode: int, + softmax_scale: float, + window_size: Optional[int], + compress_ratio: int, + max_compressed_len: Optional[int], + rms_norm_eps: float, + rope_dim: Optional[int], + out: Optional[torch.Tensor], +) -> torch.Tensor: + q_flat = q.reshape(-1, *q.shape[2:]) + q_decode = q_flat[:num_decode] + kv_decode = _flatten_decode_tokens(kv, num_decode) + topk_decode = _flatten_decode_tokens(topk_idxs, num_decode) + del slot_idx + seq_idx_decode = torch.arange(num_decode, dtype=torch.long, device=input_pos.device) + input_pos_decode = input_pos.reshape(-1)[:num_decode].to(torch.long) + position_ids_decode = position_ids.reshape(-1)[:num_decode].to(torch.long) + + _write_decode_cache_rows( + swa_cache, kv_decode, seq_idx_decode, input_pos_decode, cu_num_pages, cache_loc + ) + if compress_ratio: + assert window_size is not None + assert max_compressed_len is not None + assert rope_dim is not None + compressor_kv_decode = _flatten_decode_tokens(compressor_kv, num_decode) + compressor_gate_decode = _flatten_decode_tokens(compressor_gate, num_decode) + _update_decode_compressed_caches( + compressor_kv_decode, + compressor_gate_decode, + position_ids_decode, + compressor_ape, + compressor_norm_weight, + cos_table, + sin_table, + seq_idx_decode, + input_pos_decode, + cu_num_pages, + cache_loc, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + rms_norm_eps, + rope_dim, + compress_ratio, + max_compressed_len, + ) + mode = _compression_mode(compress_ratio) + if mode.uses_indexer: + indexer_compressor_kv_decode = _flatten_decode_tokens(indexer_compressor_kv, num_decode) + indexer_compressor_gate_decode = _flatten_decode_tokens( + indexer_compressor_gate, num_decode + ) + _write_decode_cache_rows( + indexer_compressor_kv_cache, + indexer_compressor_kv_decode, + seq_idx_decode, + input_pos_decode, + cu_num_pages, + cache_loc, + ) + _write_decode_cache_rows( + indexer_compressor_gate_cache, + indexer_compressor_gate_decode, + seq_idx_decode, + input_pos_decode, + cu_num_pages, + cache_loc, + ) + indexer_q_decode = _flatten_decode_tokens(indexer_q, num_decode) + indexer_weights_decode = _flatten_decode_tokens(indexer_weights, num_decode) + decode_output = _decode_compressed_cache_attention( + q_decode, + attn_sink, + topk_decode, + indexer_q_decode, + indexer_weights_decode, + swa_cache, + mhc_cache, + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + indexer_compressor_ape, + indexer_compressor_norm_weight, + cos_table, + sin_table, + seq_idx_decode, + input_pos_decode, + cu_num_pages, + cache_loc, + position_ids_decode, + window_size, + compress_ratio, + max_compressed_len, + softmax_scale, + rms_norm_eps, + rope_dim, + ) + else: + decode_output = _decode_topk_cache_attention( + q_decode, + attn_sink, + topk_decode, + swa_cache, + seq_idx_decode, + input_pos_decode, + cu_num_pages, + cache_loc, + softmax_scale, + window_size, + ) + + output_flat = torch.zeros_like(q_flat) + output_flat[:num_decode].copy_(decode_output) + output = output_flat.view_as(q) + if out is not None: + out.copy_(output) + return out.new_empty(0) + return output + + +def _flash_mla_deepseek_v4_sparse_attention_decode_with_cache( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + compressor_kv: torch.Tensor, + compressor_gate: torch.Tensor, + compressor_ape: torch.Tensor, + compressor_norm_weight: torch.Tensor, + indexer_q: torch.Tensor, + indexer_weights: torch.Tensor, + indexer_compressor_kv: torch.Tensor, + indexer_compressor_gate: torch.Tensor, + indexer_compressor_ape: torch.Tensor, + indexer_compressor_norm_weight: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + position_ids: torch.Tensor, + input_pos: torch.Tensor, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + swa_cache: torch.Tensor, + mhc_cache: torch.Tensor, + compressor_kv_cache: torch.Tensor, + compressor_gate_cache: torch.Tensor, + indexer_compressor_kv_cache: torch.Tensor, + indexer_compressor_gate_cache: torch.Tensor, + num_decode: int, + softmax_scale: float, + window_size: Optional[int], + compress_ratio: int, + max_compressed_len: Optional[int], + rms_norm_eps: float, + rope_dim: Optional[int], + out: Optional[torch.Tensor], + token_start: int = 0, + seq_start: int = 0, +) -> torch.Tensor: + q_flat = q.reshape(-1, *q.shape[2:]) + q_decode = _flatten_decode_tokens(q, num_decode, token_start, seq_start) + kv_decode = _flatten_decode_tokens(kv, num_decode, token_start, seq_start) + topk_decode = _flatten_decode_tokens(topk_idxs, num_decode, token_start, seq_start) + seq_idx_decode = torch.arange( + seq_start, seq_start + num_decode, dtype=torch.long, device=input_pos.device + ) + input_pos_decode = input_pos.reshape(-1)[seq_start : seq_start + num_decode].to(torch.long) + position_ids_decode = _flatten_decode_position_ids( + position_ids, num_decode, token_start, seq_start + ) + + _write_decode_cache_rows( + swa_cache, + kv_decode, + seq_idx_decode, + input_pos_decode, + cu_num_pages, + cache_loc, + ) + + extra_indices = None + extra_topk_length = None + if compress_ratio: + assert window_size is not None + assert max_compressed_len is not None + assert rope_dim is not None + compressor_kv_decode = _flatten_decode_tokens( + compressor_kv, num_decode, token_start, seq_start + ) + compressor_gate_decode = _flatten_decode_tokens( + compressor_gate, num_decode, token_start, seq_start + ) + _update_decode_compressed_caches( + compressor_kv_decode, + compressor_gate_decode, + position_ids_decode, + compressor_ape, + compressor_norm_weight, + cos_table, + sin_table, + seq_idx_decode, + input_pos_decode, + cu_num_pages, + cache_loc, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + rms_norm_eps, + rope_dim, + compress_ratio, + max_compressed_len, + ) + mode = _compression_mode(compress_ratio) + if mode.uses_indexer: + indexer_compressor_kv_decode = _flatten_decode_tokens( + indexer_compressor_kv, + num_decode, + token_start, + seq_start, + ) + indexer_compressor_gate_decode = _flatten_decode_tokens( + indexer_compressor_gate, + num_decode, + token_start, + seq_start, + ) + _write_decode_cache_rows( + indexer_compressor_kv_cache, + indexer_compressor_kv_decode, + seq_idx_decode, + input_pos_decode, + cu_num_pages, + cache_loc, + ) + _write_decode_cache_rows( + indexer_compressor_gate_cache, + indexer_compressor_gate_decode, + seq_idx_decode, + input_pos_decode, + cu_num_pages, + cache_loc, + ) + + indexer_q_decode = _flatten_decode_tokens(indexer_q, num_decode, token_start, seq_start) + indexer_weights_decode = _flatten_decode_tokens( + indexer_weights, num_decode, token_start, seq_start + ) + extra_indices, extra_topk_length = _build_flash_mla_compressed_decode_indices( + topk_decode, + indexer_q_decode, + indexer_weights_decode, + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + indexer_compressor_ape, + indexer_compressor_norm_weight, + cos_table, + sin_table, + seq_idx_decode, + input_pos_decode, + cu_num_pages, + cache_loc, + position_ids_decode, + mhc_cache, + window_size, + compress_ratio, + max_compressed_len, + rms_norm_eps, + rope_dim, + ) + + if window_size is not None: + indices, topk_length = _build_flash_mla_local_decode_indices( + swa_cache, + seq_idx_decode, + input_pos_decode, + cu_num_pages, + cache_loc, + window_size, + ) + else: + indices, topk_length = _build_flash_mla_topk_decode_indices( + swa_cache, + seq_idx_decode, + topk_decode, + cu_num_pages, + cache_loc, + ) + + decode_output = _flash_mla_sparse_decode( + q_decode, + swa_cache, + attn_sink, + indices, + topk_length, + softmax_scale, + extra_k_cache=mhc_cache if extra_indices is not None else None, + extra_indices=extra_indices, + extra_topk_length=extra_topk_length, + ) + + output_flat = torch.zeros_like(q_flat) + output_flat[token_start : token_start + num_decode].copy_(decode_output) + output = output_flat.view_as(q) + if out is not None: + out.copy_(output) + return out.new_empty(0) + return output + + +def _cached_decode_topk_positions( + topk_seq: torch.Tensor, + input_pos: int, + window_size: Optional[int], + compress_ratio: int, +) -> torch.Tensor: + if compress_ratio == 0 or window_size is None: + return topk_seq + + local_window_cols = min(window_size, topk_seq.shape[-1]) + if local_window_cols == 0: + return topk_seq + + token_offsets = torch.arange(topk_seq.shape[0], device=topk_seq.device).unsqueeze(1) + query_positions = input_pos + token_offsets + window_offsets = torch.arange(local_window_cols, device=topk_seq.device) + local_topk = query_positions - local_window_cols + 1 + window_offsets + local_topk = torch.where(local_topk < 0, -1, local_topk) + local_topk = local_topk.to(topk_seq.dtype) + if local_window_cols == topk_seq.shape[-1]: + return local_topk + return torch.cat((local_topk, topk_seq[..., local_window_cols:]), dim=-1) + + +def _sparse_attention_query_chunk_size( + num_tokens: int, + num_heads: int, + head_dim: int, + k_select: int, + compute_dtype: torch.dtype, +) -> int: + compute_element_size = torch.empty((), dtype=compute_dtype).element_size() + logits_element_size = torch.empty((), dtype=torch.float32).element_size() + bytes_per_token = ( + k_select * head_dim * compute_element_size + + 3 * num_heads * (k_select + 1) * logits_element_size + + num_heads * head_dim * compute_element_size + ) + if bytes_per_token <= 0: + return 1 + chunk_size = _SPARSE_ATTENTION_CHUNK_TARGET_BYTES // bytes_per_token + chunk_size = max(1, int(chunk_size)) + chunk_size = min(chunk_size, _SPARSE_ATTENTION_MAX_CHUNK_TOKENS) + return min(num_tokens, chunk_size) + + +def _validate_flash_mla_sparse_prefill_contract(op: Callable) -> None: + supports_dsv4_contract = getattr(op, "_supports_dsv4_contract", None) + try: + parameters = inspect.signature(op).parameters + except (TypeError, ValueError) as exc: + raise RuntimeError( + "FlashMLA sparse prefill must expose the DeepSeek V4 contract. " + "Rebuild the pinned vLLM FlashMLA dependency." + ) from exc + + required_kwargs = {"attn_sink", "topk_length", "out"} + missing_kwargs = required_kwargs.difference(parameters) + if supports_dsv4_contract is False or missing_kwargs: + missing_str = ", ".join(sorted(missing_kwargs)) or "_supports_dsv4_contract" + raise RuntimeError( + "FlashMLA sparse prefill is available, but it does not expose the " + f"DeepSeek V4 contract ({missing_str} missing). Rebuild the pinned " + "vLLM FlashMLA dependency." + ) + + +@lru_cache(maxsize=1) +def _get_flash_mla_sparse_prefill() -> Optional[Callable]: + try: + from tensorrt_llm.flash_mla import flash_mla_sparse_fwd + except (ImportError, ModuleNotFoundError, OSError, RuntimeError): + return None + _validate_flash_mla_sparse_prefill_contract(flash_mla_sparse_fwd) + return flash_mla_sparse_fwd + + +def _flash_mla_sparse_prefill_head_count( + num_heads: int, + capability: tuple[int, int], +) -> Optional[int]: + if capability[0] >= 10: + return 128 if num_heads <= 128 else None + if capability[0] == 9: + return ((num_heads + 63) // 64) * 64 + return None + + +def _prepare_flash_mla_sparse_prefill_inputs( + q: torch.Tensor, + kv: torch.Tensor, + topk_idxs: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size, seq_len, num_heads, head_dim = q.shape + kv_rows = kv.shape[1] + k_select = topk_idxs.shape[-1] + + q_flat = q.reshape(batch_size * seq_len, num_heads, head_dim).contiguous() + kv_flat = kv.reshape(batch_size * kv_rows, 1, head_dim).contiguous() + + batch_offsets = torch.arange(batch_size, dtype=torch.long, device=q.device) + batch_offsets = (batch_offsets * kv_rows).view(batch_size, 1, 1) + topk_long = topk_idxs.to(torch.long) + valid_topk = (topk_long >= 0) & (topk_long < kv_rows) + indices = torch.where(topk_long >= 0, topk_long + batch_offsets, topk_long) + indices = torch.where(valid_topk, indices, torch.full_like(indices, -1)) + indices = indices.reshape(batch_size * seq_len, 1, k_select).to(torch.int32).contiguous() + + padded_topk = ( + (k_select + _FLASH_MLA_PREFILL_TOPK_MULTIPLE - 1) // _FLASH_MLA_PREFILL_TOPK_MULTIPLE + ) * _FLASH_MLA_PREFILL_TOPK_MULTIPLE + if padded_topk != k_select: + pad = torch.full( + (indices.shape[0], 1, padded_topk - k_select), + -1, + dtype=torch.int32, + device=q.device, + ) + indices = torch.cat((indices, pad), dim=-1).contiguous() + + topk_length = torch.full( + (batch_size * seq_len,), + k_select, + dtype=torch.int32, + device=q.device, + ) + return q_flat, kv_flat, indices, topk_length + + +def _flash_mla_sparse_prefill( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + flash_mla_sparse_fwd = _get_flash_mla_sparse_prefill() + if flash_mla_sparse_fwd is None: + raise RuntimeError( + "FlashMLA DeepSeek V4 sparse backend requires the FlashMLA sparse prefill kernel." + ) + if not q.is_cuda or not kv.is_cuda: + raise RuntimeError("FlashMLA DeepSeek V4 sparse backend requires CUDA q and kv tensors.") + if q.dtype != torch.bfloat16 or kv.dtype != torch.bfloat16: + raise RuntimeError( + "FlashMLA DeepSeek V4 sparse backend requires bfloat16 q and kv tensors." + ) + if q.shape[-1] != _FLASH_MLA_PREFILL_HEAD_DIM: + raise RuntimeError( + "FlashMLA DeepSeek V4 sparse backend requires head_dim=" + f"{_FLASH_MLA_PREFILL_HEAD_DIM}, got {q.shape[-1]}." + ) + if topk_idxs.shape[-1] <= 0: + raise RuntimeError("FlashMLA DeepSeek V4 sparse backend requires at least one top-k row.") + if not torch.cuda.is_available(): + raise RuntimeError("FlashMLA DeepSeek V4 sparse backend requires CUDA availability.") + + capability = torch.cuda.get_device_capability(q.device) + padded_heads = _flash_mla_sparse_prefill_head_count(q.shape[2], capability) + if padded_heads is None: + raise RuntimeError( + "FlashMLA DeepSeek V4 sparse backend supports Hopper and Blackwell head layouts; " + f"got compute capability sm_{capability[0]}{capability[1]} and {q.shape[2]} heads." + ) + + q_flat, kv_flat, indices, topk_length = _prepare_flash_mla_sparse_prefill_inputs( + q, kv, topk_idxs + ) + num_heads = q_flat.shape[1] + if padded_heads != num_heads: + q_padded = q_flat.new_zeros(q_flat.shape[0], padded_heads, q_flat.shape[2]) + q_padded[:, :num_heads].copy_(q_flat) + q_flat = q_padded + + attn_sink_padded = torch.zeros( + padded_heads, + dtype=torch.float32, + device=attn_sink.device, + ) + attn_sink_padded[:num_heads].copy_(attn_sink.to(torch.float32)) + attn_sink_for_kernel = attn_sink_padded + else: + attn_sink_for_kernel = attn_sink.to(torch.float32).contiguous() + + output, _, _ = flash_mla_sparse_fwd( + q_flat, + kv_flat, + indices, + softmax_scale, + d_v=_FLASH_MLA_PREFILL_HEAD_DIM, + attn_sink=attn_sink_for_kernel, + topk_length=topk_length, + out=None, + ) + output = output[:, :num_heads].contiguous() + return output.view_as(q) + + +@lru_cache(maxsize=1) +def _get_flash_mla_decode() -> Optional[tuple[Callable, Callable]]: + try: + from tensorrt_llm.flash_mla import flash_mla_with_kvcache, get_mla_metadata + except (ImportError, ModuleNotFoundError, OSError, RuntimeError): + return None + return flash_mla_with_kvcache, get_mla_metadata + + +def _flash_mla_sparse_decode_head_count( + num_heads: int, + capability: tuple[int, int], +) -> Optional[int]: + if capability[0] < 9: + return None + if num_heads <= 64: + return 64 + if num_heads <= 128: + return 128 + return None + + +def _decode_positions_to_flash_mla_indices( + cache: torch.Tensor, + seq_idx: torch.Tensor, + positions: torch.Tensor, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + valid: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + page_ids, page_offsets, page_valid = _decode_page_ids_and_offsets( + cache, + seq_idx, + positions, + cu_num_pages, + cache_loc, + ) + positions_valid = positions.to(torch.long) >= 0 + if valid is not None: + positions_valid = positions_valid & valid + positions_valid = positions_valid & page_valid + block_size = int(cache.shape[1]) + indices = page_ids * block_size + page_offsets + indices = torch.where(positions_valid, indices, torch.full_like(indices, -1)) + return indices.to(torch.int32).contiguous(), positions_valid + + +def _pad_flash_mla_decode_indices(indices: torch.Tensor) -> torch.Tensor: + width = int(indices.shape[-1]) + padded_width = ( + (width + _FLASH_MLA_DECODE_TOPK_MULTIPLE - 1) // _FLASH_MLA_DECODE_TOPK_MULTIPLE + ) * _FLASH_MLA_DECODE_TOPK_MULTIPLE + if padded_width == width: + return indices.contiguous() + pad = torch.full( + (*indices.shape[:-1], padded_width - width), + -1, + dtype=indices.dtype, + device=indices.device, + ) + return torch.cat((indices, pad), dim=-1).contiguous() + + +def _build_flash_mla_local_decode_indices( + cache: torch.Tensor, + seq_idx: torch.Tensor, + input_pos: torch.Tensor, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + window_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + offsets = torch.arange(window_size, dtype=torch.long, device=input_pos.device) + input_pos_long = input_pos.to(torch.long) + lengths = torch.minimum(input_pos_long + 1, torch.full_like(input_pos_long, window_size)) + lengths = lengths.clamp(min=0) + start_positions = input_pos_long - lengths + 1 + positions = start_positions.unsqueeze(1) + offsets.view(1, -1) + valid = offsets.view(1, -1) < lengths.unsqueeze(1) + indices, valid = _decode_positions_to_flash_mla_indices( + cache, + seq_idx, + positions, + cu_num_pages, + cache_loc, + valid, + ) + topk_length = valid.sum(dim=-1).to(torch.int32).contiguous() + return _pad_flash_mla_decode_indices(indices).unsqueeze(1), topk_length + + +def _build_flash_mla_topk_decode_indices( + cache: torch.Tensor, + seq_idx: torch.Tensor, + topk_decode: torch.Tensor, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + positions = topk_decode.to(torch.long) + indices, valid = _decode_positions_to_flash_mla_indices( + cache, + seq_idx, + positions, + cu_num_pages, + cache_loc, + ) + topk_length = valid.sum(dim=-1).to(torch.int32).contiguous() + return _pad_flash_mla_decode_indices(indices).unsqueeze(1), topk_length + + +def _build_flash_mla_compressed_decode_indices( + topk_decode: torch.Tensor, + indexer_q_decode: torch.Tensor, + indexer_weights_decode: torch.Tensor, + indexer_compressor_kv_cache: torch.Tensor, + indexer_compressor_gate_cache: torch.Tensor, + indexer_compressor_ape: torch.Tensor, + indexer_compressor_norm_weight: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + seq_idx: torch.Tensor, + input_pos: torch.Tensor, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + position_ids_decode: torch.Tensor, + mhc_cache: torch.Tensor, + window_size: int, + compress_ratio: int, + max_compressed_len: int, + rms_norm_eps: float, + rope_dim: int, +) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + mode = _compression_mode(compress_ratio) + if mode.uses_indexer: + index_topk = max(int(topk_decode.shape[-1]) - int(window_size), 0) + if index_topk <= 0: + return None, None + selected_rows, compressed_valid = _select_decode_ratio4_indexer_rows( + indexer_q_decode, + indexer_weights_decode, + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + seq_idx, + input_pos, + position_ids_decode, + index_topk, + cu_num_pages, + cache_loc, + indexer_compressor_ape, + indexer_compressor_norm_weight, + cos_table, + sin_table, + rms_norm_eps, + rope_dim, + max_compressed_len, + ) + compressed_positions = selected_rows.clamp(min=0) * compress_ratio + compressed_valid = compressed_valid & (selected_rows >= 0) + else: + selected_rows = torch.arange( + max_compressed_len, + dtype=torch.long, + device=topk_decode.device, + ) + selected_rows = selected_rows.view(1, -1).expand(topk_decode.shape[0], -1) + compressed_len = ((input_pos + 1) // compress_ratio).clamp(max=max_compressed_len) + compressed_valid = selected_rows < compressed_len.unsqueeze(1) + compressed_positions = selected_rows * compress_ratio + + indices, valid = _decode_positions_to_flash_mla_indices( + mhc_cache, + seq_idx, + compressed_positions, + cu_num_pages, + cache_loc, + compressed_valid, + ) + topk_length = valid.sum(dim=-1).to(torch.int32).contiguous() + return _pad_flash_mla_decode_indices(indices).unsqueeze(1), topk_length + + +def _flash_mla_sparse_decode( + q_decode: torch.Tensor, + swa_cache: torch.Tensor, + attn_sink: torch.Tensor, + indices: torch.Tensor, + topk_length: torch.Tensor, + softmax_scale: float, + extra_k_cache: Optional[torch.Tensor] = None, + extra_indices: Optional[torch.Tensor] = None, + extra_topk_length: Optional[torch.Tensor] = None, +) -> torch.Tensor: + flash_mla_decode = _get_flash_mla_decode() + if flash_mla_decode is None: + raise RuntimeError( + "FlashMLA DeepSeek V4 sparse backend requires the FlashMLA decode kernel." + ) + flash_mla_with_kvcache, get_mla_metadata = flash_mla_decode + if q_decode.dtype != torch.bfloat16: + raise RuntimeError( + f"FlashMLA DeepSeek V4 sparse decode requires bfloat16 q tensors, got {q_decode.dtype}." + ) + if q_decode.shape[-1] != _FLASH_MLA_MODEL1_HEAD_DIM: + raise RuntimeError( + "FlashMLA DeepSeek V4 sparse decode requires head_dim=" + f"{_FLASH_MLA_MODEL1_HEAD_DIM}, got {q_decode.shape[-1]}." + ) + if not _is_flash_mla_model1_cache(swa_cache): + raise RuntimeError("FlashMLA sparse decode requires a packed uint8 SWA cache.") + if extra_k_cache is not None and not _is_flash_mla_model1_cache(extra_k_cache): + raise RuntimeError("FlashMLA sparse decode requires a packed uint8 extra cache.") + if indices.shape[-1] <= 0: + raise RuntimeError("FlashMLA sparse decode requires at least one primary index.") + if not q_decode.is_cuda or not swa_cache.is_cuda: + raise RuntimeError("FlashMLA DeepSeek V4 sparse decode requires CUDA tensors.") + if not torch.cuda.is_available(): + raise RuntimeError("FlashMLA DeepSeek V4 sparse decode requires CUDA availability.") + + capability = torch.cuda.get_device_capability(q_decode.device) + padded_heads = _flash_mla_sparse_decode_head_count(q_decode.shape[1], capability) + if padded_heads is None: + raise RuntimeError( + "FlashMLA DeepSeek V4 sparse decode supports Hopper and Blackwell with " + f"64 or 128 query heads; got sm_{capability[0]}{capability[1]} and " + f"{q_decode.shape[1]} heads." + ) + + num_heads = q_decode.shape[1] + if padded_heads != num_heads: + q_padded = q_decode.new_zeros(q_decode.shape[0], padded_heads, q_decode.shape[2]) + q_padded[:, :num_heads].copy_(q_decode) + q_decode = q_padded + attn_sink_padded = torch.zeros(padded_heads, dtype=torch.float32, device=attn_sink.device) + attn_sink_padded[:num_heads].copy_(attn_sink.to(torch.float32)) + attn_sink = attn_sink_padded + else: + attn_sink = attn_sink.to(torch.float32).contiguous() + + q_decode = q_decode.unsqueeze(1).contiguous() + sched_meta = get_mla_metadata()[0] + output, _ = flash_mla_with_kvcache( + q=q_decode, + k_cache=swa_cache.unsqueeze(-2), + block_table=None, + cache_seqlens=None, + head_dim_v=_FLASH_MLA_MODEL1_HEAD_DIM, + tile_scheduler_metadata=sched_meta, + is_fp8_kvcache=True, + indices=indices.contiguous(), + topk_length=topk_length.contiguous(), + softmax_scale=softmax_scale, + attn_sink=attn_sink, + extra_k_cache=extra_k_cache.unsqueeze(-2) if extra_k_cache is not None else None, + extra_indices_in_kvcache=extra_indices.contiguous() if extra_indices is not None else None, + extra_topk_length=extra_topk_length.contiguous() if extra_topk_length is not None else None, + ) + return output[:, 0, :num_heads].contiguous() + + +def _deepseek_v4_sparse_attention( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + """Reference DeepSeek V4 sparse attention implementation. + + Args: + q: Query states with shape ``[batch, seq_len, num_heads, head_dim]``. + kv: Shared sparse key/value rows with shape ``[batch, kv_rows, head_dim]``. + attn_sink: Per-head sink logits with shape ``[num_heads]``. + topk_idxs: Selected row indices into ``kv`` with shape + ``[batch, seq_len, k_select]``. Duplicate indices are preserved and + receive independent probability mass. Negative indices are masked + slots and receive zero probability. + softmax_scale: Scale applied to query/key logits before adding the sink + logit. + + Returns: + Attention output with shape ``[batch, seq_len, num_heads, head_dim]``. + The sink participates in softmax normalization but contributes no value + vector. + """ + _validate_deepseek_v4_sparse_attention_inputs(q, kv, attn_sink, topk_idxs) + + compute_dtype = torch.float32 if q.dtype in (torch.float16, torch.bfloat16) else q.dtype + batch_size, seq_len, num_heads, q_head_dim = q.shape + _, _, k_select = topk_idxs.shape + num_tokens = batch_size * seq_len + output = torch.empty(q.shape, dtype=q.dtype, device=q.device) + if num_tokens == 0: + return output + + chunk_size = _sparse_attention_query_chunk_size( + num_tokens, num_heads, q_head_dim, k_select, compute_dtype + ) + + q_flat = q.reshape(num_tokens, num_heads, q_head_dim) + topk_flat = topk_idxs.reshape(num_tokens, k_select) + batch_idxs = torch.arange(batch_size, device=q.device).view(batch_size, 1) + batch_idxs = batch_idxs.expand(batch_size, seq_len).reshape(num_tokens) + output_flat = output.reshape(num_tokens, num_heads, q_head_dim) + sink_logits = attn_sink.to(dtype=compute_dtype).reshape(1, num_heads, 1) + + for start in range(0, num_tokens, chunk_size): + end = min(start + chunk_size, num_tokens) + topk_chunk = topk_flat[start:end] + valid_topk = (topk_chunk >= 0) & (topk_chunk < kv.shape[1]) + selected_kv_compute = _gather_selected_kv(kv, topk_chunk, batch_idxs[start:end]).to( + compute_dtype + ) + q_compute = q_flat[start:end].to(compute_dtype) + + logits = torch.matmul(q_compute, selected_kv_compute.transpose(-1, -2)) + logits = logits * softmax_scale + logits = logits.masked_fill((~valid_topk).unsqueeze(1), float("-inf")) + chunk_sink_logits = sink_logits.expand(end - start, num_heads, 1) + logits_with_sink = torch.cat([logits, chunk_sink_logits], dim=-1) + + weights_with_sink = torch.softmax(logits_with_sink, dim=-1, dtype=torch.float32) + weights = weights_with_sink[..., :-1].to(compute_dtype) + chunk_output = torch.matmul(weights, selected_kv_compute) + output_flat[start:end].copy_(chunk_output.to(q.dtype)) + + return output + + +@torch.library.custom_op("auto_deploy::torch_deepseek_v4_sparse_attention", mutates_args=()) +def torch_deepseek_v4_sparse_attention( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + compressor_kv: torch.Tensor, + compressor_gate: torch.Tensor, + compressor_ape: torch.Tensor, + compressor_norm_weight: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + position_ids: torch.Tensor, + indexer_q: torch.Tensor, + indexer_weights: torch.Tensor, + indexer_compressor_kv: torch.Tensor, + indexer_compressor_gate: torch.Tensor, + indexer_compressor_ape: torch.Tensor, + indexer_compressor_norm_weight: torch.Tensor, + softmax_scale: float, + enable_sharding: bool = False, + layer_type: str = "mha_sparse", + layer_idx: Optional[int] = None, + window_size: Optional[int] = None, + compress_ratio: int = 0, + max_compressed_len: Optional[int] = None, + head_dim: Optional[int] = None, + rope_dim: Optional[int] = None, + rms_norm_eps: float = 1e-6, +) -> torch.Tensor: + """DeepSeek V4 sparse source op with explicit compressor projections.""" + del ( + indexer_q, + indexer_weights, + indexer_compressor_kv, + indexer_compressor_gate, + indexer_compressor_ape, + indexer_compressor_norm_weight, + enable_sharding, + layer_type, + layer_idx, + window_size, + head_dim, + ) + _validate_deepseek_v4_sparse_attention_inputs(q, kv, attn_sink, topk_idxs) + _validate_compress_ratio(compress_ratio) + if compress_ratio: + if max_compressed_len is None: + raise ValueError("max_compressed_len is required for compressed attention.") + if rope_dim is None: + raise ValueError("rope_dim is required for compressed attention.") + compressed_kv = _build_full_compressed_kv( + compressor_kv, + compressor_gate, + compressor_ape, + compressor_norm_weight, + cos_table, + sin_table, + position_ids, + rms_norm_eps, + rope_dim, + compress_ratio, + max_compressed_len, + ).to(kv.dtype) + kv = torch.cat((kv, compressed_kv), dim=1) + return _deepseek_v4_sparse_attention(q, kv, attn_sink, topk_idxs, softmax_scale) + + +@torch_deepseek_v4_sparse_attention.register_fake +def torch_deepseek_v4_sparse_attention_fake( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + compressor_kv: torch.Tensor, + compressor_gate: torch.Tensor, + compressor_ape: torch.Tensor, + compressor_norm_weight: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + position_ids: torch.Tensor, + indexer_q: torch.Tensor, + indexer_weights: torch.Tensor, + indexer_compressor_kv: torch.Tensor, + indexer_compressor_gate: torch.Tensor, + indexer_compressor_ape: torch.Tensor, + indexer_compressor_norm_weight: torch.Tensor, + softmax_scale: float, + enable_sharding: bool = False, + layer_type: str = "mha_sparse", + layer_idx: Optional[int] = None, + window_size: Optional[int] = None, + compress_ratio: int = 0, + max_compressed_len: Optional[int] = None, + head_dim: Optional[int] = None, + rope_dim: Optional[int] = None, + rms_norm_eps: float = 1e-6, +) -> torch.Tensor: + """Fake implementation for torch.export tracing.""" + del ( + softmax_scale, + enable_sharding, + layer_type, + layer_idx, + window_size, + compress_ratio, + max_compressed_len, + head_dim, + rope_dim, + rms_norm_eps, + ) + _validate_rank("q", q, 4) + _validate_rank("kv", kv, 3) + _validate_rank("attn_sink", attn_sink, 1) + _validate_rank("topk_idxs", topk_idxs, 3) + _validate_rank("compressor_kv", compressor_kv, 3) + _validate_rank("compressor_gate", compressor_gate, 3) + _validate_rank("compressor_ape", compressor_ape, 2) + _validate_rank("compressor_norm_weight", compressor_norm_weight, 1) + _validate_rank("cos_table", cos_table, 2) + _validate_rank("sin_table", sin_table, 2) + _validate_rank("position_ids", position_ids, 2) + _validate_rank("indexer_q", indexer_q, 4) + _validate_rank("indexer_weights", indexer_weights, 3) + _validate_rank("indexer_compressor_kv", indexer_compressor_kv, 3) + _validate_rank("indexer_compressor_gate", indexer_compressor_gate, 3) + _validate_rank("indexer_compressor_ape", indexer_compressor_ape, 2) + _validate_rank("indexer_compressor_norm_weight", indexer_compressor_norm_weight, 1) + return q.new_empty(q.shape).contiguous() + + +@torch.library.custom_op( + "auto_deploy::torch_deepseek_v4_sparse_attention_with_cache", + mutates_args=( + "swa_cache", + "mhc_cache", + "compressor_kv_cache", + "compressor_gate_cache", + "indexer_compressor_kv_cache", + "indexer_compressor_gate_cache", + ), +) +def torch_deepseek_v4_sparse_attention_with_cache( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + compressor_kv: torch.Tensor, + compressor_gate: torch.Tensor, + compressor_ape: torch.Tensor, + compressor_norm_weight: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + position_ids: torch.Tensor, + indexer_q: torch.Tensor, + indexer_weights: torch.Tensor, + indexer_compressor_kv: torch.Tensor, + indexer_compressor_gate: torch.Tensor, + indexer_compressor_ape: torch.Tensor, + indexer_compressor_norm_weight: torch.Tensor, + batch_info_host: torch.Tensor, + seq_len: torch.Tensor, + input_pos: torch.Tensor, + slot_idx: torch.Tensor, + cu_seqlen: torch.Tensor, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + last_page_len: torch.Tensor, + swa_cache: torch.Tensor, + mhc_cache: torch.Tensor, + compressor_kv_cache: torch.Tensor, + compressor_gate_cache: torch.Tensor, + indexer_compressor_kv_cache: torch.Tensor, + indexer_compressor_gate_cache: torch.Tensor, + softmax_scale: float, + window_size: Optional[int] = None, + compress_ratio: int = 0, + max_compressed_len: Optional[int] = None, + rms_norm_eps: float = 1e-6, + rope_dim: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Reference paged cached DeepSeek V4 sparse attention with compressor state.""" + _validate_deepseek_v4_sparse_attention_inputs(q, kv, attn_sink, topk_idxs) + _validate_swa_cache_inputs(q, kv, swa_cache) + _validate_swa_cache_inputs(q, kv, mhc_cache) + _validate_rank("compressor_kv_cache", compressor_kv_cache, 3) + _validate_rank("compressor_gate_cache", compressor_gate_cache, 3) + _validate_rank("indexer_q", indexer_q, 4) + _validate_rank("indexer_weights", indexer_weights, 3) + _validate_rank("indexer_compressor_kv", indexer_compressor_kv, 3) + _validate_rank("indexer_compressor_gate", indexer_compressor_gate, 3) + _validate_rank("indexer_compressor_ape", indexer_compressor_ape, 2) + _validate_rank("indexer_compressor_norm_weight", indexer_compressor_norm_weight, 1) + _validate_rank("indexer_compressor_kv_cache", indexer_compressor_kv_cache, 3) + _validate_rank("indexer_compressor_gate_cache", indexer_compressor_gate_cache, 3) + _validate_compress_ratio(compress_ratio) + if window_size is not None and window_size <= 0: + raise ValueError(f"window_size must be positive when provided, got {window_size}") + if compress_ratio: + if window_size is None: + raise ValueError("window_size is required for compressed cached attention.") + if max_compressed_len is None or max_compressed_len <= 0: + raise ValueError( + "max_compressed_len must be positive for compressed cached attention, " + f"got {max_compressed_len}" + ) + if rope_dim is None: + raise ValueError("rope_dim is required for compressed cached attention.") + + batch_info = BatchInfo(batch_info_host) + num_prefill, num_prefill_tokens, num_decode = batch_info.get_absorbed_info() + num_seq = num_prefill + num_decode + active_tokens = num_prefill_tokens + num_decode + q_flat = q.reshape(-1, *q.shape[2:]) + if active_tokens > q_flat.shape[0]: + raise ValueError( + f"active token count {active_tokens} exceeds flattened q tokens {q_flat.shape[0]}" + ) + if num_prefill == 0 and num_decode > 0: + return _deepseek_v4_sparse_attention_decode_with_cache( + q, + kv, + attn_sink, + topk_idxs, + compressor_kv, + compressor_gate, + compressor_ape, + compressor_norm_weight, + indexer_q, + indexer_weights, + indexer_compressor_kv, + indexer_compressor_gate, + indexer_compressor_ape, + indexer_compressor_norm_weight, + cos_table, + sin_table, + position_ids, + input_pos, + slot_idx, + cu_num_pages, + cache_loc, + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + num_decode, + softmax_scale, + window_size, + compress_ratio, + max_compressed_len, + rms_norm_eps, + rope_dim, + out, + ) + + seq_len_host = _to_host_long("seq_len", seq_len, num_seq) + input_pos_host = _to_host_long("input_pos", input_pos, num_seq) + cu_seqlen_host = _to_host_long("cu_seqlen", cu_seqlen, num_seq + 1) + cu_num_pages_host = _to_host_long("cu_num_pages", cu_num_pages, num_seq + 1) + num_page_entries = int(cu_num_pages_host[-1].item()) + cache_loc_host = _to_host_long("cache_loc", cache_loc, num_page_entries) + del slot_idx, last_page_len + + output_flat = torch.zeros_like(q_flat) + compressed_capacity = int(max_compressed_len) if compress_ratio else 0 + + for seq_idx in range(num_seq): + seq_len_i = int(seq_len_host[seq_idx].item()) + if seq_len_i == 0: + continue + flat_start = int(cu_seqlen_host[seq_idx].item()) + input_pos_i = int(input_pos_host[seq_idx].item()) + + q_seq = q_flat[flat_start : flat_start + seq_len_i] + kv_seq = _slice_sequence_tokens(kv, seq_idx, flat_start, seq_len_i) + topk_seq = _slice_sequence_tokens(topk_idxs, seq_idx, flat_start, seq_len_i) + position_ids_seq = _slice_sequence_positions(position_ids, seq_idx, flat_start, seq_len_i) + indexer_q_seq = _slice_sequence_tokens(indexer_q, seq_idx, flat_start, seq_len_i) + indexer_weights_seq = _slice_sequence_tokens( + indexer_weights, seq_idx, flat_start, seq_len_i + ) + if q_seq.shape[0] != seq_len_i: + raise ValueError( + f"Sequence {seq_idx} q slice has length {q_seq.shape[0]}, expected {seq_len_i}" + ) + if kv_seq.shape[0] != seq_len_i: + raise ValueError( + f"Sequence {seq_idx} kv slice has length {kv_seq.shape[0]}, expected {seq_len_i}" + ) + if topk_seq.shape[0] != seq_len_i: + raise ValueError( + f"Sequence {seq_idx} topk_idxs slice has length {topk_seq.shape[0]}, " + f"expected {seq_len_i}" + ) + + _write_paged_cache_rows( + kv_seq, swa_cache, seq_idx, input_pos_i, cu_num_pages_host, cache_loc_host + ) + + if compress_ratio: + compressor_kv_seq = _slice_sequence_tokens( + compressor_kv, seq_idx, flat_start, seq_len_i + ) + compressor_gate_seq = _slice_sequence_tokens( + compressor_gate, seq_idx, flat_start, seq_len_i + ) + indexer_compressor_kv_seq = _slice_sequence_tokens( + indexer_compressor_kv, seq_idx, flat_start, seq_len_i + ) + indexer_compressor_gate_seq = _slice_sequence_tokens( + indexer_compressor_gate, seq_idx, flat_start, seq_len_i + ) + _update_compressed_paged_caches( + compressor_kv_seq, + compressor_gate_seq, + position_ids_seq, + compressor_ape, + compressor_norm_weight, + cos_table, + sin_table, + seq_idx, + input_pos_i, + cu_num_pages_host, + cache_loc_host, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + rms_norm_eps, + rope_dim, + compress_ratio, + compressed_capacity, + ) + mode = _compression_mode(compress_ratio) + if mode.uses_indexer: + _update_raw_paged_caches( + indexer_compressor_kv_seq, + indexer_compressor_gate_seq, + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + seq_idx, + input_pos_i, + cu_num_pages_host, + cache_loc_host, + ) + + if input_pos_i == 0: + output_flat[flat_start : flat_start + seq_len_i] = ( + torch_deepseek_v4_sparse_attention( + q_seq.unsqueeze(0), + kv_seq.unsqueeze(0), + attn_sink, + topk_seq.unsqueeze(0), + compressor_kv_seq.unsqueeze(0), + compressor_gate_seq.unsqueeze(0), + compressor_ape, + compressor_norm_weight, + cos_table, + sin_table, + position_ids_seq, + indexer_q_seq.unsqueeze(0), + indexer_weights_seq.unsqueeze(0), + indexer_compressor_kv_seq.unsqueeze(0), + indexer_compressor_gate_seq.unsqueeze(0), + indexer_compressor_ape, + indexer_compressor_norm_weight, + softmax_scale, + window_size=window_size, + compress_ratio=compress_ratio, + max_compressed_len=max_compressed_len, + rope_dim=rope_dim, + rms_norm_eps=rms_norm_eps, + ).squeeze(0) + ) + else: + output_flat[flat_start : flat_start + seq_len_i] = _cached_compressed_attention( + q_seq, + attn_sink, + swa_cache, + mhc_cache, + seq_idx, + input_pos_i, + cu_num_pages_host, + cache_loc_host, + position_ids_seq, + window_size, + compress_ratio, + compressed_capacity, + softmax_scale, + topk_seq=topk_seq, + indexer_q_seq=indexer_q_seq, + indexer_weights_seq=indexer_weights_seq, + indexer_compressor_kv_cache=indexer_compressor_kv_cache, + indexer_compressor_gate_cache=indexer_compressor_gate_cache, + indexer_compressor_ape=indexer_compressor_ape, + indexer_compressor_norm_weight=indexer_compressor_norm_weight, + cos_table=cos_table, + sin_table=sin_table, + rms_norm_eps=rms_norm_eps, + rope_dim=rope_dim, + ) + elif input_pos_i == 0: + kv_source = _prefill_kv_source(kv, kv_seq, seq_idx, num_seq) + output_flat[flat_start : flat_start + seq_len_i] = _deepseek_v4_sparse_attention( + q_seq.unsqueeze(0), + kv_source, + attn_sink, + topk_seq.unsqueeze(0), + softmax_scale, + ).squeeze(0) + elif window_size is not None: + output_flat[flat_start : flat_start + seq_len_i] = _cached_local_window_attention( + q_seq, + attn_sink, + swa_cache, + seq_idx, + input_pos_i, + cu_num_pages_host, + cache_loc_host, + window_size, + softmax_scale, + ) + else: + output_flat[flat_start : flat_start + seq_len_i] = _cached_topk_attention( + q_seq, + attn_sink, + topk_seq, + swa_cache, + seq_idx, + cu_num_pages_host, + cache_loc_host, + softmax_scale, + ) + + if active_tokens < q_flat.shape[0]: + output_flat[active_tokens:].zero_() + output = output_flat.view_as(q) + if out is not None: + out.copy_(output) + return out.new_empty(0) + return output + + +@torch_deepseek_v4_sparse_attention_with_cache.register_fake +def torch_deepseek_v4_sparse_attention_with_cache_fake( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + compressor_kv: torch.Tensor, + compressor_gate: torch.Tensor, + compressor_ape: torch.Tensor, + compressor_norm_weight: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + position_ids: torch.Tensor, + indexer_q: torch.Tensor, + indexer_weights: torch.Tensor, + indexer_compressor_kv: torch.Tensor, + indexer_compressor_gate: torch.Tensor, + indexer_compressor_ape: torch.Tensor, + indexer_compressor_norm_weight: torch.Tensor, + batch_info_host: torch.Tensor, + seq_len: torch.Tensor, + input_pos: torch.Tensor, + slot_idx: torch.Tensor, + cu_seqlen: torch.Tensor, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + last_page_len: torch.Tensor, + swa_cache: torch.Tensor, + mhc_cache: torch.Tensor, + compressor_kv_cache: torch.Tensor, + compressor_gate_cache: torch.Tensor, + indexer_compressor_kv_cache: torch.Tensor, + indexer_compressor_gate_cache: torch.Tensor, + softmax_scale: float, + window_size: Optional[int] = None, + compress_ratio: int = 0, + max_compressed_len: Optional[int] = None, + rms_norm_eps: float = 1e-6, + rope_dim: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if out is not None: + return out.new_empty(0) + _validate_compress_ratio(compress_ratio) + _validate_rank("q", q, 4) + _validate_rank("kv", kv, 3) + _validate_rank("attn_sink", attn_sink, 1) + _validate_rank("topk_idxs", topk_idxs, 3) + _validate_rank("compressor_kv", compressor_kv, 3) + _validate_rank("compressor_gate", compressor_gate, 3) + _validate_rank("compressor_ape", compressor_ape, 2) + _validate_rank("compressor_norm_weight", compressor_norm_weight, 1) + _validate_rank("cos_table", cos_table, 2) + _validate_rank("sin_table", sin_table, 2) + _validate_rank("position_ids", position_ids, 2) + _validate_rank("indexer_q", indexer_q, 4) + _validate_rank("indexer_weights", indexer_weights, 3) + _validate_rank("indexer_compressor_kv", indexer_compressor_kv, 3) + _validate_rank("indexer_compressor_gate", indexer_compressor_gate, 3) + _validate_rank("indexer_compressor_ape", indexer_compressor_ape, 2) + _validate_rank("indexer_compressor_norm_weight", indexer_compressor_norm_weight, 1) + _validate_rank("swa_cache", swa_cache, 3) + _validate_rank("mhc_cache", mhc_cache, 3) + _validate_rank("compressor_kv_cache", compressor_kv_cache, 3) + _validate_rank("compressor_gate_cache", compressor_gate_cache, 3) + _validate_rank("indexer_compressor_kv_cache", indexer_compressor_kv_cache, 3) + _validate_rank("indexer_compressor_gate_cache", indexer_compressor_gate_cache, 3) + del ( + kv, + attn_sink, + topk_idxs, + compressor_kv, + compressor_gate, + compressor_ape, + compressor_norm_weight, + cos_table, + sin_table, + position_ids, + indexer_q, + indexer_weights, + indexer_compressor_kv, + indexer_compressor_gate, + indexer_compressor_ape, + indexer_compressor_norm_weight, + batch_info_host, + seq_len, + input_pos, + slot_idx, + cu_seqlen, + cu_num_pages, + cache_loc, + last_page_len, + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + softmax_scale, + window_size, + compress_ratio, + max_compressed_len, + rms_norm_eps, + rope_dim, + out, + ) + return q.new_empty(q.shape).contiguous() + + +@torch.library.custom_op( + "auto_deploy::flashmla_deepseek_v4_sparse_attention_with_cache", + mutates_args=( + "swa_cache", + "mhc_cache", + "compressor_kv_cache", + "compressor_gate_cache", + "indexer_compressor_kv_cache", + "indexer_compressor_gate_cache", + ), +) +def flashmla_deepseek_v4_sparse_attention_with_cache( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + compressor_kv: torch.Tensor, + compressor_gate: torch.Tensor, + compressor_ape: torch.Tensor, + compressor_norm_weight: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + position_ids: torch.Tensor, + indexer_q: torch.Tensor, + indexer_weights: torch.Tensor, + indexer_compressor_kv: torch.Tensor, + indexer_compressor_gate: torch.Tensor, + indexer_compressor_ape: torch.Tensor, + indexer_compressor_norm_weight: torch.Tensor, + batch_info_host: torch.Tensor, + seq_len: torch.Tensor, + input_pos: torch.Tensor, + slot_idx: torch.Tensor, + cu_seqlen: torch.Tensor, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + last_page_len: torch.Tensor, + swa_cache: torch.Tensor, + mhc_cache: torch.Tensor, + compressor_kv_cache: torch.Tensor, + compressor_gate_cache: torch.Tensor, + indexer_compressor_kv_cache: torch.Tensor, + indexer_compressor_gate_cache: torch.Tensor, + softmax_scale: float, + window_size: Optional[int] = None, + compress_ratio: int = 0, + max_compressed_len: Optional[int] = None, + rms_norm_eps: float = 1e-6, + rope_dim: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """FlashMLA-backed paged cached DeepSeek V4 sparse attention. + + This backend op is selected by AD graph lowering. It intentionally does not + fall back to the torch reference path: unsupported batch shapes should fail + at the backend boundary so the semantic source op remains backend-neutral. + """ + del slot_idx, last_page_len + _validate_deepseek_v4_sparse_attention_inputs(q, kv, attn_sink, topk_idxs) + _validate_flash_mla_model1_cache_inputs(q, kv, "swa_cache", swa_cache) + _validate_flash_mla_model1_cache_inputs(q, kv, "mhc_cache", mhc_cache) + _validate_rank("compressor_kv_cache", compressor_kv_cache, 3) + _validate_rank("compressor_gate_cache", compressor_gate_cache, 3) + _validate_rank("indexer_compressor_kv", indexer_compressor_kv, 3) + _validate_rank("indexer_compressor_gate", indexer_compressor_gate, 3) + _validate_rank("indexer_compressor_ape", indexer_compressor_ape, 2) + _validate_rank("indexer_compressor_norm_weight", indexer_compressor_norm_weight, 1) + _validate_rank("indexer_compressor_kv_cache", indexer_compressor_kv_cache, 3) + _validate_rank("indexer_compressor_gate_cache", indexer_compressor_gate_cache, 3) + _validate_compress_ratio(compress_ratio) + if window_size is not None and window_size <= 0: + raise ValueError(f"window_size must be positive when provided, got {window_size}") + if compress_ratio: + if window_size is None: + raise ValueError("window_size is required for compressed cached attention.") + if max_compressed_len is None or max_compressed_len <= 0: + raise ValueError( + "max_compressed_len must be positive for compressed cached attention, " + f"got {max_compressed_len}" + ) + if rope_dim is None: + raise ValueError("rope_dim is required for compressed cached attention.") + + batch_info = BatchInfo(batch_info_host) + num_prefill, num_prefill_tokens, num_decode = batch_info.get_absorbed_info() + num_seq = num_prefill + num_decode + active_tokens = num_prefill_tokens + num_decode + q_flat = q.reshape(-1, *q.shape[2:]) + if active_tokens > q_flat.shape[0]: + raise ValueError( + f"active token count {active_tokens} exceeds flattened q tokens {q_flat.shape[0]}" + ) + if num_prefill == 0 and num_decode > 0: + return _flash_mla_deepseek_v4_sparse_attention_decode_with_cache( + q, + kv, + attn_sink, + topk_idxs, + compressor_kv, + compressor_gate, + compressor_ape, + compressor_norm_weight, + indexer_q, + indexer_weights, + indexer_compressor_kv, + indexer_compressor_gate, + indexer_compressor_ape, + indexer_compressor_norm_weight, + cos_table, + sin_table, + position_ids, + input_pos, + cu_num_pages, + cache_loc, + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + num_decode, + softmax_scale, + window_size, + compress_ratio, + max_compressed_len, + rms_norm_eps, + rope_dim, + out, + ) + + seq_len_host = _to_host_long("seq_len", seq_len, num_seq) + input_pos_host = _to_host_long("input_pos", input_pos, num_seq) + cu_seqlen_host = _to_host_long("cu_seqlen", cu_seqlen, num_seq + 1) + cu_num_pages_host = _to_host_long("cu_num_pages", cu_num_pages, num_seq + 1) + num_page_entries = int(cu_num_pages_host[-1].item()) + cache_loc_host = _to_host_long("cache_loc", cache_loc, num_page_entries) + + output_flat = torch.zeros_like(q_flat) + compressed_capacity = int(max_compressed_len) if compress_ratio else 0 + + for seq_idx in range(num_prefill): + seq_len_i = int(seq_len_host[seq_idx].item()) + if seq_len_i == 0: + continue + input_pos_i = int(input_pos_host[seq_idx].item()) + if input_pos_i != 0: + raise RuntimeError( + "FlashMLA DeepSeek V4 sparse backend currently supports fresh prefill " + f"spans only; sequence {seq_idx} starts at input_pos={input_pos_i}." + ) + + flat_start = int(cu_seqlen_host[seq_idx].item()) + q_seq = q_flat[flat_start : flat_start + seq_len_i] + kv_seq = _slice_sequence_tokens(kv, seq_idx, flat_start, seq_len_i) + topk_seq = _slice_sequence_tokens(topk_idxs, seq_idx, flat_start, seq_len_i) + position_ids_seq = _slice_sequence_positions(position_ids, seq_idx, flat_start, seq_len_i) + if q_seq.shape[0] != seq_len_i: + raise ValueError( + f"Sequence {seq_idx} q slice has length {q_seq.shape[0]}, expected {seq_len_i}" + ) + if kv_seq.shape[0] != seq_len_i: + raise ValueError( + f"Sequence {seq_idx} kv slice has length {kv_seq.shape[0]}, expected {seq_len_i}" + ) + if topk_seq.shape[0] != seq_len_i: + raise ValueError( + f"Sequence {seq_idx} topk_idxs slice has length {topk_seq.shape[0]}, " + f"expected {seq_len_i}" + ) + + _write_paged_cache_rows( + kv_seq, swa_cache, seq_idx, input_pos_i, cu_num_pages_host, cache_loc_host + ) + + if compress_ratio: + assert max_compressed_len is not None + assert rope_dim is not None + compressor_kv_seq = _slice_sequence_tokens( + compressor_kv, seq_idx, flat_start, seq_len_i + ) + compressor_gate_seq = _slice_sequence_tokens( + compressor_gate, seq_idx, flat_start, seq_len_i + ) + indexer_compressor_kv_seq = _slice_sequence_tokens( + indexer_compressor_kv, seq_idx, flat_start, seq_len_i + ) + indexer_compressor_gate_seq = _slice_sequence_tokens( + indexer_compressor_gate, seq_idx, flat_start, seq_len_i + ) + _update_compressed_paged_caches( + compressor_kv_seq, + compressor_gate_seq, + position_ids_seq, + compressor_ape, + compressor_norm_weight, + cos_table, + sin_table, + seq_idx, + input_pos_i, + cu_num_pages_host, + cache_loc_host, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + rms_norm_eps, + rope_dim, + compress_ratio, + compressed_capacity, + ) + mode = _compression_mode(compress_ratio) + if mode.uses_indexer: + _update_raw_paged_caches( + indexer_compressor_kv_seq, + indexer_compressor_gate_seq, + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + seq_idx, + input_pos_i, + cu_num_pages_host, + cache_loc_host, + ) + + compressed_kv = _build_full_compressed_kv( + compressor_kv_seq.unsqueeze(0), + compressor_gate_seq.unsqueeze(0), + compressor_ape, + compressor_norm_weight, + cos_table, + sin_table, + position_ids_seq, + rms_norm_eps, + rope_dim, + compress_ratio, + max_compressed_len, + ).to(kv.dtype) + kv_source = torch.cat((kv_seq.unsqueeze(0), compressed_kv), dim=1) + else: + kv_source = _prefill_kv_source(kv, kv_seq, seq_idx, num_seq) + + output_flat[flat_start : flat_start + seq_len_i] = _flash_mla_sparse_prefill( + q_seq.unsqueeze(0), + kv_source, + attn_sink, + topk_seq.unsqueeze(0), + softmax_scale, + ).squeeze(0) + + if num_decode > 0: + decode_output = _flash_mla_deepseek_v4_sparse_attention_decode_with_cache( + q, + kv, + attn_sink, + topk_idxs, + compressor_kv, + compressor_gate, + compressor_ape, + compressor_norm_weight, + indexer_q, + indexer_weights, + indexer_compressor_kv, + indexer_compressor_gate, + indexer_compressor_ape, + indexer_compressor_norm_weight, + cos_table, + sin_table, + position_ids, + input_pos, + cu_num_pages, + cache_loc, + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + num_decode, + softmax_scale, + window_size, + compress_ratio, + max_compressed_len, + rms_norm_eps, + rope_dim, + None, + token_start=num_prefill_tokens, + seq_start=num_prefill, + ) + output_flat[num_prefill_tokens:active_tokens].copy_( + decode_output.reshape(-1, *q.shape[2:])[num_prefill_tokens:active_tokens] + ) + + if active_tokens < q_flat.shape[0]: + output_flat[active_tokens:].zero_() + output = output_flat.view_as(q) + if out is not None: + out.copy_(output) + return out.new_empty(0) + return output + + +@flashmla_deepseek_v4_sparse_attention_with_cache.register_fake +def flashmla_deepseek_v4_sparse_attention_with_cache_fake( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + compressor_kv: torch.Tensor, + compressor_gate: torch.Tensor, + compressor_ape: torch.Tensor, + compressor_norm_weight: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + position_ids: torch.Tensor, + indexer_q: torch.Tensor, + indexer_weights: torch.Tensor, + indexer_compressor_kv: torch.Tensor, + indexer_compressor_gate: torch.Tensor, + indexer_compressor_ape: torch.Tensor, + indexer_compressor_norm_weight: torch.Tensor, + batch_info_host: torch.Tensor, + seq_len: torch.Tensor, + input_pos: torch.Tensor, + slot_idx: torch.Tensor, + cu_seqlen: torch.Tensor, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + last_page_len: torch.Tensor, + swa_cache: torch.Tensor, + mhc_cache: torch.Tensor, + compressor_kv_cache: torch.Tensor, + compressor_gate_cache: torch.Tensor, + indexer_compressor_kv_cache: torch.Tensor, + indexer_compressor_gate_cache: torch.Tensor, + softmax_scale: float, + window_size: Optional[int] = None, + compress_ratio: int = 0, + max_compressed_len: Optional[int] = None, + rms_norm_eps: float = 1e-6, + rope_dim: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch_deepseek_v4_sparse_attention_with_cache_fake( + q, + kv, + attn_sink, + topk_idxs, + compressor_kv, + compressor_gate, + compressor_ape, + compressor_norm_weight, + cos_table, + sin_table, + position_ids, + indexer_q, + indexer_weights, + indexer_compressor_kv, + indexer_compressor_gate, + indexer_compressor_ape, + indexer_compressor_norm_weight, + batch_info_host, + seq_len, + input_pos, + slot_idx, + cu_seqlen, + cu_num_pages, + cache_loc, + last_page_len, + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + softmax_scale, + window_size, + compress_ratio, + max_compressed_len, + rms_norm_eps, + rope_dim, + out, + ) + + +@AttentionRegistry.register("deepseek_v4_sparse") +class DeepSeekV4SparseAttention(AttentionDescriptor): + """Cached DeepSeek V4 sparse attention descriptor for reference validation.""" + + @classmethod + def get_attention_layout(cls) -> AttentionLayout: + return "bsnd" + + @classmethod + def get_num_qkv_args(cls) -> int: + return len(_SOURCE_TENSOR_ARG_NAMES) + + @classmethod + def get_source_attention_op(cls) -> OpOverloadPacket: + return torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention + + @classmethod + def get_cached_attention_op(cls) -> MHACallable: + return torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention_with_cache.default + + @classmethod + def get_standard_metadata_args(cls) -> list[str]: + return [ + "batch_info_host", + "seq_len", + "input_pos", + "slot_idx", + "cu_seqlen", + "cu_num_pages", + "cache_loc", + "last_page_len", + ] + + @classmethod + def get_cache_initializers( + cls, source_attn_node: Node, cache_config: KvCacheConfig + ) -> ResourceHandlerDict: + kv_node, compressor_kv_node, indexer_compressor_kv_node = extract_op_args( + source_attn_node, + "kv", + "compressor_kv", + "indexer_compressor_kv", + ) + kv_fake: FakeTensor = kv_node.meta["val"] + head_dim = int(kv_fake.shape[-1]) + compressor_kv_fake: FakeTensor = compressor_kv_node.meta["val"] + compressor_state_dim = int(compressor_kv_fake.shape[-1]) + if compressor_state_dim <= 0: + compressor_state_dim = head_dim + indexer_compressor_kv_fake: FakeTensor = indexer_compressor_kv_node.meta["val"] + indexer_compressor_state_dim = int(indexer_compressor_kv_fake.shape[-1]) + dtype = cls.resolve_cache_dtype(cache_config.dtype, kv_fake.dtype) + return { + "swa_cache": PagedResourceHandler( + head_dim, + dtype=dtype, + ), + "mhc_cache": PagedResourceHandler(head_dim, dtype=dtype), + "compressor_kv_cache": PagedResourceHandler( + compressor_state_dim, + dtype=torch.float32, + ), + "compressor_gate_cache": PagedResourceHandler( + compressor_state_dim, + dtype=torch.float32, + ), + "indexer_compressor_kv_cache": PagedResourceHandler( + indexer_compressor_state_dim, + dtype=torch.float32, + ), + "indexer_compressor_gate_cache": PagedResourceHandler( + indexer_compressor_state_dim, + dtype=torch.float32, + ), + } + + @classmethod + def get_constants(cls, source_attn_node: Node) -> list[Constant]: + softmax_scale, window_size, compress_ratio, max_compressed_len, rms_norm_eps, rope_dim = ( + extract_op_args( + source_attn_node, + "softmax_scale", + "window_size", + "compress_ratio", + "max_compressed_len", + "rms_norm_eps", + "rope_dim", + ) + ) + if not isinstance(softmax_scale, float): + raise RuntimeError( + "DeepSeek V4 sparse attention source node must carry a literal " + f"float softmax_scale, got {softmax_scale!r}." + ) + if window_size is not None and not isinstance(window_size, int): + raise RuntimeError( + "DeepSeek V4 sparse attention source node must carry a literal " + f"int window_size or None, got {window_size!r}." + ) + if not isinstance(compress_ratio, int): + raise RuntimeError( + "DeepSeek V4 sparse attention source node must carry a literal " + f"int compress_ratio, got {compress_ratio!r}." + ) + try: + _compression_mode(compress_ratio) + except ValueError as exc: + raise RuntimeError( + "DeepSeek V4 sparse attention cache insertion supports " + f"compress_ratio in {_SUPPORTED_COMPRESS_RATIOS}, got {compress_ratio}." + ) from exc + if max_compressed_len is not None and not isinstance(max_compressed_len, int): + raise RuntimeError( + "DeepSeek V4 sparse attention source node must carry a literal " + f"int max_compressed_len or None, got {max_compressed_len!r}." + ) + if not isinstance(rms_norm_eps, float): + raise RuntimeError( + "DeepSeek V4 sparse attention source node must carry a literal " + f"float rms_norm_eps, got {rms_norm_eps!r}." + ) + if rope_dim is not None and not isinstance(rope_dim, int): + raise RuntimeError( + "DeepSeek V4 sparse attention source node must carry a literal " + f"int rope_dim or None, got {rope_dim!r}." + ) + return [ + softmax_scale, + window_size, + compress_ratio, + max_compressed_len, + rms_norm_eps, + rope_dim, + ] + + +@AttentionRegistry.register("flashmla_deepseek_v4_sparse") +class FlashMLADeepSeekV4SparseAttention(DeepSeekV4SparseAttention): + """DeepSeek V4 sparse attention descriptor lowered to the FlashMLA backend op.""" + + @classmethod + def get_cached_attention_op(cls) -> MHACallable: + return torch.ops.auto_deploy.flashmla_deepseek_v4_sparse_attention_with_cache.default + + @classmethod + def get_cache_initializers( + cls, source_attn_node: Node, cache_config: KvCacheConfig + ) -> ResourceHandlerDict: + handlers = super().get_cache_initializers(source_attn_node, cache_config) + kv_node = extract_op_args(source_attn_node, "kv")[0] + kv_fake: FakeTensor = kv_node.meta["val"] + head_dim = int(kv_fake.shape[-1]) + if head_dim != _FLASH_MLA_MODEL1_HEAD_DIM: + raise RuntimeError( + "FlashMLA DeepSeek V4 sparse backend requires kv head dimension " + f"{_FLASH_MLA_MODEL1_HEAD_DIM}, got {head_dim}." + ) + handlers["swa_cache"] = PagedResourceHandler( + _FLASH_MLA_MODEL1_BYTES_PER_TOKEN, + dtype=torch.uint8, + ) + handlers["mhc_cache"] = PagedResourceHandler( + _FLASH_MLA_MODEL1_BYTES_PER_TOKEN, + dtype=torch.uint8, + ) + return handlers diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_attention.py index cb2ec53a0125..e2b8c8216858 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_attention.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_attention.py @@ -218,7 +218,7 @@ def torch_attention( sinks = torch.exp(sinks_expanded - logits_max) unnormalized_scores = torch.exp(attn_scores - logits_max) normalizer = unnormalized_scores.sum(dim=-1, keepdim=True) + sinks - scores = unnormalized_scores / normalizer + scores = (unnormalized_scores / normalizer).to(value_t.dtype) # Use only the non-sink portion for computing output # We added exactly 1 column, so remove exactly 1 column attn_out = torch.matmul(scores, value_t) # [b, n_heads, s_q, v_head_dim] 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 e789a23236ea..b2022f3f2ed7 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py @@ -2051,6 +2051,50 @@ def allocate(self, sequence_info: SequenceInfo) -> torch.Tensor: raise ValueError(f"Invalid kv_layout: {self.kv_layout}") +class PagedResourceHandler(ResourceHandler): + """Handler for local per-token paged resources. + + These resources have shape ``[num_blocks, tokens_per_block, *token_shape]``. They + participate in paged cache memory sizing, but are allocated locally instead of being + managed as a standard K/V pair by ``KVCacheManager.get_buffers()``. + """ + + @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 PagedResourceHandler. + + 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 __eq__(self, other: Optional[ResourceHandler]) -> bool: + """Check compatibility for locally allocated paged resources.""" + if type(other) is not type(self): + return False + return self.token_shape == other.token_shape and self.dtype == other.dtype + + def _get_bytes_per_token(self) -> int: + """The size of the resource per token in bytes.""" + return math.prod(self.token_shape) * self.dtype.itemsize + + def allocate(self, sequence_info: SequenceInfo) -> torch.Tensor: + """Allocate a local paged resource for the given sequence info.""" + return torch.empty( + sequence_info.num_blocks, + sequence_info.tokens_per_block, + *self.token_shape, + device=sequence_info.device, + dtype=self.dtype, + ) + + class StateResourceHandler(ResourceHandler): """Handler for per-sequence state resources (e.g., Mamba SSM/conv states). diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py index 909e2eac91f9..afeddb79f130 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py @@ -13,34 +13,99 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Triton-kernels-based MXFP4 MoE ops (GPT-OSS style) with routing, swizzling, and fused activation +# MXFP4 MoE ops with packed expert decode and SwiGLU activation. import weakref from collections import OrderedDict from collections.abc import Callable +from dataclasses import dataclass import torch import torch.nn.functional as F -from triton_kernels.matmul_ogs import ( - FlexCtx, - FnSpecs, - FusedActivation, - GatherIndx, - PrecisionConfig, - RoutingData, - ScatterIndx, - matmul_ogs, -) -from triton_kernels.numerics import InFlexData -from triton_kernels.swiglu import swiglu_fn -from triton_kernels.target_info import cuda_capability_geq -from triton_kernels.tensor import FP4, Tensor, convert_layout, wrap_torch_tensor -from triton_kernels.tensor_details import layout -from triton_kernels.tensor_details.layout import StridedLayout -from tensorrt_llm._torch.modules.fused_moe.fused_moe_triton import TritonEPRouter +# NOTE: ``triton_kernels`` is an optional dependency. The torch-reference MXFP4 ops +# (``torch_mxfp4_moe`` etc.) must remain importable and usable without it, so all +# ``triton_kernels`` symbols are imported lazily inside the functions that need them. +try: + import triton + import triton.language as tl +except ImportError: + triton = None + tl = None + +if triton is not None: + + @triton.jit + def _fill_int_kernel(tensor, n_elements, value, BLOCK: tl.constexpr): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + tl.store(tensor + offsets, value, mask=offsets < n_elements) + + @triton.jit + def _routing_col_sum_kernel(expert_indices, col_sum, n_indices, BLOCK: tl.constexpr): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + cols = tl.load(expert_indices + offsets, mask=offsets < n_indices, other=-1) + valid = (offsets < n_indices) & (cols >= 0) + tl.atomic_add(col_sum + cols, 1, sem="relaxed", mask=valid) + + @triton.jit + def _routing_prefix_sum_kernel( + col_sum, + col_offsets, + col_counters, + n_cols, + BLOCK: tl.constexpr, + ): + offsets = tl.arange(0, BLOCK) + values = tl.load(col_sum + offsets, mask=offsets < n_cols, other=0) + offsets_exclusive = tl.cumsum(values, 0) - values + tl.store(col_offsets + offsets, offsets_exclusive, mask=offsets < n_cols) + tl.store(col_counters + offsets, offsets_exclusive, mask=offsets < n_cols) + + @triton.jit + def _routing_sorted_indices_kernel( + expert_indices, + col_counters, + col_sorted_indx, + row_sorted_indx, + n_indices, + BLOCK: tl.constexpr, + ): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + cols = tl.load(expert_indices + offsets, mask=offsets < n_indices, other=-1) + valid = (offsets < n_indices) & (cols >= 0) + sorted_offsets = tl.atomic_add(col_counters + cols, 1, sem="relaxed", mask=valid) + tl.store(col_sorted_indx + sorted_offsets, offsets, mask=valid) + tl.store(row_sorted_indx + offsets, sorted_offsets, mask=valid) + -PreparedWeights = tuple[Tensor, Tensor, Tensor, Tensor] +_E2M1_VALUES = ( + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, +) +_E2M1_PACKED_VALUES = tuple( + (_E2M1_VALUES[byte & 0x0F], _E2M1_VALUES[(byte >> 4) & 0x0F]) for byte in range(256) +) +_E8M0_EXPONENT_BIAS = 127 +_MXFP4_BLOCK_SIZE = 32 +_TORCH_MXFP4_ROUTED_MOE_TOKEN_CHUNK = 16 + +# Prepared (swizzled) triton_kernels tensors; typed as ``object`` so the module +# imports without ``triton_kernels``. +PreparedWeights = tuple[object, object, object, object] TensorCacheKey = tuple[ str, str, @@ -53,11 +118,129 @@ ] WeightCacheKey = tuple[object, ...] + +@dataclass +class _RoutingSparseMatrix: + vals: torch.Tensor + indx: torch.Tensor + mask: object + mask_metadata: object + + +def _is_power_of_two(value: int) -> bool: + return value > 0 and (value & (value - 1)) == 0 + + +def _next_power_of_two(value: int) -> int: + return 1 << (value - 1).bit_length() + + +def _make_routing_metadata_from_indices(expert_indices: torch.Tensor, num_experts: int) -> object: + if triton is None: + raise ImportError("triton is required for Triton MXFP4 routed MoE metadata.") + + from triton_kernels.tensor_details.bitmatrix import BitmatrixMetadata + + expert_indices = expert_indices.to(torch.int32).contiguous() + n_indices = expert_indices.numel() + device = expert_indices.device + col_sum = torch.empty((num_experts,), dtype=torch.int32, device=device) + col_offsets = torch.empty((num_experts,), dtype=torch.int32, device=device) + col_counters = torch.empty((num_experts,), dtype=torch.int32, device=device) + col_sorted_indx = torch.empty((n_indices,), dtype=torch.int32, device=device) + row_sorted_indx = torch.empty((n_indices,), dtype=torch.int32, device=device) + + block_indices = 256 + block_cols = _next_power_of_two(num_experts) + _fill_int_kernel[(triton.cdiv(num_experts, block_indices),)]( + col_sum, + num_experts, + 0, + BLOCK=block_indices, + ) + _fill_int_kernel[(triton.cdiv(n_indices, block_indices),)]( + col_sorted_indx, + n_indices, + -1, + BLOCK=block_indices, + ) + _fill_int_kernel[(triton.cdiv(n_indices, block_indices),)]( + row_sorted_indx, + n_indices, + -1, + BLOCK=block_indices, + ) + _routing_col_sum_kernel[(triton.cdiv(n_indices, block_indices),)]( + expert_indices, + col_sum, + n_indices, + BLOCK=block_indices, + ) + _routing_prefix_sum_kernel[(1,)]( + col_sum, + col_offsets, + col_counters, + num_experts, + BLOCK=block_cols, + ) + _routing_sorted_indices_kernel[(triton.cdiv(n_indices, block_indices),)]( + expert_indices, + col_counters, + col_sorted_indx, + row_sorted_indx, + n_indices, + BLOCK=block_indices, + ) + return BitmatrixMetadata( + col_sum=col_sum, + col_sorted_indx=col_sorted_indx, + row_sorted_indx=row_sorted_indx, + ) + + +def _make_routing_bitmatrix( + selected: torch.Tensor, + num_experts: int, + valid_mask: torch.Tensor | None = None, +) -> object: + from triton_kernels.tensor import BIT, Bitmatrix + + num_tokens, top_k = selected.shape + + def cdiv(x, y): + return (x + y - 1) // y + + bitmatrix_data = torch.zeros( + (cdiv(num_experts, 32), cdiv(num_tokens, 32) * 32), + dtype=torch.int32, + device=selected.device, + ) + bitmatrix_data = bitmatrix_data.t()[:num_tokens] + rows = ( + torch.arange(num_tokens, device=selected.device).unsqueeze(1).expand(-1, top_k).reshape(-1) + ) + cols = selected.reshape(-1) + if valid_mask is not None: + flat_valid_mask = valid_mask.reshape(-1) + rows = rows[flat_valid_mask] + cols = cols[flat_valid_mask] + word_idx = torch.div(cols, 32, rounding_mode="floor") + bit_idx = cols % 32 + masks = torch.ones_like(bit_idx, dtype=torch.int32) << bit_idx + bitmatrix_data.index_put_((rows, word_idx), masks, accumulate=True) + return Bitmatrix( + bitmatrix_data.view(torch.uint32), + dtype=BIT, + shape=[num_tokens, num_experts], + ) + + # ``convert_layout`` swizzles the packed expert weights. Cache the result by # underlying storage/view metadata so decode steps do not repeat that work. _MXFP4_WEIGHT_CACHE: OrderedDict[WeightCacheKey, tuple[PreparedWeights, list[weakref.finalize]]] = ( OrderedDict() ) +_E2M1_PACKED_TABLE_CACHE: dict[tuple[str, torch.dtype], torch.Tensor] = {} # copied from transformers.integrations.mxfp4::swizzle_mxfp4 with minor modification @@ -65,19 +248,28 @@ def _mxfp4_value_layout(mx_axis: int): # Blackwell's default value layout is only supported by the persistent TMA # kernel. GPT-OSS MoE can select the non-persistent kernel for small shapes, # where unswizzled values use the native MXFP4 dot_scaled path. + from triton_kernels.target_info import cuda_capability_geq + from triton_kernels.tensor_details import layout + from triton_kernels.tensor_details.layout import StridedLayout + if cuda_capability_geq(10): return StridedLayout, {} return layout.make_default_matmul_mxfp4_w_layout(mx_axis=mx_axis) def _swizzle_mxfp4(w, w_scale): + from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor + from triton_kernels.tensor_details.layout import StridedLayout + value_layout, value_layout_opts = _mxfp4_value_layout(mx_axis=1) w = convert_layout(wrap_torch_tensor(w, dtype=FP4), value_layout, **value_layout_opts) w_scale = convert_layout(wrap_torch_tensor(w_scale), StridedLayout) return w, w_scale -RouteFn = Callable[[torch.Tensor], tuple[RoutingData, GatherIndx, ScatterIndx]] +# route_fn produces ``(routing_data, gather_idx, scatter_idx)`` from triton_kernels; +# typed as ``object`` so the module imports without ``triton_kernels``. +RouteFn = Callable[[torch.Tensor], tuple[object, object, object]] def _mxfp4_layout_cache_key() -> tuple[object, ...]: @@ -153,6 +345,264 @@ def _register_cache_finalizers( return finalizers +def _as_uint8(tensor: torch.Tensor, name: str) -> torch.Tensor: + if tensor.dtype == torch.uint8: + return tensor + if tensor.dtype == torch.int8: + return tensor.view(torch.uint8) + raise TypeError(f"{name} should contain packed MXFP4 bytes with dtype uint8 or int8.") + + +def _decode_e8m0_scales(scales: torch.Tensor) -> torch.Tensor: + scales_u8 = _as_uint8(scales, "scales") + exponents = scales_u8.to(torch.int32) - _E8M0_EXPONENT_BIAS + return torch.ldexp(torch.ones_like(exponents, dtype=torch.float32), exponents) + + +def _get_e2m1_packed_table(device: torch.device, dtype: torch.dtype) -> torch.Tensor: + key = (str(device), dtype) + table = _E2M1_PACKED_TABLE_CACHE.get(key) + if table is not None: + return table + if device.type == "cuda" and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "MXFP4 E2M1 decode table should be initialized before CUDA graph capture." + ) + table = torch.tensor(_E2M1_PACKED_VALUES, device=device, dtype=dtype) + _E2M1_PACKED_TABLE_CACHE[key] = table + return table + + +def _decode_mxfp4_blocks(blocks: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + blocks_u8 = _as_uint8(blocks, "blocks") + if blocks_u8.dim() < 2: + raise ValueError("MXFP4 blocks should have at least block and packed-byte dimensions.") + if blocks_u8.shape[-1] * 2 != _MXFP4_BLOCK_SIZE: + raise ValueError( + f"MXFP4 blocks should pack {_MXFP4_BLOCK_SIZE} values per scale block, " + f"got last dimension {blocks_u8.shape[-1]}." + ) + if scales.shape != blocks_u8.shape[:-1]: + raise ValueError( + f"MXFP4 scales shape {tuple(scales.shape)} should match blocks shape " + f"{tuple(blocks_u8.shape[:-1])} without the packed-byte dimension." + ) + + table = _get_e2m1_packed_table(blocks.device, torch.float32) + values = table[blocks_u8.to(torch.int64)].reshape(*blocks_u8.shape[:-1], _MXFP4_BLOCK_SIZE) + decoded = values * _decode_e8m0_scales(scales).unsqueeze(-1) + return decoded.reshape(*blocks_u8.shape[:-2], blocks_u8.shape[-2] * _MXFP4_BLOCK_SIZE) + + +def _split_gate_up( + gate_up: torch.Tensor, + gate_up_order: str, +) -> tuple[torch.Tensor, torch.Tensor]: + if gate_up_order in ("interleaved", "gpt_oss_interleaved"): + return gate_up[..., ::2], gate_up[..., 1::2] + + gate_up_size = gate_up.shape[-1] + if gate_up_size % 2 != 0: + raise ValueError(f"gate_up output size should be even, got {gate_up_size}.") + midpoint = gate_up_size // 2 + first, second = gate_up[..., :midpoint], gate_up[..., midpoint:] + if gate_up_order in ("up_gate", "w3_w1"): + return second, first + if gate_up_order in ("gate_up", "w1_w3"): + return first, second + raise ValueError(f"Unsupported gate_up_order: {gate_up_order}.") + + +def _apply_swiglu( + gate_up: torch.Tensor, + alpha: float, + limit: float, + gate_up_order: str, + swiglu_mode: str, +) -> torch.Tensor: + gate, up = _split_gate_up(gate_up, gate_up_order) + if swiglu_mode == "gpt_oss" or limit > 0.0: + gate = gate.clamp(max=float(limit)) + up = up.clamp(min=-float(limit), max=float(limit)) + gate = gate * torch.sigmoid(gate * float(alpha)) + if swiglu_mode == "deepseek": + return gate * up + if swiglu_mode == "gpt_oss": + return gate * (up + 1.0) + raise ValueError(f"Unsupported swiglu_mode: {swiglu_mode}.") + + +def _router_topk( + hidden_states: torch.Tensor, + router_weight: torch.Tensor, + router_bias: torch.Tensor, + top_k: int, +) -> tuple[torch.Tensor, torch.Tensor]: + if top_k <= 0: + raise ValueError(f"top_k should be positive, got {top_k}.") + + bias = router_bias.to(torch.float32) if router_bias.numel() > 0 else None + router_logits = F.linear(hidden_states.to(torch.float32), router_weight.to(torch.float32), bias) + if top_k > router_logits.shape[-1]: + raise ValueError(f"top_k={top_k} should be <= number of experts {router_logits.shape[-1]}.") + router_top_value, router_indices = torch.topk(router_logits, top_k, dim=-1) + routing_weights = torch.softmax(router_top_value, dim=-1, dtype=torch.float32) + return router_indices, routing_weights + + +def _split_range_last_remainder(num_experts: int, world_size: int, rank: int) -> tuple[int, int]: + if world_size <= 0: + raise ValueError(f"ep_size should be positive, got {world_size}.") + if rank < 0 or rank >= world_size: + raise ValueError(f"ep_rank should be in [0, {world_size}), got {rank}.") + base = num_experts // world_size + lo = base * rank + hi = num_experts if rank == world_size - 1 else base * (rank + 1) + return lo, hi + + +def _run_torch_mxfp4_from_routing_slots( + x: torch.Tensor, + leading_shape: torch.Size, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + gate_up_weight: torch.Tensor, + gate_up_bias: torch.Tensor, + down_weight: torch.Tensor, + down_bias: torch.Tensor, + alpha: float, + limit: float, + expert_start: int, + gate_up_order: str, + swiglu_mode: str, + output_dtype: torch.dtype, +) -> torch.Tensor: + hidden_size = x.shape[-1] + output = torch.zeros((x.shape[0], hidden_size), device=x.device, dtype=torch.float32) + local_experts = gate_up_weight.shape[0] + if local_experts <= 0: + raise ValueError("MXFP4 MoE requires at least one local expert.") + + local_expert_idx = selected_experts - int(expert_start) + valid_route = (local_expert_idx >= 0) & (local_expert_idx < local_experts) + local_expert_idx = local_expert_idx.clamp(0, local_experts - 1).to(torch.int64) + + x_for_bmm = x.unsqueeze(-1) + for route_idx in range(local_expert_idx.shape[1]): + for start in range(0, x.shape[0], _TORCH_MXFP4_ROUTED_MOE_TOKEN_CHUNK): + end = min(start + _TORCH_MXFP4_ROUTED_MOE_TOKEN_CHUNK, x.shape[0]) + token_slice = slice(start, end) + expert_idx = local_expert_idx[token_slice, route_idx] + gate_up = torch.bmm( + gate_up_weight.index_select(0, expert_idx), x_for_bmm[token_slice] + ).squeeze(-1) + gate_up = gate_up + gate_up_bias.index_select(0, expert_idx).to(torch.float32) + inter = _apply_swiglu(gate_up, alpha, limit, gate_up_order, swiglu_mode) + expert_output = torch.bmm( + down_weight.index_select(0, expert_idx), inter.unsqueeze(-1) + ).squeeze(-1) + expert_output = expert_output + down_bias.index_select(0, expert_idx).to(torch.float32) + route_scale = routing_weights[token_slice, route_idx, None] * valid_route[ + token_slice, route_idx, None + ].to(torch.float32) + output[token_slice] = output[token_slice] + expert_output * route_scale + + return output.reshape(*leading_shape, hidden_size).to(output_dtype) + + +def _run_torch_mxfp4_mlp_core( + hidden_states: torch.Tensor, + router_weight: torch.Tensor, + router_bias: torch.Tensor, + top_k: int, + gate_up_blocks: torch.Tensor, + gate_up_bias: torch.Tensor, + gate_up_scales: torch.Tensor, + alpha: float, + limit: float, + down_blocks: torch.Tensor, + down_bias: torch.Tensor, + down_scales: torch.Tensor, + expert_start: int = 0, +) -> torch.Tensor: + selected_experts, routing_weights = _router_topk( + hidden_states.reshape(-1, hidden_states.shape[-1]), + router_weight, + router_bias, + top_k, + ) + return _run_torch_mxfp4_from_routing_core( + hidden_states, + selected_experts, + routing_weights, + gate_up_blocks, + gate_up_bias, + gate_up_scales, + alpha, + limit, + down_blocks, + down_bias, + down_scales, + expert_start=expert_start, + gate_up_order="interleaved", + swiglu_mode="gpt_oss", + ) + + +def _run_torch_mxfp4_from_routing_core( + hidden_states: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + gate_up_blocks: torch.Tensor, + gate_up_bias: torch.Tensor, + gate_up_scales: torch.Tensor, + alpha: float, + limit: float, + down_blocks: torch.Tensor, + down_bias: torch.Tensor, + down_scales: torch.Tensor, + expert_start: int = 0, + gate_up_order: str = "up_gate", + swiglu_mode: str = "deepseek", +) -> torch.Tensor: + leading_shape = hidden_states.shape[:-1] + hidden_size = hidden_states.shape[-1] + x = hidden_states.reshape(-1, hidden_size).to(torch.float32) + + selected_experts = selected_experts.reshape(x.shape[0], -1).to(torch.int64) + routing_weights = routing_weights.reshape_as(selected_experts).to(torch.float32) + gate_up_weight = _decode_mxfp4_blocks(gate_up_blocks, gate_up_scales) + down_weight = _decode_mxfp4_blocks(down_blocks, down_scales) + + if gate_up_weight.shape[-1] != hidden_size: + raise ValueError( + f"Decoded gate/up MXFP4 weight hidden dimension {gate_up_weight.shape[-1]} " + f"does not match hidden states dimension {hidden_size}." + ) + if down_weight.shape[-2] != hidden_size: + raise ValueError( + f"Decoded down MXFP4 weight output dimension {down_weight.shape[-2]} " + f"does not match hidden states dimension {hidden_size}." + ) + + return _run_torch_mxfp4_from_routing_slots( + x, + leading_shape, + selected_experts, + routing_weights, + gate_up_weight, + gate_up_bias, + down_weight, + down_bias, + alpha, + limit, + expert_start, + gate_up_order, + swiglu_mode, + hidden_states.dtype, + ) + + def _prepare_weights_scales( hidden_size: int, gate_up_blocks: torch.Tensor, # [E_local, 2I, H//32, 16] in unit8 @@ -233,6 +683,16 @@ def _run_mxfp4_mlp_core( Shared core for both triton_mxfp4_moe and triton_mxfp4_moe_ep. - route_fn encapsulates the only difference: how we produce (routing_data, gather_idx, scatter_idx). """ + from triton_kernels.matmul_ogs import ( + FlexCtx, + FnSpecs, + FusedActivation, + PrecisionConfig, + matmul_ogs, + ) + from triton_kernels.numerics import InFlexData + from triton_kernels.swiglu import swiglu_fn + leading_shape = hidden_states.shape[:-1] hidden_size = hidden_states.shape[-1] x = hidden_states.reshape(-1, hidden_size) @@ -290,6 +750,164 @@ def _run_mxfp4_mlp_core( return y +def _triton_route_from_selected_experts( + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + num_experts_total: int, + expert_start: int, + num_local_experts: int, +): + """Build Triton OGS routing metadata from externally computed top-k routing.""" + from triton_kernels.matmul_ogs import GatherIndx, RoutingData, ScatterIndx + from triton_kernels.tensor import make_bitmatrix_metadata, make_ragged_tensor_metadata + + from tensorrt_llm._torch.modules.fused_moe.fused_moe_triton import TritonEPRouter + + selected = selected_experts.to(torch.int32).contiguous() + weights = routing_weights.to(torch.float32).contiguous() + num_tokens, top_k = selected.shape + + expert_stop = expert_start + num_local_experts + is_ep = expert_start != 0 or num_local_experts != num_experts_total + if _is_power_of_two(top_k): + mask = _make_routing_bitmatrix(selected, num_experts_total) + sparse_matrix = _RoutingSparseMatrix( + vals=weights, + indx=selected, + mask=mask, + mask_metadata=make_bitmatrix_metadata(selected, mask), + ) + expt_scal = sparse_matrix.vals + expt_indx = sparse_matrix.indx + if is_ep: + triton_ep_router = TritonEPRouter() + expt_scal, expt_indx, sparse_matrix = triton_ep_router.prune_routing_ep( + expt_scal, + expt_indx, + sparse_matrix, + num_experts_total, + expert_start, + expert_stop, + ) + metadata = make_bitmatrix_metadata(expt_indx, sparse_matrix.mask) + else: + metadata = sparse_matrix.mask_metadata + else: + if is_ep: + keep = (selected >= expert_start) & (selected < expert_stop) + expt_indx = torch.where(keep, selected - expert_start, -1) + expt_scal = torch.where(keep, weights, -1.0) + metadata = _make_routing_metadata_from_indices(expt_indx, num_local_experts) + else: + expt_scal = weights + expt_indx = selected + metadata = _make_routing_metadata_from_indices(expt_indx, num_experts_total) + + expt_data = make_ragged_tensor_metadata(metadata.col_sum, num_tokens * top_k) + gate_scal_sorted = expt_scal.reshape(-1)[metadata.col_sorted_indx] + + rdata = RoutingData( + gate_scal=gate_scal_sorted, + expt_hist=metadata.col_sum, + n_expts_tot=num_local_experts, + n_expts_act=top_k, + expt_data=expt_data, + ) + gather_idx = GatherIndx( + src_indx=metadata.col_sorted_indx, + dst_indx=metadata.row_sorted_indx, + ) + scatter_idx = ScatterIndx( + src_indx=metadata.row_sorted_indx, + dst_indx=metadata.col_sorted_indx, + ) + return rdata, gather_idx, scatter_idx + + +def _run_triton_mxfp4_from_routing_core( + hidden_states: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + gate_up_blocks: torch.Tensor, + gate_up_bias: torch.Tensor, + gate_up_scales: torch.Tensor, + alpha: float, + limit: float, + down_blocks: torch.Tensor, + down_bias: torch.Tensor, + down_scales: torch.Tensor, + gate_up_order: str, + swiglu_mode: str, + expert_start: int = 0, + num_experts_total: int = -1, +) -> torch.Tensor: + from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig, matmul_ogs + from triton_kernels.numerics import InFlexData + + leading_shape = hidden_states.shape[:-1] + hidden_size = hidden_states.shape[-1] + x = hidden_states.reshape(-1, hidden_size) + selected = selected_experts.reshape(x.shape[0], -1) + weights = routing_weights.reshape_as(selected) + + num_local_experts = int(gate_up_blocks.shape[0]) + if num_experts_total < 0: + num_experts_total = expert_start + num_local_experts + if num_experts_total < expert_start + num_local_experts: + raise ValueError( + "num_experts_total should cover the local expert range, got " + f"num_experts_total={num_experts_total}, expert_start={expert_start}, " + f"num_local_experts={num_local_experts}." + ) + + with torch.cuda.device(x.device): + routing_data, gather_idx, scatter_idx = _triton_route_from_selected_experts( + selected, + weights, + int(num_experts_total), + int(expert_start), + num_local_experts, + ) + + ( + triton_gate_up_w, + gate_up_w_scale_raw, + triton_down_w, + down_w_scale_raw, + ) = _prepare_weights_scales_cached( + hidden_size, gate_up_blocks, gate_up_scales, down_blocks, down_scales + ) + + gate_pc = PrecisionConfig( + weight_scale=gate_up_w_scale_raw, flex_ctx=FlexCtx(rhs_data=InFlexData()) + ) + down_pc = PrecisionConfig( + weight_scale=down_w_scale_raw, flex_ctx=FlexCtx(rhs_data=InFlexData()) + ) + + gate_up = matmul_ogs( + x, + triton_gate_up_w, + gate_up_bias.to(torch.float32), + routing_data, + gather_indx=gather_idx, + precision_config=gate_pc, + gammas=None, + ) + inter = _apply_swiglu(gate_up, alpha, limit, gate_up_order, swiglu_mode) + y = matmul_ogs( + inter, + triton_down_w, + down_bias.to(torch.float32), + routing_data, + scatter_indx=scatter_idx, + precision_config=down_pc, + gammas=routing_data.gate_scal, + ) + + return y.reshape(*leading_shape, hidden_size) + + @torch.library.custom_op("auto_deploy::triton_mxfp4_moe", mutates_args=()) def triton_mxfp4_moe( hidden_states: torch.Tensor, # [B, S, H] or [B*S, H] @@ -309,6 +927,8 @@ def triton_mxfp4_moe( down_scales: torch.Tensor, # [E, H, I//32] in uint8 layer_type: str = "moe", ) -> torch.Tensor: + from tensorrt_llm._torch.modules.fused_moe.fused_moe_triton import TritonEPRouter + def _global_route_fn(logits: torch.Tensor): # routing() removed in triton_kernels 3.6.0 # TritonEPRouter(ep=1) is equivalent @@ -330,6 +950,64 @@ def _global_route_fn(logits: torch.Tensor): ) +@torch.library.custom_op("auto_deploy::triton_mxfp4_moe_from_routing", mutates_args=()) +def triton_mxfp4_moe_from_routing( + hidden_states: torch.Tensor, # [B, S, H] or [B*S, H] + selected_experts: torch.Tensor, # [B*S, top_k] + routing_weights: torch.Tensor, # [B*S, top_k] + gate_up_blocks: torch.Tensor, # [E, 2I, H//32, 16] in uint8 + gate_up_bias: torch.Tensor, # [E, 2I] + gate_up_scales: torch.Tensor, # [E, 2I, H//32] in uint8 + alpha: float, + limit: float, + down_blocks: torch.Tensor, # [E, H, I//32, 16] in uint8 + down_bias: torch.Tensor, # [E, H] + down_scales: torch.Tensor, # [E, H, I//32] in uint8 + gate_up_order: str = "up_gate", + swiglu_mode: str = "deepseek", + layer_type: str = "moe", + num_experts_total: int = -1, +) -> torch.Tensor: + return _run_triton_mxfp4_from_routing_core( + hidden_states, + selected_experts, + routing_weights, + gate_up_blocks, + gate_up_bias, + gate_up_scales, + alpha, + limit, + down_blocks, + down_bias, + down_scales, + gate_up_order, + swiglu_mode, + num_experts_total=num_experts_total, + ) + + +@triton_mxfp4_moe_from_routing.register_fake +def _triton_mxfp4_mlp_from_routing_fake( + hidden_states: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + gate_up_blocks: torch.Tensor, + gate_up_bias: torch.Tensor, + gate_up_scales: torch.Tensor, + alpha: float, + limit: float, + down_blocks: torch.Tensor, + down_bias: torch.Tensor, + down_scales: torch.Tensor, + gate_up_order: str = "up_gate", + swiglu_mode: str = "deepseek", + layer_type: str = "moe", + num_experts_total: int = -1, +): + del num_experts_total + return torch.empty_like(hidden_states) + + @triton_mxfp4_moe.register_fake def _mxfp4_mlp_fake( hidden_states: torch.Tensor, @@ -349,6 +1027,120 @@ def _mxfp4_mlp_fake( return torch.empty_like(hidden_states) +@torch.library.custom_op("auto_deploy::torch_mxfp4_moe", mutates_args=()) +def torch_mxfp4_moe( + hidden_states: torch.Tensor, # [B, S, H] or [B*S, H] + # router + router_weight: torch.Tensor, # [E, H] + router_bias: torch.Tensor, # [E] + top_k: int, + # gate_up path + gate_up_blocks: torch.Tensor, # [E, 2I, H//32, 16] in uint8 + gate_up_bias: torch.Tensor, # [E, 2I] + gate_up_scales: torch.Tensor, # [E, 2I, H//32] in uint8 + alpha: float, + limit: float, + # down path + down_blocks: torch.Tensor, # [E, H, I//32, 16] in uint8 + down_bias: torch.Tensor, # [E, H] + down_scales: torch.Tensor, # [E, H, I//32] in uint8 + layer_type: str = "moe", +) -> torch.Tensor: + return _run_torch_mxfp4_mlp_core( + hidden_states, + router_weight, + router_bias, + top_k, + gate_up_blocks, + gate_up_bias, + gate_up_scales, + alpha, + limit, + down_blocks, + down_bias, + down_scales, + ) + + +@torch_mxfp4_moe.register_fake +def _torch_mxfp4_mlp_fake( + hidden_states: torch.Tensor, + router_weight: torch.Tensor, + router_bias: torch.Tensor, + top_k: int, + gate_up_blocks: torch.Tensor, + gate_up_bias: torch.Tensor, + gate_up_scales: torch.Tensor, + alpha: float, + limit: float, + down_blocks: torch.Tensor, + down_bias: torch.Tensor, + down_scales: torch.Tensor, + layer_type: str = "moe", +): + return torch.empty_like(hidden_states) + + +@torch.library.custom_op("auto_deploy::torch_mxfp4_moe_from_routing", mutates_args=()) +def torch_mxfp4_moe_from_routing( + hidden_states: torch.Tensor, # [B, S, H] or [B*S, H] + selected_experts: torch.Tensor, # [B*S, top_k] + routing_weights: torch.Tensor, # [B*S, top_k] + # gate_up path + gate_up_blocks: torch.Tensor, # [E, 2I, H//32, 16] in uint8 + gate_up_bias: torch.Tensor, # [E, 2I] + gate_up_scales: torch.Tensor, # [E, 2I, H//32] in uint8 + alpha: float, + limit: float, + # down path + down_blocks: torch.Tensor, # [E, H, I//32, 16] in uint8 + down_bias: torch.Tensor, # [E, H] + down_scales: torch.Tensor, # [E, H, I//32] in uint8 + gate_up_order: str = "up_gate", + swiglu_mode: str = "deepseek", + layer_type: str = "moe", + num_experts_total: int = -1, +) -> torch.Tensor: + del num_experts_total + return _run_torch_mxfp4_from_routing_core( + hidden_states, + selected_experts, + routing_weights, + gate_up_blocks, + gate_up_bias, + gate_up_scales, + alpha, + limit, + down_blocks, + down_bias, + down_scales, + gate_up_order=gate_up_order, + swiglu_mode=swiglu_mode, + ) + + +@torch_mxfp4_moe_from_routing.register_fake +def _torch_mxfp4_mlp_from_routing_fake( + hidden_states: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + gate_up_blocks: torch.Tensor, + gate_up_bias: torch.Tensor, + gate_up_scales: torch.Tensor, + alpha: float, + limit: float, + down_blocks: torch.Tensor, + down_bias: torch.Tensor, + down_scales: torch.Tensor, + gate_up_order: str = "up_gate", + swiglu_mode: str = "deepseek", + layer_type: str = "moe", + num_experts_total: int = -1, +): + del num_experts_total + return torch.empty_like(hidden_states) + + @torch.library.custom_op("auto_deploy::triton_mxfp4_moe_ep", mutates_args=()) def triton_mxfp4_moe_ep( hidden_states: torch.Tensor, # [B, S, H] or [B*S, H] @@ -370,6 +1162,8 @@ def triton_mxfp4_moe_ep( ep_rank: int, layer_type: str = "moe", ) -> torch.Tensor: + from tensorrt_llm._torch.modules.fused_moe.fused_moe_triton import TritonEPRouter + triton_ep_router = TritonEPRouter() def _ep_route_fn(logits: torch.Tensor): @@ -410,3 +1204,185 @@ def _mxfp4_mlp_ep_fake( layer_type: str = "moe", ): return torch.empty_like(hidden_states) + + +@torch.library.custom_op("auto_deploy::triton_mxfp4_moe_from_routing_ep", mutates_args=()) +def triton_mxfp4_moe_from_routing_ep( + hidden_states: torch.Tensor, # [B, S, H] or [B*S, H] + selected_experts: torch.Tensor, # [B*S, top_k] + routing_weights: torch.Tensor, # [B*S, top_k] + gate_up_blocks: torch.Tensor, # [E_local, 2I, H//32, 16] in uint8 + gate_up_bias: torch.Tensor, # [E_local, 2I] + gate_up_scales: torch.Tensor, # [E_local, 2I, H//32] in uint8 + alpha: float, + limit: float, + down_blocks: torch.Tensor, # [E_local, H, I//32, 16] in uint8 + down_bias: torch.Tensor, # [E_local, H] + down_scales: torch.Tensor, # [E_local, H, I//32] in uint8 + gate_up_order: str = "up_gate", + swiglu_mode: str = "deepseek", + layer_type: str = "moe", + expert_start: int = 0, + num_experts_total: int = -1, +) -> torch.Tensor: + return _run_triton_mxfp4_from_routing_core( + hidden_states, + selected_experts, + routing_weights, + gate_up_blocks, + gate_up_bias, + gate_up_scales, + alpha, + limit, + down_blocks, + down_bias, + down_scales, + gate_up_order, + swiglu_mode, + expert_start=expert_start, + num_experts_total=num_experts_total, + ) + + +@triton_mxfp4_moe_from_routing_ep.register_fake +def _triton_mxfp4_mlp_from_routing_ep_fake( + hidden_states: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + gate_up_blocks: torch.Tensor, + gate_up_bias: torch.Tensor, + gate_up_scales: torch.Tensor, + alpha: float, + limit: float, + down_blocks: torch.Tensor, + down_bias: torch.Tensor, + down_scales: torch.Tensor, + gate_up_order: str = "up_gate", + swiglu_mode: str = "deepseek", + layer_type: str = "moe", + expert_start: int = 0, + num_experts_total: int = -1, +): + return torch.empty_like(hidden_states) + + +@torch.library.custom_op("auto_deploy::torch_mxfp4_moe_from_routing_ep", mutates_args=()) +def torch_mxfp4_moe_from_routing_ep( + hidden_states: torch.Tensor, # [B, S, H] or [B*S, H] + selected_experts: torch.Tensor, # [B*S, top_k] + routing_weights: torch.Tensor, # [B*S, top_k] + # expert params (already sharded along dim 0) + gate_up_blocks: torch.Tensor, # [E_local, 2I, H//32, 16] in uint8 + gate_up_bias: torch.Tensor, # [E_local, 2I] + gate_up_scales: torch.Tensor, # [E_local, 2I, H//32] in uint8 + alpha: float, + limit: float, + down_blocks: torch.Tensor, # [E_local, H, I//32, 16] in uint8 + down_bias: torch.Tensor, # [E_local, H] + down_scales: torch.Tensor, # [E_local, H, I//32] in uint8 + gate_up_order: str = "up_gate", + swiglu_mode: str = "deepseek", + layer_type: str = "moe", + expert_start: int = 0, + num_experts_total: int = -1, +) -> torch.Tensor: + del num_experts_total + return _run_torch_mxfp4_from_routing_core( + hidden_states, + selected_experts, + routing_weights, + gate_up_blocks, + gate_up_bias, + gate_up_scales, + alpha, + limit, + down_blocks, + down_bias, + down_scales, + expert_start=expert_start, + gate_up_order=gate_up_order, + swiglu_mode=swiglu_mode, + ) + + +@torch_mxfp4_moe_from_routing_ep.register_fake +def _torch_mxfp4_mlp_from_routing_ep_fake( + hidden_states: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + gate_up_blocks: torch.Tensor, + gate_up_bias: torch.Tensor, + gate_up_scales: torch.Tensor, + alpha: float, + limit: float, + down_blocks: torch.Tensor, + down_bias: torch.Tensor, + down_scales: torch.Tensor, + gate_up_order: str = "up_gate", + swiglu_mode: str = "deepseek", + layer_type: str = "moe", + expert_start: int = 0, + num_experts_total: int = -1, +): + del num_experts_total + return torch.empty_like(hidden_states) + + +@torch.library.custom_op("auto_deploy::torch_mxfp4_moe_ep", mutates_args=()) +def torch_mxfp4_moe_ep( + hidden_states: torch.Tensor, # [B, S, H] or [B*S, H] + # router (replicated across EP) + router_weight: torch.Tensor, # [E_total, H] + router_bias: torch.Tensor, # [E_total] + top_k: int, + # expert params (already sharded along dim 0) + gate_up_blocks: torch.Tensor, # [E_local, 2I, H//32, 16] in uint8 + gate_up_bias: torch.Tensor, # [E_local, 2I] + gate_up_scales: torch.Tensor, # [E_local, 2I, H//32] in uint8 + alpha: float, + limit: float, + down_blocks: torch.Tensor, # [E_local, H, I//32, 16] in uint8 + down_bias: torch.Tensor, # [E_local, H] + down_scales: torch.Tensor, # [E_local, H, I//32] in uint8 + # EP topology + ep_size: int, + ep_rank: int, + layer_type: str = "moe", +) -> torch.Tensor: + expert_start, _ = _split_range_last_remainder(router_weight.shape[0], ep_size, ep_rank) + return _run_torch_mxfp4_mlp_core( + hidden_states, + router_weight, + router_bias, + top_k, + gate_up_blocks, + gate_up_bias, + gate_up_scales, + alpha, + limit, + down_blocks, + down_bias, + down_scales, + expert_start=expert_start, + ) + + +@torch_mxfp4_moe_ep.register_fake +def _torch_mxfp4_mlp_ep_fake( + hidden_states: torch.Tensor, + router_weight: torch.Tensor, + router_bias: torch.Tensor, + top_k: int, + gate_up_blocks: torch.Tensor, + gate_up_bias: torch.Tensor, + gate_up_scales: torch.Tensor, + alpha: float, + limit: float, + down_blocks: torch.Tensor, + down_bias: torch.Tensor, + down_scales: torch.Tensor, + ep_size: int, + ep_rank: int, + layer_type: str = "moe", +): + return torch.empty_like(hidden_states) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/prepare_trtllm_gen_moe_mxfp4_weights.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/prepare_trtllm_gen_moe_mxfp4_weights.py index cc1f75e9d863..ee70b3818393 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/prepare_trtllm_gen_moe_mxfp4_weights.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/prepare_trtllm_gen_moe_mxfp4_weights.py @@ -237,20 +237,50 @@ def _deinterleave_gate_up( gu_3d: torch.Tensor, # [E, 2I, H/2] gate_up_scales: torch.Tensor, # [E, 2I, H/32] gate_up_bias: torch.Tensor, # [E, 2I] + gate_up_order: str, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Split the HF interleaved 2I axis (gate at even rows, up at odd rows) into separate halves. + """Split the HF 2I axis into separate gate/up halves. Kept separate so downstream row-padding puts the zero-pad rows INSIDE each half before re-concat as ``[up | gate]`` — matches PT's ``dst_w3 = up`` / ``dst_w1 = gate`` chunk layout. """ - return ( - gu_3d[:, 0::2, :].contiguous(), # gate_rows_w - gu_3d[:, 1::2, :].contiguous(), # up_rows_w - gate_up_scales[:, 0::2, :].contiguous(), # gate_rows_s - gate_up_scales[:, 1::2, :].contiguous(), # up_rows_s - gate_up_bias[:, 0::2].contiguous(), # gate_b - gate_up_bias[:, 1::2].contiguous(), # up_b - ) + if gate_up_order in ("interleaved", "gpt_oss_interleaved"): + return ( + gu_3d[:, 0::2, :].contiguous(), # gate_rows_w + gu_3d[:, 1::2, :].contiguous(), # up_rows_w + gate_up_scales[:, 0::2, :].contiguous(), # gate_rows_s + gate_up_scales[:, 1::2, :].contiguous(), # up_rows_s + gate_up_bias[:, 0::2].contiguous(), # gate_b + gate_up_bias[:, 1::2].contiguous(), # up_b + ) + + gate_up_size = gu_3d.shape[1] + if gate_up_size % 2 != 0: + raise ValueError(f"gate_up 2I axis should be even, got {gate_up_size}.") + midpoint = gate_up_size // 2 + first_w, second_w = gu_3d[:, :midpoint, :], gu_3d[:, midpoint:, :] + first_s, second_s = gate_up_scales[:, :midpoint, :], gate_up_scales[:, midpoint:, :] + first_b, second_b = gate_up_bias[:, :midpoint], gate_up_bias[:, midpoint:] + + if gate_up_order in ("up_gate", "w3_w1"): + return ( + second_w.contiguous(), + first_w.contiguous(), + second_s.contiguous(), + first_s.contiguous(), + second_b.contiguous(), + first_b.contiguous(), + ) + if gate_up_order in ("gate_up", "w1_w3"): + return ( + first_w.contiguous(), + second_w.contiguous(), + first_s.contiguous(), + second_s.contiguous(), + first_b.contiguous(), + second_b.contiguous(), + ) + raise ValueError(f"Unsupported gate_up_order: {gate_up_order}.") def _pad_and_slice_axis(t: torch.Tensor, dim: int, target: int, lo: int, hi: int) -> torch.Tensor: @@ -511,6 +541,7 @@ def prepare_trtllm_gen_moe_mxfp4_weights( tp_size: int = 1, tp_rank: int = 0, scratch: MXFP4PrepScratch | None = None, + gate_up_order: str = "interleaved", ) -> TRTLLMGenMXFP4MoEWeights: """Convert HF on-disk MXFP4 expert weights to the trtllm-gen kernel layout. @@ -542,11 +573,11 @@ def prepare_trtllm_gen_moe_mxfp4_weights( assert down_blocks.size(0) == gate_up_blocks.size(0) - # 1. Flatten blocks ([..., 16] → flattened) and de-interleave gate/up halves. + # 1. Flatten blocks ([..., 16] → flattened) and split gate/up halves. gu_3d = _flatten_block_dim(gate_up_blocks) # [E, 2I, H/2] dn_3d = _flatten_block_dim(down_blocks) # [E, H, I/2] gate_rows_w, up_rows_w, gate_rows_s, up_rows_s, gate_b, up_b = _deinterleave_gate_up( - gu_3d, gate_up_scales, gate_up_bias + gu_3d, gate_up_scales, gate_up_bias, gate_up_order ) # 2. TP slicing on the intermediate axis (no-op for tp_size == 1). diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py index bf648f0a0526..f268948b65ba 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py @@ -189,6 +189,18 @@ def relu2(x: torch.Tensor) -> torch.Tensor: return torch_fn +def _gated_mlp_product( + gate_out: torch.Tensor, + up_out: torch.Tensor, + torch_act_fn: Callable[[torch.Tensor], torch.Tensor], + swiglu_limit: float, +) -> torch.Tensor: + if swiglu_limit > 0.0: + gate_out = gate_out.clamp(max=swiglu_limit) + up_out = up_out.clamp(min=-swiglu_limit, max=swiglu_limit) + return torch_act_fn(gate_out) * up_out + + def _template_moe( x: torch.Tensor, selected_experts: torch.Tensor, @@ -283,6 +295,7 @@ def torch_moe( w3_weight: List[torch.Tensor], is_gated_mlp: bool = True, act_fn: int = int(ActivationType.Silu), + swiglu_limit: float = 0.0, mapping_config: str = "", max_num_tokens: int = 0, apply_routing_on_input: bool = False, @@ -306,6 +319,7 @@ def torch_moe( is_gated_mlp: If True, use a gated MLP. If False, use a simple MLP. act_fn: Activation function applied inside the expert MLP. Supported: ActivationType.Silu (default), ActivationType.Relu2 (ReLU then square). + swiglu_limit: If positive, clamp gated MLP gate and up projections before the activation/product. apply_routing_on_input: If True, multiply routing weights with INPUT before MLP This means: silu(input * routing_weight) If False, multiply routing weights with OUTPUT after MLP @@ -326,7 +340,13 @@ def make_mlp(i: int): W2 = w2_weight[i] # (H, I) W3 = w3_weight[i] # (I, H) return lambda inp: F.linear( - torch_act_fn(F.linear(inp.to(W1.dtype), W1)) * F.linear(inp.to(W3.dtype), W3), W2 + _gated_mlp_product( + F.linear(inp.to(W1.dtype), W1), + F.linear(inp.to(W3.dtype), W3), + torch_act_fn, + swiglu_limit, + ), + W2, ) mlps = [make_mlp(i) for i in range(len(w1_weight))] @@ -361,6 +381,7 @@ def torch_moe_fake( w3_weight: List[torch.Tensor], is_gated_mlp: bool = True, act_fn: int = int(ActivationType.Silu), + swiglu_limit: float = 0.0, mapping_config: str = "", max_num_tokens: int = 0, apply_routing_on_input: bool = False, diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py index 6e10bf5d1230..61188f62fe3f 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py @@ -1510,3 +1510,151 @@ def trtllm_quant_mxfp4_trtllm_gen_moe_fused_fake( out_shape = list(x.shape) out_shape[-1] = valid_hidden_size return x.new_empty(out_shape, dtype=x.dtype) + + +@torch.library.custom_op( + "auto_deploy::trtllm_quant_mxfp4_trtllm_gen_moe_from_routing_fused", + mutates_args=(), +) +def trtllm_quant_mxfp4_trtllm_gen_moe_from_routing_fused( + x: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + fc1_weights_mxfp4: torch.Tensor, + fc2_weights_mxfp4: torch.Tensor, + fc1_weights_scale_ue8m0: torch.Tensor, + fc2_weights_scale_ue8m0: torch.Tensor, + fc1_bias_f32: torch.Tensor, + fc2_bias_f32: torch.Tensor, + swiglu_alpha: torch.Tensor, + swiglu_beta: torch.Tensor, + swiglu_limit: torch.Tensor, + valid_hidden_size: int, + valid_intermediate_size: int, + act_dtype: str, + num_experts_total: int, + local_expert_offset: int = 0, + local_num_experts: int = -1, + routing_method_type: int = int(RoutingMethodType.Renormalize), +) -> torch.Tensor: + """TensorRT-LLM Gen MXFP4 MoE using externally computed top-k routing. + + DeepSeek V4 computes hash/sqrt-softplus routing in Python modeling code, so + this op consumes ``selected_experts`` and already-scaled ``routing_weights`` + directly instead of deriving top-k from router logits. + """ + x_shape = x.shape + x2d = x.view(-1, x_shape[-1]) + selected_2d = selected_experts.reshape(x2d.shape[0], -1).to(torch.int32).contiguous() + weights_2d = routing_weights.reshape_as(selected_2d).to(torch.bfloat16).contiguous() + top_k = int(selected_2d.shape[-1]) + + expected_hidden = int(fc1_weights_mxfp4.shape[-1] * 2) + pad_size = expected_hidden - int(x2d.shape[-1]) + if pad_size > 0: + x2d = torch.nn.functional.pad(x2d, (0, pad_size)) + + if local_num_experts < 0: + local_num_experts = int(fc1_weights_mxfp4.shape[0]) + intermediate_size_padded = int(fc1_weights_mxfp4.shape[1] // 2) + + if act_dtype == "mxfp8": + x_mxfp8, x_scale = torch.ops.trtllm.mxfp8_quantize( + x2d, + False, + alignment=512, + ) + result = torch.ops.trtllm.mxe4m3_mxe2m1_block_scale_moe_runner( + None, # routing_logits + None, # routing_bias + x_mxfp8, + x_scale, + fc1_weights_mxfp4, + fc1_weights_scale_ue8m0, + fc1_bias_f32, + swiglu_alpha, + swiglu_beta, + swiglu_limit, + fc2_weights_mxfp4, + fc2_weights_scale_ue8m0, + fc2_bias_f32, + int(num_experts_total), + top_k, + None, # n_group + None, # topk_group + intermediate_size_padded, + valid_hidden_size, + valid_intermediate_size, + int(local_expert_offset), + int(local_num_experts), + None, # routed_scaling_factor: routing_weights are already scaled + int(routing_method_type), + int(ActivationType.Swiglu), + topk_weights=weights_2d, + topk_ids=selected_2d, + ) + elif act_dtype == "bf16": + result = torch.ops.trtllm.bf16_mxe2m1_block_scale_moe_runner( + None, # routing_logits + None, # routing_bias + x2d, + fc1_weights_mxfp4, + fc1_weights_scale_ue8m0, + fc1_bias_f32, + swiglu_alpha, + swiglu_beta, + swiglu_limit, + fc2_weights_mxfp4, + fc2_weights_scale_ue8m0, + fc2_bias_f32, + int(num_experts_total), + top_k, + None, # n_group + None, # topk_group + intermediate_size_padded, + valid_hidden_size, + valid_intermediate_size, + int(local_expert_offset), + int(local_num_experts), + None, # routed_scaling_factor: routing_weights are already scaled + int(routing_method_type), + int(ActivationType.Swiglu), + topk_weights=weights_2d, + topk_ids=selected_2d, + ) + else: + raise ValueError( + "trtllm_quant_mxfp4_trtllm_gen_moe_from_routing_fused: " + f"act_dtype must be 'bf16' or 'mxfp8', got {act_dtype!r}." + ) + + if result.shape[-1] > valid_hidden_size: + result = result[..., :valid_hidden_size].contiguous() + return result.view(*x_shape[:-1], valid_hidden_size) + + +@trtllm_quant_mxfp4_trtllm_gen_moe_from_routing_fused.register_fake +def trtllm_quant_mxfp4_trtllm_gen_moe_from_routing_fused_fake( + x: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + fc1_weights_mxfp4: torch.Tensor, + fc2_weights_mxfp4: torch.Tensor, + fc1_weights_scale_ue8m0: torch.Tensor, + fc2_weights_scale_ue8m0: torch.Tensor, + fc1_bias_f32: torch.Tensor, + fc2_bias_f32: torch.Tensor, + swiglu_alpha: torch.Tensor, + swiglu_beta: torch.Tensor, + swiglu_limit: torch.Tensor, + valid_hidden_size: int, + valid_intermediate_size: int, + act_dtype: str, + num_experts_total: int, + local_expert_offset: int = 0, + local_num_experts: int = -1, + routing_method_type: int = int(RoutingMethodType.Renormalize), +) -> torch.Tensor: + out_shape = list(x.shape) + out_shape[-1] = valid_hidden_size + return x.new_empty(out_shape, dtype=x.dtype) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/linear.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/linear.py index 61bea3a3d6ae..d1d1a1bf8bf6 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/linear.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/linear.py @@ -100,3 +100,48 @@ def simple_fake( ): """Fake implementation of simple_linear.""" return torch.ops.aten.linear(input, weight, bias) + + +def _grouped_linear(input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor]): + output = torch.matmul(input.unsqueeze(-2), weight.transpose(-1, -2)).squeeze(-2) + if bias is not None: + output = output + bias.reshape(weight.shape[0], weight.shape[1]).to(output.dtype) + return output.flatten(-2) + + +@torch.library.custom_op("auto_deploy::torch_grouped_linear", mutates_args=()) +def grouped( + input: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + tp_mode: str = "none", + output_sizes: Optional[List[int]] = None, + tp_min_local_shape: int = 1, + layer_type: str = "unknown", +) -> torch.Tensor: + """Grouped linear projection over ``input[..., group, in_features]``. + + ``weight`` is shaped ``[group, out_features_per_group, in_features]``. + The result is flattened over group and per-group output dimensions, yielding + ``input.shape[:-2] + [group * out_features_per_group]``. + + Sharding hint arguments follow :func:`simple`. + """ + del tp_mode, output_sizes, tp_min_local_shape, layer_type + return _grouped_linear(input, weight, bias) + + +@grouped.register_fake +def grouped_fake( + input, + weight, + bias, + tp_mode="none", + output_sizes=None, + tp_min_local_shape=1, + layer_type="unknown", +): + """Fake implementation of grouped linear.""" + del bias, tp_mode, output_sizes, tp_min_local_shape, layer_type + out_features = weight.shape[0] * weight.shape[1] + return torch.empty((*input.shape[:-2], out_features), dtype=input.dtype, device=input.device) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/quantization/torch_quant.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/quantization/torch_quant.py index 4eccebb41c93..dbfcf41c7b91 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/quantization/torch_quant.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/quantization/torch_quant.py @@ -23,6 +23,7 @@ import triton import triton.language as tl +from ...utils.fp8_dequant import dequant_fp8_weight_two_dim_block_grid from ...utils.quantization_utils import ( cutlass_fp4_scale_to_modelopt_fp4_scale, unpack_uint8_to_int4_weight_2d, @@ -57,6 +58,18 @@ def _dequant_weight_fp8( return weight_fp8.to(dtype) * weight_scale +def _dequant_block_fp8_weight( + weight_fp8: torch.Tensor, + weight_scale: torch.Tensor, + block_n: int, + block_k: int, + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + return dequant_fp8_weight_two_dim_block_grid( + weight_fp8, weight_scale, block_n, block_k, dtype=dtype + ) + + # The NVFP4 helpers below are adapted from modelopt.torch.quantization.qtensor.nvfp4_tensor.NVFP4QTensor def _nvfp4_get_weights_scaling_factor( input: torch.Tensor, @@ -478,7 +491,7 @@ def torch_fake_quant_int4_gptq_linear_fake( @triton.jit -def _act_quant_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr): +def _act_quant_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr, ROUND_SCALE: tl.constexpr): """Block-wise FP8 activation quantization, safe for all-zero blocks. Identical to HuggingFace's act_quant_kernel except that the per-block scale @@ -488,16 +501,22 @@ def _act_quant_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr): pid = tl.program_id(axis=0) offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) x = tl.load(x_ptr + offs).to(tl.float32) - s = tl.max(tl.abs(x)) / 448.0 - # Clamp scale so that all-zero blocks produce 0/eps = 0 instead of 0/0 = NaN. - s = tl.maximum(s, 1e-12) + amax = tl.max(tl.abs(x)) + if ROUND_SCALE: + amax = tl.maximum(amax, 1e-4) + s = amax / 448.0 + s = tl.exp2(tl.ceil(tl.log2(s))) + else: + s = amax / 448.0 + # Clamp scale so that all-zero blocks produce 0/eps = 0 instead of 0/0 = NaN. + s = tl.maximum(s, 1e-12) y = x / s y = y.to(y_ptr.dtype.element_ty) tl.store(y_ptr + offs, y) tl.store(s_ptr + pid, s) -def _safe_act_quant(x: torch.Tensor, block_size: int = 128) -> tuple: +def _safe_act_quant(x: torch.Tensor, block_size: int = 128, input_scale_fmt: str = "") -> tuple: """Block-wise FP8 activation quantization (CUDA-graph safe). Drop-in replacement for ``transformers.integrations.finegrained_fp8.act_quant`` @@ -513,7 +532,8 @@ def _safe_act_quant(x: torch.Tensor, block_size: int = 128) -> tuple: s = x.new_empty(*x.shape[:-1], x.shape[-1] // block_size, dtype=x.dtype) grid = lambda meta: (triton.cdiv(x.numel(), meta["BLOCK_SIZE"]),) # noqa: E731 - _act_quant_kernel[grid](x, y, s, BLOCK_SIZE=block_size) + round_scale = input_scale_fmt.lower() == "ue8m0" + _act_quant_kernel[grid](x, y, s, BLOCK_SIZE=block_size, ROUND_SCALE=round_scale) return y, s @@ -675,6 +695,7 @@ def torch_fake_quant_finegrained_fp8_linear( output_sizes: Optional[List[int]] = None, tp_min_local_shape: int = 1, layer_type: str = "unknown", + input_scale_fmt: str = "", ) -> torch.Tensor: """FineGrainedFP8 linear operation. - weight_scale[0] = weight_scale_inv (per-block weight scale) @@ -691,7 +712,7 @@ def torch_fake_quant_finegrained_fp8_linear( block_k = triton.cdiv(K, scale_k) block_size = [block_n, block_k] - qinput, scale = _safe_act_quant(input, block_size[1]) + qinput, scale = _safe_act_quant(input, block_size[1], input_scale_fmt) output = _w8a8_block_fp8_matmul_triton( qinput, weight_quantized, @@ -720,7 +741,102 @@ def _torch_fake_quant_finegrained_fp8_linear_fake( output_sizes: Optional[List[int]] = None, tp_min_local_shape: int = 1, layer_type: str = "unknown", + input_scale_fmt: str = "", ) -> torch.Tensor: """Fake implementation for torch.export tracing.""" out_features = weight_quantized.shape[0] return torch.empty((*input.shape[:-1], out_features), dtype=input.dtype, device=input.device) + + +@torch.library.custom_op( + "auto_deploy::torch_fake_quant_grouped_finegrained_fp8_linear", + mutates_args=(), +) +def torch_fake_quant_grouped_finegrained_fp8_linear( + input: torch.Tensor, # [..., G, K] + weight_quantized: torch.Tensor, # [G * R, K] float8_e4m3fn + bias: Optional[torch.Tensor], # [G * R], [G, R], or None + input_scale: List[torch.Tensor], # unused for FineGrained FP8 (input quantized on the fly) + weight_scale: List[torch.Tensor], # [weight_scale_inv] + input_zp: List[torch.Tensor], # unused + weight_zp: List[torch.Tensor], # unused + tp_mode: str = "none", + output_sizes: Optional[List[int]] = None, + tp_min_local_shape: int = 1, + layer_type: str = "unknown", + input_scale_fmt: str = "", +) -> torch.Tensor: + """Grouped FineGrainedFP8 projection for flattened checkpoint weights. + + Consumes a flattened checkpoint weight layout ``[G * R, K]`` and matching + per-block ``weight_scale_inv`` buffers. + """ + del input_scale, input_zp, weight_zp, tp_mode, output_sizes, tp_min_local_shape, layer_type + if input.dim() < 2: + raise ValueError(f"input must have at least grouped and K dimensions, got {input.shape}") + if weight_quantized.dim() != 2: + raise ValueError(f"weight must have shape [G * R, K], got {weight_quantized.shape}") + if weight_quantized.dtype != torch.float8_e4m3fn: + raise TypeError("Grouped FineGrained FP8 path requires float8_e4m3fn weight") + + weight_scale_inv = _expect_single_scale(weight_scale, "weight_scale") + num_groups = input.shape[-2] + in_features = input.shape[-1] + out_rows, weight_in_features = weight_quantized.shape + if weight_in_features != in_features: + raise ValueError(f"weight K ({weight_in_features}) must match input K ({in_features})") + if out_rows % num_groups != 0: + raise ValueError(f"weight rows ({out_rows}) must be divisible by groups ({num_groups})") + + scale_n, scale_k = weight_scale_inv.shape + if scale_n == 0 or scale_k == 0: + raise ValueError(f"weight_scale has zero dimension {tuple(weight_scale_inv.shape)}") + block_n = triton.cdiv(out_rows, scale_n) + block_k = triton.cdiv(in_features, scale_k) + + input_contiguous = input.contiguous() + qinput, input_scales = _safe_act_quant(input_contiguous, block_k, input_scale_fmt) + qinput_blocks = qinput.reshape(*input_contiguous.shape[:-1], -1, block_k) + input_dequant = (qinput_blocks.to(input.dtype) * input_scales.unsqueeze(-1)).reshape_as( + input_contiguous + ) + + weight_dequant = _dequant_block_fp8_weight( + weight_quantized, + weight_scale_inv, + block_n, + block_k, + dtype=input.dtype, + ) + rank = out_rows // num_groups + weight_grouped = weight_dequant.view(num_groups, rank, in_features) + output = torch.matmul( + input_dequant.unsqueeze(-2), + weight_grouped.transpose(-1, -2), + ).squeeze(-2) + output = output.flatten(-2) + if bias is not None: + output = output + bias.reshape(out_rows).to(output.dtype) + return output.to(dtype=input.dtype) + + +@torch_fake_quant_grouped_finegrained_fp8_linear.register_fake +def _torch_fake_quant_grouped_finegrained_fp8_linear_fake( + input: torch.Tensor, + weight_quantized: torch.Tensor, + bias: Optional[torch.Tensor], + input_scale: List[torch.Tensor], + weight_scale: List[torch.Tensor], + input_zp: List[torch.Tensor], + weight_zp: List[torch.Tensor], + tp_mode: str = "none", + output_sizes: Optional[List[int]] = None, + tp_min_local_shape: int = 1, + layer_type: str = "unknown", + input_scale_fmt: str = "", +) -> torch.Tensor: + """Fake implementation for torch.export tracing.""" + del bias, input_scale, weight_scale, input_zp, weight_zp + del tp_mode, output_sizes, tp_min_local_shape, layer_type, input_scale_fmt + out_features = weight_quantized.shape[0] + return torch.empty((*input.shape[:-2], out_features), dtype=input.dtype, device=input.device) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/sharding_ops.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/sharding_ops.py index 5f5c71102687..926d17e693e4 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/sharding_ops.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/sharding_ops.py @@ -32,6 +32,7 @@ def view( shape: List[int], tp_scaled_dim: int = -1, layer_type: str = "unknown", + tp_min_local_shape: int = 1, ) -> torch.Tensor: """Sharding-aware view/reshape. @@ -49,10 +50,12 @@ def view( layer_type: Layer classification for selective sharding via ``shard_layers`` config. Values: ``"mha"``, ``"mla"``, ``"mlp"``, ``"moe"``, ``"ssm"``, ``"delta"``, ``"unknown"``. + tp_min_local_shape: Minimum chunk size used when the view input is a + parameter that should be TP-sharded along the scaled dimension. Sharding hint arguments (graph-level metadata for ``apply_sharding_hints``): - ``tp_scaled_dim`` and ``layer_type`` are hints only; they do not change the - unsharded reshape result. + ``tp_scaled_dim``, ``layer_type``, and ``tp_min_local_shape`` are hints only; + they do not change the unsharded reshape result. Returns: Reshaped tensor, same values as ``x.reshape(shape)`` up to clone semantics. @@ -66,6 +69,7 @@ def _view_fake( shape: List[int], tp_scaled_dim: int = -1, layer_type: str = "unknown", + tp_min_local_shape: int = 1, ) -> torch.Tensor: return x.reshape(shape).clone() diff --git a/tensorrt_llm/_torch/auto_deploy/models/checkpoint_metadata.py b/tensorrt_llm/_torch/auto_deploy/models/checkpoint_metadata.py new file mode 100644 index 000000000000..ff490e389609 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/models/checkpoint_metadata.py @@ -0,0 +1,109 @@ +# 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. + +"""Safetensors checkpoint metadata helpers.""" + +from __future__ import annotations + +import json +import struct +from collections.abc import Mapping +from pathlib import Path + +from .quant_checkpoint_layout import QuantizedCheckpointLayoutError + + +def has_safetensors_metadata(ckpt_dir: str | Path) -> bool: + ckpt_path = Path(ckpt_dir) + index_path = ckpt_path / "model.safetensors.index.json" + if index_path.exists(): + index = _read_safetensors_index(index_path) + weight_map = index.get("weight_map") + if not isinstance(weight_map, Mapping) or not weight_map: + raise QuantizedCheckpointLayoutError( + f"safetensors index is missing a non-empty weight_map: {index_path}" + ) + return any((ckpt_path / str(filename)).exists() for filename in set(weight_map.values())) + if not ckpt_path.is_dir(): + return False + return any(ckpt_path.glob("*.safetensors")) + + +def read_safetensors_metadata(ckpt_dir: str | Path) -> dict[str, dict[str, object]]: + ckpt_path = Path(ckpt_dir) + index_path = ckpt_path / "model.safetensors.index.json" + if index_path.exists(): + index = _read_safetensors_index(index_path) + weight_map = index.get("weight_map") + if not isinstance(weight_map, Mapping) or not weight_map: + raise QuantizedCheckpointLayoutError( + f"safetensors index is missing a non-empty weight_map: {index_path}" + ) + safetensors_files = sorted({str(filename) for filename in weight_map.values()}) + else: + safetensors_files = sorted(path.name for path in ckpt_path.glob("*.safetensors")) + + if not safetensors_files: + raise QuantizedCheckpointLayoutError( + f"Quantized checkpoint layout requires safetensors metadata in {ckpt_path}" + ) + + tensor_metadata: dict[str, dict[str, object]] = {} + for filename in safetensors_files: + tensor_metadata.update(_read_safetensors_header(ckpt_path / filename)) + return tensor_metadata + + +def _read_safetensors_index(path: Path) -> Mapping[str, object]: + try: + with path.open("r", encoding="utf-8") as f: + index = json.load(f) + except json.JSONDecodeError as error: + raise QuantizedCheckpointLayoutError(f"Invalid safetensors index JSON: {path}") from error + if not isinstance(index, Mapping): + raise QuantizedCheckpointLayoutError(f"Invalid safetensors index JSON: {path}") + return index + + +def _read_safetensors_header(path: Path) -> dict[str, dict[str, object]]: + if not path.exists(): + raise QuantizedCheckpointLayoutError(f"safetensors file not found: {path}") + + with path.open("rb") as f: + header_size_bytes = f.read(8) + if len(header_size_bytes) != 8: + raise QuantizedCheckpointLayoutError(f"Invalid safetensors header in {path}") + header_size = struct.unpack(" object: + if isinstance(token, Mapping): + return token.get("content") + return token + + +def _message_content_text(message: Mapping[str, Any]) -> str: + content = message.get("content", "") + if isinstance(content, str): + return content + if isinstance(content, Sequence): + parts: list[str] = [] + for block in content: + if isinstance(block, Mapping) and block.get("type") == "text": + parts.append(str(block.get("text", ""))) + return "\n\n".join(parts) + return str(content) + + +def _format_deepseek_v4_chat_messages(messages: Sequence[Mapping[str, Any]]) -> str: + """Render the official DeepSeek V4 chat prompt for text messages.""" + + rendered = _BOS_MARKER + last_user_idx = -1 + for idx in range(len(messages) - 1, -1, -1): + if messages[idx].get("role") in ("user", "developer"): + last_user_idx = idx + break + + for idx, message in enumerate(messages): + role = message.get("role") + content = _message_content_text(message) + if role == "system": + rendered += content + elif role in ("user", "developer"): + rendered += _USER_MARKER + content + elif role == "latest_reminder": + rendered += _LATEST_REMINDER_MARKER + content + elif role == "assistant": + rendered += content + if not message.get("wo_eos", False): + rendered += _EOS_MARKER + else: + raise NotImplementedError(f"Unsupported DeepSeek V4 chat role: {role}") + + has_next = idx + 1 < len(messages) + if has_next and messages[idx + 1].get("role") not in ("assistant", "latest_reminder"): + continue + if role in ("user", "developer"): + rendered += _ASSISTANT_MARKER + if idx >= last_user_idx and message.get("thinking_mode") == "thinking": + rendered += _THINKING_START_MARKER + else: + rendered += _THINKING_END_MARKER + + return rendered + + +class ADDeepseekV4Tokenizer(PreTrainedTokenizerFast): + """Tokenizer wrapper that exposes DeepSeek V4's official chat encoding.""" + + vocab_files_names = {"tokenizer_file": _TOKENIZER_FILE} + model_input_names = ["input_ids", "attention_mask"] + slow_tokenizer_class = None + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: str | Path, + *inputs, + **kwargs, + ) -> "ADDeepseekV4Tokenizer": + del inputs + for key in ("_from_auto", "_commit_hash", "trust_remote_code"): + kwargs.pop(key, None) + + config_path = cached_file(pretrained_model_name_or_path, _TOKENIZER_CONFIG_FILE, **kwargs) + assert config_path is not None + config = json.loads(Path(config_path).read_text()) + + tokenizer_file = cached_file(pretrained_model_name_or_path, _TOKENIZER_FILE, **kwargs) + assert tokenizer_file is not None + + tokenizer = cls( + tokenizer_object=Tokenizer.from_file(tokenizer_file), + name_or_path=str(pretrained_model_name_or_path), + bos_token=_token_content(config.get("bos_token")), + eos_token=_token_content(config.get("eos_token")), + unk_token=_token_content(config.get("unk_token")), + pad_token=_token_content(config.get("pad_token")), + clean_up_tokenization_spaces=config.get("clean_up_tokenization_spaces", False), + model_max_length=config.get("model_max_length"), + padding_side=config.get("padding_side", "left"), + truncation_side=config.get("truncation_side", "left"), + is_local=True, + fix_mistral_regex=False, + ) + tokenizer.chat_template = _DSV4_CHAT_TEMPLATE_MARKER + return tokenizer + + def apply_chat_template( + self, + conversation: Sequence[Mapping[str, Any]], + tokenize: bool = True, + return_dict: bool = False, + return_tensors: Optional[str] = None, + return_attention_mask: bool = True, + **kwargs, + ): + prompt = _format_deepseek_v4_chat_messages(conversation) + if not tokenize: + return prompt + + kwargs.pop("add_generation_prompt", None) + kwargs["add_special_tokens"] = False + encoded = self( + prompt, + return_tensors=return_tensors, + return_attention_mask=return_attention_mask, + **kwargs, + ) + if return_dict: + return encoded + return encoded["input_ids"] + + +def _normalize_deepseek_v4_rope_scaling( + rope_scaling: Optional[Mapping[str, Any]], + max_position_embeddings: int, +) -> dict[str, Any]: + """Return a YaRN config accepted by current transformers validators.""" + + normalized = ( + dict(rope_scaling) + if rope_scaling is not None + else { + "type": "yarn", + "rope_type": "yarn", + "factor": 1.0, + "original_max_position_embeddings": max_position_embeddings, + "beta_fast": 32.0, + "beta_slow": 1.0, + } + ) + rope_type = normalized.get("rope_type", normalized.get("type", "yarn")) + normalized["rope_type"] = rope_type + normalized["type"] = normalized.get("type", rope_type) + normalized["factor"] = float(normalized.get("factor", 1.0)) + normalized["beta_fast"] = float(normalized.get("beta_fast", 32.0)) + normalized["beta_slow"] = float(normalized.get("beta_slow", 1.0)) + original_seq_len = int(normalized.get("original_max_position_embeddings", 0) or 0) + if original_seq_len <= 0: + original_seq_len = max_position_embeddings + normalized["original_max_position_embeddings"] = original_seq_len + return normalized + + +class DeepseekV4Config(PretrainedConfig): + """Minimal local config for DeepSeek V4 when transformers lacks one.""" + + model_type = "deepseek_v4" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size: int = 129280, + hidden_size: int = 4096, + num_hidden_layers: int = 43, + num_attention_heads: int = 64, + num_key_value_heads: int = 1, + head_dim: int = 512, + q_lora_rank: int = 1024, + qk_rope_head_dim: int = 64, + o_lora_rank: int = 1024, + o_groups: int = 8, + sliding_window: int = 128, + compress_ratios: Optional[Tuple[int, ...]] = None, + compress_rope_theta: float = 160000.0, + index_n_heads: int = 64, + index_head_dim: int = 128, + index_topk: int = 512, + moe_intermediate_size: int = 2048, + n_routed_experts: int = 256, + n_shared_experts: int = 1, + num_experts_per_tok: int = 6, + num_hash_layers: int = 3, + scoring_func: str = "sqrtsoftplus", + routed_scaling_factor: float = 1.5, + norm_topk_prob: bool = True, + swiglu_limit: float = 10.0, + hidden_act: str = "silu", + max_position_embeddings: int = 1048576, + rope_theta: float = 10000.0, + rope_scaling: Optional[dict] = None, + rms_norm_eps: float = 1e-6, + hc_mult: int = 4, + hc_sinkhorn_iters: int = 20, + hc_eps: float = 1e-6, + ad_rope_cache_len: Optional[int] = None, + ad_compress_max_seq_len: Optional[int] = None, + skip_mtp: bool = True, + use_cache: bool = False, + attention_bias: bool = False, + attention_dropout: float = 0.0, + initializer_range: float = 0.02, + tie_word_embeddings: bool = False, + pad_token_id: Optional[int] = None, + bos_token_id: int = 0, + eos_token_id: int = 1, + **kwargs, + ) -> None: + ratios = tuple(int(ratio) for ratio in (compress_ratios or (0,) * num_hidden_layers)) + if len(ratios) < num_hidden_layers: + ratios = ratios + (0,) * (num_hidden_layers - len(ratios)) + + self.vocab_size = vocab_size + self.hidden_size = hidden_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.head_dim = head_dim + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.o_lora_rank = o_lora_rank + self.o_groups = o_groups + self.sliding_window = sliding_window + self.compress_ratios = ratios[:num_hidden_layers] + self.compress_rope_theta = compress_rope_theta + self.index_n_heads = index_n_heads + self.index_head_dim = index_head_dim + self.index_topk = index_topk + self.moe_intermediate_size = moe_intermediate_size + self.n_routed_experts = n_routed_experts + self.n_shared_experts = n_shared_experts + self.num_experts_per_tok = num_experts_per_tok + self.num_hash_layers = num_hash_layers + self.scoring_func = scoring_func + self.routed_scaling_factor = routed_scaling_factor + self.norm_topk_prob = norm_topk_prob + self.swiglu_limit = swiglu_limit + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.rope_theta = rope_theta + self.rope_scaling = _normalize_deepseek_v4_rope_scaling( + rope_scaling, max_position_embeddings + ) + self.rms_norm_eps = rms_norm_eps + self.hc_mult = hc_mult + self.hc_sinkhorn_iters = hc_sinkhorn_iters + self.hc_eps = hc_eps + self.ad_rope_cache_len = ad_rope_cache_len or min(max_position_embeddings, 4096) + self.ad_compress_max_seq_len = ad_compress_max_seq_len or self.ad_rope_cache_len + self.skip_mtp = skip_mtp + self.use_cache = use_cache + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.initializer_range = initializer_range + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +try: + AutoConfig.register(DeepseekV4Config.model_type, DeepseekV4Config, exist_ok=True) +except TypeError: + try: + AutoConfig.register(DeepseekV4Config.model_type, DeepseekV4Config) + except ValueError: + pass + + +_DEEPSEEK_V4_FP8_TARGET_PATTERNS = ( + r"layers\.\d+\.(?:" + r"attn\.(?:wq_a|wq_b|wkv|wo_a|wo_b)|" + r"attn\.indexer\.wq_b|" + r"ffn\.shared_experts\.w[123]" + r")", +) +_DEEPSEEK_V4_EXCLUDED_MODULES = ( + "embed", + "head", + "*.ffn.gate", + "*.attn.compressor", + "*.attn.indexer.compressor", + "*.attn.indexer.weights_proj", + "*.norm", + "*.hc_*", + "*.attn_sink", + "mtp.*", +) + +_DEEPSEEK_V4_MXFP4_EXPERT_RE = re.compile( + r"layers\.(?P\d+)\.ffn\.experts\.(?P\d+)\." + r"(?Pw[123])\.(?Pweight|scale)" +) +_DEEPSEEK_V4_MXFP4_KEY_TEMPLATE = "layers.{layer}.ffn.experts.{expert}.{projection}.{kind}" +_DEEPSEEK_V4_MXFP4_LAYER_NAME_RE = re.compile(r"(?:^|[._])layers[._](?P\d+)[._]ffn[._]") +_DEEPSEEK_V4_MXFP4_PROJECTIONS = ("w1", "w2", "w3") +_DEEPSEEK_V4_MXFP4_GATE_UP_ORDER = ("w3", "w1") +_DEEPSEEK_V4_MXFP4_DOWN_PROJECTION = "w2" + +DeepseekV4PackedMxfp4ExpertsCheckpointLayout = PackedMXFP4ExpertLayout + + +def build_deepseek_v4_packed_mxfp4_experts_layout() -> PackedMXFP4ExpertLayout: + return PackedMXFP4ExpertLayout( + key_pattern=_DEEPSEEK_V4_MXFP4_EXPERT_RE, + key_template=_DEEPSEEK_V4_MXFP4_KEY_TEMPLATE, + layer_name_pattern=_DEEPSEEK_V4_MXFP4_LAYER_NAME_RE, + projections=_DEEPSEEK_V4_MXFP4_PROJECTIONS, + gate_up_order=_DEEPSEEK_V4_MXFP4_GATE_UP_ORDER, + down_projection=_DEEPSEEK_V4_MXFP4_DOWN_PROJECTION, + ) + + +def _deepseek_v4_scale_fmt(qconf: Mapping[str, object]) -> str: + scale_fmt = qconf.get("scale_fmt") + if not isinstance(scale_fmt, str) or not scale_fmt: + raise QuantizedCheckpointLayoutError("DeepSeek V4 quantization_config requires scale_fmt.") + scale_fmt = scale_fmt.lower() + if scale_fmt != "ue8m0": + raise QuantizedCheckpointLayoutError( + "DeepSeek V4 quantized checkpoint layout supports " + f"scale_fmt='ue8m0', got '{scale_fmt}'." + ) + return scale_fmt + + +def _deepseek_v4_weight_block_size(qconf: Mapping[str, object]) -> tuple[int, int]: + block_size = qconf.get("weight_block_size") + if not isinstance(block_size, Sequence) or isinstance(block_size, str): + raise QuantizedCheckpointLayoutError( + "DeepSeek V4 quantization_config requires weight_block_size " + "as a two-element integer sequence." + ) + try: + parsed = tuple(_deepseek_v4_positive_int("weight_block_size", dim) for dim in block_size) + except TypeError as error: + raise QuantizedCheckpointLayoutError( + "DeepSeek V4 quantization_config requires weight_block_size " + f"as a two-element integer sequence, got {block_size}." + ) from error + if len(parsed) != 2: + raise QuantizedCheckpointLayoutError( + "DeepSeek V4 quantization_config requires weight_block_size " + f"with two dimensions, got {block_size}." + ) + return parsed + + +def _deepseek_v4_positive_int(name: str, value: object) -> int: + if isinstance(value, bool): + raise TypeError(f"{name} should contain integers, got {value}.") + parsed = operator.index(value) + if parsed <= 0: + raise QuantizedCheckpointLayoutError( + f"DeepSeek V4 quantization_config {name} values should be positive, got {value}." + ) + return parsed + + +@QuantCheckpointLayoutRegistry.register("deepseek_v4") +def _build_deepseek_v4_checkpoint_layout( + config: Mapping[str, object], +) -> QuantizedCheckpointLayout | None: + qconf = config.get("quantization_config") + if not isinstance(qconf, Mapping): + return None + if str(qconf.get("quant_method", "")).lower() != "fp8": + return None + scale_fmt = _deepseek_v4_scale_fmt(qconf) + weight_block_size = _deepseek_v4_weight_block_size(qconf) + expert_layout = build_deepseek_v4_packed_mxfp4_experts_layout() + + return QuantizedCheckpointLayout( + finegrained_fp8=FineGrainedFP8CheckpointLayout( + weight_name_patterns=tuple( + pattern + r"\.weight$" for pattern in _DEEPSEEK_V4_FP8_TARGET_PATTERNS + ), + exclude_patterns=_DEEPSEEK_V4_EXCLUDED_MODULES, + weight_block_size=weight_block_size, + scale_suffix="scale", + runtime_scale_name="weight_scale_inv", + scale_fmt=scale_fmt, + ), + checkpoint_consumers=(expert_layout,), + extra_quant_config={ + "expert_quant_method": expert_layout.quant_method, + "expert_block_size": expert_layout.expert_block_size, + }, + ) + + +_EXPERT_ALIAS_TO_WEIGHT = { + "gate_proj": "w1", + "up_proj": "w3", + "down_proj": "w2", +} + + +def _rename_deepseek_v4_checkpoint_key(key: str) -> str: + if key.startswith("model."): + key = key.removeprefix("model.") + if key == "embed_tokens.weight": + return "embed.weight" + if key == "lm_head.weight": + return "head.weight" + + def expert_alias(match: re.Match) -> str: + return f"{match.group(1)}{_EXPERT_ALIAS_TO_WEIGHT[match.group(2)]}" + + key = re.sub( + r"(\.ffn\.experts\.\d+\.)(gate_proj|up_proj|down_proj)(?=\.)", + expert_alias, + key, + ) + return re.sub( + r"(\.ffn\.shared_experts\.)(gate_proj|up_proj|down_proj)(?=\.)", + expert_alias, + key, + ) + + +def _remap_deepseek_v4_checkpoint_keys(state_dict: dict[str, torch.Tensor]) -> None: + removals = [] + updates = {} + for key, tensor in list(state_dict.items()): + new_key = _rename_deepseek_v4_checkpoint_key(key) + if new_key.startswith("mtp.0."): + removals.append(key) + continue + + if new_key != key: + removals.append(key) + updates[new_key] = tensor + + for key in removals: + state_dict.pop(key, None) + for key, tensor in updates.items(): + state_dict.setdefault(key, tensor) + + +def _linear( + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + tp_mode: str = "none", + layer_type: str = "unknown", + tp_min_local_shape: int = 1, +) -> torch.Tensor: + return torch.ops.auto_deploy.torch_linear_simple( + x, + weight, + bias, + tp_mode=tp_mode, + layer_type=layer_type, + tp_min_local_shape=tp_min_local_shape, + ) + + +def _linear_module( + x: torch.Tensor, + module: nn.Linear, + tp_mode: str = "none", + layer_type: str = "unknown", + tp_min_local_shape: int = 1, +) -> torch.Tensor: + return _linear( + x, + module.weight, + module.bias, + tp_mode=tp_mode, + layer_type=layer_type, + tp_min_local_shape=tp_min_local_shape, + ) + + +def _rope_scaling_value(config: PretrainedConfig, key: str, default): + rope_scaling = getattr(config, "rope_scaling", None) or {} + return rope_scaling.get(key, default) + + +def _yarn_find_correction_dim( + num_rotations: float, + dim: int, + base: float, + max_position_embeddings: int, +) -> float: + return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / ( + 2 * math.log(base) + ) + + +def _yarn_find_correction_range( + low_rot: int, + high_rot: int, + dim: int, + base: float, + max_position_embeddings: int, +) -> Tuple[int, int]: + low = math.floor(_yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)) + high = math.ceil(_yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)) + return max(low, 0), min(high, dim - 1) + + +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) + + +def _build_rope_tables( + rope_head_dim: int, + max_seq_len: int, + base: float, + original_seq_len: int, + factor: float, + beta_fast: int, + beta_slow: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + freqs = 1.0 / (base ** (torch.arange(0, rope_head_dim, 2, dtype=torch.float32) / rope_head_dim)) + if original_seq_len > 0: + low, high = _yarn_find_correction_range( + beta_fast, beta_slow, rope_head_dim, base, original_seq_len + ) + smooth = 1 - _yarn_linear_ramp_mask(low, high, rope_head_dim // 2) + freqs = freqs / factor * (1 - smooth) + freqs * smooth + + positions = torch.arange(max_seq_len, dtype=torch.float32) + freqs = torch.outer(positions, freqs) + return freqs.cos(), freqs.sin() + + +def _apply_interleaved_rope( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + inverse: bool = False, +) -> torch.Tensor: + if inverse: + sin = -sin + x_even = x[..., 0::2] + x_odd = x[..., 1::2] + out_even = x_even * cos - x_odd * sin + out_odd = x_even * sin + x_odd * cos + return torch.stack((out_even, out_odd), dim=-1).flatten(-2).to(x.dtype) + + +def _window_topk_idxs( + window_size: int, + batch_size: int, + seq_len: int, + device: torch.device, +) -> torch.Tensor: + query_positions = torch.arange(seq_len, device=device).unsqueeze(1) + key_positions = query_positions - window_size + 1 + torch.arange(window_size, device=device) + key_positions = torch.where( + (key_positions < 0) | (key_positions > query_positions), + -1, + key_positions, + ) + return key_positions.unsqueeze(0).expand(batch_size, -1, -1) + + +def _compress_topk_idxs( + ratio: int, + batch_size: int, + seq_len: int, + offset: int, + max_compressed_len: int, + device: torch.device, +) -> torch.Tensor: + compressed_positions = torch.arange(max_compressed_len, device=device) + valid_lengths = torch.arange(1, seq_len + 1, device=device).unsqueeze(1) // ratio + topk_idxs = compressed_positions.unsqueeze(0).expand(seq_len, -1) + topk_idxs = torch.where(topk_idxs < valid_lengths, topk_idxs + offset, -1) + return topk_idxs.unsqueeze(0).expand(batch_size, -1, -1) + + +def _overlap_transform(tensor: torch.Tensor, head_dim: int, value: float) -> torch.Tensor: + batch_size, compressed_len, ratio, _ = tensor.shape + previous = tensor[:, :, :, :head_dim] + current = tensor[:, :, :, head_dim:] + prefix = tensor.new_full((batch_size, 1, ratio, head_dim), value) + previous = torch.cat((prefix, previous[:, :-1]), dim=1) + return torch.cat((previous, current), dim=2) + + +def _sparse_attention( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + compressor_kv: torch.Tensor, + compressor_gate: torch.Tensor, + compressor_ape: torch.Tensor, + compressor_norm_weight: torch.Tensor, + cos_compress_table: torch.Tensor, + sin_compress_table: torch.Tensor, + position_ids: torch.Tensor, + indexer_q: torch.Tensor, + indexer_weights: torch.Tensor, + indexer_compressor_kv: torch.Tensor, + indexer_compressor_gate: torch.Tensor, + indexer_compressor_ape: torch.Tensor, + indexer_compressor_norm_weight: torch.Tensor, + softmax_scale: float, + layer_idx: int, + window_size: int, + compress_ratio: int, + max_compressed_len: Optional[int], + rope_dim: Optional[int], + rms_norm_eps: float, +) -> torch.Tensor: + return torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention( + q, + kv, + attn_sink, + topk_idxs, + compressor_kv, + compressor_gate, + compressor_ape, + compressor_norm_weight, + cos_compress_table, + sin_compress_table, + position_ids, + indexer_q, + indexer_weights, + indexer_compressor_kv, + indexer_compressor_gate, + indexer_compressor_ape, + indexer_compressor_norm_weight, + softmax_scale, + enable_sharding=True, + layer_idx=layer_idx, + layer_type="mla", + window_size=window_size, + compress_ratio=compress_ratio, + max_compressed_len=max_compressed_len, + head_dim=kv.shape[-1], + rope_dim=rope_dim, + rms_norm_eps=rms_norm_eps, + ) + + +def _hc_split_sinkhorn( + mixes: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + hc_mult: int, + sinkhorn_iters: int, + eps: float, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + pre_logits = mixes[..., :hc_mult] * hc_scale[0] + hc_base[:hc_mult] + post_logits = mixes[..., hc_mult : 2 * hc_mult] * hc_scale[1] + hc_base[hc_mult : 2 * hc_mult] + comb_logits = mixes[..., 2 * hc_mult :] * hc_scale[2] + hc_base[2 * hc_mult :] + + pre = torch.sigmoid(pre_logits) + eps + post = 2.0 * torch.sigmoid(post_logits) + comb = comb_logits.view(*comb_logits.shape[:-1], hc_mult, hc_mult) + comb = comb.softmax(dim=-1) + eps + comb = comb / (comb.sum(dim=-2, keepdim=True) + eps) + for _ in range(sinkhorn_iters - 1): + comb = comb / (comb.sum(dim=-1, keepdim=True) + eps) + comb = comb / (comb.sum(dim=-2, keepdim=True) + eps) + return pre, post, comb + + +class DeepseekV4RMSNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=torch.float32)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.auto_deploy.torch_rmsnorm(x, self.weight, self.eps) + + +class DeepseekV4RotaryEmbedding(nn.Module): + def __init__(self, config: DeepseekV4Config) -> None: + super().__init__() + factor = float(_rope_scaling_value(config, "factor", 1.0)) + original_seq_len = int(_rope_scaling_value(config, "original_max_position_embeddings", 0)) + beta_fast = int(_rope_scaling_value(config, "beta_fast", 32)) + beta_slow = int(_rope_scaling_value(config, "beta_slow", 1)) + + cos_base, sin_base = _build_rope_tables( + config.qk_rope_head_dim, + config.ad_rope_cache_len, + config.rope_theta, + 0, + 1.0, + beta_fast, + beta_slow, + ) + cos_compress, sin_compress = _build_rope_tables( + config.qk_rope_head_dim, + config.ad_rope_cache_len, + config.compress_rope_theta, + original_seq_len, + factor, + beta_fast, + beta_slow, + ) + self.register_buffer("_ad_cos_base", cos_base, persistent=False) + self.register_buffer("_ad_sin_base", sin_base, persistent=False) + self.register_buffer("_ad_cos_compress", cos_compress, persistent=False) + self.register_buffer("_ad_sin_compress", sin_compress, persistent=False) + + def forward(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return ( + self._ad_cos_base, + self._ad_sin_base, + self._ad_cos_compress, + self._ad_sin_compress, + ) + + +class DeepseekV4MLP(nn.Module): + """SwiGLU body with checkpoint names w1=gate, w3=up, w2=down.""" + + def __init__(self, hidden_size: int, intermediate_size: int, swiglu_limit: float = 0.0) -> None: + super().__init__() + self.w1 = nn.Linear(hidden_size, intermediate_size, bias=False) + self.w2 = nn.Linear(intermediate_size, hidden_size, bias=False) + self.w3 = nn.Linear(hidden_size, intermediate_size, bias=False) + self.swiglu_limit = swiglu_limit + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate = _linear_module(x, self.w1, tp_mode="colwise", layer_type="moe") + up = _linear_module(x, self.w3, tp_mode="colwise", layer_type="moe") + if self.swiglu_limit > 0: + gate = torch.clamp(gate, max=self.swiglu_limit) + up = torch.clamp(up, min=-self.swiglu_limit, max=self.swiglu_limit) + hidden = (F.silu(gate.float()) * up.float()).to(self.w2.weight.dtype) + down = _linear_module(hidden, self.w2, tp_mode="rowwise", layer_type="moe") + down = torch.ops.auto_deploy.all_reduce(down, layer_type="moe") + return down.to(x.dtype) + + +class DeepseekV4MoEGate(nn.Module): + def __init__(self, config: DeepseekV4Config, layer_idx: int) -> None: + 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.norm_topk_prob = config.norm_topk_prob + self.score_func = config.scoring_func + self.hash_routing = layer_idx < config.num_hash_layers + self.weight = nn.Parameter( + torch.empty(config.n_routed_experts, config.hidden_size, dtype=torch.float32) + ) + if self.hash_routing: + self.register_buffer( + "tid2eid", + torch.zeros(config.vocab_size, self.top_k, dtype=torch.long), + persistent=True, + ) + self.register_parameter("bias", None) + else: + self.register_buffer("tid2eid", None, persistent=False) + self.bias = nn.Parameter(torch.zeros(config.n_routed_experts, dtype=torch.float32)) + + def forward( + self, + hidden_states_flat: torch.Tensor, + input_ids_flat: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if self.score_func != "sqrtsoftplus": + raise ValueError(f"Unsupported DeepSeek V4 scoring_func: {self.score_func}") + + router_logits = _linear(hidden_states_flat.to(self.weight.dtype), self.weight).float() + scores = F.softplus(router_logits).sqrt() + + if self.hash_routing: + selected_experts = self.tid2eid[input_ids_flat.to(torch.long)].to(torch.int64) + else: + selected_experts = (scores + self.bias).topk(self.top_k, dim=-1).indices + + routing_weights = scores.gather(1, selected_experts) + if self.norm_topk_prob: + routing_weights = routing_weights / (routing_weights.sum(dim=-1, keepdim=True) + 1e-20) + routing_weights = routing_weights * self.routed_scaling_factor + return selected_experts, routing_weights + + +class DeepseekV4MoE(nn.Module): + def __init__(self, config: DeepseekV4Config, layer_idx: int) -> None: + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.intermediate_size = config.moe_intermediate_size + self.n_routed_experts = config.n_routed_experts + self.swiglu_limit = config.swiglu_limit + self.gate = DeepseekV4MoEGate(config, layer_idx) + self.experts = nn.Module() + self._register_mxfp4_runtime_buffers() + self._register_load_state_dict_pre_hook(self._load_mxfp4_checkpoint_experts) + shared_intermediate_size = config.moe_intermediate_size * config.n_shared_experts + self.shared_experts = DeepseekV4MLP( + config.hidden_size, + shared_intermediate_size, + swiglu_limit=self.swiglu_limit, + ) + + def _register_mxfp4_runtime_buffers(self) -> None: + layout = build_deepseek_v4_packed_mxfp4_experts_layout() + block_size = layout.expert_block_size + packed_block_width = block_size // layout.packed_values_per_byte + if self.hidden_size % block_size != 0 or self.intermediate_size % block_size != 0: + raise ValueError( + "DeepSeek V4 MXFP4 experts require hidden_size and " + f"moe_intermediate_size to be divisible by {block_size}." + ) + self.experts.register_buffer( + "gate_up_proj_blocks", + torch.zeros( + self.n_routed_experts, + 2 * self.intermediate_size, + self.hidden_size // block_size, + packed_block_width, + dtype=torch.uint8, + ), + persistent=True, + ) + self.experts.register_buffer( + "gate_up_proj_scales", + torch.zeros( + self.n_routed_experts, + 2 * self.intermediate_size, + self.hidden_size // block_size, + dtype=torch.uint8, + ), + persistent=True, + ) + self.experts.register_buffer( + "down_proj_blocks", + torch.zeros( + self.n_routed_experts, + self.hidden_size, + self.intermediate_size // block_size, + packed_block_width, + dtype=torch.uint8, + ), + persistent=True, + ) + self.experts.register_buffer( + "down_proj_scales", + torch.zeros( + self.n_routed_experts, + self.hidden_size, + self.intermediate_size // block_size, + dtype=torch.uint8, + ), + persistent=True, + ) + self.experts.register_buffer( + "gate_up_proj_bias", + torch.zeros( + self.n_routed_experts, + 2 * self.intermediate_size, + dtype=torch.float32, + ), + persistent=False, + ) + self.experts.register_buffer( + "down_proj_bias", + torch.zeros(self.n_routed_experts, self.hidden_size, dtype=torch.float32), + persistent=False, + ) + + def _load_mxfp4_checkpoint_experts( + self, + state_dict: dict[str, torch.Tensor], + prefix: str, + *args, + ) -> None: + del args + layout = build_deepseek_v4_packed_mxfp4_experts_layout() + target_names = { + "gate_up_blocks": f"{prefix}experts.gate_up_proj_blocks", + "gate_up_scales": f"{prefix}experts.gate_up_proj_scales", + "down_blocks": f"{prefix}experts.down_proj_blocks", + "down_scales": f"{prefix}experts.down_proj_scales", + } + packed = load_packed_mxfp4_expert_tensors( + state_dict, + "", + checkpoint_layout=layout, + target_names=target_names, + layer=self.layer_idx, + num_experts=self.n_routed_experts, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + ) + if packed is None: + return + + self._resize_mxfp4_runtime_buffer("gate_up_proj_blocks", packed.gate_up_blocks) + self._resize_mxfp4_runtime_buffer("gate_up_proj_scales", packed.gate_up_scales) + self._resize_mxfp4_runtime_buffer("down_proj_blocks", packed.down_blocks) + self._resize_mxfp4_runtime_buffer("down_proj_scales", packed.down_scales) + + def _resize_mxfp4_runtime_buffer(self, name: str, loaded: torch.Tensor) -> None: + current = getattr(self.experts, name) + if current.shape == loaded.shape and current.dtype == loaded.dtype: + return + setattr( + self.experts, + name, + torch.empty(loaded.shape, dtype=loaded.dtype, device=current.device), + ) + + def forward(self, hidden_states: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor: + original_shape = hidden_states.shape + hidden_states_flat = hidden_states.view(-1, original_shape[-1]) + selected_experts, routing_weights = self.gate(hidden_states_flat, input_ids.reshape(-1)) + routed = torch.ops.auto_deploy.torch_mxfp4_moe_from_routing( + hidden_states_flat, + selected_experts, + routing_weights.to(hidden_states_flat.dtype), + self.experts.gate_up_proj_blocks, + self.experts.gate_up_proj_bias, + self.experts.gate_up_proj_scales, + 1.0, + float(self.swiglu_limit), + self.experts.down_proj_blocks, + self.experts.down_proj_bias, + self.experts.down_proj_scales, + "up_gate", + "deepseek", + "moe", + self.n_routed_experts, + ) + return routed.view(*original_shape).to(hidden_states.dtype) + self.shared_experts( + hidden_states + ) + + +class DeepseekV4Compressor(nn.Module): + def __init__( + self, + config: DeepseekV4Config, + compress_ratio: int, + head_dim: int, + rotate: bool = False, + ) -> None: + super().__init__() + assert compress_ratio > 0, "DeepSeek V4 compressor requires a positive ratio" + self.hidden_size = config.hidden_size + self.head_dim = head_dim + self.rope_head_dim = config.qk_rope_head_dim + self.compress_ratio = compress_ratio + self.rotate = rotate + self.overlap = compress_ratio == 4 + channels = 2 if self.overlap else 1 + self.max_compressed_len = max( + 1, + (int(config.ad_compress_max_seq_len) + compress_ratio - 1) // compress_ratio, + ) + self.max_compressed_tokens = self.max_compressed_len * compress_ratio + + self.ape = nn.Parameter( + torch.zeros(compress_ratio, channels * self.head_dim, dtype=torch.float32) + ) + self.wkv = nn.Linear(self.hidden_size, channels * self.head_dim, bias=False) + self.wgate = nn.Linear(self.hidden_size, channels * self.head_dim, bias=False) + self.norm = DeepseekV4RMSNorm(self.head_dim, config.rms_norm_eps) + + def project(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + kv_all = _linear_module(hidden_states, self.wkv, layer_type="mla").float() + score_all = _linear_module(hidden_states, self.wgate, layer_type="mla").float() + return kv_all, score_all + + def compress_projected( + self, + kv_all: torch.Tensor, + score_all: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + position_ids: torch.Tensor, + output_dtype: torch.dtype, + ) -> torch.Tensor: + batch_size, seq_len, _ = kv_all.shape + ratio = self.compress_ratio + torch._check( + seq_len <= self.max_compressed_tokens, + lambda: "DeepSeek V4 compressor sequence length exceeds ad_compress_max_seq_len", + ) + + row_offsets = torch.arange(self.max_compressed_len, device=kv_all.device) + token_offsets = torch.arange(ratio, device=kv_all.device) + gather_idxs = row_offsets.unsqueeze(1) * ratio + token_offsets + valid = gather_idxs < seq_len + gather_idxs = torch.where(valid, gather_idxs, torch.zeros_like(gather_idxs)) + flat_idxs = gather_idxs.reshape(-1) + + kv = kv_all[:, flat_idxs].view(batch_size, self.max_compressed_len, ratio, -1) + score = score_all[:, flat_idxs].view(batch_size, self.max_compressed_len, ratio, -1) + score = score + self.ape + score = torch.where( + valid.view(1, self.max_compressed_len, ratio, 1), + score, + score.new_full((), -1.0e20), + ) + if self.overlap: + kv = _overlap_transform(kv, self.head_dim, 0.0) + score = _overlap_transform(score, self.head_dim, -1.0e20) + + compressed = (kv * score.softmax(dim=2)).sum(dim=2).to(output_dtype) + compressed = self.norm(compressed) + + row_start = row_offsets * ratio + row_start = torch.minimum(row_start, torch.full_like(row_start, seq_len - 1)) + compressed_position_ids = position_ids[:, row_start] + cos = cos_table[compressed_position_ids] + sin = sin_table[compressed_position_ids] + + nope, pe = torch.split( + compressed, + [self.head_dim - self.rope_head_dim, self.rope_head_dim], + dim=-1, + ) + pe = _apply_interleaved_rope(pe, cos, sin) + compressed = torch.cat((nope, pe), dim=-1) + if self.rotate: + return _fake_fp4_act_quant(_hadamard_rotate(compressed), block_size=32) + + nope, pe = torch.split( + compressed, + [self.head_dim - self.rope_head_dim, self.rope_head_dim], + dim=-1, + ) + nope = _fake_fp8_act_quant(nope, block_size=64) + return torch.cat((nope, pe), dim=-1) + + def forward( + self, + hidden_states: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + kv_all, score_all = self.project(hidden_states) + return self.compress_projected( + kv_all, + score_all, + cos_table, + sin_table, + position_ids, + hidden_states.dtype, + ) + + +class DeepseekV4Indexer(nn.Module): + def __init__(self, config: DeepseekV4Config, compress_ratio: int) -> None: + super().__init__() + self.index_n_heads = config.index_n_heads + self.index_head_dim = config.index_head_dim + self.rope_head_dim = config.qk_rope_head_dim + self.index_topk = config.index_topk + self.compress_ratio = compress_ratio + self.softmax_scale = self.index_head_dim**-0.5 + self.wq_b = nn.Linear( + config.q_lora_rank, config.index_n_heads * config.index_head_dim, bias=False + ) + self.weights_proj = nn.Linear(config.hidden_size, config.index_n_heads, bias=False) + self.compressor = DeepseekV4Compressor( + config, + compress_ratio, + config.index_head_dim, + rotate=True, + ) + + def project( + self, + hidden_states: torch.Tensor, + q_lora: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size, seq_len, _ = hidden_states.shape + q = _linear_module( + q_lora, + self.wq_b, + tp_mode="colwise", + layer_type="mla", + tp_min_local_shape=self.index_head_dim, + ) + q = torch.ops.auto_deploy.view( + q, + [batch_size, seq_len, self.index_n_heads, self.index_head_dim], + tp_scaled_dim=2, + layer_type="mla", + ) + q_nope, q_pe = torch.split( + q, + [self.index_head_dim - self.rope_head_dim, self.rope_head_dim], + dim=-1, + ) + q_pe = _apply_interleaved_rope(q_pe, cos.unsqueeze(2), sin.unsqueeze(2)) + q = _fake_fp4_act_quant(_hadamard_rotate(torch.cat((q_nope, q_pe), dim=-1)), block_size=32) + + weights = _linear_module( + hidden_states, + self.weights_proj, + tp_mode="colwise", + layer_type="mla", + ) + weights = weights.float() * (self.softmax_scale * self.index_n_heads**-0.5) + compressor_kv, compressor_gate = self.compressor.project(hidden_states) + return q, weights, compressor_kv, compressor_gate + + def select_topk( + self, + q: torch.Tensor, + index_k: torch.Tensor, + weights: torch.Tensor, + seq_len: int, + offset: int, + ) -> torch.Tensor: + index_score = torch.matmul( + q.transpose(1, 2), + index_k.transpose(1, 2).unsqueeze(1), + ).transpose(1, 2) + index_score = index_score.float() + index_score = (index_score.relu() * weights.unsqueeze(-1)).sum(dim=2) + index_score = torch.ops.auto_deploy.all_reduce(index_score, layer_type="mla") + + batch_size = q.shape[0] + compressed_positions = torch.arange( + self.compressor.max_compressed_len, + device=q.device, + ) + valid_lengths = torch.arange(1, seq_len + 1, device=q.device).unsqueeze(1) + valid_lengths = valid_lengths // self.compress_ratio + index_score = index_score.masked_fill( + (compressed_positions.unsqueeze(0) >= valid_lengths).unsqueeze(0), + -1.0e20, + ) + + if self.index_topk == 0: + return torch.empty(batch_size, seq_len, 0, device=q.device, dtype=torch.int64) + + topk_count = min(self.index_topk, self.compressor.max_compressed_len) + topk_idxs = index_score.topk(topk_count, dim=-1).indices + invalid = topk_idxs >= valid_lengths.unsqueeze(0) + topk_idxs = torch.where(invalid, -1, topk_idxs + offset) + if topk_count < self.index_topk: + pad = torch.full( + (batch_size, seq_len, self.index_topk - topk_count), + -1, + device=q.device, + dtype=topk_idxs.dtype, + ) + topk_idxs = torch.cat((topk_idxs, pad), dim=-1) + return topk_idxs.to(torch.int64) + + def forward( + self, + hidden_states: torch.Tensor, + q_lora: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + position_ids: torch.Tensor, + offset: int, + ) -> torch.Tensor: + q, weights, compressor_kv, compressor_gate = self.project(hidden_states, q_lora, cos, sin) + index_k = self.compressor.compress_projected( + compressor_kv, + compressor_gate, + cos_table, + sin_table, + position_ids, + hidden_states.dtype, + ) + return self.select_topk(q, index_k, weights, hidden_states.shape[1], offset) + + +class DeepseekV4Attention(nn.Module): + def __init__(self, config: DeepseekV4Config, layer_idx: int) -> None: + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.n_heads = config.num_attention_heads + self.head_dim = config.head_dim + self.rope_head_dim = config.qk_rope_head_dim + self.nope_head_dim = self.head_dim - self.rope_head_dim + self.q_lora_rank = config.q_lora_rank + self.o_lora_rank = config.o_lora_rank + self.n_groups = config.o_groups + self.window_size = config.sliding_window + self.compress_ratio = int(config.compress_ratios[layer_idx]) + self.rms_eps = config.rms_norm_eps + self.softmax_scale = self.head_dim**-0.5 + assert config.num_key_value_heads == 1, "DeepSeek V4 semantic model expects one KV head" + assert self.n_heads % self.n_groups == 0, "o_groups must divide num_attention_heads" + self.group_head_width = (self.n_heads * self.head_dim) // self.n_groups + + self.wq_a = nn.Linear(self.hidden_size, self.q_lora_rank, bias=False) + self.q_norm = DeepseekV4RMSNorm(self.q_lora_rank, self.rms_eps) + self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False) + self.wkv = nn.Linear(self.hidden_size, self.head_dim, bias=False) + self.kv_norm = DeepseekV4RMSNorm(self.head_dim, self.rms_eps) + self.wo_a = nn.Linear(self.group_head_width, self.n_groups * self.o_lora_rank, bias=False) + self.wo_b = nn.Linear(self.n_groups * self.o_lora_rank, self.hidden_size, bias=False) + self.attn_sink = nn.Parameter(torch.zeros(self.n_heads, dtype=torch.float32)) + if self.compress_ratio: + self.compressor = DeepseekV4Compressor(config, self.compress_ratio, self.head_dim) + self.indexer = ( + DeepseekV4Indexer(config, self.compress_ratio) if self.compress_ratio == 4 else None + ) + else: + self.compressor = None + self.indexer = None + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ], + position_ids: torch.Tensor, + ) -> torch.Tensor: + batch_size, seq_len, _ = hidden_states.shape + cos_base_table, sin_base_table, cos_compress_table, sin_compress_table = position_embeddings + cos_base = cos_base_table[position_ids] + sin_base = sin_base_table[position_ids] + cos_compress = cos_compress_table[position_ids] + sin_compress = sin_compress_table[position_ids] + cos = cos_compress if self.compress_ratio else cos_base + sin = sin_compress if self.compress_ratio else sin_base + + q_lora = _linear_module(hidden_states, self.wq_a, layer_type="mla") + q_lora = self.q_norm(q_lora) + q = _linear_module( + q_lora, + self.wq_b, + tp_mode="colwise", + layer_type="mla", + tp_min_local_shape=self.head_dim, + ) + q = torch.ops.auto_deploy.view( + q, + [batch_size, seq_len, self.n_heads, self.head_dim], + tp_scaled_dim=2, + layer_type="mla", + ) + q = q * torch.rsqrt(q.float().square().mean(-1, keepdim=True) + self.rms_eps).to(q.dtype) + + kv = _linear_module(hidden_states, self.wkv, layer_type="mla") + kv = self.kv_norm(kv) + + q_nope, q_pe = torch.split(q, [self.nope_head_dim, self.rope_head_dim], dim=-1) + kv_nope, kv_pe = torch.split(kv, [self.nope_head_dim, self.rope_head_dim], dim=-1) + q_pe = _apply_interleaved_rope(q_pe, cos.unsqueeze(2), sin.unsqueeze(2)) + kv_pe = _apply_interleaved_rope(kv_pe, cos, sin) + kv_nope = _fake_fp8_act_quant(kv_nope, block_size=64) + q = torch.cat((q_nope, q_pe), dim=-1) + kv = torch.cat((kv_nope, kv_pe), dim=-1) + + empty_compressor_state = kv.new_empty(batch_size, seq_len, 0) + empty_compressor_ape = kv.new_empty(0, 0) + empty_compressor_norm_weight = kv.new_empty(0) + empty_rope_table = kv.new_empty(0, 0) + empty_indexer_q = q.new_empty(batch_size, seq_len, 0, 0) + empty_indexer_weights = q.new_empty(batch_size, seq_len, 0) + if self.compress_ratio: + window_idxs = _window_topk_idxs( + self.window_size, batch_size, seq_len, hidden_states.device + ) + compressor_kv, compressor_gate = self.compressor.project(hidden_states) + if self.indexer is not None: + ( + indexer_q, + indexer_weights, + indexer_compressor_kv, + indexer_compressor_gate, + ) = self.indexer.project( + hidden_states, + q_lora, + cos_compress, + sin_compress, + ) + index_k = self.indexer.compressor.compress_projected( + indexer_compressor_kv, + indexer_compressor_gate, + cos_compress_table, + sin_compress_table, + position_ids, + hidden_states.dtype, + ) + compressed_idxs = self.indexer.select_topk( + indexer_q, + index_k, + indexer_weights, + seq_len, + seq_len, + ) + indexer_compressor_ape = self.indexer.compressor.ape + indexer_compressor_norm_weight = self.indexer.compressor.norm.weight + else: + indexer_q = empty_indexer_q + indexer_weights = empty_indexer_weights + indexer_compressor_kv = empty_compressor_state + indexer_compressor_gate = empty_compressor_state + indexer_compressor_ape = empty_compressor_ape + indexer_compressor_norm_weight = empty_compressor_norm_weight + compressed_idxs = _compress_topk_idxs( + self.compress_ratio, + batch_size, + seq_len, + seq_len, + self.compressor.max_compressed_len, + hidden_states.device, + ) + topk_idxs = torch.cat((window_idxs, compressed_idxs), dim=-1).to(torch.int64) + attn_output = _sparse_attention( + q, + kv, + self.attn_sink, + topk_idxs, + compressor_kv, + compressor_gate, + self.compressor.ape, + self.compressor.norm.weight, + cos_compress_table, + sin_compress_table, + position_ids, + indexer_q, + indexer_weights, + indexer_compressor_kv, + indexer_compressor_gate, + indexer_compressor_ape, + indexer_compressor_norm_weight, + self.softmax_scale, + self.layer_idx, + self.window_size, + self.compress_ratio, + self.compressor.max_compressed_len, + self.rope_head_dim, + self.rms_eps, + ) + else: + topk_idxs = _window_topk_idxs( + self.window_size, batch_size, seq_len, hidden_states.device + ).to(torch.int64) + attn_output = _sparse_attention( + q, + kv, + self.attn_sink, + topk_idxs, + empty_compressor_state, + empty_compressor_state, + empty_compressor_ape, + empty_compressor_norm_weight, + empty_rope_table, + empty_rope_table, + position_ids, + empty_indexer_q, + empty_indexer_weights, + empty_compressor_state, + empty_compressor_state, + empty_compressor_ape, + empty_compressor_norm_weight, + self.softmax_scale, + self.layer_idx, + self.window_size, + self.compress_ratio, + None, + None, + self.rms_eps, + ) + + out_nope, out_pe = torch.split( + attn_output, + [self.nope_head_dim, self.rope_head_dim], + dim=-1, + ) + out_pe = _apply_interleaved_rope(out_pe, cos.unsqueeze(2), sin.unsqueeze(2), inverse=True) + attn_output = torch.cat((out_nope, out_pe), dim=-1) + attn_output = torch.ops.auto_deploy.view( + attn_output, + [batch_size, seq_len, self.n_groups, self.group_head_width], + tp_scaled_dim=2, + layer_type="mla", + ) + wo_a = torch.ops.auto_deploy.view( + self.wo_a.weight, + [self.n_groups, self.o_lora_rank, self.group_head_width], + tp_scaled_dim=0, + layer_type="mla", + tp_min_local_shape=self.o_lora_rank, + ) + attn_output = torch.ops.auto_deploy.torch_grouped_linear( + attn_output, + wo_a, + None, + tp_mode="colwise", + layer_type="mla", + tp_min_local_shape=self.o_lora_rank, + ) + attn_output = _linear_module(attn_output, self.wo_b, tp_mode="rowwise", layer_type="mla") + attn_output = torch.ops.auto_deploy.all_reduce(attn_output, layer_type="mla") + return attn_output + + +class DeepseekV4Block(nn.Module): + def __init__(self, config: DeepseekV4Config, layer_idx: int) -> None: + super().__init__() + self.hc_mult = config.hc_mult + self.hc_sinkhorn_iters = config.hc_sinkhorn_iters + self.hc_eps = config.hc_eps + self.norm_eps = config.rms_norm_eps + self.attn = DeepseekV4Attention(config, layer_idx) + self.ffn = DeepseekV4MoE(config, layer_idx) + self.attn_norm = DeepseekV4RMSNorm(config.hidden_size, self.norm_eps) + self.ffn_norm = DeepseekV4RMSNorm(config.hidden_size, self.norm_eps) + + mix_hc = (2 + self.hc_mult) * self.hc_mult + hc_dim = self.hc_mult * config.hidden_size + self.hc_attn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim, dtype=torch.float32)) + self.hc_ffn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim, dtype=torch.float32)) + self.hc_attn_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32)) + self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32)) + self.hc_attn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32)) + self.hc_ffn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32)) + self._reset_hc_parameters() + + def _reset_hc_parameters(self) -> None: + nn.init.normal_(self.hc_attn_fn, mean=0.0, std=0.02) + nn.init.normal_(self.hc_ffn_fn, mean=0.0, std=0.02) + nn.init.zeros_(self.hc_attn_base) + nn.init.zeros_(self.hc_ffn_base) + nn.init.ones_(self.hc_attn_scale) + nn.init.ones_(self.hc_ffn_scale) + + def _hc_pre( + self, + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + original_shape = x.shape + flat = x.flatten(2).float() + rsqrt = torch.rsqrt(flat.square().mean(-1, keepdim=True) + self.norm_eps) + mixes = _linear(flat, hc_fn) * rsqrt + pre, post, comb = _hc_split_sinkhorn( + mixes, + hc_scale, + hc_base, + self.hc_mult, + self.hc_sinkhorn_iters, + self.hc_eps, + ) + y = torch.sum(pre.unsqueeze(-1) * flat.view(original_shape), dim=2) + return y.to(x.dtype), post, comb + + @staticmethod + def _hc_post( + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + ) -> torch.Tensor: + y = post.unsqueeze(-1) * x.unsqueeze(-2) + y = y + torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim=2) + return y.to(x.dtype) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ], + position_ids: torch.Tensor, + input_ids: torch.Tensor, + ) -> torch.Tensor: + residual = hidden_states + hidden_states, post, comb = self._hc_pre( + hidden_states, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + ) + hidden_states = self.attn_norm(hidden_states) + hidden_states = self.attn(hidden_states, position_embeddings, position_ids) + hidden_states = self._hc_post(hidden_states, residual, post, comb) + + residual = hidden_states + hidden_states, post, comb = self._hc_pre( + hidden_states, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + ) + hidden_states = self.ffn_norm(hidden_states) + hidden_states = self.ffn(hidden_states, input_ids) + return self._hc_post(hidden_states, residual, post, comb) + + +class DeepseekV4PreTrainedModel(PreTrainedModel): + config_class = DeepseekV4Config + base_model_prefix = "" + _no_split_modules = ["DeepseekV4Block"] + _keys_to_ignore_on_load_unexpected = [r"^mtp\.0\."] + supports_gradient_checkpointing = False + + def __init__(self, config: PretrainedConfig) -> None: + super().__init__(config) + self.register_load_state_dict_pre_hook(self._remap_load_state_hook) + + @staticmethod + def _remap_load_state_hook( + module: nn.Module, + state_dict: dict[str, torch.Tensor], + prefix: str, + *args, + ) -> None: + del args + if prefix: + return + _remap_deepseek_v4_checkpoint_keys(state_dict) + + def _init_weights(self, module: nn.Module) -> None: + std = getattr(self.config, "initializer_range", 0.02) + 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) + elif isinstance(module, DeepseekV4RMSNorm): + module.weight.data.fill_(1.0) + elif isinstance(module, DeepseekV4MoEGate): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, DeepseekV4Compressor): + module.ape.data.normal_(mean=0.0, std=std) + elif isinstance(module, DeepseekV4Attention): + module.attn_sink.data.zero_() + elif isinstance(module, DeepseekV4Block): + module.hc_attn_fn.data.normal_(mean=0.0, std=std) + module.hc_ffn_fn.data.normal_(mean=0.0, std=std) + module.hc_attn_base.data.zero_() + module.hc_ffn_base.data.zero_() + module.hc_attn_scale.data.fill_(1.0) + module.hc_ffn_scale.data.fill_(1.0) + elif isinstance(module, DeepseekV4ForCausalLM): + module.hc_head_fn.data.normal_(mean=0.0, std=std) + module.hc_head_base.data.zero_() + module.hc_head_scale.data.fill_(1.0) + + +class DeepseekV4ForCausalLM(DeepseekV4PreTrainedModel, GenerationMixin): + def __init__(self, config: DeepseekV4Config, **kwargs) -> None: + kwargs.pop("skip_mtp", None) + if kwargs: + unexpected = ", ".join(sorted(kwargs)) + raise TypeError(f"Unexpected DeepSeek V4 model kwarg(s): {unexpected}") + super().__init__(config) + self.embed = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [DeepseekV4Block(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = DeepseekV4RMSNorm(config.hidden_size, config.rms_norm_eps) + self.head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.rotary_emb = DeepseekV4RotaryEmbedding(config) + self.hc_mult = config.hc_mult + hc_dim = config.hc_mult * config.hidden_size + self.hc_head_fn = nn.Parameter(torch.empty(config.hc_mult, hc_dim, dtype=torch.float32)) + self.hc_head_base = nn.Parameter(torch.empty(config.hc_mult, dtype=torch.float32)) + self.hc_head_scale = nn.Parameter(torch.empty(1, dtype=torch.float32)) + self.post_init() + + def get_input_embeddings(self): + return self.embed + + def set_input_embeddings(self, new_embeddings): + self.embed = new_embeddings + + def get_output_embeddings(self): + return self.head + + def set_output_embeddings(self, new_embeddings): + self.head = new_embeddings + + def get_decoder(self): + return self + + def _hc_head(self, hidden_states: torch.Tensor) -> torch.Tensor: + original_dtype = hidden_states.dtype + original_shape = hidden_states.shape + flat = hidden_states.flatten(2).float() + rsqrt = torch.rsqrt(flat.square().mean(-1, keepdim=True) + self.config.rms_norm_eps) + mixes = _linear(flat, self.hc_head_fn) * rsqrt + pre = torch.sigmoid(mixes * self.hc_head_scale + self.hc_head_base) + self.config.hc_eps + hidden_states = torch.sum(pre.unsqueeze(-1) * flat.view(original_shape), dim=2) + return hidden_states.to(original_dtype) + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + ) -> DeepseekV4CausalLMOutput: + assert input_ids is not None, "input_ids is required" + assert position_ids is not None, "position_ids is required" + + hidden_states = self.embed(input_ids) + position_embeddings = self.rotary_emb() + hidden_states = hidden_states.unsqueeze(2).expand(-1, -1, self.hc_mult, -1) + for layer in self.layers: + hidden_states = layer(hidden_states, position_embeddings, position_ids, input_ids) + hidden_states = self._hc_head(hidden_states) + hidden_states = self.norm(hidden_states) + logits = _linear(hidden_states.float(), self.head.weight.float(), None).float() + return DeepseekV4CausalLMOutput(logits=logits) + + +AutoModelForCausalLMFactory.register_custom_model_cls("DeepseekV4Config", DeepseekV4ForCausalLM) + + +@ModelFactoryRegistry.register("DeepseekV4AutoModelForCausalLM") +class DeepseekV4AutoModelForCausalLMFactory(AutoModelForCausalLMFactory): + """DeepSeek V4 factory for export sizing and config overrides.""" + + def init_tokenizer(self) -> Optional[Any]: + if self.tokenizer is None: + return None + return ADDeepseekV4Tokenizer.from_pretrained(self.tokenizer, **self.tokenizer_kwargs) + + def _get_model_config(self) -> Tuple[PretrainedConfig, Dict[str, Any]]: + model_config, unused_kwargs = super()._get_model_config() + if ( + "ad_rope_cache_len" in self.model_kwargs + and "ad_compress_max_seq_len" not in self.model_kwargs + and hasattr(model_config, "ad_rope_cache_len") + and hasattr(model_config, "ad_compress_max_seq_len") + ): + model_config.ad_compress_max_seq_len = model_config.ad_rope_cache_len + return model_config, unused_kwargs + + def get_example_inputs(self) -> Dict[str, torch.Tensor]: + model_config, _ = self._get_model_config() + ratios = tuple(getattr(model_config, "compress_ratios", ())) + num_layers = int(getattr(model_config, "num_hidden_layers", len(ratios))) + active_ratios = [int(ratio) for ratio in ratios[:num_layers] if int(ratio) > 0] + export_seq_len = max(4, 2 * max(active_ratios, default=0)) + export_cap = min( + self.max_seq_len, + int(getattr(model_config, "ad_rope_cache_len", self.max_seq_len)), + int(getattr(model_config, "ad_compress_max_seq_len", self.max_seq_len)), + ) + export_seq_len = max(1, min(export_cap, export_seq_len)) + input_ids = torch.ones(2, export_seq_len, dtype=torch.long) + position_ids = ( + torch.arange(export_seq_len, dtype=torch.long).unsqueeze(0).expand_as(input_ids) + ) + return {"input_ids": input_ids, "position_ids": position_ids} + + +DeepseekV4AutoModelForCausalLMFactory.register_custom_model_cls( + "DeepseekV4Config", + DeepseekV4ForCausalLM, +) diff --git a/tensorrt_llm/_torch/auto_deploy/models/quant_checkpoint_layout.py b/tensorrt_llm/_torch/auto_deploy/models/quant_checkpoint_layout.py new file mode 100644 index 000000000000..81e1f05af4de --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/models/quant_checkpoint_layout.py @@ -0,0 +1,905 @@ +# 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. + +"""Generic checkpoint-layout helpers for pre-quantized HF checkpoints.""" + +from __future__ import annotations + +import fnmatch +import math +import operator +import re +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, field +from typing import Protocol, TypeAlias, runtime_checkable + +import torch + +TensorMetadata: TypeAlias = Mapping[str, object] +CheckpointMetadata: TypeAlias = Mapping[str, TensorMetadata] + +_E8M0_DTYPE = getattr(torch, "float8_e8m0fnu", None) + + +class QuantizedCheckpointLayoutError(ValueError): + """Raised when a quantized checkpoint layout cannot consume checkpoint tensors.""" + + +@runtime_checkable +class PackedMXFP4ExpertTensors(Protocol): + """Packed MXFP4 expert tensors materialized from checkpoint tensors.""" + + gate_up_blocks: torch.Tensor + gate_up_scales: torch.Tensor + down_blocks: torch.Tensor + down_scales: torch.Tensor + expert_indices: Sequence[int] + + +@runtime_checkable +class PackedMXFP4ExpertCheckpointLayout(Protocol): + """Consumer contract for packed MXFP4 expert tensors in checkpoint layouts.""" + + quant_method: str + + def layer_from_runtime_name(self, name: str) -> int | None: + """Return the checkpoint layer index encoded in a runtime buffer name, if present.""" + ... + + def has_layer_expert_tensors(self, state_dict: Mapping[str, torch.Tensor], layer: int) -> bool: + """Return whether a state-dict fragment contains source tensors for a layer.""" + ... + + def pack_experts( + self, + state_dict: Mapping[str, torch.Tensor], + *, + layer: int, + hidden_size: int, + intermediate_size: int, + expert_indices: Sequence[int] | None = None, + num_experts: int | None = None, + ) -> PackedMXFP4ExpertTensors: + """Pack checkpoint tensors into runtime MXFP4 expert tensors.""" + ... + + def source_keys_for_packed_experts( + self, layer: int, expert_indices: Sequence[int] + ) -> Sequence[str]: + """Return checkpoint source keys consumed when packing experts.""" + ... + + +@dataclass(frozen=True) +class PackedMXFP4ExpertTensorKey: + """Parsed packed MXFP4 expert tensor key.""" + + layer: int + expert: int + projection: str + tensor_kind: str + + +@dataclass(frozen=True) +class PackedMXFP4Experts: + """Packed MXFP4 expert tensors in runtime layout.""" + + gate_up_blocks: torch.Tensor + gate_up_scales: torch.Tensor + down_blocks: torch.Tensor + down_scales: torch.Tensor + expert_indices: tuple[int, ...] + hidden_size: int + intermediate_size: int + + +@dataclass(frozen=True) +class PackedMXFP4ExpertLayout: + """Regex-driven checkpoint layout for packed MXFP4 routed expert tensors.""" + + key_pattern: re.Pattern[str] + key_template: str + layer_name_pattern: re.Pattern[str] + projections: tuple[str, ...] + gate_up_order: tuple[str, ...] + down_projection: str + expert_block_size: int = 32 + packed_values_per_byte: int = 2 + weight_dtypes: tuple[str, ...] = ("I8",) + scale_dtype: str = "F8_E8M0" + quant_method: str = "mxfp4" + + def parse_key(self, name: str) -> PackedMXFP4ExpertTensorKey | None: + match = self.key_pattern.fullmatch(name) + if match is None: + return None + projection = match.group("projection") + if projection not in self.projections: + return None + return PackedMXFP4ExpertTensorKey( + layer=int(match.group("layer")), + expert=int(match.group("expert")), + projection=projection, + tensor_kind=match.group("kind"), + ) + + def layer_from_runtime_name(self, name: str) -> int | None: + match = self.layer_name_pattern.search(name) + if match is None: + return None + return int(match.group("layer")) + + def validate_checkpoint_metadata(self, tensor_metadata: CheckpointMetadata) -> int: + parsed_metadata: dict[tuple[int, int, str], dict[str, TensorMetadata]] = {} + expert_projections: dict[tuple[int, int], set[str]] = {} + for name, metadata in tensor_metadata.items(): + parsed = self.parse_key(name) + if parsed is None: + continue + if not isinstance(metadata, Mapping): + raise QuantizedCheckpointLayoutError(f"Invalid metadata for tensor {name}.") + key = (parsed.layer, parsed.expert, parsed.projection) + parsed_metadata.setdefault(key, {})[parsed.tensor_kind] = metadata + expert_projections.setdefault((parsed.layer, parsed.expert), set()).add( + parsed.projection + ) + + for (layer, expert, projection), tensors in parsed_metadata.items(): + weight_name = self.format_key(layer, expert, projection, "weight") + scale_name = self.format_key(layer, expert, projection, "scale") + weight_metadata = tensors.get("weight") + scale_metadata = tensors.get("scale") + if not isinstance(weight_metadata, Mapping): + raise QuantizedCheckpointLayoutError( + f"{scale_name} is missing companion weight {weight_name}." + ) + if not isinstance(scale_metadata, Mapping): + raise QuantizedCheckpointLayoutError( + f"{weight_name} is missing companion scale {scale_name}." + ) + self._validate_weight_metadata(weight_name, weight_metadata) + self._validate_scale_metadata(weight_name, weight_metadata, scale_name, scale_metadata) + + for (layer, expert), projections in expert_projections.items(): + for projection in self.required_projection_names(): + if projection in projections: + continue + weight_name = self.format_key(layer, expert, projection, "weight") + scale_name = self.format_key(layer, expert, projection, "scale") + raise QuantizedCheckpointLayoutError( + "Packed MXFP4 expert metadata is missing required projection " + f"{projection} for layer {layer}, expert {expert}: " + f"{weight_name} and {scale_name}." + ) + self._validate_expert_projection_shapes(parsed_metadata, layer, expert) + + if not parsed_metadata: + raise QuantizedCheckpointLayoutError( + "Checkpoint metadata has no packed MXFP4 expert tensors." + ) + return len(parsed_metadata) + + def format_key(self, layer: int, expert: int, projection: str, tensor_kind: str) -> str: + return self.key_template.format( + layer=layer, + expert=expert, + projection=projection, + kind=tensor_kind, + tensor_kind=tensor_kind, + ) + + def has_layer_expert_tensors(self, state_dict: Mapping[str, torch.Tensor], layer: int) -> bool: + layer = _validate_nonnegative_int("layer", layer) + for name in state_dict: + parsed = self.parse_key(name) + if parsed is not None and parsed.layer == layer: + return True + return False + + def pack_experts( + self, + state_dict: Mapping[str, torch.Tensor], + *, + layer: int, + hidden_size: int, + intermediate_size: int, + expert_indices: Sequence[int] | None = None, + num_experts: int | None = None, + ) -> PackedMXFP4Experts: + layer = _validate_nonnegative_int("layer", layer) + hidden_size = _validate_positive_divisible_by( + "hidden_size", hidden_size, self.expert_block_size + ) + intermediate_size = _validate_positive_divisible_by( + "intermediate_size", intermediate_size, self.expert_block_size + ) + + expert_tensors = self._collect_expert_tensors(state_dict, layer) + expert_order = _normalize_expert_order( + expert_tensors, + expert_indices=expert_indices, + num_experts=num_experts, + ) + + gate_up_blocks = [] + gate_up_scales = [] + down_blocks = [] + down_scales = [] + for expert in expert_order: + expert_map = self._get_required_expert_map(expert_tensors, layer, expert) + gate_up_blocks.append( + torch.cat( + [ + self._load_weight_blocks( + expert_map, + layer, + expert, + projection, + expected_shape=( + intermediate_size, + hidden_size // self.packed_values_per_byte, + ), + packed_shape=self._packed_shape(intermediate_size, hidden_size), + ) + for projection in self.gate_up_order + ], + dim=0, + ) + ) + gate_up_scales.append( + torch.cat( + [ + self._load_scales( + expert_map, + layer, + expert, + projection, + expected_shape=( + intermediate_size, + hidden_size // self.expert_block_size, + ), + ) + for projection in self.gate_up_order + ], + dim=0, + ) + ) + down_blocks.append( + self._load_weight_blocks( + expert_map, + layer, + expert, + self.down_projection, + expected_shape=( + hidden_size, + intermediate_size // self.packed_values_per_byte, + ), + packed_shape=self._packed_shape(hidden_size, intermediate_size), + ) + ) + down_scales.append( + self._load_scales( + expert_map, + layer, + expert, + self.down_projection, + expected_shape=(hidden_size, intermediate_size // self.expert_block_size), + ) + ) + + return PackedMXFP4Experts( + gate_up_blocks=torch.stack(gate_up_blocks, dim=0), + gate_up_scales=torch.stack(gate_up_scales, dim=0), + down_blocks=torch.stack(down_blocks, dim=0), + down_scales=torch.stack(down_scales, dim=0), + expert_indices=expert_order, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + ) + + def source_keys_for_packed_experts( + self, layer: int, expert_indices: Sequence[int] + ) -> tuple[str, ...]: + return tuple( + self.format_key(layer, expert, projection, tensor_kind) + for expert in expert_indices + for projection in self.required_projection_names() + for tensor_kind in ("weight", "scale") + ) + + def required_projection_names(self) -> tuple[str, ...]: + return tuple(dict.fromkeys((*self.gate_up_order, self.down_projection))) + + def _collect_expert_tensors( + self, + state_dict: Mapping[str, torch.Tensor], + layer: int, + ) -> dict[int, dict[str, dict[str, torch.Tensor]]]: + expert_tensors: dict[int, dict[str, dict[str, torch.Tensor]]] = {} + for name, tensor in state_dict.items(): + parsed = self.parse_key(name) + if parsed is None or parsed.layer != layer: + continue + if not isinstance(tensor, torch.Tensor): + raise QuantizedCheckpointLayoutError(f"{name} should be a torch.Tensor.") + weight_tensors = expert_tensors.setdefault(parsed.expert, {}) + tensor_kinds = weight_tensors.setdefault(parsed.projection, {}) + tensor_kinds[parsed.tensor_kind] = tensor + return expert_tensors + + def _get_required_expert_map( + self, + expert_tensors: Mapping[int, dict[str, dict[str, torch.Tensor]]], + layer: int, + expert: int, + ) -> dict[str, dict[str, torch.Tensor]]: + expert_map = expert_tensors.get(expert) + if expert_map is None: + first_missing = self.format_key(layer, expert, self.gate_up_order[0], "weight") + raise QuantizedCheckpointLayoutError(f"Missing packed MXFP4 tensor {first_missing}.") + return expert_map + + def _load_weight_blocks( + self, + expert_map: Mapping[str, Mapping[str, torch.Tensor]], + layer: int, + expert: int, + projection: str, + *, + expected_shape: tuple[int, int], + packed_shape: tuple[int, int, int], + ) -> torch.Tensor: + name = self.format_key(layer, expert, projection, "weight") + tensor = self._get_required_tensor(expert_map, layer, expert, projection, "weight") + _validate_tensor_shape(name, tensor, expected_shape) + return self._reinterpret_weight_as_uint8(name, tensor).view(packed_shape) + + def _load_scales( + self, + expert_map: Mapping[str, Mapping[str, torch.Tensor]], + layer: int, + expert: int, + projection: str, + *, + expected_shape: tuple[int, int], + ) -> torch.Tensor: + name = self.format_key(layer, expert, projection, "scale") + tensor = self._get_required_tensor(expert_map, layer, expert, projection, "scale") + _validate_tensor_shape(name, tensor, expected_shape) + return self._reinterpret_scale_as_uint8(name, tensor).view(expected_shape) + + def _get_required_tensor( + self, + expert_map: Mapping[str, Mapping[str, torch.Tensor]], + layer: int, + expert: int, + projection: str, + tensor_kind: str, + ) -> torch.Tensor: + tensor = expert_map.get(projection, {}).get(tensor_kind) + if tensor is None: + name = self.format_key(layer, expert, projection, tensor_kind) + raise QuantizedCheckpointLayoutError(f"Missing packed MXFP4 tensor {name}.") + return tensor + + def _packed_shape(self, rows: int, logical_cols: int) -> tuple[int, int, int]: + return ( + rows, + logical_cols // self.expert_block_size, + self.expert_block_size // self.packed_values_per_byte, + ) + + def _validate_weight_metadata(self, name: str, metadata: TensorMetadata) -> None: + dtype = _get_dtype(name, metadata) + if dtype not in self.weight_dtypes: + raise QuantizedCheckpointLayoutError( + f"{name} should be packed FP4 with dtype {self.weight_dtypes}, got {dtype}." + ) + _require_2d(name, _get_shape(name, metadata)) + + def _validate_scale_metadata( + self, + weight_name: str, + weight_metadata: TensorMetadata, + scale_name: str, + scale_metadata: TensorMetadata, + ) -> None: + dtype = _get_dtype(scale_name, scale_metadata) + if dtype != self.scale_dtype: + raise QuantizedCheckpointLayoutError( + f"{scale_name} should be {self.scale_dtype} scale, got {dtype}." + ) + weight_shape = _get_shape(weight_name, weight_metadata) + _require_2d(weight_name, weight_shape) + expected = [ + weight_shape[0], + math.ceil(weight_shape[1] * self.packed_values_per_byte / self.expert_block_size), + ] + shape = _get_shape(scale_name, scale_metadata) + if shape != expected: + raise QuantizedCheckpointLayoutError( + f"{scale_name} has shape {shape}, expected {expected} for packed FP4 {weight_name}." + ) + + def _validate_expert_projection_shapes( + self, + parsed_metadata: Mapping[tuple[int, int, str], Mapping[str, TensorMetadata]], + layer: int, + expert: int, + ) -> None: + weight_shapes = { + projection: _get_shape( + self.format_key(layer, expert, projection, "weight"), + parsed_metadata[(layer, expert, projection)]["weight"], + ) + for projection in self.required_projection_names() + } + gate_up_shapes = [weight_shapes[projection] for projection in self.gate_up_order] + reference_gate_up_shape = gate_up_shapes[0] + for projection, shape in zip(self.gate_up_order[1:], gate_up_shapes[1:]): + if shape != reference_gate_up_shape: + raise QuantizedCheckpointLayoutError( + "Packed MXFP4 expert metadata has inconsistent gate/up shapes for " + f"layer {layer}, expert {expert}: {self.gate_up_order[0]} " + f"{reference_gate_up_shape}, {projection} {shape}." + ) + + intermediate_size = reference_gate_up_shape[0] + hidden_size = reference_gate_up_shape[1] * self.packed_values_per_byte + _validate_positive_divisible_by("hidden_size", hidden_size, self.expert_block_size) + _validate_positive_divisible_by( + "intermediate_size", intermediate_size, self.expert_block_size + ) + expected_down_shape = [hidden_size, intermediate_size // self.packed_values_per_byte] + down_shape = weight_shapes[self.down_projection] + if down_shape != expected_down_shape: + raise QuantizedCheckpointLayoutError( + "Packed MXFP4 expert metadata has inconsistent down projection shape for " + f"layer {layer}, expert {expert}: {self.down_projection} {down_shape}, " + f"expected {expected_down_shape} from gate/up shape {reference_gate_up_shape}." + ) + + def _reinterpret_weight_as_uint8(self, name: str, tensor: torch.Tensor) -> torch.Tensor: + if tensor.dtype not in (torch.int8, torch.uint8): + raise QuantizedCheckpointLayoutError( + f"{name} should be packed FP4 bytes with dtype torch.int8 or torch.uint8, " + f"got {tensor.dtype}." + ) + return _view_as_uint8(tensor) + + def _reinterpret_scale_as_uint8(self, name: str, tensor: torch.Tensor) -> torch.Tensor: + if tensor.dtype == torch.uint8 or (_E8M0_DTYPE is not None and tensor.dtype == _E8M0_DTYPE): + return _view_as_uint8(tensor) + expected = "torch.uint8" + if _E8M0_DTYPE is not None: + expected += " or torch.float8_e8m0fnu" + raise QuantizedCheckpointLayoutError( + f"{name} should contain raw E8M0 scale bytes with dtype {expected}, got {tensor.dtype}." + ) + + +def load_packed_mxfp4_expert_tensors( + state_dict: dict[str, torch.Tensor], + prefix: str, + *, + checkpoint_layout: PackedMXFP4ExpertCheckpointLayout, + target_names: Mapping[str, str], + layer: int, + num_experts: int, + hidden_size: int, + intermediate_size: int, +) -> PackedMXFP4ExpertTensors | None: + """Materialize packed MXFP4 checkpoint expert tensors into runtime buffers.""" + + prefix = prefix or "" + if prefix + target_names["gate_up_blocks"] in state_dict: + return None + source_state = _strip_prefix_from_state_dict(state_dict, prefix) + if not checkpoint_layout.has_layer_expert_tensors(source_state, layer): + return None + + packed = checkpoint_layout.pack_experts( + source_state, + layer=layer, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_experts=num_experts, + ) + for source_key in checkpoint_layout.source_keys_for_packed_experts( + layer, packed.expert_indices + ): + state_dict.pop(prefix + source_key, None) + state_dict[prefix + target_names["gate_up_blocks"]] = packed.gate_up_blocks + state_dict[prefix + target_names["gate_up_scales"]] = packed.gate_up_scales + state_dict[prefix + target_names["down_blocks"]] = packed.down_blocks + state_dict[prefix + target_names["down_scales"]] = packed.down_scales + return packed + + +def _strip_prefix_from_state_dict( + state_dict: Mapping[str, torch.Tensor], prefix: str +) -> dict[str, torch.Tensor]: + if not prefix: + return dict(state_dict) + return { + key.removeprefix(prefix): tensor + for key, tensor in state_dict.items() + if key.startswith(prefix) + } + + +class QuantCheckpointLayoutRegistry: + """Registry for model-owned quantized checkpoint layout builders.""" + + _registry: dict[str, Callable[[Mapping[str, object]], "QuantizedCheckpointLayout | None"]] = {} + + @classmethod + def register( + cls, model_type: str + ) -> Callable[ + [Callable[[Mapping[str, object]], "QuantizedCheckpointLayout | None"]], + Callable[[Mapping[str, object]], "QuantizedCheckpointLayout | None"], + ]: + def inner( + builder: Callable[[Mapping[str, object]], "QuantizedCheckpointLayout | None"], + ) -> Callable[[Mapping[str, object]], "QuantizedCheckpointLayout | None"]: + cls._registry[model_type] = builder + return builder + + return inner + + @classmethod + def build_from_config(cls, config: Mapping[str, object]) -> "QuantizedCheckpointLayout | None": + model_type = config.get("model_type") + if not isinstance(model_type, str): + return None + builder = cls._registry.get(model_type) + if builder is None: + return None + return builder(config) + + +@dataclass(frozen=True) +class FineGrainedFP8CheckpointLayout: + """Checkpoint layout for block-wise FP8 linear weights.""" + + weight_name_patterns: tuple[str, ...] + weight_block_size: tuple[int, int] + scale_suffix: str = "scale" + runtime_scale_name: str = "weight_scale_inv" + scale_fmt: str = "ue8m0" + weight_dtype: str = "F8_E4M3" + scale_dtype: str = "F8_E8M0" + exclude_patterns: tuple[str, ...] = () + quant_method: str = "finegrained_fp8" + + def is_weight_targeted( + self, + weight_name: str, + extra_exclude_patterns: Sequence[str] | None = None, + ) -> bool: + if not any(re.match(pattern, weight_name) for pattern in self.weight_name_patterns): + return False + return not self.is_weight_excluded(weight_name, extra_exclude_patterns) + + def is_weight_excluded( + self, + weight_name: str, + extra_exclude_patterns: Sequence[str] | None = None, + ) -> bool: + module_name = weight_name.removesuffix(".weight") + for pattern in (*self.exclude_patterns, *(extra_exclude_patterns or ())): + if fnmatch.fnmatchcase(weight_name, pattern) or fnmatch.fnmatchcase( + module_name, pattern + ): + return True + return False + + def scale_name_for_weight(self, weight_name: str) -> str: + module_name = weight_name.removesuffix(".weight") + return f"{module_name}.{self.scale_suffix}" + + def runtime_scale_name_for_weight(self, weight_name: str) -> str: + module_name = weight_name.rsplit(".", 1)[0] + return f"{module_name}.{self.runtime_scale_name}" + + def default_scales(self, original_weight_shape: Sequence[int]) -> dict[str, torch.Tensor]: + if len(original_weight_shape) != 2: + raise QuantizedCheckpointLayoutError( + f"FineGrained FP8 weight should be 2D, got {tuple(original_weight_shape)}." + ) + n, k = (int(dim) for dim in original_weight_shape) + block_n, block_k = self.weight_block_size + scale_shape = (math.ceil(n / block_n), math.ceil(k / block_k)) + return {self.runtime_scale_name: torch.ones(scale_shape, dtype=torch.float32)} + + def decode_scale(self, scale: torch.Tensor) -> torch.Tensor: + if self.scale_fmt.lower() == "ue8m0" and ( + scale.dtype == torch.uint8 or (_E8M0_DTYPE is not None and scale.dtype == _E8M0_DTYPE) + ): + raw_exponents = scale.contiguous().view(torch.uint8) + float_bits = raw_exponents.to(torch.int32) << 23 + float_bits = torch.where( + raw_exponents == 0, + torch.full_like(float_bits, 1 << 22), + float_bits, + ) + float_bits = torch.where( + raw_exponents == 255, + torch.full_like(float_bits, 0x7FC00000), + float_bits, + ) + return float_bits.view(torch.float32).reshape(scale.shape) + return scale.to(torch.float32) + + def validate_checkpoint_metadata(self, tensor_metadata: CheckpointMetadata) -> int: + count = 0 + for name, metadata in tensor_metadata.items(): + if not isinstance(metadata, Mapping): + raise QuantizedCheckpointLayoutError(f"Invalid metadata for tensor {name}.") + if name.endswith(".weight") and self.is_weight_targeted(name): + self._validate_weight_metadata(name, metadata) + scale_name = self.scale_name_for_weight(name) + scale_metadata = tensor_metadata.get(scale_name) + if not isinstance(scale_metadata, Mapping): + raise QuantizedCheckpointLayoutError( + f"{name} is missing companion scale {scale_name}." + ) + self._validate_scale_metadata(name, metadata, scale_name, scale_metadata) + count += 1 + elif name.endswith(f".{self.scale_suffix}"): + weight_name = name.removesuffix(f".{self.scale_suffix}") + ".weight" + if self.is_weight_targeted(weight_name): + weight_metadata = tensor_metadata.get(weight_name) + if not isinstance(weight_metadata, Mapping): + raise QuantizedCheckpointLayoutError( + f"{name} is missing companion weight {weight_name}." + ) + self._validate_scale_metadata(weight_name, weight_metadata, name, metadata) + if count == 0: + raise QuantizedCheckpointLayoutError( + "Checkpoint metadata has no FineGrained FP8 tensors." + ) + return count + + def validate_scale_shape( + self, + weight_name: str, + weight_shape: Sequence[int] | torch.Tensor, + scale_name: str, + scale_shape: Sequence[int] | torch.Tensor, + ) -> None: + if isinstance(weight_shape, torch.Tensor): + weight_shape = weight_shape.shape + if isinstance(scale_shape, torch.Tensor): + scale_shape = scale_shape.shape + _require_2d(weight_name, weight_shape) + block_n, block_k = self.weight_block_size + expected = ( + math.ceil(int(weight_shape[0]) / block_n), + math.ceil(int(weight_shape[1]) / block_k), + ) + actual = tuple(int(dim) for dim in scale_shape) + if actual != expected: + raise QuantizedCheckpointLayoutError( + f"{scale_name} shape {actual} does not match expected {expected} " + f"(expected {list(expected)}) for {weight_name}." + ) + + def _validate_weight_metadata(self, name: str, metadata: TensorMetadata) -> None: + dtype = _get_dtype(name, metadata) + if dtype != self.weight_dtype: + raise QuantizedCheckpointLayoutError( + f"{name} should be {self.weight_dtype}, got {dtype}." + ) + _require_2d(name, _get_shape(name, metadata)) + + def _validate_scale_metadata( + self, + weight_name: str, + weight_metadata: TensorMetadata, + scale_name: str, + scale_metadata: TensorMetadata, + ) -> None: + dtype = _get_dtype(scale_name, scale_metadata) + if dtype != self.scale_dtype: + raise QuantizedCheckpointLayoutError( + f"{scale_name} should be {self.scale_dtype} scale, got {dtype}." + ) + self.validate_scale_shape( + weight_name, + _get_shape(weight_name, weight_metadata), + scale_name, + _get_shape(scale_name, scale_metadata), + ) + + +@dataclass(frozen=True) +class QuantizedCheckpointLayout: + """Quantized checkpoint layout selected by a model-owned layout builder.""" + + finegrained_fp8: FineGrainedFP8CheckpointLayout | None = None + checkpoint_consumers: tuple[object, ...] = () + extra_model_kwargs: Mapping[str, object] = field(default_factory=dict) + extra_quant_config: Mapping[str, object] = field(default_factory=dict) + + def validate_checkpoint_metadata(self, tensor_metadata: CheckpointMetadata) -> None: + if self.finegrained_fp8 is not None: + self.finegrained_fp8.validate_checkpoint_metadata(tensor_metadata) + for consumer in self.checkpoint_consumers: + validate = getattr(consumer, "validate_checkpoint_metadata", None) + if validate is None: + raise QuantizedCheckpointLayoutError( + f"Checkpoint consumer {type(consumer).__name__} cannot validate metadata." + ) + validate(tensor_metadata) + + def apply_to_quant_config(self, quant_config: Mapping[str, object]) -> dict[str, object]: + normalized = dict(quant_config) + normalized["checkpoint_layout"] = self + if self.finegrained_fp8 is not None: + scale_fmt = normalized.get("scale_fmt") + if scale_fmt is not None and str(scale_fmt).lower() != self.finegrained_fp8.scale_fmt: + raise QuantizedCheckpointLayoutError( + "Quantized checkpoint layout requires " + f"scale_fmt='{self.finegrained_fp8.scale_fmt}', got '{scale_fmt}'." + ) + block_size = normalized.get("weight_block_size") + if block_size is not None: + parsed_block_size = _normalize_block_size("weight_block_size", block_size) + if parsed_block_size != self.finegrained_fp8.weight_block_size: + raise QuantizedCheckpointLayoutError( + "Quantized checkpoint layout requires " + f"weight_block_size={list(self.finegrained_fp8.weight_block_size)}, " + f"got {list(parsed_block_size)}." + ) + normalized["linear_quant_method"] = self.finegrained_fp8.quant_method + normalized["scale_fmt"] = self.finegrained_fp8.scale_fmt + normalized["weight_block_size"] = list(self.finegrained_fp8.weight_block_size) + excludes = list(normalized.get("exclude_modules") or []) + for pattern in self.finegrained_fp8.exclude_patterns: + if pattern not in excludes: + excludes.append(pattern) + normalized["exclude_modules"] = excludes + normalized.update(self.extra_quant_config) + return normalized + + +def _get_dtype(name: str, metadata: TensorMetadata) -> str: + dtype = metadata.get("dtype") + if not isinstance(dtype, str) or not dtype: + raise QuantizedCheckpointLayoutError(f"Tensor {name} is missing dtype metadata.") + return dtype.upper() + + +def _get_shape(name: str, metadata: TensorMetadata) -> list[int]: + shape = metadata.get("shape") + if not isinstance(shape, Sequence) or isinstance(shape, str): + raise QuantizedCheckpointLayoutError(f"Tensor {name} is missing shape metadata.") + try: + return [int(dim) for dim in shape] + except (TypeError, ValueError) as error: + raise QuantizedCheckpointLayoutError( + f"Tensor {name} has invalid shape metadata {shape}." + ) from error + + +def _require_2d(name: str, shape: Sequence[int]) -> None: + if len(shape) != 2: + raise QuantizedCheckpointLayoutError( + f"Tensor {name} should be 2D, got shape {list(shape)}." + ) + + +def _normalize_block_size(name: str, value: object) -> tuple[int, int]: + if not isinstance(value, Sequence) or isinstance(value, str): + raise QuantizedCheckpointLayoutError(f"{name} should be a two-element integer sequence.") + try: + normalized = tuple(int(dim) for dim in value) + except (TypeError, ValueError) as error: + raise QuantizedCheckpointLayoutError( + f"{name} should contain integer values, got {value}." + ) from error + if len(normalized) != 2 or any(dim <= 0 for dim in normalized): + raise QuantizedCheckpointLayoutError( + f"{name} should be a two-element positive integer sequence, got {value}." + ) + return normalized + + +def _validate_tensor_shape( + name: str, tensor: torch.Tensor, expected_shape: tuple[int, int] +) -> None: + actual_shape = tuple(tensor.shape) + if actual_shape != expected_shape: + raise QuantizedCheckpointLayoutError( + f"{name} has shape {list(actual_shape)}, expected {list(expected_shape)}." + ) + + +def _view_as_uint8(tensor: torch.Tensor) -> torch.Tensor: + if tensor.dtype == torch.uint8 and tensor.is_contiguous(): + return tensor + return tensor.contiguous().view(torch.uint8) + + +def _normalize_expert_order( + expert_tensors: Mapping[int, object], + *, + expert_indices: Sequence[int] | None, + num_experts: int | None, +) -> tuple[int, ...]: + if num_experts is not None: + num_experts = _validate_positive_int("num_experts", num_experts) + + if expert_indices is None: + if num_experts is not None: + expert_order = tuple(range(num_experts)) + else: + expert_order = tuple(sorted(expert_tensors)) + else: + expert_order = tuple( + _validate_nonnegative_int("expert_indices", idx) for idx in expert_indices + ) + + if not expert_tensors: + raise QuantizedCheckpointLayoutError("No packed MXFP4 routed expert tensors were found.") + if not expert_order: + raise QuantizedCheckpointLayoutError("No packed MXFP4 expert indices were selected.") + if len(set(expert_order)) != len(expert_order): + raise QuantizedCheckpointLayoutError( + f"expert_indices should not contain duplicates, got {expert_order}." + ) + if num_experts is not None: + invalid_indices = [expert for expert in expert_order if expert >= num_experts] + if invalid_indices: + raise QuantizedCheckpointLayoutError( + f"expert_indices should be less than num_experts={num_experts}, got {invalid_indices}." + ) + return expert_order + + +def _validate_positive_divisible_by(name: str, value: int, divisor: int) -> int: + value = _validate_positive_int(name, value) + if value % divisor != 0: + raise QuantizedCheckpointLayoutError( + f"{name} should be divisible by {divisor}, got {value}." + ) + return value + + +def _validate_positive_int(name: str, value: int) -> int: + value = _validate_int(name, value) + if value <= 0: + raise QuantizedCheckpointLayoutError(f"{name} should be positive, got {value}.") + return value + + +def _validate_nonnegative_int(name: str, value: int) -> int: + value = _validate_int(name, value) + if value < 0: + raise QuantizedCheckpointLayoutError(f"{name} should be non-negative, got {value}.") + return value + + +def _validate_int(name: str, value: int) -> int: + if isinstance(value, bool): + raise QuantizedCheckpointLayoutError(f"{name} should be an integer, got {value}.") + try: + return operator.index(value) + except TypeError as error: + raise QuantizedCheckpointLayoutError( + f"{name} should be an integer, got {type(value).__name__}." + ) from error diff --git a/tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py b/tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py index bcffb6a7c339..37b79b24ffc8 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py +++ b/tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py @@ -23,10 +23,17 @@ import json import os from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import Any, Callable, Dict, Optional, Tuple, Type from .._compat import is_modelopt_quant_config, read_modelopt_quant_config from ..utils.logger import ad_logger +from .checkpoint_metadata import has_safetensors_metadata, read_safetensors_metadata +from .quant_checkpoint_layout import ( + QuantCheckpointLayoutRegistry, + QuantizedCheckpointLayout, + QuantizedCheckpointLayoutError, +) class QuantConfigReader(ABC): @@ -56,7 +63,9 @@ def read_config(self, config: Dict) -> Dict: @classmethod @abstractmethod - def from_file(cls, file_path: str) -> Optional[Tuple["QuantConfigReader", Dict[str, Any]]]: + def from_file( + cls, file_path: str, require_checkpoint_metadata: bool = True + ) -> Optional[Tuple["QuantConfigReader", Dict[str, Any]]]: """ Load and parse a quantization config file from disk. @@ -64,6 +73,8 @@ def from_file(cls, file_path: str) -> Optional[Tuple["QuantConfigReader", Dict[s Args: file_path: Path to the quant config JSON file. + require_checkpoint_metadata: Whether checkpoint metadata required by the reader must + be present during detection. Returns: A (reader, extra_model_kwargs) tuple, or None if the file doesn't exist. @@ -185,17 +196,20 @@ def _handle_kv_cache(self, quant_config: Dict) -> None: @classmethod def from_file( - cls, ckpt_dir: str + cls, ckpt_dir: str, require_checkpoint_metadata: bool = True ) -> Optional[Tuple["ModelOPTQuantConfigReader", Dict[str, Any]]]: """ Load and parse a modelopt-style quantization config from a checkpoint directory. Args: ckpt_dir: Path to the root directory containing the checkpoint. + require_checkpoint_metadata: Ignored; ModelOPT quant config detection does not read + checkpoint metadata. Returns: An initialized ModelOPTQuantConfigReader instance, or None if the file doesn't exist. """ + _ = require_checkpoint_metadata quant_file = os.path.join(ckpt_dir, "hf_quant_config.json") if not os.path.exists(quant_file): return None @@ -226,6 +240,18 @@ def read_config(self, config: Dict) -> Dict: if not qconf: raise ValueError("HF quantization_config not found.") + layout = self._get_checkpoint_layout(config) + if layout is not None: + tensor_metadata = config.get("checkpoint_tensor_metadata") + if not isinstance(tensor_metadata, Mapping): + raise QuantizedCheckpointLayoutError( + "HF quantized checkpoint layout requires checkpoint_tensor_metadata." + ) + layout.validate_checkpoint_metadata(tensor_metadata) + self._quant_config = layout.apply_to_quant_config(qconf) + self._hf_quantizer = None + return dict(layout.extra_model_kwargs) + # Inject default exclusion, add "model.embed_tokens" for "tie_word_embedding:true" case excludes = qconf.get("exclude_modules", []) qconf["exclude_modules"] = excludes + [n for n in self._ALWAYS_EXCLUDE if n not in excludes] @@ -240,8 +266,17 @@ def read_config(self, config: Dict) -> Dict: return {} + @staticmethod + def _get_checkpoint_layout(config: Mapping[str, object]) -> QuantizedCheckpointLayout | None: + layout = config.get("checkpoint_layout") + if isinstance(layout, QuantizedCheckpointLayout): + return layout + return QuantCheckpointLayoutRegistry.build_from_config(config) + @classmethod - def from_file(cls, ckpt_dir: str) -> Optional[Tuple["HFQuantConfigReader", Dict[str, Any]]]: + def from_file( + cls, ckpt_dir: str, require_checkpoint_metadata: bool = True + ) -> Optional[Tuple["HFQuantConfigReader", Dict[str, Any]]]: config_file = os.path.join(ckpt_dir, "config.json") if not os.path.exists(config_file): return None @@ -257,6 +292,17 @@ def from_file(cls, ckpt_dir: str) -> Optional[Tuple["HFQuantConfigReader", Dict[ if quant_method not in cls._SUPPORTED_QUANT_METHODS: return None + layout = QuantCheckpointLayoutRegistry.build_from_config(raw) + if layout is not None: + if not require_checkpoint_metadata and not has_safetensors_metadata(ckpt_dir): + return None + raw = dict(raw) + raw["checkpoint_layout"] = layout + raw["checkpoint_tensor_metadata"] = read_safetensors_metadata(ckpt_dir) + reader = cls() + extra_model_kwargs = reader.read_config(raw) + return reader, extra_model_kwargs + # Validate GPTQ config: currently only INT4 with group_size=128 is supported if quant_method == "gptq": bits = qconf.get("bits") @@ -287,15 +333,22 @@ def post_process_model(self, model, model_config): def autodetect_quant_config_reader( fetched_dir: str, + require_checkpoint_metadata: bool = True, ) -> Optional[Tuple["QuantConfigReader", Dict[str, Any]]]: """Try ModelOPT first; if not found, fall back to HF. Returns (reader, extra_kwargs) or None.""" reader_cls = QuantConfigReaderRegistry.get("modelopt") - result = reader_cls.from_file(fetched_dir) + result = reader_cls.from_file( + fetched_dir, require_checkpoint_metadata=require_checkpoint_metadata + ) # Fallback to HF reader if ModelOPT not present if result is None: hf_cls = QuantConfigReaderRegistry.get("hf") try: - result = hf_cls.from_file(fetched_dir) + result = hf_cls.from_file( + fetched_dir, require_checkpoint_metadata=require_checkpoint_metadata + ) + except QuantizedCheckpointLayoutError: + raise except Exception: # Skip HF reader if it errors out during probing result = None diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe_mxfp4.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe_mxfp4.py index e5782caaf53c..efeb0fe8bbdf 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe_mxfp4.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe_mxfp4.py @@ -12,6 +12,8 @@ # 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. +import operator +from functools import partial from typing import Literal, Optional, Tuple, Type import torch @@ -20,9 +22,13 @@ from torch.fx import GraphModule, Node from ..._compat import get_sm_version +from ...models.quant_checkpoint_layout import ( + PackedMXFP4ExpertCheckpointLayout, + load_packed_mxfp4_expert_tensors, +) from ...utils.logger import ad_logger from ...utils.module import get_submodule_of_param -from ...utils.node_utils import is_op +from ...utils.node_utils import extract_op_args, is_op from ...utils.pattern_matcher import ADPatternMatcherPass, register_ad_pattern from ..interface import BaseTransform, TransformConfig, TransformInfo, TransformRegistry @@ -31,6 +37,21 @@ # in ``QuantizeMXFP4MOE._apply_trtllm``. _MXFP4_SCALING_VECTOR_SIZE = 32 _WEIGHT_ALIGNMENT = 128 +_MXFP4_DEFAULT_EXPERT_BLOCK_SIZE = 32 +_MXFP4_SUPPORTED_EXPERT_BLOCK_SIZE = 32 +_MXFP4_VALUES_PER_BYTE = 2 +_MXFP4_LAYOUT_ARG_NAMES = ( + "gate_up_blocks", + "gate_up_scales", + "down_blocks", + "down_scales", +) +_MXFP4_RUNTIME_OPS = ( + torch.ops.auto_deploy.torch_mxfp4_moe, + torch.ops.auto_deploy.triton_mxfp4_moe, + torch.ops.auto_deploy.torch_mxfp4_moe_from_routing, + torch.ops.auto_deploy.triton_mxfp4_moe_from_routing, +) # Backend selection for MXFP4 MoE quantization. # - "triton": use the triton_mxfp4_moe kernel (Ampere/Hopper compatible). @@ -169,12 +190,211 @@ def _get_topk_from_router(node: Node) -> int: return int(node.args[3]) if len(node.args) >= 4 else 2 +def _get_packed_mxfp4_expert_layout( + qcfg: dict, +) -> PackedMXFP4ExpertCheckpointLayout | None: + checkpoint_layout = qcfg.get("checkpoint_layout") + if checkpoint_layout is None: + return None + for consumer in getattr(checkpoint_layout, "checkpoint_consumers", ()): + if ( + isinstance(consumer, PackedMXFP4ExpertCheckpointLayout) + and consumer.quant_method == "mxfp4" + ): + return consumer + return None + + +def _normalize_mxfp4_expert_block_size(source: str, value: object) -> int: + if isinstance(value, bool): + raise ValueError(f"MXFP4 {source} should be an integer, got {value}.") + try: + block_size = operator.index(value) + except TypeError as error: + raise ValueError( + f"MXFP4 {source} should be an integer, got {type(value).__name__}." + ) from error + if block_size <= 0: + raise ValueError(f"MXFP4 {source} should be positive, got {block_size}.") + return block_size + + +def _resolve_mxfp4_expert_block_size( + qcfg: dict, + checkpoint_layout: object | None, +) -> int: + candidates = [] + if qcfg.get("expert_block_size") is not None: + candidates.append(("quant config expert_block_size", qcfg["expert_block_size"])) + + layout_block_size = ( + getattr(checkpoint_layout, "expert_block_size", None) + if checkpoint_layout is not None + else None + ) + if layout_block_size is not None: + candidates.append(("checkpoint layout expert_block_size", layout_block_size)) + + if not candidates: + block_size = _MXFP4_DEFAULT_EXPERT_BLOCK_SIZE + else: + source, value = candidates[0] + block_size = _normalize_mxfp4_expert_block_size(source, value) + for other_source, other_value in candidates[1:]: + other_block_size = _normalize_mxfp4_expert_block_size(other_source, other_value) + if other_block_size != block_size: + raise ValueError( + "MXFP4 expert_block_size mismatch: " + f"{source} is {block_size}, but {other_source} is {other_block_size}." + ) + + if block_size != _MXFP4_SUPPORTED_EXPERT_BLOCK_SIZE: + raise ValueError( + "MXFP4 MoE supports expert_block_size=32 because the kernels decode " + f"32 values per scale block, got {block_size}." + ) + return block_size + + +def _layout_layer_from_names( + checkpoint_layout: PackedMXFP4ExpertCheckpointLayout, + source_names: dict[str, str], +) -> int | None: + for source_name in source_names.values(): + layer = checkpoint_layout.layer_from_runtime_name(source_name) + if layer is not None: + return layer + return None + + +def _get_attr_tensor(gm: GraphModule, target: str) -> torch.Tensor: + mod_name, _, attr_name = target.rpartition(".") + submod = gm.get_submodule(mod_name) + tensor = getattr(submod, attr_name) + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"Expected get_attr target {target} to be a tensor, got {type(tensor)}.") + return tensor + + +def _mxfp4_target_names_from_node(node: Node) -> dict[str, str] | None: + if not is_op(node, _MXFP4_RUNTIME_OPS): + return None + + expert_args = extract_op_args(node, *_MXFP4_LAYOUT_ARG_NAMES) + target_names: dict[str, str] = {} + for name, arg in zip(_MXFP4_LAYOUT_ARG_NAMES, expert_args): + if not isinstance(arg, Node) or arg.op != "get_attr": + return None + target_names[name] = str(arg.target) + return target_names + + +def _load_mxfp4_expert_layout_hook( + state_dict, + prefix, + *args, + checkpoint_layout: PackedMXFP4ExpertCheckpointLayout, + target_names: dict[str, str], + layer: int, + num_experts: int, + hidden_size: int, + intermediate_size: int, +) -> None: + if prefix + target_names["gate_up_blocks"] in state_dict: + return + load_packed_mxfp4_expert_tensors( + state_dict, + prefix, + checkpoint_layout=checkpoint_layout, + target_names=target_names, + layer=layer, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + ) + + +def _register_mxfp4_expert_layout_hook( + gm: GraphModule, + checkpoint_layout: PackedMXFP4ExpertCheckpointLayout, + target_names: dict[str, str], + *, + registered_targets: set[tuple[str, str, str, str]] | None = None, +) -> bool: + target_key = ( + target_names["gate_up_blocks"], + target_names["gate_up_scales"], + target_names["down_blocks"], + target_names["down_scales"], + ) + if registered_targets is not None: + if target_key in registered_targets: + return False + registered_targets.add(target_key) + + layer = _layout_layer_from_names(checkpoint_layout, target_names) + if layer is None: + return False + gate_up_blocks = _get_attr_tensor(gm, target_names["gate_up_blocks"]) + down_blocks = _get_attr_tensor(gm, target_names["down_blocks"]) + if gate_up_blocks.dim() < 4 or down_blocks.dim() < 4: + raise ValueError( + "MXFP4 expert runtime buffers should be at least rank 4, got " + f"{gate_up_blocks.shape} and {down_blocks.shape}." + ) + gm._register_load_state_dict_pre_hook( + partial( + _load_mxfp4_expert_layout_hook, + checkpoint_layout=checkpoint_layout, + target_names=target_names, + layer=layer, + num_experts=int(gate_up_blocks.shape[0]), + hidden_size=int(down_blocks.shape[1]), + intermediate_size=int(gate_up_blocks.shape[1] // 2), + ) + ) + return True + + +def _register_existing_mxfp4_expert_layout_hooks( + gm: GraphModule, + checkpoint_layout: PackedMXFP4ExpertCheckpointLayout, + *, + registered_targets: set[tuple[str, str, str, str]], +) -> int: + num_hooks = 0 + for node in gm.graph.nodes: + target_names = _mxfp4_target_names_from_node(node) + if target_names is None: + continue + if _register_mxfp4_expert_layout_hook( + gm, + checkpoint_layout, + target_names, + registered_targets=registered_targets, + ): + num_hooks += 1 + return num_hooks + + +def _mxfp4_block_count(name: str, dim: int, expert_block_size: int) -> int: + if dim <= 0: + raise ValueError(f"MXFP4 expert {name} should be positive, got {dim}.") + if dim % expert_block_size != 0: + raise ValueError( + f"MXFP4 expert {name} should be divisible by expert_block_size=" + f"{expert_block_size}, got {dim}." + ) + return dim // expert_block_size + + def _register_mxfp4_expert_params( gm: GraphModule, gate_up_w_name: str, gate_up_b_name: str, down_w_name: str, down_b_name: str, + expert_block_size: int = _MXFP4_DEFAULT_EXPERT_BLOCK_SIZE, ) -> Tuple[str, str, str, str]: """Create (if missing) the four MXFP4 params under the experts module and return their full names. @@ -198,9 +418,9 @@ def _register_mxfp4_expert_params( # Fallback: use down bias last dim H = int(dn_b.shape[1]) - # Compute block dims (assume divisible; zero-init anyway) - H_blk = max(1, H // 32) - I_blk = max(1, In // 32) + packed_block_width = expert_block_size // _MXFP4_VALUES_PER_BYTE + H_blk = _mxfp4_block_count("hidden_size", H, expert_block_size) + I_blk = _mxfp4_block_count("intermediate_size", In, expert_block_size) experts_mod, experts_path, _ = get_submodule_of_param(gm, gate_up_w_name) @@ -215,9 +435,13 @@ def _register_mxfp4_expert_params( # (meta in the normal meta-device build) so we don't materialize giant CPU # buffers before load. param_device = gu_w.device - gu_blocks = torch.empty((E, 2 * In, H_blk, 16), dtype=torch.uint8, device=param_device) + gu_blocks = torch.empty( + (E, 2 * In, H_blk, packed_block_width), dtype=torch.uint8, device=param_device + ) gu_scales = torch.empty((E, 2 * In, H_blk), dtype=torch.uint8, device=param_device) - dn_blocks = torch.empty((E, H, I_blk, 16), dtype=torch.uint8, device=param_device) + dn_blocks = torch.empty( + (E, H, I_blk, packed_block_width), dtype=torch.uint8, device=param_device + ) dn_scales = torch.empty((E, H, I_blk), dtype=torch.uint8, device=param_device) experts_mod.register_parameter(gu_blocks_name, nn.Parameter(gu_blocks, requires_grad=False)) @@ -511,28 +735,59 @@ def _apply( 2. Resolves the backend (``triton`` | ``trtllm``) and dispatches. """ qcfg = factory.get_quant_config() - if not qcfg or qcfg.get("quant_method", "") != self.algo_name: + if not qcfg: return gm, TransformInfo( skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True ) + checkpoint_layout = _get_packed_mxfp4_expert_layout(qcfg) + expert_block_size = _resolve_mxfp4_expert_block_size(qcfg, checkpoint_layout) + num_existing_hooks = 0 + if checkpoint_layout is not None: + num_existing_hooks = _register_existing_mxfp4_expert_layout_hooks( + gm, + checkpoint_layout, + registered_targets=set(), + ) + + if qcfg.get("quant_method", "") != self.algo_name: + return gm, TransformInfo( + skipped=num_existing_hooks == 0, + num_matches=num_existing_hooks, + is_clean=num_existing_hooks == 0, + has_valid_shapes=num_existing_hooks == 0, + ) backend = self._resolve_backend() ad_logger.info(f"quantize_mxfp4_moe: dispatching to backend={backend!r}") if backend == "triton": - return self._apply_triton(gm, cm, factory, shared_config) + gm, info = self._apply_triton( + gm, cm, factory, shared_config, expert_block_size=expert_block_size + ) elif backend == "trtllm": - return self._apply_trtllm(gm, cm, factory, shared_config) + gm, info = self._apply_trtllm(gm, cm, factory, shared_config) else: # _resolve_backend should only return "triton" or "trtllm". raise ValueError(f"Unexpected backend resolved: {backend!r}") + if num_existing_hooks == 0: + return gm, info + num_matches = info.num_matches + num_existing_hooks + return gm, TransformInfo( + skipped=False, + num_matches=num_matches, + is_clean=info.is_clean and num_matches == 0, + has_valid_shapes=info.has_valid_shapes and num_matches == 0, + ) + def _apply_triton( self, gm: GraphModule, cm, factory, shared_config, + *, + expert_block_size: int = _MXFP4_DEFAULT_EXPERT_BLOCK_SIZE, ) -> Tuple[GraphModule, TransformInfo]: """Triton backend: graph rewrite to ``triton_mxfp4_moe``. @@ -586,7 +841,14 @@ def _apply_triton( # Register MXFP4 params on experts gu_blocks_name, gu_scales_name, dn_blocks_name, dn_scales_name = ( - _register_mxfp4_expert_params(gm, gu_w_name, gu_b_name, dn_w_name, dn_b_name) + _register_mxfp4_expert_params( + gm, + gu_w_name, + gu_b_name, + dn_w_name, + dn_b_name, + expert_block_size=expert_block_size, + ) ) # Alpha/limit (from dense call) @@ -1059,6 +1321,294 @@ class FuseMXFP4Moe(BaseTransform): def get_config_class(cls) -> Type[TransformConfig]: return FuseMXFP4MoeConfig + @staticmethod + def _graph_uses_attr(gm: GraphModule, target: str) -> bool: + return any(n.op == "get_attr" and str(n.target) == target for n in gm.graph.nodes) + + @staticmethod + def _register_prepared_parameter(module: nn.Module, name: str, tensor: torch.Tensor) -> None: + _delete_module_attr(module, name) + module.register_parameter( + name, + nn.Parameter(tensor.detach().clone(), requires_grad=False), + ) + + @staticmethod + def _routed_target_for_sm(sm_version: int): + if sm_version in (100, 103): + return ( + torch.ops.auto_deploy.trtllm_quant_mxfp4_trtllm_gen_moe_from_routing_fused.default + ) + if sm_version == 90: + return torch.ops.auto_deploy.triton_mxfp4_moe_from_routing.default + return None + + @staticmethod + def _routed_ep_target_for_sm(sm_version: int): + if sm_version in (100, 103): + return ( + torch.ops.auto_deploy.trtllm_quant_mxfp4_trtllm_gen_moe_from_routing_fused.default + ) + if sm_version == 90: + return torch.ops.auto_deploy.triton_mxfp4_moe_from_routing_ep.default + return None + + @staticmethod + def _routed_source_kind(node: Node) -> str | None: + if is_op(node, torch.ops.auto_deploy.torch_mxfp4_moe_from_routing): + return "base" + if is_op(node, torch.ops.auto_deploy.torch_mxfp4_moe_from_routing_ep): + return "ep" + return None + + def _apply_routed_from_routing( + self, + gm: GraphModule, + *, + prepare_trtllm_gen_moe_mxfp4_weights, + ) -> int: + """Rewrite DeepSeek-style externally-routed MXFP4 MoE to the arch fast path. + + DeepSeek V4 computes selected experts and routing weights in modeling + code. The arch policy here is intentionally narrow: + + * SM100/SM103: TRTLLMGen W4A8 (MXFP8 activations) with external top-k. + * SM90: Triton W4A16 (BF16/FP16 activations) with external top-k. + * Otherwise: leave the torch reference op unchanged. + """ + routed_nodes = [ + (n, source_kind) + for n in list(gm.graph.nodes) + if (source_kind := self._routed_source_kind(n)) is not None + ] + if not routed_nodes: + return 0 + + sm_version = int(get_sm_version()) + num_matches = 0 + + for n, source_kind in routed_nodes: + target_op = ( + self._routed_target_for_sm(sm_version) + if source_kind == "base" + else self._routed_ep_target_for_sm(sm_version) + ) + if target_op is None: + continue + + ( + hidden_node, + selected_experts_node, + routing_weights_node, + gate_up_blocks_node, + gate_up_bias_node, + gate_up_scales_node, + alpha, + limit, + down_blocks_node, + down_bias_node, + down_scales_node, + gate_up_order, + swiglu_mode, + layer_type, + expert_start, + num_experts_total, + ) = extract_op_args( + n, + "hidden_states", + "selected_experts", + "routing_weights", + "gate_up_blocks", + "gate_up_bias", + "gate_up_scales", + "alpha", + "limit", + "down_blocks", + "down_bias", + "down_scales", + "gate_up_order", + "swiglu_mode", + "layer_type", + "expert_start", + "num_experts_total", + ) + + if gate_up_order not in ("up_gate", "w3_w1") or swiglu_mode != "deepseek": + continue + + raw_get_attrs = ( + gate_up_blocks_node, + down_blocks_node, + gate_up_scales_node, + down_scales_node, + gate_up_bias_node, + down_bias_node, + ) + if not all(isinstance(a, Node) and a.op == "get_attr" for a in raw_get_attrs): + continue + + raw_names = tuple(str(a.target) for a in raw_get_attrs) + gate_up_blocks = _get_attr_tensor(gm, raw_names[0]) + down_blocks = _get_attr_tensor(gm, raw_names[1]) + gate_up_scales = _get_attr_tensor(gm, raw_names[2]) + down_scales = _get_attr_tensor(gm, raw_names[3]) + gate_up_bias = _get_attr_tensor(gm, raw_names[4]) + down_bias = _get_attr_tensor(gm, raw_names[5]) + + num_local_experts = int(gate_up_blocks.shape[0]) + hidden_size = int(down_blocks.shape[1]) + intermediate_size = int(gate_up_blocks.shape[1] // 2) + expert_start = 0 if expert_start is None else int(expert_start) + if num_experts_total is None or int(num_experts_total) < 0: + num_experts_total = expert_start + num_local_experts + else: + num_experts_total = int(num_experts_total) + + if sm_version == 90: + n.target = target_op + if source_kind == "base": + n.args = ( + hidden_node, + selected_experts_node, + routing_weights_node, + gate_up_blocks_node, + gate_up_bias_node, + gate_up_scales_node, + float(alpha), + float(limit), + down_blocks_node, + down_bias_node, + down_scales_node, + gate_up_order, + swiglu_mode, + layer_type, + num_experts_total, + ) + else: + n.args = ( + hidden_node, + selected_experts_node, + routing_weights_node, + gate_up_blocks_node, + gate_up_bias_node, + gate_up_scales_node, + float(alpha), + float(limit), + down_blocks_node, + down_bias_node, + down_scales_node, + gate_up_order, + swiglu_mode, + layer_type, + expert_start, + num_experts_total, + ) + n.kwargs = {} + num_matches += 1 + continue + + experts_mod, experts_path, _ = get_submodule_of_param(gm, raw_names[0]) + prep = prepare_trtllm_gen_moe_mxfp4_weights( + gate_up_blocks, + gate_up_scales, + gate_up_bias, + down_blocks, + down_scales, + down_bias, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + tp_size=1, + tp_rank=0, + gate_up_order="up_gate", + ) + + self._register_prepared_parameter(experts_mod, "fc1_w_trtllm", prep.fc1_weights_mxfp4) + self._register_prepared_parameter( + experts_mod, "fc1_w_scale_trtllm", prep.fc1_weights_scale_ue8m0 + ) + self._register_prepared_parameter(experts_mod, "fc1_bias_trtllm", prep.fc1_bias_f32) + self._register_prepared_parameter(experts_mod, "fc2_w_trtllm", prep.fc2_weights_mxfp4) + self._register_prepared_parameter( + experts_mod, "fc2_w_scale_trtllm", prep.fc2_weights_scale_ue8m0 + ) + self._register_prepared_parameter(experts_mod, "fc2_bias_trtllm", prep.fc2_bias_f32) + + constants_device = prep.fc1_bias_f32.device + swiglu_alpha = torch.full( + (num_local_experts,), + float(alpha), + dtype=torch.float32, + device=constants_device, + ) + swiglu_beta = torch.zeros( + (num_local_experts,), + dtype=torch.float32, + device=constants_device, + ) + swiglu_limit = torch.full( + (num_local_experts,), + float(limit), + dtype=torch.float32, + device=constants_device, + ) + self._register_prepared_parameter(experts_mod, "swiglu_alpha_trtllm", swiglu_alpha) + self._register_prepared_parameter(experts_mod, "swiglu_beta_trtllm", swiglu_beta) + self._register_prepared_parameter(experts_mod, "swiglu_limit_trtllm", swiglu_limit) + + prefix_path = (experts_path + ".") if experts_path else "" + with gm.graph.inserting_before(n): + fc1_w_attr = gm.graph.create_node("get_attr", prefix_path + "fc1_w_trtllm") + fc2_w_attr = gm.graph.create_node("get_attr", prefix_path + "fc2_w_trtllm") + fc1_s_attr = gm.graph.create_node("get_attr", prefix_path + "fc1_w_scale_trtllm") + fc2_s_attr = gm.graph.create_node("get_attr", prefix_path + "fc2_w_scale_trtllm") + fc1_b_attr = gm.graph.create_node("get_attr", prefix_path + "fc1_bias_trtllm") + fc2_b_attr = gm.graph.create_node("get_attr", prefix_path + "fc2_bias_trtllm") + sa_attr = gm.graph.create_node("get_attr", prefix_path + "swiglu_alpha_trtllm") + sb_attr = gm.graph.create_node("get_attr", prefix_path + "swiglu_beta_trtllm") + sl_attr = gm.graph.create_node("get_attr", prefix_path + "swiglu_limit_trtllm") + + n.target = target_op + n.args = ( + hidden_node, + selected_experts_node, + routing_weights_node, + fc1_w_attr, + fc2_w_attr, + fc1_s_attr, + fc2_s_attr, + fc1_b_attr, + fc2_b_attr, + sa_attr, + sb_attr, + sl_attr, + prep.valid_hidden_size, + prep.valid_intermediate_size, + "mxfp8", + num_experts_total, + expert_start, + num_local_experts, + 1, # routing_method_type = RoutingMethodType.Renormalize + ) + n.kwargs = {} + + for stale_node in raw_get_attrs: + if len(stale_node.users) == 0: + gm.graph.erase_node(stale_node) + for raw_name in raw_names: + if self._graph_uses_attr(gm, raw_name): + continue + _, _, attr_short = raw_name.rpartition(".") + _delete_module_attr(experts_mod, attr_short) + + num_matches += 1 + + if num_matches: + ad_logger.info( + f"fuse_mxfp4_moe: rewrote {num_matches} routed MXFP4 MoE node(s) " + f"for sm{sm_version}." + ) + return num_matches + def _apply( self, gm: GraphModule, @@ -1103,6 +1653,11 @@ def _apply( # Single MXFP4 trtllm-gen MoE op (act_dtype="bf16" or "mxfp8"). target_op = torch.ops.auto_deploy.trtllm_quant_mxfp4_trtllm_gen_moe_fused.default + routed_matches = self._apply_routed_from_routing( + gm, + prepare_trtllm_gen_moe_mxfp4_weights=prepare_trtllm_gen_moe_mxfp4_weights, + ) + # ---- Pass 1: collect MoE node info, validate consistent shape ---- # Arg index layout from ``_apply_trtllm`` (kept in sync; comment # block there documents the full slot list). @@ -1166,7 +1721,12 @@ def _apply( num_matches = len(layer_infos) if num_matches == 0: - info = TransformInfo(skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True) + info = TransformInfo( + skipped=(routed_matches == 0), + num_matches=routed_matches, + is_clean=(routed_matches == 0), + has_valid_shapes=True, + ) return gm, info # ---- Pass 2: allocate scratch ONCE ---- @@ -1294,9 +1854,10 @@ def _apply( f"with shared scratch (E={e_local_g}, I={per_rank_i_g}, H={H_g})" ) + total_matches = num_matches + routed_matches info = TransformInfo( skipped=False, - num_matches=num_matches, + num_matches=total_matches, is_clean=False, has_valid_shapes=True, ) diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py b/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py index e0eae9ae01cd..46fda88b3260 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py @@ -662,6 +662,26 @@ class InsertCachedAttention(_InsertCachedOperator): """A transform to insert cached attention into the graph module.""" +@TransformRegistry.register("insert_cached_deepseek_v4_sparse_attention") +class InsertCachedDeepSeekV4SparseAttention(_InsertCachedOperator): + """Opt-in transform for DeepSeek V4 sparse cached attention backends.""" + + _SUPPORTED_BACKENDS = { + "deepseek_v4_sparse", + "flashmla_deepseek_v4_sparse", + } + + def _apply(self, *args, **kwargs): + if self.config.backend is None: + self.config.backend = "deepseek_v4_sparse" + elif self.config.backend not in self._SUPPORTED_BACKENDS: + raise ValueError( + "insert_cached_deepseek_v4_sparse_attention only supports " + f"backends {sorted(self._SUPPORTED_BACKENDS)}, got {self.config.backend!r}." + ) + return super()._apply(*args, **kwargs) + + @TransformRegistry.register("insert_cached_mla_attention") class InsertCachedMLAAttention(_InsertCachedOperator): """A transform to insert cached MLA attention into the graph module.""" diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py b/tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py index 98bd87fbe358..e03ccd28ff1f 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py @@ -13,8 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. import math +from collections.abc import Sequence +from dataclasses import dataclass from functools import partial -from typing import Dict, List, Tuple +from typing import Dict, List, Optional, Tuple import torch import torch.nn as nn @@ -34,10 +36,12 @@ from ...utils.node_utils import ( WeightBiasInfoCache, extract_op_args, + extract_weight_name, extract_weight_nodes, get_quantization_params_from_linear_node, is_bmm_op, is_linear_op, + is_op, ) from ...utils.quantization_utils import ( fp4_global_scale, @@ -59,6 +63,63 @@ float4_sf_dtype = None +def _is_view_or_reshape_node(node: Node) -> bool: + if not isinstance(node, Node): + return False + if node.op == "call_method": + return node.target in {"view", "reshape"} + return is_op( + node, + [ + torch.ops.aten.view, + torch.ops.aten.reshape, + torch.ops.auto_deploy.view, + ], + ) + + +def _view_base_and_shape(node: Node) -> Tuple[object, Tuple[object, ...]]: + if node.op == "call_method": + base = node.args[0] + shape_args = node.args[1:] + elif is_op(node, torch.ops.auto_deploy.view): + base = node.args[0] + shape_args = (node.args[1],) + else: + base = node.args[0] + shape_args = node.args[1:] + if len(shape_args) == 1 and isinstance(shape_args[0], (list, tuple)): + shape_args = tuple(shape_args[0]) + return base, tuple(shape_args) + + +def _shape_from_view_or_meta(node: Node) -> Tuple[object, ...] | None: + if _is_view_or_reshape_node(node): + _, view_shape = _view_base_and_shape(node) + return view_shape + meta_val = node.meta.get("val") if isinstance(node, Node) else None + if meta_val is not None and hasattr(meta_val, "shape"): + return tuple(meta_val.shape) + return None + + +def _is_static_dim_value(value: object) -> bool: + return isinstance(value, int) and not isinstance(value, bool) + + +def _extract_layer_type_hint(node: Node, default: str = "unknown") -> str: + if not isinstance(node, Node): + return default + if node.op == "call_function": + try: + [layer_type] = extract_op_args(node, "layer_type") + except RuntimeError: + layer_type = None + if layer_type is not None: + return layer_type + return node.kwargs.get("layer_type", default) + + class Quantization(BaseTransform): """Abstract base for config-driven quantization of a single algorithm/op-kind. @@ -174,6 +235,8 @@ def _insert_quantized_linear( gm: GraphModule, node: Node, is_quantized_graph: bool = False, + load_hook_kwargs: Optional[Dict[str, object]] = None, + custom_kwargs: Optional[Dict[str, object]] = None, ): """Replaces the matmul node with a new custom quantized linear node. @@ -218,7 +281,11 @@ def _insert_quantized_linear( lin_weight.submod.register_buffer(scale_name, scale) gm._register_load_state_dict_pre_hook( - partial(self.load_hook, weight_name=lin_weight.node_key) + partial( + self.load_hook, + weight_name=lin_weight.node_key, + **(load_hook_kwargs or {}), + ) ) post_load_hook = getattr(type(self), "post_load_hook", None) if post_load_hook is not None and post_load_hook is not Quantization.post_load_hook: @@ -246,6 +313,7 @@ def _insert_quantized_linear( "output_sizes": output_sizes, "tp_min_local_shape": tp_min_local_shape, "layer_type": layer_type, + **(custom_kwargs or {}), } def _insert_quantized_bmm( @@ -847,6 +915,14 @@ def load_hook(state_dict, prefix, *args, weight_name: str): del state_dict[qweight_ckpt] +@dataclass(frozen=True) +class _FineGrainedFP8LinearContext: + weight_block_size: Tuple[int, int] + runtime_scale_name: str + default_scales_layout: Optional[object] + input_scale_fmt: str + + @TransformRegistry.register("quantize_finegrained_fp8_linear_from_config") class FineGrainedFP8LinearQuantization(Quantization): """Quantization transform for FineGrainedFP8 (block-wise FP8) models. @@ -868,41 +944,303 @@ class FineGrainedFP8LinearQuantization(Quantization): def target_op(self): return torch.ops.auto_deploy.torch_fake_quant_finegrained_fp8_linear.default + def grouped_target_op(self): + return torch.ops.auto_deploy.torch_fake_quant_grouped_finegrained_fp8_linear.default + def quantize_weight(self, w: torch.Tensor) -> torch.Tensor: return torch.empty_like(w, dtype=torch.float8_e4m3fn, device=w.device) - def scale_names(self) -> List[str]: - return ["weight_scale_inv"] + def scale_names(self, context: Optional[_FineGrainedFP8LinearContext] = None) -> List[str]: + runtime_scale_name = context.runtime_scale_name if context else "weight_scale_inv" + return [runtime_scale_name] + + def default_scales( + self, + original_weight_shape: Tuple, + context: Optional[_FineGrainedFP8LinearContext] = None, + ) -> Dict[str, torch.Tensor]: + if context and context.default_scales_layout is not None: + return context.default_scales_layout.default_scales(original_weight_shape) - def default_scales(self, original_weight_shape: Tuple) -> Dict[str, torch.Tensor]: - # Default block size is 128x128 for FineGrained FP8 N, K = original_weight_shape - block_n, block_k = 128, 128 + block_n, block_k = context.weight_block_size if context else (128, 128) # Use ceil to handle dimensions smaller than or not divisible by block size # (e.g. after TP sharding or small projection weights). scale_shape = (math.ceil(N / block_n), math.ceil(K / block_k)) - return {"weight_scale_inv": torch.ones(scale_shape, dtype=torch.bfloat16)} + return {self.scale_names(context)[0]: torch.ones(scale_shape, dtype=torch.float32)} - def build_custom_args_for_linear(self, scales: Dict[str, Node]) -> Tuple: - return ([], [scales["weight_scale_inv"]], [], []) + def build_custom_args_for_linear( + self, + scales: Dict[str, Node], + context: Optional[_FineGrainedFP8LinearContext] = None, + ) -> Tuple: + return ([], [scales[self.scale_names(context)[0]]], [], []) + + def _insert_quantized_linear( + self, + gm: GraphModule, + node: Node, + is_quantized_graph: bool = False, + load_hook_kwargs: Optional[Dict[str, object]] = None, + custom_kwargs: Optional[Dict[str, object]] = None, + context: Optional[_FineGrainedFP8LinearContext] = None, + ): + """Replaces a linear node with a fine-grained FP8 custom op.""" + weight_nodes = extract_weight_nodes(node) + if len(weight_nodes.weights) == 0: + raise ValueError(f"Linear node {node.name} has no weight") + lin_weight = weight_nodes.weights[0] - def load_hook(self, state_dict, prefix, *args, weight_name: str): + new_param = nn.Parameter(self.quantize_weight(lin_weight.tensor), requires_grad=False) + modname, _, attrname = lin_weight.node_key.rpartition(".") + setattr(lin_weight.submod, attrname, new_param) + + if is_quantized_graph: + input_params, weight_params, _output_params = get_quantization_params_from_linear_node( + node + ) + node.args = (input_params.input_node, weight_params.input_node, *node.args[2:]) + + user = list(node.users.keys())[0] + if len(node.users) == 1 and is_quantized_op(user): + user.replace_all_uses_with(node) + + input_scale_name = self.scale_names(context)[0] + gm._register_load_state_dict_pre_hook( + partial( + self.convert_amax_hook, + scale_name=modname + "." + input_scale_name, + amax_name=input_params.amax.target, + ) + ) + + for scale_name, scale in self.default_scales(lin_weight.tensor.shape, context).items(): + lin_weight.submod.register_buffer(scale_name, scale) + + gm._register_load_state_dict_pre_hook( + partial( + self.load_hook, + weight_name=lin_weight.node_key, + **(load_hook_kwargs or {}), + ) + ) + if self.post_load_hook: + gm.register_load_state_dict_post_hook( + partial(self.post_load_hook, weight_name=lin_weight.node_key) + ) + + with gm.graph.inserting_before(node): + scales = {} + for scale_name in self.scale_names(context): + scales[scale_name] = gm.graph.create_node("get_attr", modname + "." + scale_name) + + custom_args = self.build_custom_args_for_linear(scales, context) + + [tp_mode, output_sizes, tp_min_local_shape, layer_type] = extract_op_args( + node, "tp_mode", "output_sizes", "tp_min_local_shape", "layer_type" + ) + [inp, weight, bias] = extract_op_args(node, "input", "weight", "bias") + node.target = self.target_op() + node.args = (inp, weight, bias, *custom_args) + node.kwargs = { + **node.kwargs, + "tp_mode": tp_mode, + "output_sizes": output_sizes, + "tp_min_local_shape": tp_min_local_shape, + "layer_type": layer_type, + **(custom_kwargs or {}), + } + + def _extract_grouped_linear_match( + self, + node: Node, + checkpoint_layout: object, + excluded: List[str], + ) -> Tuple[Node, Node, object, Node, str, int] | None: + if not is_op(node, torch.ops.auto_deploy.torch_grouped_linear): + return None + + input_node, weight_view, bias_node = extract_op_args(node, "input", "weight", "bias") + if not isinstance(input_node, Node) or not isinstance(weight_view, Node): + return None + if bias_node is not None and not isinstance(bias_node, Node): + return None + if not _is_view_or_reshape_node(weight_view): + return None + + weight_base, weight_shape = _view_base_and_shape(weight_view) + if not isinstance(weight_base, Node) or weight_base.op != "get_attr": + return None + if not isinstance(weight_base.target, str): + return None + weight_name = weight_base.target + if not checkpoint_layout.is_weight_targeted(weight_name, excluded): + return None + + input_shape = _shape_from_view_or_meta(input_node) + if input_shape is None or len(input_shape) != 4 or len(weight_shape) != 3: + return None + num_groups = weight_shape[0] + rank = weight_shape[1] + if not _is_static_dim_value(num_groups) or not _is_static_dim_value(rank): + return None + if _is_static_dim_value(input_shape[2]) and input_shape[2] != num_groups: + return None + + return node, input_node, bias_node, weight_base, weight_name, rank + + def _insert_grouped_quantized_linear( + self, + gm: GraphModule, + source_node: Node, + input_node: Node, + bias_node: object, + weight_node: Node, + weight_name: str, + rank: int, + load_hook_kwargs: Optional[Dict[str, object]] = None, + context: Optional[_FineGrainedFP8LinearContext] = None, + ) -> None: + original_weight = gm.get_parameter(weight_name) + new_param = nn.Parameter(self.quantize_weight(original_weight), requires_grad=False) + modname, _, attrname = weight_name.rpartition(".") + submod = gm.get_submodule(modname) + setattr(submod, attrname, new_param) + weight_node.meta["val"] = new_param.detach() + + for scale_name, scale in self.default_scales(original_weight.shape, context).items(): + if scale_name in submod._buffers: + submod._buffers[scale_name] = scale + else: + submod.register_buffer(scale_name, scale) + + gm._register_load_state_dict_pre_hook( + partial( + self.load_hook, + weight_name=weight_name, + **(load_hook_kwargs or {}), + ) + ) + + scale_name = self.scale_names(context)[0] + scale_target = f"{modname}.{scale_name}" if modname else scale_name + layer_type = _extract_layer_type_hint(source_node, _extract_layer_type_hint(input_node)) + + with gm.graph.inserting_before(source_node): + scale_node = gm.graph.create_node("get_attr", scale_target) + scale_node.meta["val"] = getattr(submod, scale_name).detach() + grouped_node = gm.graph.call_function( + self.grouped_target_op(), + args=(input_node, weight_node, bias_node, [], [scale_node], [], []), + kwargs={ + "tp_mode": "colwise", + "output_sizes": None, + "tp_min_local_shape": rank, + "layer_type": layer_type, + "input_scale_fmt": context.input_scale_fmt if context else "", + }, + ) + + source_node.replace_all_uses_with(grouped_node) + gm.graph.erase_node(source_node) + + @staticmethod + def _get_finegrained_fp8_layout(qcfg: Dict): + checkpoint_layout = qcfg.get("checkpoint_layout") + if checkpoint_layout is None: + return None + return getattr(checkpoint_layout, "finegrained_fp8", None) + + @staticmethod + def _normalize_weight_block_size(block_size: object) -> Tuple[int, int]: + if ( + isinstance(block_size, Sequence) + and not isinstance(block_size, str) + and len(block_size) == 2 + ): + return int(block_size[0]), int(block_size[1]) + raise ValueError(f"FineGrained FP8 weight_block_size must have two dims, got {block_size}") + + @staticmethod + def _resolve_weight_block_size(qcfg: Dict, checkpoint_layout: object) -> Tuple[int, int]: + if checkpoint_layout is not None: + return FineGrainedFP8LinearQuantization._normalize_weight_block_size( + checkpoint_layout.weight_block_size + ) + if qcfg.get("weight_block_size") is not None: + return FineGrainedFP8LinearQuantization._normalize_weight_block_size( + qcfg["weight_block_size"] + ) + return 128, 128 + + @staticmethod + def _resolve_input_scale_fmt(qcfg: Dict, checkpoint_layout: object) -> str: + scale_fmt = getattr(checkpoint_layout, "scale_fmt", None) + if scale_fmt is None: + scale_fmt = qcfg.get("scale_fmt") + if scale_fmt is None: + return "" + return str(scale_fmt).lower() + + @staticmethod + def _add_prefix(prefix: str, name: str) -> str: + if prefix and not name.startswith(prefix): + return prefix + name + return name + + def load_hook(self, state_dict, prefix, *args, weight_name: str, checkpoint_layout=None): """Load hook to handle FineGrainedFP8 checkpoint format. FineGrained FP8 checkpoints store: - weight: float8_e4m3fn tensor - weight_scale_inv: per-block scale tensor """ - if weight_name not in state_dict: + prefix = prefix or "" + weight_key = prefix + weight_name + if weight_key not in state_dict and weight_name in state_dict: + weight_key = weight_name + if weight_key not in state_dict: return - weight = state_dict[weight_name] - if weight.dtype == torch.float8_e4m3fn: - scale_inv_name = weight_name + "_scale_inv" + weight = state_dict[weight_key] + if weight.dtype != torch.float8_e4m3fn: + return + + active_prefix = prefix if weight_key == prefix + weight_name else "" + mod_prefix = weight_key.rsplit(".", 1)[0] + if checkpoint_layout is None: + scale_inv_name = weight_key + "_scale_inv" if scale_inv_name in state_dict: - # Rename to match our buffer name - mod_prefix = weight_name.rsplit(".", 1)[0] + # Rename to match our buffer name. state_dict[mod_prefix + ".weight_scale_inv"] = state_dict[scale_inv_name] + return + + layout_weight_name = weight_name + if weight_key != weight_name and not checkpoint_layout.is_weight_targeted( + layout_weight_name + ): + layout_weight_name = weight_key + + source_name = checkpoint_layout.scale_name_for_weight(layout_weight_name) + target_name = checkpoint_layout.runtime_scale_name_for_weight(layout_weight_name) + scale_key = self._add_prefix(active_prefix, source_name) + if scale_key not in state_dict and source_name in state_dict: + scale_key = source_name + if scale_key not in state_dict: + return + + target_key = self._add_prefix(active_prefix, target_name) + scale = state_dict[scale_key] + checkpoint_layout.validate_scale_shape( + weight_key, + weight.shape, + scale_key, + scale.shape, + ) + decoded_scale = checkpoint_layout.decode_scale(scale) + if scale_key != target_key: + state_dict.pop(scale_key) + state_dict[target_key] = decoded_scale # NOTE: post_load_hook intentionally inherited as None from the base # `Quantization`. UE8M0 conversion + TMA col-major layout for DeepGEMM is @@ -935,28 +1273,94 @@ def _apply( skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True ) + checkpoint_layout = self._get_finegrained_fp8_layout(qcfg) quant_method = str(qcfg.get("quant_method", "")).lower() - if quant_method != self.algo_name: + if quant_method != self.algo_name and checkpoint_layout is None: return gm, TransformInfo( skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True ) - if qcfg.get("weight_block_size") is None: + if qcfg.get("weight_block_size") is None and checkpoint_layout is None: return gm, TransformInfo( skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True ) - excluded = qcfg.get("modules_to_not_convert", []) + quant_context = _FineGrainedFP8LinearContext( + weight_block_size=self._resolve_weight_block_size(qcfg, checkpoint_layout), + runtime_scale_name=( + checkpoint_layout.runtime_scale_name + if checkpoint_layout is not None + else "weight_scale_inv" + ), + default_scales_layout=checkpoint_layout, + input_scale_fmt=self._resolve_input_scale_fmt(qcfg, checkpoint_layout), + ) + if checkpoint_layout is not None: + excluded = list( + dict.fromkeys( + [ + *(qcfg.get("exclude_modules") or []), + *(qcfg.get("modules_to_not_convert") or []), + ] + ) + ) + else: + excluded = qcfg.get("modules_to_not_convert") or [] cnt = 0 with WeightBiasInfoCache(): for n in gm.graph.nodes: if not is_linear_op(n): continue - if should_skip_quantization(n, excluded): - continue - self._insert_quantized_linear(gm, n, is_quantized_graph=False) + load_hook_kwargs = None + if checkpoint_layout is not None: + weight_name = extract_weight_name(n) + if not isinstance(weight_name, str): + continue + if not checkpoint_layout.is_weight_targeted(weight_name, excluded): + continue + load_hook_kwargs = {"checkpoint_layout": checkpoint_layout} + else: + if should_skip_quantization(n, excluded): + continue + self._insert_quantized_linear( + gm, + n, + is_quantized_graph=False, + load_hook_kwargs=load_hook_kwargs, + custom_kwargs={"input_scale_fmt": quant_context.input_scale_fmt}, + context=quant_context, + ) cnt += 1 + if checkpoint_layout is not None: + load_hook_kwargs = {"checkpoint_layout": checkpoint_layout} + for n in list(gm.graph.nodes): + grouped_match = self._extract_grouped_linear_match( + n, checkpoint_layout, excluded + ) + if grouped_match is None: + continue + source_node, input_node, bias_node, weight_node, weight_name, rank = ( + grouped_match + ) + self._insert_grouped_quantized_linear( + gm, + source_node, + input_node, + bias_node, + weight_node, + weight_name, + rank, + load_hook_kwargs=load_hook_kwargs, + context=quant_context, + ) + cnt += 1 + + if cnt: + gm.graph.eliminate_dead_code() + gm.graph.lint() + gm.recompile() + return gm, TransformInfo( skipped=False, num_matches=cnt, is_clean=False, has_valid_shapes=(cnt == 0) ) diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py b/tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py index 443fd5591dc0..9748592f8692 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py @@ -14,7 +14,7 @@ # limitations under the License. """Graph transform to optimize RMSNorm execution using FlashInfer.""" -from typing import Tuple, Type +from typing import Optional, Tuple, Type import torch from pydantic import Field @@ -42,6 +42,30 @@ } +def _extract_dtype_from_meta(node: Node) -> Optional[torch.dtype]: + val = node.meta.get("val") + if hasattr(val, "dtype"): + return val.dtype + + tensor_meta = node.meta.get("tensor_meta") + if hasattr(tensor_meta, "dtype"): + return tensor_meta.dtype + + return None + + +def _copy_meta_with_dtype(dst: Node, src: Node, dtype: torch.dtype) -> None: + dst.meta.update(src.meta) + + val = dst.meta.get("val") + if hasattr(val, "to"): + dst.meta["val"] = val.to(dtype=dtype) + + tensor_meta = dst.meta.get("tensor_meta") + if hasattr(tensor_meta, "_replace") and hasattr(tensor_meta, "dtype"): + dst.meta["tensor_meta"] = tensor_meta._replace(dtype=dtype) + + def _rms_norm_pattern(data: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: """Implements the RMSNorm pattern for pattern matching. @@ -375,11 +399,31 @@ def _apply( # Replace torch_rmsnorm ops with the selected backend for node in list(graph.nodes): if is_op(node, torch.ops.auto_deploy.torch_rmsnorm): + args = node.args + if backend == "flashinfer": + data, weight, eps = node.args + if isinstance(data, Node) and isinstance(weight, Node): + original_weight = weight + data_dtype = _extract_dtype_from_meta(data) + weight_dtype = _extract_dtype_from_meta(weight) + if ( + data_dtype is not None + and weight_dtype is not None + and data_dtype != weight_dtype + ): + with graph.inserting_before(node): + weight = graph.call_function( + torch.ops.aten.to.dtype, + args=(weight, data_dtype), + ) + _copy_meta_with_dtype(weight, original_weight, data_dtype) + args = (data, weight, eps) + # Replace with the selected backend op with graph.inserting_after(node): new_node: Node = graph.call_function( target_op, - args=node.args, + args=args, kwargs=node.kwargs, ) # Preserve metadata (including val/tensor_meta) for downstream transforms. diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py index f2bf0d307910..987702a57917 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py @@ -74,6 +74,7 @@ def is_trtllm_op_available(): filtered_nodes, get_all_layer_subgraphs, get_all_weights_in_subgraph, + get_op_schema, is_any_conv_op, is_any_delta_op, is_any_lin_op, @@ -984,7 +985,13 @@ class MXFP4EPShardingInfo(EPShardingInfo): def validate(self, gm: GraphModule = None, node: Node = None) -> bool: """Validate the transformation configuration.""" - if not is_op(node, torch.ops.auto_deploy.triton_mxfp4_moe): + if not is_op( + node, + [ + torch.ops.auto_deploy.torch_mxfp4_moe, + torch.ops.auto_deploy.triton_mxfp4_moe, + ], + ): ad_logger.warning(f"EP sharding is only supported for MOE nodes. Skipping {self}.") return False return True @@ -1062,6 +1069,7 @@ def apply(self, gm: GraphModule, node: Node) -> None: FineGrainedFP8EPShardingInfo, ), (lambda n: is_op(n, torch.ops.auto_deploy.torch_moe), EPShardingInfo), + (lambda n: is_op(n, torch.ops.auto_deploy.torch_mxfp4_moe), MXFP4EPShardingInfo), (lambda n: is_op(n, torch.ops.auto_deploy.triton_mxfp4_moe), MXFP4EPShardingInfo), ] @@ -1676,8 +1684,11 @@ def shard_weight_tensor( Tuple of (sharded_tensor, sharded_shape) """ + if custom_shard_fn is not None: + f_split = custom_shard_fn + # Handle fused weights - if fused_weight_dims is not None: + elif fused_weight_dims is not None: def f_split( t: torch.Tensor, @@ -2241,9 +2252,9 @@ def _insert_sharded_mxfp4_mlp_ep( config: ShardingTransformConfig, ): """ - Transform a call to auto_deploy::triton_mxfp4_moe into: + Transform a call to auto_deploy::torch_mxfp4_moe or auto_deploy::triton_mxfp4_moe into: - sharded expert parameters along dim 0 (this rank's slice), - - call to auto_deploy::triton_mxfp4_moe_ep(..., local_lo, local_hi), + - call to the matching *_mxfp4_moe_ep op, - followed by torch_dist_all_reduce. Expects the original op signature: @@ -2254,31 +2265,43 @@ def _insert_sharded_mxfp4_mlp_ep( down_blocks, down_bias, down_scales) """ - IDX_GATE_UP_BLOCKS = 4 - IDX_GATE_UP_BIAS = 5 - IDX_GATE_UP_SCALES = 6 - IDX_DOWN_BLOCKS = 9 - IDX_DOWN_BIAS = 10 - IDX_DOWN_SCALES = 11 - - gate_up_blocks_node = node.args[IDX_GATE_UP_BLOCKS] + expert_arg_names = ( + "gate_up_blocks", + "gate_up_bias", + "gate_up_scales", + "down_blocks", + "down_bias", + "down_scales", + ) + expert_args = extract_op_args(node, *expert_arg_names) + gate_up_blocks_node = expert_args[0] num_experts = shape(gate_up_blocks_node)[0] rank, world_size = config.rank, config.world_size local_lo, local_hi = _split_range_last_remainder(num_experts, world_size, rank) - # Prepare new args with slices for this rank + sliced_expert_args = { + name: _slice_expert_dim(gm, arg, local_lo, local_hi) + for name, arg in zip(expert_arg_names, expert_args) + } + set_op_args(node, **sliced_expert_args) + args = list(node.args) - args[IDX_GATE_UP_BLOCKS] = _slice_expert_dim(gm, args[IDX_GATE_UP_BLOCKS], local_lo, local_hi) - args[IDX_GATE_UP_BIAS] = _slice_expert_dim(gm, args[IDX_GATE_UP_BIAS], local_lo, local_hi) - args[IDX_GATE_UP_SCALES] = _slice_expert_dim(gm, args[IDX_GATE_UP_SCALES], local_lo, local_hi) - args[IDX_DOWN_BLOCKS] = _slice_expert_dim(gm, args[IDX_DOWN_BLOCKS], local_lo, local_hi) - args[IDX_DOWN_BIAS] = _slice_expert_dim(gm, args[IDX_DOWN_BIAS], local_lo, local_hi) - args[IDX_DOWN_SCALES] = _slice_expert_dim(gm, args[IDX_DOWN_SCALES], local_lo, local_hi) - - args_ep = tuple(args) + (int(world_size), int(rank)) - node.target = torch.ops.auto_deploy.triton_mxfp4_moe_ep.default - node.args = args_ep + kwargs = dict(node.kwargs or {}) + base_arg_names = [arg.name for arg in get_op_schema(node.target).arguments] + if "layer_type" in base_arg_names: + layer_type_pos = base_arg_names.index("layer_type") + if len(args) > layer_type_pos and "layer_type" not in kwargs: + kwargs["layer_type"] = args[layer_type_pos] + args = args[:layer_type_pos] + node.args = tuple(args) + node.kwargs = kwargs + + if is_op(node, torch.ops.auto_deploy.torch_mxfp4_moe): + node.target = torch.ops.auto_deploy.torch_mxfp4_moe_ep.default + else: + node.target = torch.ops.auto_deploy.triton_mxfp4_moe_ep.default + set_op_args(node, ep_size=int(world_size), ep_rank=int(rank)) # Add a dist all-reduce after the op (sum partial results across EP ranks) _, all_reduce_op = _get_dist_ops(config.dist_backend) diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py index f191a817c13e..768d049fcdb2 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py @@ -27,7 +27,7 @@ import operator from abc import ABC, abstractmethod from functools import partial -from typing import Any, Dict, List, Optional, Tuple, Type +from typing import Any, Dict, List, Literal, NamedTuple, Optional, Tuple, Type import torch from pydantic import Field, field_validator @@ -46,6 +46,7 @@ _get_op_schema, extract_op_args, extract_weight_nodes, + get_param_or_buffer, invalidate_weight_node_cache, is_any_lin_op, set_op_args, @@ -148,6 +149,80 @@ def _fp4_weight_scale_pipeline_cache_spec( } +def _assert_divisible(value: int, divisor: int, msg: str) -> None: + assert value % divisor == 0, f"{msg}: {value} is not divisible by {divisor}" + + +def _contiguous_partition(total: int, world_size: int, rank: int) -> Tuple[int, int]: + _assert_divisible(total, world_size, "contiguous partition") + part = total // world_size + return rank * part, (rank + 1) * part + + +def _split_flattened_groups( + tensor: torch.Tensor, + *, + num_groups: int, + rank: int, + world_size: int, +) -> torch.Tensor: + """Split a row-major ``[num_groups * rows_per_group, ...]`` tensor by group.""" + _assert_divisible(num_groups, world_size, "group count") + _assert_divisible(int(tensor.shape[0]), num_groups, "flattened grouped rows") + group_lo, group_hi = _contiguous_partition(num_groups, world_size, rank) + rows_per_group = int(tensor.shape[0]) // num_groups + grouped = tensor.reshape(num_groups, rows_per_group, *tensor.shape[1:]) + return grouped[group_lo:group_hi].reshape(-1, *tensor.shape[1:]).contiguous() + + +def _split_grouped_fp8_scale( + tensor: torch.Tensor, + *, + num_groups: int, + rank: int, + world_size: int, +) -> torch.Tensor: + if int(tensor.shape[0]) % num_groups == 0: + return _split_flattened_groups( + tensor, + num_groups=num_groups, + rank=rank, + world_size=world_size, + ) + return _split_fp8_block_scale(tensor, dim=0, rank=rank, world_size=world_size) + + +def _slice_group_range(tensor: torch.Tensor, group_lo: int, group_hi: int) -> torch.Tensor: + return tensor[group_lo:group_hi].contiguous() + + +def _get_attr_tensor(gm: GraphModule, node: Node, arg_name: str) -> Tuple[str, torch.Tensor]: + assert isinstance(node, Node) and node.op == "get_attr" and isinstance(node.target, str), ( + f"Expected {arg_name} to be a get_attr node, got {node!r}" + ) + return node.target, get_param_or_buffer(node.target, gm) + + +def _weight_node_from_attr(gm: GraphModule, node: Node, tensor: torch.Tensor) -> WeightNode: + node_key = str(node.target) + return WeightNode( + node=node, + node_key=node_key, + tensor=tensor, + submod=gm.get_submodule(node_key.rpartition(".")[0]), + ) + + +def _set_node_meta_shape(node: Node, new_shape: Tuple[int, ...]) -> None: + meta_val = node.meta.get("val") + if isinstance(meta_val, torch.Tensor): + node.meta["val"] = torch.empty( + new_shape, + dtype=meta_val.dtype, + device=meta_val.device, + ) + + _SHARDING_HINT_NAMES = frozenset( { "tp_mode", @@ -385,6 +460,133 @@ def _shard_scales(self, gm, dc, weight_nodes, dim, min_shape=1, fused=None): ) +@ShardableNode.register(torch.ops.auto_deploy.torch_fake_quant_grouped_finegrained_fp8_linear) +class GroupedFineGrainedFP8LinearShardableNode(ShardableNode): + """Grouped FineGrained FP8 linear: shard by whole output group boundaries.""" + + @staticmethod + def _input_is_tp_scaled_group_view(input_node: Node, group_dim: int) -> bool: + if not ( + isinstance(input_node, Node) + and input_node.op == "call_function" + and input_node.target == torch.ops.auto_deploy.view.default + ): + return False + + [tp_scaled_dim, view_shape] = extract_op_args(input_node, "tp_scaled_dim", "shape") + if tp_scaled_dim == -1: + return False + if tp_scaled_dim < 0: + tp_scaled_dim = len(view_shape) + tp_scaled_dim + return tp_scaled_dim == group_dim + + def apply(self, gm: GraphModule, dc: DistConfig, max_num_tokens: int = 0) -> int: + if dc.tp_size <= 1: + return 0 + + [input_node, weight_node, bias_node, weight_scale, tp_min_local_shape] = extract_op_args( + self.node, + "input", + "weight_quantized", + "bias", + "weight_scale", + "tp_min_local_shape", + ) + assert isinstance(weight_scale, (list, tuple)) and len(weight_scale) == 1, ( + "Grouped FineGrained FP8 linear expects a single weight_scale_inv tensor." + ) + scale_node = weight_scale[0] + weight_name, weight_tensor = _get_attr_tensor(gm, weight_node, "weight_quantized") + _, scale_tensor = _get_attr_tensor(gm, scale_node, "weight_scale") + + rank = int(tp_min_local_shape) if tp_min_local_shape else 1 + _assert_divisible(int(weight_tensor.shape[0]), rank, "grouped FP8 weight rows") + num_groups = int(weight_tensor.shape[0]) // rank + _assert_divisible(num_groups, dc.tp_size, "grouped FP8 output groups") + local_groups = num_groups // dc.tp_size + group_lo, group_hi = _contiguous_partition(num_groups, dc.tp_size, dc.tp_rank) + + input_shape = shape(input_node) + assert input_shape is not None and len(input_shape) >= 2, ( + f"Cannot determine grouped FP8 input shape for node {self.node.name}." + ) + group_dim = len(input_shape) - 2 + observed_groups = int(input_shape[group_dim]) + assert observed_groups in (num_groups, local_groups), ( + "Grouped FineGrained FP8 input group count must match either global groups " + f"({num_groups}) or TP-local groups ({local_groups}), got {observed_groups}." + ) + + split_groups = partial( + _split_flattened_groups, + num_groups=num_groups, + rank=dc.tp_rank, + world_size=dc.tp_size, + ) + shard_weight_tensor( + gm=gm, + weight_tensor=weight_tensor, + param_key=weight_name, + dim=SplitDimension.COLUMN, + rank=dc.tp_rank, + world_size=dc.tp_size, + custom_shard_fn=split_groups, + ) + + if isinstance(bias_node, Node): + bias_name, bias_tensor = _get_attr_tensor(gm, bias_node, "bias") + bias_split = ( + split_groups + if bias_tensor.ndim == 1 + else partial(_slice_group_range, group_lo=group_lo, group_hi=group_hi) + ) + shard_weight_tensor( + gm=gm, + weight_tensor=bias_tensor, + param_key=bias_name, + dim=SplitDimension.COLUMN, + rank=dc.tp_rank, + world_size=dc.tp_size, + custom_shard_fn=bias_split, + ) + + scale_split = partial( + _split_grouped_fp8_scale, + num_groups=num_groups, + rank=dc.tp_rank, + world_size=dc.tp_size, + ) + scale_weight_node = _weight_node_from_attr(gm, scale_node, scale_tensor) + _shard_scale_and_hook(gm, scale_weight_node, scale_split(scale_tensor), scale_split) + + view_localizes_groups = self._input_is_tp_scaled_group_view(input_node, group_dim) + if view_localizes_groups: + local_input_shape = list(input_shape) + local_input_shape[group_dim] = local_groups + _set_node_meta_shape(input_node, tuple(local_input_shape)) + elif observed_groups == num_groups: + with gm.graph.inserting_before(self.node): + local_input = gm.graph.call_function( + torch.ops.aten.slice.Tensor, + args=(input_node, group_dim, group_lo, group_hi, 1), + ) + local_shape = list(input_shape) + local_shape[group_dim] = local_groups + _set_node_meta_shape(local_input, tuple(local_shape)) + set_op_args(self.node, input=local_input) + + output_shape = shape(self.node) + if output_shape is not None: + local_output_shape = list(output_shape) + local_output_shape[-1] = local_groups * rank + _set_node_meta_shape(self.node, tuple(local_output_shape)) + + ad_logger.debug( + f" sharded grouped FP8 linear groups [{group_lo}:{group_hi}] / {num_groups}" + ) + return 1 + + @ShardableNode.register( torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear, torch.ops.auto_deploy.torch_quant_nvfp4_linear, @@ -431,19 +633,58 @@ def _strip_node_hints(cls, node: Node) -> bool: return True def apply(self, gm: GraphModule, dc: DistConfig, max_num_tokens: int = 0) -> int: - [tp_scaled_dim, view_shape] = extract_op_args(self.node, "tp_scaled_dim", "shape") + [tp_scaled_dim, view_shape, tp_min_local_shape] = extract_op_args( + self.node, "tp_scaled_dim", "shape", "tp_min_local_shape" + ) if tp_scaled_dim == -1: return 0 view_shape = list(view_shape) if tp_scaled_dim < 0: tp_scaled_dim = len(view_shape) + tp_scaled_dim + modified = self._shard_parameter_input(gm, dc, tp_scaled_dim, tp_min_local_shape) if tp_scaled_dim < len(view_shape) and isinstance(view_shape[tp_scaled_dim], int): view_shape[tp_scaled_dim] = -1 set_op_args(self.node, shape=view_shape) ad_logger.debug(f" updated view shape at dim {tp_scaled_dim} to -1 (inferred)") - return 1 - return 0 + modified = 1 + return modified + + def _shard_parameter_input( + self, + gm: GraphModule, + dc: DistConfig, + tp_scaled_dim: int, + tp_min_local_shape: Optional[int], + ) -> int: + """Shard get_attr tensors that are reshaped by a sharding-aware view.""" + view_input = self.node.args[0] + if not isinstance(view_input, Node) or view_input.op != "get_attr": + return 0 + if tp_scaled_dim != 0: + raise RuntimeError( + "Parameter sharding through auto_deploy.view currently supports " + f"tp_scaled_dim=0, got {tp_scaled_dim} on {self.node.name}." + ) + + weight_nodes = extract_weight_nodes(view_input) + shardable = weight_nodes.weights + weight_nodes.biases + if not shardable: + return 0 + + min_shape = tp_min_local_shape if tp_min_local_shape else 1 + for wn in shardable: + shard_weight_tensor( + gm=gm, + weight_tensor=wn.tensor, + param_key=wn.node_key, + dim=0, + rank=dc.tp_rank, + world_size=dc.tp_size, + min_local_shape=min_shape, + ) + ad_logger.debug(f" sharded {len(shardable)} parameter(s) feeding view {self.node.name}") + return 1 @ShardableNode.register(torch.ops.auto_deploy.split_with_sizes) @@ -494,7 +735,7 @@ def apply(self, gm: GraphModule, dc: DistConfig, max_num_tokens: int = 0) -> int if dc.tp_size <= 1: return 0 - _, all_reduce_op = _get_dist_ops("auto") + _, all_reduce_op = _get_dist_ops(dc.dist_backend) [x] = extract_op_args(self.node, "x") self.node.target = all_reduce_op self.node.args = (x, dc.allreduce_strategy) @@ -595,6 +836,87 @@ def apply(self, gm: GraphModule, dc: DistConfig, max_num_tokens: int = 0) -> int return 1 if count > 0 else 0 +@ShardableNode.register(torch.ops.auto_deploy.torch_attention) +class AttentionSinkShardableNode(ShardableNode): + """Attention ops with per-head sink tensors: shard sinks with TP heads.""" + + def apply(self, gm: GraphModule, dc: DistConfig, max_num_tokens: int = 0) -> int: + if dc.tp_size <= 1: + return 0 + + [enable_sharding] = extract_op_args(self.node, "enable_sharding") + if enable_sharding is False: + return 0 + + [sinks_node] = extract_op_args(self.node, "sinks") + if not isinstance(sinks_node, Node): + return 0 + + weight_nodes = extract_weight_nodes(sinks_node) + count = 0 + for wn in weight_nodes.weights: + if wn.tensor.ndim != 1: + continue + shard_weight_tensor( + gm=gm, + weight_tensor=wn.tensor, + param_key=wn.node_key, + dim=0, + rank=dc.tp_rank, + world_size=dc.tp_size, + ) + count += 1 + + if count: + ad_logger.debug(f" sharded {count} attention sink tensor(s)") + return 1 if count > 0 else 0 + + +class AttentionSinkArgShardableNode(ShardableNode): + """Attention op with an ``attn_sink`` argument: shard only the per-head sink tensor.""" + + def apply(self, gm: GraphModule, dc: DistConfig, max_num_tokens: int = 0) -> int: + if dc.tp_size <= 1: + return 0 + + [enable_sharding] = extract_op_args(self.node, "enable_sharding") + if enable_sharding is False: + return 0 + + [attn_sink_node] = extract_op_args(self.node, "attn_sink") + if not isinstance(attn_sink_node, Node): + return 0 + + weight_nodes = extract_weight_nodes(attn_sink_node) + count = 0 + for wn in weight_nodes.weights: + shard_weight_tensor( + gm=gm, + weight_tensor=wn.tensor, + param_key=wn.node_key, + dim=0, + rank=dc.tp_rank, + world_size=dc.tp_size, + ) + count += 1 + + if count: + ad_logger.debug(f" sharded {count} attention sink tensor(s)") + return 1 if count > 0 else 0 + + +for _sparse_attention_op_name in ( + "torch_deepseek_v4_sparse_attention", + "torch_deepseek_v4_sparse_attention_with_cache", + "flashmla_deepseek_v4_sparse_attention_with_cache", +): + try: + _sparse_attention_op = getattr(torch.ops.auto_deploy, _sparse_attention_op_name) + except AttributeError: + continue + ShardableNode.register(_sparse_attention_op)(AttentionSinkArgShardableNode) + + @ShardableNode.register(*_auto_deploy_ops("torch_rmsnorm_gated", "triton_rmsnorm_gated")) class NormShardableNode(ShardableNode): """Gated RMSNorm op: shard weight parameter.""" @@ -876,6 +1198,31 @@ def _localize_expert_indices( ) +class _StackedMoESchema(NamedTuple): + """Argument layout for stacked-expert MoE custom ops.""" + + expert_arg_names: Tuple[str, ...] + optional_trailing_arg_names: Tuple[str, ...] + + +_STACKED_MOE_EXPERT_ARGS = ( + "gate_up_blocks", + "gate_up_bias", + "gate_up_scales", + "down_blocks", + "down_bias", + "down_scales", +) +_ROUTER_STACKED_MOE_SCHEMA = _StackedMoESchema( + expert_arg_names=_STACKED_MOE_EXPERT_ARGS, + optional_trailing_arg_names=("layer_type",), +) +_ROUTING_DRIVEN_STACKED_MOE_SCHEMA = _StackedMoESchema( + expert_arg_names=_STACKED_MOE_EXPERT_ARGS, + optional_trailing_arg_names=("gate_up_order", "swiglu_mode", "layer_type", "num_experts_total"), +) + + class StackedMoEShardableNode(ShardableNode): """Stacked-tensor MoE EP sharding: slice along the expert dimension and rewrite. @@ -884,18 +1231,188 @@ class StackedMoEShardableNode(ShardableNode): stacked into 3-D tensors (``Tensor[num_experts, ...]``). Sharding slices along dim 0 to select the local expert partition. - Currently the only registered variant is ``triton_mxfp4_moe`` (MXFP4 - quantized), but the approach generalises to any stacked-tensor MoE op. - The op is rewritten to ``triton_mxfp4_moe_ep`` with an explicit + Registered MXFP4 variants carry explicit argument schemas. This covers both + router-style ops that compute top-k internally and routing-driven ops that + receive ``selected_experts``/``routing_weights`` from a separate router. + Each variant rewrites to its matching ``*_ep`` op with an explicit ``all_reduce`` after the node. """ - _IDX_GATE_UP_BLOCKS = 4 - _IDX_GATE_UP_BIAS = 5 - _IDX_GATE_UP_SCALES = 6 - _IDX_DOWN_BLOCKS = 9 - _IDX_DOWN_BIAS = 10 - _IDX_DOWN_SCALES = 11 + _EP_TARGET_BY_BASE: Dict[Any, OpOverload] = {} + _SCHEMA_BY_BASE: Dict[Any, _StackedMoESchema] = {} + + @classmethod + def register_variant( + cls, + base_target, + ep_target: OpOverload, + schema: Optional[_StackedMoESchema] = None, + ) -> None: + """Register one stacked MXFP4 MoE op and remember its matching EP op.""" + schema = schema or cls._infer_schema(base_target) + ShardableNode.register(base_target)(cls) + cls._EP_TARGET_BY_BASE[base_target] = ep_target + cls._SCHEMA_BY_BASE[base_target] = schema + if isinstance(base_target, OpOverloadPacket): + for overload_name in base_target.overloads(): + overload = getattr(base_target, overload_name) + cls._EP_TARGET_BY_BASE[overload] = ep_target + cls._SCHEMA_BY_BASE[overload] = schema + + @classmethod + def _arg_positions(cls, target) -> Dict[str, int]: + overload = ( + getattr(target, "default", target) if isinstance(target, OpOverloadPacket) else target + ) + schema = getattr(overload, "_schema", None) + if schema is None: + return {} + return {arg.name: idx for idx, arg in enumerate(schema.arguments)} + + @classmethod + def _infer_schema(cls, target) -> _StackedMoESchema: + arg_positions = cls._arg_positions(target) + if not arg_positions: + raise RuntimeError(f"Cannot infer stacked MoE schema for op {target}") + + if not all(arg_name in arg_positions for arg_name in _STACKED_MOE_EXPERT_ARGS): + raise RuntimeError(f"Op {target} does not match the stacked MXFP4 MoE expert schema") + + if {"selected_experts", "routing_weights"}.issubset(arg_positions): + return _ROUTING_DRIVEN_STACKED_MOE_SCHEMA + if {"router_weight", "router_bias", "top_k"}.issubset(arg_positions): + return _ROUTER_STACKED_MOE_SCHEMA + raise RuntimeError(f"Op {target} does not expose a supported stacked MXFP4 routing schema") + + @classmethod + def _ep_target_for(cls, target) -> Optional[OpOverload]: + ep_target = cls._EP_TARGET_BY_BASE.get(target) + if ep_target is None and isinstance(target, OpOverloadPacket): + ep_target = cls._EP_TARGET_BY_BASE.get(getattr(target, "default", None)) + return ep_target + + @classmethod + def _schema_for(cls, target) -> Optional[_StackedMoESchema]: + schema = cls._SCHEMA_BY_BASE.get(target) + if schema is None and isinstance(target, OpOverloadPacket): + schema = cls._SCHEMA_BY_BASE.get(getattr(target, "default", None)) + return schema + + @classmethod + def _move_optional_args_to_kwargs( + cls, + args: List[Any], + kwargs: Dict[str, Any], + base_target, + ep_target, + arg_names: Tuple[str, ...], + ) -> Tuple[List[Any], Dict[str, Any]]: + base_positions = cls._arg_positions(base_target) + ep_positions = cls._arg_positions(ep_target) + ordered_names = sorted( + arg_names, + key=lambda name: base_positions.get(name, -1), + reverse=True, + ) + for arg_name in ordered_names: + base_pos = base_positions.get(arg_name) + if base_pos is None: + continue + if arg_name not in ep_positions: + if arg_name in kwargs or base_pos < len(args): + raise RuntimeError( + f"Cannot preserve argument '{arg_name}' when rewriting " + f"{base_target} to {ep_target}: EP op does not accept it." + ) + continue + if arg_name in kwargs: + continue + if base_pos < len(args): + kwargs[arg_name] = args.pop(base_pos) + return args, kwargs + + @classmethod + def _ep_topology_args( + cls, + ep_target: OpOverload, + expert_start: int, + ep_size: int, + ep_rank: int, + num_experts_total: Optional[int] = None, + ) -> Dict[str, int]: + ep_positions = cls._arg_positions(ep_target) + topology_args: Dict[str, int] = {} + if "expert_start" in ep_positions: + topology_args["expert_start"] = int(expert_start) + if "num_experts_total" in ep_positions: + if num_experts_total is None: + raise RuntimeError( + f"EP op {ep_target} accepts num_experts_total but the source op " + "did not expose a global expert count." + ) + topology_args["num_experts_total"] = int(num_experts_total) + + has_ep_size = "ep_size" in ep_positions + has_ep_rank = "ep_rank" in ep_positions + if has_ep_size or has_ep_rank: + if not (has_ep_size and has_ep_rank): + raise RuntimeError(f"EP op {ep_target} must accept both ep_size and ep_rank") + topology_args["ep_size"] = int(ep_size) + topology_args["ep_rank"] = int(ep_rank) + + if not topology_args: + raise RuntimeError( + f"EP op {ep_target} must accept expert_start or ep_size/ep_rank topology args" + ) + return topology_args + + @staticmethod + def _slice_experts(tensor: torch.Tensor, lo: int, hi: int) -> torch.Tensor: + return tensor[lo:hi] + + @classmethod + def _shard_get_attr_expert_arg(cls, gm: GraphModule, arg: Node, lo: int, hi: int) -> Node: + mod_name, _, attr_name = arg.target.rpartition(".") + submod = gm.get_submodule(mod_name) + tensor = getattr(submod, attr_name) + local_tensor = cls._slice_experts(tensor, lo, hi).detach().clone() + + if attr_name in submod._buffers: + persistent = attr_name not in submod._non_persistent_buffers_set + submod.register_buffer(attr_name, local_tensor, persistent=persistent) + elif attr_name in submod._parameters: + requires_grad = submod._parameters[attr_name].requires_grad + setattr( + submod, + attr_name, + torch.nn.Parameter(local_tensor, requires_grad=requires_grad), + ) + else: + setattr(submod, attr_name, local_tensor) + + arg.meta["val"] = local_tensor + f_split = partial(cls._slice_experts, lo=lo, hi=hi) + submod._register_load_state_dict_pre_hook( + partial( + _load_hook, + f_split=f_split, + param_key=attr_name, + param_shape=local_tensor.shape, + ) + ) + invalidate_weight_node_cache(gm) + return arg + + @classmethod + def _local_expert_arg(cls, gm: GraphModule, node: Node, arg: Any, lo: int, hi: int) -> Any: + if isinstance(arg, Node) and arg.op == "get_attr": + return cls._shard_get_attr_expert_arg(gm, arg, lo, hi) + + with gm.graph.inserting_before(node): + return gm.graph.call_function( + torch.ops.aten.slice.Tensor, + args=(arg, 0, lo, hi, 1), + ) def apply(self, gm: GraphModule, dc: DistConfig, max_num_tokens: int = 0) -> int: ep_size = dc.moe_ep_size @@ -903,10 +1420,19 @@ def apply(self, gm: GraphModule, dc: DistConfig, max_num_tokens: int = 0) -> int if ep_size <= 1: return 0 - - expert_shape = shape(self.node.args[self._IDX_GATE_UP_BLOCKS]) + ep_target = self._ep_target_for(self.node.target) + if ep_target is None: + raise RuntimeError(f"No EP target registered for stacked MoE op {self.node.target}") + schema = self._schema_for(self.node.target) + if schema is None: + raise RuntimeError(f"No stacked MoE schema registered for op {self.node.target}") + + expert_args = extract_op_args(self.node, *schema.expert_arg_names) + first_expert_name = schema.expert_arg_names[0] + first_expert_arg = expert_args[0] + expert_shape = shape(first_expert_arg) assert expert_shape is not None, ( - f"Cannot determine num_experts: gate_up_blocks arg has no shape metadata " + f"Cannot determine num_experts: {first_expert_name} arg has no shape metadata " f"(node: {self.node.name})" ) num_experts = expert_shape[0] @@ -914,25 +1440,36 @@ def apply(self, gm: GraphModule, dc: DistConfig, max_num_tokens: int = 0) -> int lo = base * ep_rank hi = num_experts if ep_rank == ep_size - 1 else base * (ep_rank + 1) + sliced_expert_args = {} + for arg_name, arg in zip(schema.expert_arg_names, expert_args): + sliced_expert_args[arg_name] = self._local_expert_arg(gm, self.node, arg, lo, hi) + set_op_args(self.node, **sliced_expert_args) + args = list(self.node.args) - for idx in ( - self._IDX_GATE_UP_BLOCKS, - self._IDX_GATE_UP_BIAS, - self._IDX_GATE_UP_SCALES, - self._IDX_DOWN_BLOCKS, - self._IDX_DOWN_BIAS, - self._IDX_DOWN_SCALES, - ): - with gm.graph.inserting_after(args[idx]): - args[idx] = gm.graph.call_function( - torch.ops.aten.slice.Tensor, - args=(args[idx], 0, lo, hi, 1), - ) + kwargs = dict(self.node.kwargs or {}) + args, kwargs = self._move_optional_args_to_kwargs( + args, + kwargs, + self.node.target, + ep_target, + schema.optional_trailing_arg_names, + ) - self.node.target = torch.ops.auto_deploy.triton_mxfp4_moe_ep.default - self.node.args = tuple(args) + (int(ep_size), int(ep_rank)) + self.node.target = ep_target + self.node.args = tuple(args) + self.node.kwargs = kwargs + set_op_args( + self.node, + **self._ep_topology_args( + ep_target=ep_target, + expert_start=lo, + ep_size=ep_size, + ep_rank=ep_rank, + num_experts_total=num_experts, + ), + ) - _, all_reduce_op = _get_dist_ops("auto") + _, all_reduce_op = _get_dist_ops(dc.dist_backend) with gm.graph.inserting_after(self.node): red = gm.graph.call_function( all_reduce_op, @@ -947,13 +1484,26 @@ def apply(self, gm: GraphModule, dc: DistConfig, max_num_tokens: int = 0) -> int return 1 -# MXFP4 MoE ops depend on triton_kernels + TRT-LLM internals that aren't in the -# standalone package, so the op may not be registered at import time. Register -# only when the op is actually available. -try: - ShardableNode.register(torch.ops.auto_deploy.triton_mxfp4_moe)(StackedMoEShardableNode) -except AttributeError: - pass +def _register_stacked_mxfp4_moe_variant(base_name: str, ep_name: str) -> None: + """Register an optional stacked MXFP4 MoE op pair when both schemas exist.""" + try: + base_target = getattr(torch.ops.auto_deploy, base_name) + except AttributeError: + return + try: + ep_target = getattr(torch.ops.auto_deploy, ep_name).default + except AttributeError: + return + StackedMoEShardableNode.register_variant(base_target, ep_target) + + +for _base_name, _ep_name in ( + ("triton_mxfp4_moe", "triton_mxfp4_moe_ep"), + ("triton_mxfp4_moe_from_routing", "triton_mxfp4_moe_from_routing_ep"), + ("torch_mxfp4_moe", "torch_mxfp4_moe_ep"), + ("torch_mxfp4_moe_from_routing", "torch_mxfp4_moe_from_routing_ep"), +): + _register_stacked_mxfp4_moe_variant(_base_name, _ep_name) # ============================================================================= @@ -980,6 +1530,7 @@ class IRShardingConfig(TransformConfig): description="When set, only shard nodes whose layer_type hint is in this list.", ) enable_attention_dp: bool = Field(default=False) + dist_backend: Literal["auto", "torch", "trtllm"] = Field(default="auto") dist_mapping: dict[str, int] = Field(default_factory=dict) dist_config: DistConfig = Field(default_factory=DistConfig) @@ -1003,6 +1554,7 @@ def _init_dist_config(self, rank: int, world_size: int): dist_mapping=self.dist_mapping, enable_attention_dp=self.enable_attention_dp, allreduce_strategy=self.allreduce_strategy.name, + dist_backend=self.dist_backend, ) @@ -1017,7 +1569,7 @@ def _log_sharding_prelude(dc: DistConfig) -> None: ad_logger.info( f"apply_sharding_hints{skip}: tp_size={dc.tp_size}, tp_rank={dc.tp_rank}, " f"moe grid: [ep x tp] = [{dc.moe_ep_size} x {dc.moe_tp_size}], " - f"strategy={dc.allreduce_strategy}" + f"strategy={dc.allreduce_strategy}, backend={dc.dist_backend}" ) @@ -1167,6 +1719,7 @@ def _apply( self.config._init_dist_config(shared_config.local_rank, shared_config.world_size) dc = self.config.dist_config + dc.dist_backend = self.config.dist_backend _log_sharding_prelude(dc) if shared_config.world_size < 2: diff --git a/tensorrt_llm/_torch/auto_deploy/utils/dist_config.py b/tensorrt_llm/_torch/auto_deploy/utils/dist_config.py index 060ff28c1d40..ccf1b5311e16 100644 --- a/tensorrt_llm/_torch/auto_deploy/utils/dist_config.py +++ b/tensorrt_llm/_torch/auto_deploy/utils/dist_config.py @@ -22,7 +22,7 @@ """ import json -from typing import Any +from typing import Any, Literal from pydantic import BaseModel, Field, model_validator @@ -41,6 +41,7 @@ class DistConfig(BaseModel): moe_cluster_size: int = Field(default=1, ge=1) enable_attention_dp: bool = Field(default=False) allreduce_strategy: str = Field(default="NCCL") + dist_backend: Literal["auto", "torch", "trtllm"] = Field(default="auto") @model_validator(mode="after") def _validate_grid(self) -> "DistConfig": @@ -93,6 +94,7 @@ def to_dict(self) -> dict: "moe_cluster_size": self.moe_cluster_size, "enable_attention_dp": self.enable_attention_dp, "allreduce_strategy": self.allreduce_strategy, + "dist_backend": self.dist_backend, } @classmethod @@ -123,6 +125,7 @@ def from_mapping(mapping: Any) -> "DistConfig": moe_ep_size=mapping.moe_ep_size, moe_cluster_size=mapping.moe_cluster_size, enable_attention_dp=mapping.enable_attention_dp, + dist_backend=getattr(mapping, "dist_backend", "auto"), ) @classmethod @@ -134,6 +137,7 @@ def from_sharding_params( dist_mapping: dict, enable_attention_dp: bool = False, allreduce_strategy: str = "NCCL", + dist_backend: Literal["auto", "torch", "trtllm"] = "auto", ) -> "DistConfig": """Build ``DistConfig`` from sharding-transform YAML inputs + runtime MPI info. @@ -151,6 +155,7 @@ def from_sharding_params( moe_cluster_size=dist_mapping.get("moe_cluster", 1), enable_attention_dp=enable_attention_dp, allreduce_strategy=allreduce_strategy, + dist_backend=dist_backend, ) def to_mapping(self) -> Any: diff --git a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py index 17b4012a83e4..5af8d0eb852a 100644 --- a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py +++ b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py @@ -665,6 +665,7 @@ def is_finegrained_fp8_linear_op(node: Node) -> bool: node, [ torch.ops.auto_deploy.torch_fake_quant_finegrained_fp8_linear, + torch.ops.auto_deploy.torch_fake_quant_grouped_finegrained_fp8_linear, torch.ops.auto_deploy.trtllm_finegrained_fp8_linear, ], ) @@ -680,6 +681,8 @@ def is_any_moe_op(node: Node) -> bool: _auto_deploy_op("torch_quant_fp8_moe"), _auto_deploy_op("torch_quant_nvfp4_moe"), _auto_deploy_op("torch_quant_finegrained_fp8_moe"), + _auto_deploy_op("torch_mxfp4_moe"), + _auto_deploy_op("torch_mxfp4_moe_from_routing"), _auto_deploy_op("triton_mxfp4_moe"), _auto_deploy_op("torch_moe_fused"), _auto_deploy_op("torch_moe_dense_mlp"), diff --git a/tensorrt_llm/_torch/auto_deploy/utils/quantization_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/quantization_utils.py index 6dd2fc4c19e5..37b2c110a9cf 100644 --- a/tensorrt_llm/_torch/auto_deploy/utils/quantization_utils.py +++ b/tensorrt_llm/_torch/auto_deploy/utils/quantization_utils.py @@ -115,6 +115,64 @@ def fp8_scale(input: torch.Tensor) -> torch.Tensor: return torch.max(torch.abs(input).to(torch.float)) / FP8_MAX +def ceil_pow2_scale(amax: torch.Tensor, max_value: float, min_value: float) -> torch.Tensor: + return torch.pow(2.0, torch.ceil(torch.log2(amax.clamp_min(min_value) / max_value))) + + +def fake_fp8_act_quant(x: torch.Tensor, block_size: int = 64) -> torch.Tensor: + dim = x.shape[-1] + if dim == 0 or dim % block_size != 0: + return x + + dtype = x.dtype + x_float = x.float() + grouped = x_float.reshape(-1, dim // block_size, block_size) + scale = ceil_pow2_scale(grouped.abs().amax(dim=-1, keepdim=True), 448.0, 1.0e-4) + quant = torch.clamp(grouped / scale, -448.0, 448.0).to(dtype).float() + return (quant * scale).reshape_as(x_float).to(dtype) + + +def fake_fp4_act_quant(x: torch.Tensor, block_size: int = 32) -> torch.Tensor: + dim = x.shape[-1] + if dim == 0 or dim % block_size != 0: + return x + + dtype = x.dtype + x_float = x.float() + grouped = x_float.reshape(-1, dim // block_size, block_size) + scale = ceil_pow2_scale(grouped.abs().amax(dim=-1, keepdim=True), 6.0, 6.0 * 2.0**-126) + normalized = torch.clamp(grouped / scale, -6.0, 6.0) + abs_normalized = normalized.abs() + quant_abs = torch.zeros_like(abs_normalized) + quant_abs = torch.where(abs_normalized > 0.25, torch.full_like(quant_abs, 0.5), quant_abs) + quant_abs = torch.where(abs_normalized > 0.75, torch.full_like(quant_abs, 1.0), quant_abs) + quant_abs = torch.where(abs_normalized > 1.25, torch.full_like(quant_abs, 1.5), quant_abs) + quant_abs = torch.where(abs_normalized > 1.75, torch.full_like(quant_abs, 2.0), quant_abs) + quant_abs = torch.where(abs_normalized > 2.5, torch.full_like(quant_abs, 3.0), quant_abs) + quant_abs = torch.where(abs_normalized > 3.5, torch.full_like(quant_abs, 4.0), quant_abs) + quant_abs = torch.where(abs_normalized > 5.0, torch.full_like(quant_abs, 6.0), quant_abs) + quant = quant_abs * normalized.sign() + return (quant * scale).reshape_as(x_float).to(dtype) + + +def hadamard_rotate(x: torch.Tensor) -> torch.Tensor: + dim = x.shape[-1] + if dim <= 1: + return x + if dim & (dim - 1): + raise ValueError(f"Hadamard rotation requires power-of-two dimension, got {dim}.") + + out = x.reshape(-1, dim).float() + width = 1 + while width < dim: + out = out.reshape(-1, dim // (2 * width), 2, width) + left = out[..., 0, :] + right = out[..., 1, :] + out = torch.cat((left + right, left - right), dim=-1).flatten(-2) + width *= 2 + return (out * (dim**-0.5)).reshape_as(x).to(x.dtype) + + def is_quantized_graph(gm: GraphModule): """Check if the graph is quantized by modelopt.""" for n in gm.graph.nodes: diff --git a/tests/integration/defs/accuracy/configs/deepseek_v4_flash_4gpu_smoke.yaml b/tests/integration/defs/accuracy/configs/deepseek_v4_flash_4gpu_smoke.yaml new file mode 100644 index 000000000000..8ff5985a6d58 --- /dev/null +++ b/tests/integration/defs/accuracy/configs/deepseek_v4_flash_4gpu_smoke.yaml @@ -0,0 +1,22 @@ +# 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. + +transforms: + apply_sharding_hints: + dist_mapping: + tp: 4 + moe_ep: 4 + moe_tp: 1 + moe_cluster: 1 diff --git a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py index 902f7addf53a..a6153c00b1e6 100644 --- a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py +++ b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @@ -24,6 +24,7 @@ from test_common.llm_data import hf_id_to_local_model_dir from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM +from tensorrt_llm.evaluate import GSM8K as GSM8KEvaluator from tensorrt_llm.llmapi import Eagle3DecodingConfig from tensorrt_llm.quantization import QuantAlgo from tensorrt_llm.sampling_params import SamplingParams @@ -35,6 +36,7 @@ 'model_registry' / 'configs') _AD_MODEL_REGISTRY_DIR = Path( get_llm_root()) / 'examples' / 'auto_deploy' / 'model_registry' +_ACCURACY_CONFIGS_DIR = Path(__file__).resolve().parent / "configs" def _load_ad_config(config_name): @@ -1592,6 +1594,60 @@ def test_autodeploy_from_registry(self, model_name, config_overrides, tasks, raise type(e)(f"[{task_cls.__name__}] {e}") from None +class TestDeepSeekV4Flash(LlmapiAccuracyTestHarness): + MODEL_NAME = "deepseek-ai/DeepSeek-V4-Flash" + MODEL_PATH = hf_id_to_local_model_dir(MODEL_NAME) + WORLD_SIZE = 4 + YAML_EXTRA = [ + str(_AD_CONFIGS_DIR / "dashboard_default.yaml"), + str(_AD_CONFIGS_DIR / "world_size_4.yaml"), + str(_AD_CONFIGS_DIR / "deepseek_v4_flash.yaml"), + str(_ACCURACY_CONFIGS_DIR / "deepseek_v4_flash_4gpu_smoke.yaml"), + ] + GSM8K_NUM_SAMPLES = 15 + GSM8K_NUM_FEWSHOT = 0 + GSM8K_MAX_INPUT_LEN = 1024 + GSM8K_MAX_OUTPUT_LEN = 128 + GSM8K_MIN_ACCURACY = 40.0 + + def get_default_sampling_params(self): + return SamplingParams(end_id=None, + pad_id=None, + max_tokens=self.GSM8K_MAX_OUTPUT_LEN, + n=1, + use_beam_search=False) + + @pytest.mark.skip_less_device(4) + @pytest.mark.skip_less_device_memory(80000) + def test_gsm8k_smoke(self): + if get_device_count() < self.WORLD_SIZE: + pytest.skip( + f"DeepSeek V4 Flash smoke requires {self.WORLD_SIZE} GPUs") + + with AutoDeployLLM(model=self.MODEL_PATH, + tokenizer=self.MODEL_PATH, + world_size=self.WORLD_SIZE, + yaml_extra=self.YAML_EXTRA, + max_seq_len=self.GSM8K_MAX_INPUT_LEN + + self.GSM8K_MAX_OUTPUT_LEN, + trust_remote_code=True) as llm: + task = GSM8KEvaluator(dataset_path=GSM8K.DATASET_DIR, + num_samples=self.GSM8K_NUM_SAMPLES, + random_seed=0) + for task_obj in task.task_dict.values(): + task_obj.set_config("num_fewshot", self.GSM8K_NUM_FEWSHOT) + score = task.evaluate( + llm, + sampling_params=self.get_default_sampling_params(), + scores_filter="exact_match,flexible-extract", + ) + + assert score >= self.GSM8K_MIN_ACCURACY, ( + f"DeepSeek V4 Flash GSM8K smoke accuracy {score:.2f} is below " + f"{self.GSM8K_MIN_ACCURACY:.2f} on {self.GSM8K_NUM_SAMPLES} samples" + ) + + # ============================================================================= # IR Sharding Path Tests # ============================================================================= diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index db3730b3f9ed..569d88e9b99d 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -383,6 +383,7 @@ l0_dgx_h100: - disaggregated/test_ad_disagg.py::test_async_generation_no_overlap_matches_aggregate - disaggregated/test_ad_disagg.py::test_async_sharded_generation_handoff - disaggregated/test_ad_disagg.py::test_async_eagle3_full_model_handoff + - accuracy/test_llm_api_autodeploy.py::TestDeepSeekV4Flash::test_gsm8k_smoke # ------------- AutoDeploy Backend Stages L1 / Nightly only --------------- - condition: ranges: diff --git a/tests/test_common/llm_data.py b/tests/test_common/llm_data.py index 19cd069153dd..71cbd8dd605b 100644 --- a/tests/test_common/llm_data.py +++ b/tests/test_common/llm_data.py @@ -41,6 +41,7 @@ "deepseek-ai/DeepSeek-V3": "DeepSeek-V3", "deepseek-ai/DeepSeek-R1": "DeepSeek-R1/DeepSeek-R1", "deepseek-ai/DeepSeek-R1-0528": "DeepSeek-R1/DeepSeek-R1-0528", + "deepseek-ai/DeepSeek-V4-Flash": "DeepSeek-V4-Flash", "ibm-ai-platform/Bamba-9B-v2": "Bamba-9B-v2", "nvidia/NVIDIA-Nemotron-Nano-12B-v2": "NVIDIA-Nemotron-Nano-12B-v2", "nvidia/NVIDIA-Nemotron-Nano-31B-A3-v3": "NVIDIA-Nemotron-Nano-31B-A3-v3", diff --git a/tests/unittest/_torch/auto_deploy/unit/models/test_deepseek_v4_quant_checkpoint_layout.py b/tests/unittest/_torch/auto_deploy/unit/models/test_deepseek_v4_quant_checkpoint_layout.py new file mode 100644 index 000000000000..28f2499ec25a --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/models/test_deepseek_v4_quant_checkpoint_layout.py @@ -0,0 +1,108 @@ +# 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. + +import json +import struct +from pathlib import Path + +from tensorrt_llm._torch.auto_deploy.models.checkpoint_metadata import has_safetensors_metadata +from tensorrt_llm._torch.auto_deploy.models.quant_config_reader import HFQuantConfigReader + + +def _write_deepseek_v4_checkpoint_fixture(tmp_path: Path) -> None: + config = { + "model_type": "deepseek_v4", + "quantization_config": { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + "scale_fmt": "ue8m0", + "weight_block_size": [128, 128], + }, + } + tensor_metadata = { + "layers.0.attn.wq_a.weight": {"dtype": "F8_E4M3", "shape": [1024, 4096]}, + "layers.0.attn.wq_a.scale": {"dtype": "F8_E8M0", "shape": [8, 32]}, + "layers.0.ffn.experts.0.w1.weight": {"dtype": "I8", "shape": [2048, 2048]}, + "layers.0.ffn.experts.0.w1.scale": {"dtype": "F8_E8M0", "shape": [2048, 128]}, + "layers.0.ffn.experts.0.w2.weight": {"dtype": "I8", "shape": [4096, 1024]}, + "layers.0.ffn.experts.0.w2.scale": {"dtype": "F8_E8M0", "shape": [4096, 64]}, + "layers.0.ffn.experts.0.w3.weight": {"dtype": "I8", "shape": [2048, 2048]}, + "layers.0.ffn.experts.0.w3.scale": {"dtype": "F8_E8M0", "shape": [2048, 128]}, + } + safetensors_name = "model-00001-of-00001.safetensors" + header = { + name: { + "dtype": metadata["dtype"], + "shape": metadata["shape"], + "data_offsets": [0, 0], + } + for name, metadata in tensor_metadata.items() + } + header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8") + + (tmp_path / "config.json").write_text(json.dumps(config), encoding="utf-8") + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps( + { + "metadata": {"total_size": 0}, + "weight_map": {name: safetensors_name for name in tensor_metadata}, + } + ), + encoding="utf-8", + ) + (tmp_path / safetensors_name).write_bytes(struct.pack(" None: + _write_deepseek_v4_checkpoint_fixture(tmp_path) + + result = HFQuantConfigReader.from_file(str(tmp_path)) + + assert result is not None + reader, extra_model_kwargs = result + qcfg = reader.get_config() + assert extra_model_kwargs == {} + assert qcfg["checkpoint_layout"] is not None + assert qcfg["expert_quant_method"] == "mxfp4" + assert qcfg["expert_block_size"] == 32 + assert qcfg["scale_fmt"] == "ue8m0" + assert qcfg["weight_block_size"] == [128, 128] + assert "*.ffn.gate" in qcfg["exclude_modules"] + assert "mtp.*" in qcfg["exclude_modules"] + + +def test_non_deepseek_fp8_config_uses_generic_hf_behavior() -> None: + reader = HFQuantConfigReader() + reader.read_config( + { + "model_type": "llama", + "quantization_config": { + "quant_method": "fp8", + "exclude_modules": ["custom"], + }, + } + ) + + qcfg = reader.get_config() + assert qcfg["quant_method"] == "fp8" + assert qcfg["exclude_modules"] == ["custom", "lm_head", "model.embed_tokens"] + assert "checkpoint_layout" not in qcfg + + +def test_safetensors_metadata_probe_tolerates_missing_directory(tmp_path: Path) -> None: + assert not has_safetensors_metadata(tmp_path / "missing") diff --git a/tests/unittest/_torch/auto_deploy/unit/quantization/test_deepseek_v4_finegrained_fp8_linear.py b/tests/unittest/_torch/auto_deploy/unit/quantization/test_deepseek_v4_finegrained_fp8_linear.py new file mode 100644 index 000000000000..2bdc5cda6052 --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/quantization/test_deepseek_v4_finegrained_fp8_linear.py @@ -0,0 +1,494 @@ +# 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. + +import pytest +import torch +import torch.nn as nn +from torch.fx import Graph, GraphModule, Node + +import tensorrt_llm._torch.auto_deploy.custom_ops.quantization.torch_quant # noqa: F401 +from tensorrt_llm._torch.auto_deploy.models.quant_checkpoint_layout import ( + FineGrainedFP8CheckpointLayout, + QuantCheckpointLayoutRegistry, + QuantizedCheckpointLayout, +) +from tensorrt_llm._torch.auto_deploy.transform.interface import ( + SharedConfig, + Stages, + TransformConfig, +) +from tensorrt_llm._torch.auto_deploy.transform.library.quantization import ( + FineGrainedFP8LinearQuantization, +) +from tensorrt_llm._torch.auto_deploy.transform.library.sharding_ir import ShardableNode +from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op + +_FP8_DTYPE = getattr(torch, "float8_e4m3fn", None) +_E8M0_DTYPE = getattr(torch, "float8_e8m0fnu", None) + +pytestmark = pytest.mark.skipif(_FP8_DTYPE is None, reason="Requires torch.float8_e4m3fn") + +_ALLOWED_WEIGHT_NAMES = ( + "layers.0.attn.wq_a.weight", + "layers.0.attn.wq_b.weight", + "layers.0.attn.wkv.weight", + "layers.0.attn.wo_a.weight", + "layers.0.attn.wo_b.weight", + "layers.0.attn.indexer.wq_b.weight", + "layers.0.ffn.shared_experts.w1.weight", + "layers.0.ffn.shared_experts.w2.weight", + "layers.0.ffn.shared_experts.w3.weight", +) +_SKIPPED_WEIGHT_NAMES = ( + "layers.0.attn.compressor.wkv.weight", + "layers.0.attn.indexer.compressor.wkv.weight", + "layers.0.attn.indexer.weights_proj.weight", + "layers.0.ffn.experts.0.w1.weight", + "layers.0.ffn.experts.0.w2.weight", + "layers.0.ffn.experts.0.w3.weight", + "layers.0.ffn.gate.weight", + "head.weight", +) + + +class _Factory: + def __init__(self, quant_config: dict[str, object]) -> None: + self._quant_config = quant_config + + def get_quant_config(self) -> dict[str, object]: + return self._quant_config + + +class _WeightOnlyLinear(nn.Module): + def __init__(self, out_features: int = 129, in_features: int = 257) -> None: + super().__init__() + self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=torch.float16)) + + +class _DeepSeekLinearSurface(nn.Module): + def __init__(self) -> None: + super().__init__() + self.layers = nn.ModuleList([nn.Module()]) + layer = self.layers[0] + + layer.attn = nn.Module() + for name in ("wq_a", "wq_b", "wkv", "wo_a", "wo_b"): + setattr(layer.attn, name, _WeightOnlyLinear()) + layer.attn.compressor = nn.Module() + layer.attn.compressor.wkv = _WeightOnlyLinear() + layer.attn.indexer = nn.Module() + layer.attn.indexer.wq_b = _WeightOnlyLinear() + layer.attn.indexer.weights_proj = _WeightOnlyLinear() + layer.attn.indexer.compressor = nn.Module() + layer.attn.indexer.compressor.wkv = _WeightOnlyLinear() + + layer.ffn = nn.Module() + layer.ffn.shared_experts = nn.Module() + for name in ("w1", "w2", "w3"): + setattr(layer.ffn.shared_experts, name, _WeightOnlyLinear()) + layer.ffn.experts = nn.ModuleList([nn.Module()]) + for name in ("w1", "w2", "w3"): + setattr(layer.ffn.experts[0], name, _WeightOnlyLinear()) + layer.ffn.gate = _WeightOnlyLinear() + + self.head = _WeightOnlyLinear() + + +def _transform() -> FineGrainedFP8LinearQuantization: + return FineGrainedFP8LinearQuantization(TransformConfig(stage=Stages.PATTERN_MATCHER)) + + +def _deepseek_v4_checkpoint_layout() -> QuantizedCheckpointLayout: + checkpoint_layout = QuantCheckpointLayoutRegistry.build_from_config( + { + "model_type": "deepseek_v4", + "quantization_config": { + "quant_method": "fp8", + "scale_fmt": "ue8m0", + "weight_block_size": [128, 128], + }, + } + ) + assert isinstance(checkpoint_layout, QuantizedCheckpointLayout) + return checkpoint_layout + + +def _finegrained_fp8_layout() -> FineGrainedFP8CheckpointLayout: + checkpoint_layout = _deepseek_v4_checkpoint_layout() + assert isinstance(checkpoint_layout.finegrained_fp8, FineGrainedFP8CheckpointLayout) + return checkpoint_layout.finegrained_fp8 + + +def _fp8_weight(shape: tuple[int, int]) -> torch.Tensor: + return torch.empty(shape, dtype=_FP8_DTYPE) + + +def _run_load_hook( + state_dict: dict[str, torch.Tensor], + weight_name: str, + *, + checkpoint_layout: object | None = None, +) -> None: + load_hook_kwargs = {} + if checkpoint_layout is not None: + load_hook_kwargs["checkpoint_layout"] = checkpoint_layout + _transform().load_hook( + state_dict, + "", + weight_name=weight_name, + **load_hook_kwargs, + ) + + +def _build_graph_module(weight_names: tuple[str, ...]) -> GraphModule: + model = _DeepSeekLinearSurface() + params = dict(model.named_parameters()) + graph = Graph() + x = graph.placeholder("x") + outputs = [] + + for weight_name in weight_names: + weight = graph.get_attr(weight_name) + weight.meta["val"] = params[weight_name] + outputs.append(graph.call_function(torch.ops.aten.linear.default, (x, weight, None))) + + graph.output(tuple(outputs)) + return GraphModule(model, graph) + + +def _build_grouped_linear_graph_module(weight_name: str) -> GraphModule: + batch_size, seq_len, num_groups, rank, group_width = 2, 3, 4, 8, 16 + model = _DeepSeekLinearSurface() + modname, _, attrname = weight_name.rpartition(".") + submod = model.get_submodule(modname) + setattr( + submod, + attrname, + nn.Parameter(torch.empty(num_groups * rank, group_width, dtype=torch.float16)), + ) + params = dict(model.named_parameters()) + + graph = Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty(batch_size, seq_len, num_groups, group_width, dtype=torch.float16) + input_view = graph.call_function( + torch.ops.auto_deploy.view.default, + args=(x, [batch_size, seq_len, num_groups, group_width]), + kwargs={"tp_scaled_dim": 2, "layer_type": "mla"}, + ) + weight = graph.get_attr(weight_name) + weight.meta["val"] = params[weight_name] + weight_view = graph.call_function( + torch.ops.auto_deploy.view.default, + args=(weight, [num_groups, rank, group_width]), + kwargs={ + "tp_scaled_dim": 0, + "layer_type": "mla", + "tp_min_local_shape": rank, + }, + ) + output = graph.call_function( + torch.ops.auto_deploy.torch_grouped_linear.default, + args=(input_view, weight_view, None), + kwargs={ + "tp_mode": "colwise", + "layer_type": "mla", + "tp_min_local_shape": rank, + }, + ) + graph.output(output) + return GraphModule(model, graph) + + +def _linear_weight_names(gm: GraphModule, op: object) -> set[str]: + return {node.args[1].target for node in gm.graph.nodes if is_op(node, op)} + + +def _linear_nodes(gm: GraphModule, op: object) -> list[Node]: + return [node for node in gm.graph.nodes if is_op(node, op)] + + +def _call_function_nodes(gm: GraphModule, op: object) -> list[Node]: + return [node for node in gm.graph.nodes if is_op(node, op)] + + +def test_layout_scale_alias_loads_weight_scale_inv_and_removes_alias() -> None: + weight_name = "layers.0.attn.wq_a.weight" + scale_name = "layers.0.attn.wq_a.scale" + scale = torch.arange(6, dtype=torch.float16).reshape(2, 3) + 1 + state_dict = { + weight_name: _fp8_weight((129, 257)), + scale_name: scale, + } + + _run_load_hook(state_dict, weight_name, checkpoint_layout=_finegrained_fp8_layout()) + + assert scale_name not in state_dict + loaded_scale = state_dict["layers.0.attn.wq_a.weight_scale_inv"] + assert loaded_scale.dtype == torch.float32 + torch.testing.assert_close(loaded_scale, scale.to(torch.float32)) + + +@pytest.mark.skipif(_E8M0_DTYPE is None, reason="Requires torch.float8_e8m0fnu") +def test_layout_e8m0_scale_alias_decodes_raw_exponent_bytes() -> None: + weight_name = "layers.0.attn.wq_a.weight" + scale_name = "layers.0.attn.wq_a.scale" + raw_exponents = torch.tensor([[0, 1, 255], [128, 126, 130]], dtype=torch.uint8) + try: + scale = raw_exponents.view(_E8M0_DTYPE) + except RuntimeError as error: + pytest.skip(f"torch.float8_e8m0fnu view is unavailable: {error}") + expected_bits = raw_exponents.to(torch.int32) << 23 + expected_bits[raw_exponents == 0] = 1 << 22 + expected_bits[raw_exponents == 255] = 0x7FC00000 + expected = expected_bits.view(torch.float32) + state_dict = { + weight_name: _fp8_weight((129, 257)), + scale_name: scale, + } + + _run_load_hook(state_dict, weight_name, checkpoint_layout=_finegrained_fp8_layout()) + + loaded_scale = state_dict["layers.0.attn.wq_a.weight_scale_inv"] + assert loaded_scale.dtype == torch.float32 + torch.testing.assert_close(loaded_scale, expected, equal_nan=True) + loaded_bits = loaded_scale.contiguous().view(torch.int32) + torch.testing.assert_close(loaded_bits, expected_bits) + assert loaded_bits[0, 0].item() == 1 << 22 + assert torch.isnan(loaded_scale[0, 2]) + assert loaded_bits[0, 2].item() == 0x7FC00000 + + +@pytest.mark.parametrize( + ("weight_shape", "scale_shape"), + ( + ((1, 1), (1, 1)), + ((128, 128), (1, 1)), + ((129, 257), (2, 3)), + ), +) +def test_layout_scale_alias_accepts_ceil_shape( + weight_shape: tuple[int, int], + scale_shape: tuple[int, int], +) -> None: + weight_name = "layers.0.attn.wq_a.weight" + scale_name = "layers.0.attn.wq_a.scale" + state_dict = { + weight_name: _fp8_weight(weight_shape), + scale_name: torch.ones(scale_shape, dtype=torch.float32), + } + + _run_load_hook(state_dict, weight_name, checkpoint_layout=_finegrained_fp8_layout()) + + assert state_dict["layers.0.attn.wq_a.weight_scale_inv"].shape == scale_shape + + +def test_layout_scale_alias_rejects_bad_ceil_shape() -> None: + weight_name = "layers.0.attn.wq_a.weight" + scale_name = "layers.0.attn.wq_a.scale" + state_dict = { + weight_name: _fp8_weight((129, 257)), + scale_name: torch.ones((1, 3), dtype=torch.float32), + } + + with pytest.raises(ValueError, match=r"expected \(2, 3\)"): + _run_load_hook(state_dict, weight_name, checkpoint_layout=_finegrained_fp8_layout()) + + assert scale_name in state_dict + assert "layers.0.attn.wq_a.weight_scale_inv" not in state_dict + + +def test_layout_config_quantizes_only_direct_finegrained_fp8_linear_paths() -> None: + gm = _build_graph_module(_ALLOWED_WEIGHT_NAMES + _SKIPPED_WEIGHT_NAMES) + modules_to_not_convert_weight = "layers.0.attn.wq_b.weight" + qcfg = { + "quant_method": "fp8", + "weight_block_size": [128, 128], + "scale_fmt": "ue8m0", + "checkpoint_layout": _deepseek_v4_checkpoint_layout(), + "exclude_modules": [ + "embed", + "head", + "*.ffn.gate", + "*.attn.compressor", + "*.attn.indexer.compressor", + "*.attn.indexer.weights_proj", + "*.norm", + ], + "modules_to_not_convert": ["*.attn.wq_b"], + } + quant_op = torch.ops.auto_deploy.torch_fake_quant_finegrained_fp8_linear.default + expected_quantized = set(_ALLOWED_WEIGHT_NAMES) - {modules_to_not_convert_weight} + expected_skipped = set(_SKIPPED_WEIGHT_NAMES) | {modules_to_not_convert_weight} + + gm, info = _transform()._apply(gm, None, _Factory(qcfg), SharedConfig()) + + assert info.num_matches == len(expected_quantized) + assert _linear_weight_names(gm, quant_op) == expected_quantized + assert _linear_weight_names(gm, torch.ops.aten.linear) == expected_skipped + for node in _linear_nodes(gm, quant_op): + assert node.kwargs["input_scale_fmt"] == "ue8m0" + + buffers = dict(gm.named_buffers()) + for weight_name in expected_quantized: + scale_name = weight_name.removesuffix(".weight") + ".weight_scale_inv" + assert buffers[scale_name].shape == (2, 3) + assert buffers[scale_name].dtype == torch.float32 + + for weight_name in expected_skipped: + scale_name = weight_name.removesuffix(".weight") + ".weight_scale_inv" + assert scale_name not in buffers + + +def test_layout_config_quantizes_targeted_grouped_linear_with_ue8m0_scale_fmt() -> None: + weight_name = "layers.0.attn.wo_a.weight" + gm = _build_grouped_linear_graph_module(weight_name) + qcfg = { + "quant_method": "fp8", + "weight_block_size": [128, 128], + "scale_fmt": "ue8m0", + "checkpoint_layout": _deepseek_v4_checkpoint_layout(), + } + grouped_op = torch.ops.auto_deploy.torch_fake_quant_grouped_finegrained_fp8_linear.default + + gm, info = _transform()._apply(gm, None, _Factory(qcfg), SharedConfig()) + + grouped_nodes = _call_function_nodes(gm, grouped_op) + assert info.num_matches == 1 + assert len(grouped_nodes) == 1 + grouped_node = grouped_nodes[0] + assert grouped_node.args[1].target == weight_name + assert grouped_node.kwargs["input_scale_fmt"] == "ue8m0" + assert grouped_node.kwargs["tp_mode"] == "colwise" + assert grouped_node.kwargs["tp_min_local_shape"] == 8 + assert ShardableNode.from_node(grouped_node) is not None + assert not _call_function_nodes(gm, torch.ops.auto_deploy.torch_grouped_linear) + + buffers = dict(gm.named_buffers()) + scale_name = weight_name.removesuffix(".weight") + ".weight_scale_inv" + assert buffers[scale_name].shape == (1, 1) + assert buffers[scale_name].dtype == torch.float32 + + +def test_layout_config_preserves_exported_positional_view_layer_type_on_grouped_linear() -> None: + class _ExportedGroupedLinear(nn.Module): + def __init__(self) -> None: + super().__init__() + self.layers = nn.ModuleList([nn.Module()]) + self.layers[0].attn = nn.Module() + self.layers[0].attn.wo_a = _WeightOnlyLinear(out_features=32, in_features=16) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = torch.ops.auto_deploy.view(x, [2, 3, 4, 16], 2, "mla") + weight = torch.ops.auto_deploy.view( + self.layers[0].attn.wo_a.weight, + [4, 8, 16], + 0, + "mla", + 8, + ) + return torch.ops.auto_deploy.torch_grouped_linear( + x, weight, None, "colwise", None, 8, "mla" + ) + + gm = torch.export.export(_ExportedGroupedLinear(), (torch.randn(2, 3, 4, 16),)).module() + qcfg = { + "quant_method": "fp8", + "weight_block_size": [128, 128], + "scale_fmt": "ue8m0", + "checkpoint_layout": _deepseek_v4_checkpoint_layout(), + } + grouped_op = torch.ops.auto_deploy.torch_fake_quant_grouped_finegrained_fp8_linear.default + + gm, info = _transform()._apply(gm, None, _Factory(qcfg), SharedConfig()) + + grouped_nodes = _call_function_nodes(gm, grouped_op) + assert info.num_matches == 1 + assert len(grouped_nodes) == 1 + assert grouped_nodes[0].kwargs["layer_type"] == "mla" + + +def test_layout_config_leaves_non_targeted_grouped_linear_untouched() -> None: + gm = _build_grouped_linear_graph_module("layers.0.attn.compressor.wkv.weight") + qcfg = { + "quant_method": "fp8", + "weight_block_size": [128, 128], + "scale_fmt": "ue8m0", + "checkpoint_layout": _deepseek_v4_checkpoint_layout(), + } + grouped_op = torch.ops.auto_deploy.torch_fake_quant_grouped_finegrained_fp8_linear.default + + gm, info = _transform()._apply(gm, None, _Factory(qcfg), SharedConfig()) + + assert info.num_matches == 0 + assert not _call_function_nodes(gm, grouped_op) + assert len(_call_function_nodes(gm, torch.ops.auto_deploy.torch_grouped_linear)) == 1 + assert "layers.0.attn.compressor.wkv.weight_scale_inv" not in dict(gm.named_buffers()) + + +def test_layout_config_skips_linear_without_extractable_weight() -> None: + graph = Graph() + x = graph.placeholder("x") + weight = graph.placeholder("weight") + linear = graph.call_function(torch.ops.aten.linear.default, (x, weight, None)) + graph.output(linear) + gm = GraphModule(nn.Module(), graph) + qcfg = { + "quant_method": "fp8", + "weight_block_size": [128, 128], + "checkpoint_layout": _deepseek_v4_checkpoint_layout(), + } + quant_op = torch.ops.auto_deploy.torch_fake_quant_finegrained_fp8_linear.default + + gm, info = _transform()._apply(gm, None, _Factory(qcfg), SharedConfig()) + + assert info.num_matches == 0 + assert _linear_weight_names(gm, quant_op) == set() + assert any(is_op(node, torch.ops.aten.linear) for node in gm.graph.nodes) + + +def test_generic_config_quantized_linear_uses_default_input_scale_fmt() -> None: + gm = _build_graph_module(("layers.0.attn.wq_a.weight",)) + qcfg = { + "quant_method": "fp8", + "weight_block_size": [128, 128], + } + quant_op = torch.ops.auto_deploy.torch_fake_quant_finegrained_fp8_linear.default + + gm, info = _transform()._apply(gm, None, _Factory(qcfg), SharedConfig()) + + assert info.num_matches == 1 + quant_nodes = _linear_nodes(gm, quant_op) + assert len(quant_nodes) == 1 + assert quant_nodes[0].kwargs["input_scale_fmt"] == "" + + +def test_generic_fp8_uses_scale_inv_and_does_not_consume_scale_alias() -> None: + weight_name = "layers.0.attn.wq_a.weight" + alias_name = "layers.0.attn.wq_a.scale" + scale_inv_name = weight_name + "_scale_inv" + scale_inv = torch.full((2, 3), 3.0, dtype=torch.float16) + alias = torch.full((2, 3), 7.0, dtype=torch.float16) + state_dict = { + weight_name: _fp8_weight((129, 257)), + scale_inv_name: scale_inv, + alias_name: alias, + } + + _run_load_hook(state_dict, weight_name) + + assert alias_name in state_dict + assert state_dict[alias_name] is alias + assert state_dict["layers.0.attn.wq_a.weight_scale_inv"] is scale_inv diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_deepseek_v4_modeling.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_deepseek_v4_modeling.py new file mode 100644 index 000000000000..c5f8268f88e3 --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_deepseek_v4_modeling.py @@ -0,0 +1,1533 @@ +# 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. + +"""Hierarchical DeepSeek V4 Flash semantic tests for AutoDeploy.""" + +from __future__ import annotations + +import inspect + +import pytest +import torch +import torch.nn.functional as F +from torch import nn +from torch.export import Dim + +import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 +import tensorrt_llm._torch.auto_deploy.transform.library # noqa: F401 +from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm +from tensorrt_llm._torch.auto_deploy.models import custom as custom_models +from tensorrt_llm._torch.auto_deploy.models.custom import modeling_deepseek_v4 as dsv4 +from tensorrt_llm._torch.auto_deploy.models.custom.modeling_deepseek_v4 import ( + DeepseekV4Attention, + DeepseekV4Compressor, + DeepseekV4Config, + DeepseekV4ForCausalLM, + DeepseekV4Indexer, + DeepseekV4MLP, + DeepseekV4MoE, + DeepseekV4RMSNorm, +) +from tensorrt_llm._torch.auto_deploy.models.factory import ModelFactoryRegistry +from tensorrt_llm._torch.auto_deploy.models.hf import AutoModelForCausalLMFactory +from tensorrt_llm._torch.auto_deploy.transform.interface import SharedConfig +from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer +from tensorrt_llm._torch.auto_deploy.utils.dist_config import DistConfig +from tensorrt_llm._torch.auto_deploy.utils.node_utils import extract_op_args, is_op + +DeepseekV4Router = ( + getattr(dsv4, "DeepseekV4Router", None) + or getattr(dsv4, "DeepseekV4MoEGate", None) + or getattr(dsv4, "DeepseekV4Gate", None) +) + + +@pytest.fixture(autouse=True) +def _set_seed() -> None: + torch.manual_seed(1234) + + +def _small_config(**overrides) -> DeepseekV4Config: + values = { + "vocab_size": 32, + "hidden_size": 32, + "num_hidden_layers": 3, + "num_attention_heads": 4, + "num_key_value_heads": 1, + "head_dim": 4, + "q_lora_rank": 8, + "qk_rope_head_dim": 2, + "o_groups": 2, + "o_lora_rank": 8, + "sliding_window": 3, + "compress_ratios": (0, 4, 128), + "compress_rope_theta": 16000.0, + "index_n_heads": 2, + "index_head_dim": 4, + "index_topk": 2, + "moe_intermediate_size": 32, + "n_routed_experts": 4, + "n_shared_experts": 1, + "num_experts_per_tok": 2, + "num_hash_layers": 1, + "scoring_func": "sqrtsoftplus", + "routed_scaling_factor": 1.25, + "norm_topk_prob": True, + "swiglu_limit": 0.5, + "hidden_act": "silu", + "max_position_embeddings": 256, + "rope_theta": 10000.0, + "rope_scaling": { + "type": "yarn", + "rope_type": "yarn", + "factor": 1.0, + "original_max_position_embeddings": 256, + "beta_fast": 32.0, + "beta_slow": 1.0, + }, + "rms_norm_eps": 1e-6, + "hc_mult": 2, + "hc_sinkhorn_iters": 3, + "hc_eps": 1e-6, + "ad_rope_cache_len": 256, + "ad_compress_max_seq_len": 256, + "attention_bias": False, + "tie_word_embeddings": False, + } + values.update(overrides) + return DeepseekV4Config(**values) + + +def _make_apply_sharding_hints_optimizer( + world_size: int, + rank: int, + dist_backend: str | None = None, +) -> InferenceOptimizer: + apply_config = {"stage": "sharding"} + if dist_backend is not None: + apply_config["dist_backend"] = dist_backend + + opt = InferenceOptimizer( + factory=None, + config={"apply_sharding_hints": apply_config}, + ) + opt.shared_config = SharedConfig( + local_rank=rank, + world_size=world_size, + dist_config=DistConfig( + world_size=world_size, + rank=rank, + tp_size=world_size, + moe_ep_size=world_size, + ), + ) + return opt + + +def _call_nodes(gm: torch.fx.GraphModule, op) -> list[torch.fx.Node]: + return [node for node in gm.graph.nodes if is_op(node, op)] + + +def _position_ids(batch: int, seq: int, device: torch.device | str) -> torch.Tensor: + return torch.arange(seq, device=device).unsqueeze(0).expand(batch, -1) + + +def _output_logits(output) -> torch.Tensor: + if isinstance(output, torch.Tensor): + return output + if isinstance(output, tuple): + return output[0] + if isinstance(output, dict): + return output["logits"] + return output.logits + + +def assert_rmse_close( + actual: torch.Tensor, + expected: torch.Tensor, + rmse_ratio_tol: float, + msg: str = "", +) -> None: + diff = actual.float() - expected.float() + rmse_diff = torch.sqrt(torch.mean(diff**2)) + rmse_ref = torch.sqrt(torch.mean(expected.float() ** 2)) + ratio = (rmse_diff / rmse_ref).item() + assert ratio < rmse_ratio_tol, ( + f"{msg}RMSE ratio {ratio:.6f} exceeds tolerance {rmse_ratio_tol}. " + f"(rmse_diff={rmse_diff.item():.6f}, rmse_ref={rmse_ref.item():.6f})" + ) + + +def _get_linear(module: nn.Module, *names: str) -> nn.Module: + for name in names: + if hasattr(module, name): + return getattr(module, name) + raise AssertionError(f"{type(module).__name__} is missing one of {names}") + + +def _linear_ref(x: torch.Tensor, linear: nn.Module) -> torch.Tensor: + bias = getattr(linear, "bias", None) + return F.linear(x.float(), linear.weight.float(), None if bias is None else bias.float()).to( + x.dtype + ) + + +def _copy_deterministic_linear(linear: nn.Module, offset: float) -> None: + values = torch.linspace(-0.7 + offset, 0.7 + offset, linear.weight.numel()) + with torch.no_grad(): + linear.weight.copy_(values.view_as(linear.weight)) + bias = getattr(linear, "bias", None) + if bias is not None: + bias.zero_() + + +def _sequential_tensor(shape: torch.Size | tuple[int, ...], offset: float = 0.0) -> torch.Tensor: + numel = 1 + for dim in shape: + numel *= int(dim) + return (torch.arange(numel, dtype=torch.float32) + offset).view(shape) + + +def _set_mlp_weights(module: nn.Module, offset: float = 0.0) -> None: + _copy_deterministic_linear(_get_linear(module, "gate_proj", "w1"), offset) + _copy_deterministic_linear(_get_linear(module, "up_proj", "w3"), offset + 0.1) + _copy_deterministic_linear(_get_linear(module, "down_proj", "w2"), offset + 0.2) + + +def _ref_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + y = x.float() * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + eps) + return (y * weight.float()).to(x.dtype) + + +def _ref_ceil_pow2_scale(amax: torch.Tensor, max_value: float, min_value: float) -> torch.Tensor: + return torch.pow(2.0, torch.ceil(torch.log2(amax.clamp_min(min_value) / max_value))) + + +def _ref_fake_fp8_act_quant(x: torch.Tensor, block_size: int = 64) -> torch.Tensor: + dim = x.shape[-1] + if dim == 0 or dim % block_size != 0: + return x + + dtype = x.dtype + x_float = x.float() + grouped = x_float.reshape(*x_float.shape[:-1], dim // block_size, block_size) + scale = _ref_ceil_pow2_scale(grouped.abs().amax(dim=-1, keepdim=True), 448.0, 1.0e-4) + quant = torch.clamp(grouped / scale, -448.0, 448.0).to(dtype).float() + return (quant * scale).reshape_as(x_float).to(dtype) + + +def _ref_fake_fp4_act_quant(x: torch.Tensor, block_size: int = 32) -> torch.Tensor: + dim = x.shape[-1] + if dim == 0 or dim % block_size != 0: + return x + + dtype = x.dtype + x_float = x.float() + grouped = x_float.reshape(*x_float.shape[:-1], dim // block_size, block_size) + scale = _ref_ceil_pow2_scale(grouped.abs().amax(dim=-1, keepdim=True), 6.0, 6.0 * 2.0**-126) + normalized = torch.clamp(grouped / scale, -6.0, 6.0) + levels = normalized.new_tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) + level_idx = (normalized.abs().unsqueeze(-1) - levels).abs().argmin(dim=-1) + quant = levels[level_idx] * normalized.sign() + return (quant * scale).reshape_as(x_float).to(dtype) + + +def _ref_hadamard_rotate(x: torch.Tensor) -> torch.Tensor: + dim = x.shape[-1] + if dim <= 1: + return x + if dim & (dim - 1): + raise ValueError(f"Hadamard rotation requires power-of-two dimension, got {dim}.") + + original_shape = x.shape + out = x.float() + width = 1 + while width < dim: + out = out.reshape(*out.shape[:-1], dim // (2 * width), 2, width) + left = out[..., 0, :] + right = out[..., 1, :] + out = torch.cat((left + right, left - right), dim=-1).reshape(original_shape) + width *= 2 + return (out * (dim**-0.5)).to(x.dtype) + + +def _ref_overlap_transform(tensor: torch.Tensor, head_dim: int, value: float) -> torch.Tensor: + batch_size, compressed_len, ratio, _ = tensor.shape + previous = tensor[:, :, :, :head_dim] + current = tensor[:, :, :, head_dim:] + prefix = tensor.new_full((batch_size, 1, ratio, head_dim), value) + previous = torch.cat((prefix, previous[:, :-1]), dim=1) + return torch.cat((previous, current), dim=2) + + +def _ref_mlp(module: nn.Module, x: torch.Tensor, swiglu_limit: float | None = None) -> torch.Tensor: + gate_proj = _get_linear(module, "gate_proj", "w1") + up_proj = _get_linear(module, "up_proj", "w3") + down_proj = _get_linear(module, "down_proj", "w2") + gate = _linear_ref(x, gate_proj).float() + up = _linear_ref(x, up_proj).float() + limit = getattr(module, "swiglu_limit", 0.0) if swiglu_limit is None else swiglu_limit + if limit > 0: + gate = torch.clamp(gate, max=limit) + up = torch.clamp(up, min=-limit, max=limit) + hidden = F.silu(gate) * up + return _linear_ref(hidden.to(x.dtype), down_proj) + + +def _set_router_weights(router: nn.Module) -> None: + values = torch.linspace(-0.4, 0.5, router.weight.numel()) + with torch.no_grad(): + router.weight.copy_(values.view_as(router.weight)) + bias = getattr(router, "bias", None) + if bias is not None: + bias.copy_(torch.tensor([0.0, 0.4, -0.3, 0.8], dtype=bias.dtype)) + tid2eid = getattr(router, "tid2eid", None) + if tid2eid is not None: + pattern = torch.tensor([[0, 2], [3, 1], [1, 0], [2, 3]], dtype=tid2eid.dtype) + repeats = (tid2eid.shape[0] + pattern.shape[0] - 1) // pattern.shape[0] + tid2eid.copy_(pattern.repeat(repeats, 1)[: tid2eid.shape[0]]) + + +def _ref_router( + router: nn.Module, + hidden_states: torch.Tensor, + input_ids: torch.Tensor, + config: DeepseekV4Config, + hash_routing: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + scores = F.softplus(F.linear(hidden_states, router.weight.float())).sqrt() + if hash_routing: + tid2eid = getattr(router, "tid2eid") + selected = tid2eid[input_ids].to(torch.long) + selected = selected.clamp(0, config.n_routed_experts - 1) + else: + bias = getattr(router, "bias", None) + assert bias is not None + selected = (scores + bias.float()).topk(config.num_experts_per_tok, dim=-1).indices + weights = scores.gather(1, selected) + if getattr(config, "norm_topk_prob", True): + weights = weights / (weights.sum(dim=-1, keepdim=True) + 1e-20) + return selected, weights * config.routed_scaling_factor + + +def _interleaved_rope( + x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, inverse: bool = False +) -> torch.Tensor: + if inverse: + sin = -sin + x_pair = x.float().reshape(*x.shape[:-1], -1, 2) + cos = cos.float().unsqueeze(-1) + sin = sin.float().unsqueeze(-1) + first = x_pair[..., 0:1] * cos - x_pair[..., 1:2] * sin + second = x_pair[..., 1:2] * cos + x_pair[..., 0:1] * sin + return torch.cat([first, second], dim=-1).flatten(-2).to(x.dtype) + + +def _rope_cos_sin( + position_ids: torch.Tensor, rope_dim: int, base: float +) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq = 1.0 / ( + base ** (torch.arange(0, rope_dim, 2, device=position_ids.device).float() / rope_dim) + ) + phase = position_ids.float().unsqueeze(-1) * inv_freq + return phase.cos(), phase.sin() + + +def _position_embeddings( + config: DeepseekV4Config, position_ids: torch.Tensor +) -> tuple[torch.Tensor, ...]: + rotary_cls = getattr(dsv4, "DeepseekV4RotaryEmbedding", None) + if rotary_cls is not None: + try: + return rotary_cls(config)() + except TypeError: + return rotary_cls(config)(position_ids) + + table_positions = torch.arange(config.ad_rope_cache_len, device=position_ids.device) + cos, sin = _rope_cos_sin(table_positions, config.qk_rope_head_dim, config.rope_theta) + cos_comp, sin_comp = _rope_cos_sin( + table_positions, config.qk_rope_head_dim, config.compress_rope_theta + ) + return cos, sin, cos_comp, sin_comp + + +def _ref_compressor( + compressor: DeepseekV4Compressor, + hidden_states: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + position_ids: torch.Tensor, +) -> torch.Tensor: + batch_size, seq_len, _ = hidden_states.shape + ratio = compressor.compress_ratio + max_compressed_len = compressor.max_compressed_len + row_offsets = torch.arange(max_compressed_len, device=hidden_states.device) + token_offsets = torch.arange(ratio, device=hidden_states.device) + gather_idxs = row_offsets.unsqueeze(1) * ratio + token_offsets + valid = gather_idxs < seq_len + gather_idxs = torch.where(valid, gather_idxs, torch.zeros_like(gather_idxs)) + flat_idxs = gather_idxs.reshape(-1) + + kv_all = _linear_ref(hidden_states, compressor.wkv).float() + score_all = _linear_ref(hidden_states, compressor.wgate).float() + kv = kv_all[:, flat_idxs].view(batch_size, max_compressed_len, ratio, -1) + score = score_all[:, flat_idxs].view(batch_size, max_compressed_len, ratio, -1) + score = score + compressor.ape.float() + score = torch.where( + valid.view(1, max_compressed_len, ratio, 1), + score, + score.new_full((), -1.0e20), + ) + if compressor.overlap: + kv = _ref_overlap_transform(kv, compressor.head_dim, 0.0) + score = _ref_overlap_transform(score, compressor.head_dim, -1.0e20) + + compressed = (kv * score.softmax(dim=2)).sum(dim=2).to(hidden_states.dtype) + compressed = _ref_rmsnorm(compressed, compressor.norm.weight, compressor.norm.eps) + + row_start = row_offsets * ratio + row_start = torch.minimum(row_start, torch.full_like(row_start, seq_len - 1)) + compressed_position_ids = position_ids[:, row_start] + cos = cos_table[compressed_position_ids] + sin = sin_table[compressed_position_ids] + nope_dim = compressor.head_dim - compressor.rope_head_dim + nope, pe = torch.split(compressed, [nope_dim, compressor.rope_head_dim], dim=-1) + pe = _interleaved_rope(pe, cos, sin) + compressed = torch.cat((nope, pe), dim=-1) + + if compressor.rotate: + return _ref_fake_fp4_act_quant(_ref_hadamard_rotate(compressed), block_size=32) + + nope, pe = torch.split(compressed, [nope_dim, compressor.rope_head_dim], dim=-1) + nope = _ref_fake_fp8_act_quant(nope, block_size=64) + return torch.cat((nope, pe), dim=-1) + + +def _ref_indexer( + indexer: DeepseekV4Indexer, + hidden_states: torch.Tensor, + q_lora: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + position_ids: torch.Tensor, + offset: int, +) -> torch.Tensor: + batch_size, seq_len, _ = hidden_states.shape + q = _linear_ref(q_lora, indexer.wq_b) + q = q.view(batch_size, seq_len, indexer.index_n_heads, indexer.index_head_dim) + q_nope, q_pe = torch.split( + q, + [indexer.index_head_dim - indexer.rope_head_dim, indexer.rope_head_dim], + dim=-1, + ) + q_pe = _interleaved_rope(q_pe, cos.unsqueeze(2), sin.unsqueeze(2)) + q = _ref_fake_fp4_act_quant(_ref_hadamard_rotate(torch.cat((q_nope, q_pe), dim=-1))) + + index_k = _ref_compressor(indexer.compressor, hidden_states, cos_table, sin_table, position_ids) + weights = _linear_ref(hidden_states, indexer.weights_proj).float() + weights = weights * (indexer.softmax_scale * indexer.index_n_heads**-0.5) + + index_score = torch.matmul( + q.transpose(1, 2), + index_k.transpose(1, 2).unsqueeze(1), + ).transpose(1, 2) + index_score = index_score.float() + index_score = (index_score.relu() * weights.unsqueeze(-1)).sum(dim=2) + + compressed_positions = torch.arange( + indexer.compressor.max_compressed_len, + device=hidden_states.device, + ) + valid_lengths = torch.arange(1, seq_len + 1, device=hidden_states.device).unsqueeze(1) + valid_lengths = valid_lengths // indexer.compress_ratio + index_score = index_score.masked_fill( + (compressed_positions.unsqueeze(0) >= valid_lengths).unsqueeze(0), + -1.0e20, + ) + + if indexer.index_topk == 0: + return torch.empty(batch_size, seq_len, 0, device=hidden_states.device, dtype=torch.int64) + + topk_count = min(indexer.index_topk, indexer.compressor.max_compressed_len) + topk_idxs = index_score.topk(topk_count, dim=-1).indices + invalid = topk_idxs >= valid_lengths.unsqueeze(0) + topk_idxs = torch.where(invalid, -1, topk_idxs + offset) + if topk_count < indexer.index_topk: + pad = torch.full( + (batch_size, seq_len, indexer.index_topk - topk_count), + -1, + device=hidden_states.device, + dtype=topk_idxs.dtype, + ) + topk_idxs = torch.cat((topk_idxs, pad), dim=-1) + return topk_idxs.to(torch.int64) + + +def _ref_sparse_attention( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + batch_size, seq_len, num_heads, _ = q.shape + batch_idx = torch.arange(batch_size, device=q.device).view(batch_size, 1, 1) + batch_idx = batch_idx.expand(batch_size, seq_len, topk_idxs.shape[-1]) + + compute_dtype = torch.float32 if q.dtype in (torch.float16, torch.bfloat16) else q.dtype + gather_idxs = topk_idxs.to(torch.long).clamp(min=0) + selected_kv = kv[batch_idx, gather_idxs].to(compute_dtype) + logits = torch.matmul(q.to(compute_dtype), selected_kv.transpose(-1, -2)) + logits = logits * softmax_scale + logits = logits.masked_fill((topk_idxs < 0).unsqueeze(2), float("-inf")) + + sink_logits = attn_sink.to(dtype=compute_dtype).view(1, 1, num_heads, 1) + sink_logits = sink_logits.expand(batch_size, seq_len, num_heads, 1) + weights = torch.softmax(torch.cat([logits, sink_logits], dim=-1), dim=-1) + output = torch.matmul(weights[..., :-1], selected_kv) + return output.to(q.dtype) + + +def _ref_window_topk_idxs( + window_size: int, + batch_size: int, + seq_len: int, + device: torch.device, +) -> torch.Tensor: + query_positions = torch.arange(seq_len, device=device).unsqueeze(1) + key_positions = query_positions - window_size + 1 + torch.arange(window_size, device=device) + key_positions = torch.where( + (key_positions < 0) | (key_positions > query_positions), + -1, + key_positions, + ) + return key_positions.unsqueeze(0).expand(batch_size, -1, -1) + + +def _ref_compress_topk_idxs( + ratio: int, + batch_size: int, + seq_len: int, + offset: int, + max_compressed_len: int, + device: torch.device, +) -> torch.Tensor: + compressed_positions = torch.arange(max_compressed_len, device=device) + valid_lengths = torch.arange(1, seq_len + 1, device=device).unsqueeze(1) // ratio + topk_idxs = compressed_positions.unsqueeze(0).expand(seq_len, -1) + topk_idxs = torch.where(topk_idxs < valid_lengths, topk_idxs + offset, -1) + return topk_idxs.unsqueeze(0).expand(batch_size, -1, -1) + + +# `transformers.models.deepseek_v4` is unavailable in the installed transformers package. +# The HFDeepseekV4* classes below are test-only, HF-style minimal faithful copies of the +# local DeepSeek V4 `inference/model.py` semantics for equivalence checks. +class HFDeepseekV4RMSNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=torch.float32)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return _ref_rmsnorm(x, self.weight, self.eps) + + +class HFDeepseekV4Compressor(nn.Module): + def __init__( + self, + config: DeepseekV4Config, + compress_ratio: int, + head_dim: int, + rotate: bool = False, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.head_dim = head_dim + self.rope_head_dim = config.qk_rope_head_dim + self.compress_ratio = compress_ratio + self.rotate = rotate + self.overlap = compress_ratio == 4 + channels = 2 if self.overlap else 1 + self.max_compressed_len = max( + 1, + (int(config.ad_compress_max_seq_len) + compress_ratio - 1) // compress_ratio, + ) + self.max_compressed_tokens = self.max_compressed_len * compress_ratio + self.ape = nn.Parameter( + torch.zeros(compress_ratio, channels * head_dim, dtype=torch.float32) + ) + self.wkv = nn.Linear(config.hidden_size, channels * head_dim, bias=False) + self.wgate = nn.Linear(config.hidden_size, channels * head_dim, bias=False) + self.norm = HFDeepseekV4RMSNorm(head_dim, config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + return _ref_compressor(self, hidden_states, cos_table, sin_table, position_ids) + + +class HFDeepseekV4Indexer(nn.Module): + def __init__(self, config: DeepseekV4Config, compress_ratio: int) -> None: + super().__init__() + self.index_n_heads = config.index_n_heads + self.index_head_dim = config.index_head_dim + self.rope_head_dim = config.qk_rope_head_dim + self.index_topk = config.index_topk + self.compress_ratio = compress_ratio + self.softmax_scale = self.index_head_dim**-0.5 + self.wq_b = nn.Linear( + config.q_lora_rank, config.index_n_heads * config.index_head_dim, bias=False + ) + self.weights_proj = nn.Linear(config.hidden_size, config.index_n_heads, bias=False) + self.compressor = HFDeepseekV4Compressor( + config, + compress_ratio, + config.index_head_dim, + rotate=True, + ) + + def forward( + self, + hidden_states: torch.Tensor, + q_lora: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + position_ids: torch.Tensor, + offset: int, + ) -> torch.Tensor: + return _ref_indexer( + self, + hidden_states, + q_lora, + cos, + sin, + cos_table, + sin_table, + position_ids, + offset, + ) + + +class HFDeepseekV4Attention(nn.Module): + def __init__(self, config: DeepseekV4Config, layer_idx: int) -> None: + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.n_heads = config.num_attention_heads + self.head_dim = config.head_dim + self.rope_head_dim = config.qk_rope_head_dim + self.nope_head_dim = self.head_dim - self.rope_head_dim + self.q_lora_rank = config.q_lora_rank + self.o_lora_rank = config.o_lora_rank + self.n_groups = config.o_groups + self.window_size = config.sliding_window + self.compress_ratio = int(config.compress_ratios[layer_idx]) + self.rms_eps = config.rms_norm_eps + self.softmax_scale = self.head_dim**-0.5 + self.group_head_width = (self.n_heads * self.head_dim) // self.n_groups + + self.wq_a = nn.Linear(self.hidden_size, self.q_lora_rank, bias=False) + self.q_norm = HFDeepseekV4RMSNorm(self.q_lora_rank, self.rms_eps) + self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False) + self.wkv = nn.Linear(self.hidden_size, self.head_dim, bias=False) + self.kv_norm = HFDeepseekV4RMSNorm(self.head_dim, self.rms_eps) + self.wo_a = nn.Linear(self.group_head_width, self.n_groups * self.o_lora_rank, bias=False) + self.wo_b = nn.Linear(self.n_groups * self.o_lora_rank, self.hidden_size, bias=False) + self.attn_sink = nn.Parameter(torch.zeros(self.n_heads, dtype=torch.float32)) + if self.compress_ratio: + self.compressor = HFDeepseekV4Compressor(config, self.compress_ratio, self.head_dim) + self.indexer = ( + HFDeepseekV4Indexer(config, self.compress_ratio) + if self.compress_ratio == 4 + else None + ) + else: + self.compressor = None + self.indexer = None + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + position_ids: torch.Tensor, + ) -> torch.Tensor: + batch_size, seq_len, _ = hidden_states.shape + cos_base_table, sin_base_table, cos_compress_table, sin_compress_table = position_embeddings + cos_base = cos_base_table[position_ids] + sin_base = sin_base_table[position_ids] + cos_compress = cos_compress_table[position_ids] + sin_compress = sin_compress_table[position_ids] + cos = cos_compress if self.compress_ratio else cos_base + sin = sin_compress if self.compress_ratio else sin_base + + q_lora = self.q_norm(_linear_ref(hidden_states, self.wq_a)) + q = _linear_ref(q_lora, self.wq_b).view(batch_size, seq_len, self.n_heads, self.head_dim) + q = q * torch.rsqrt(q.float().square().mean(-1, keepdim=True) + self.rms_eps).to(q.dtype) + + kv = self.kv_norm(_linear_ref(hidden_states, self.wkv)) + q_nope, q_pe = torch.split(q, [self.nope_head_dim, self.rope_head_dim], dim=-1) + kv_nope, kv_pe = torch.split(kv, [self.nope_head_dim, self.rope_head_dim], dim=-1) + q_pe = _interleaved_rope(q_pe, cos.unsqueeze(2), sin.unsqueeze(2)) + kv_pe = _interleaved_rope(kv_pe, cos, sin) + kv_nope = _ref_fake_fp8_act_quant(kv_nope, block_size=64) + q = torch.cat((q_nope, q_pe), dim=-1) + kv = torch.cat((kv_nope, kv_pe), dim=-1) + + if self.compress_ratio: + window_idxs = _ref_window_topk_idxs( + self.window_size, batch_size, seq_len, hidden_states.device + ) + compressed_kv = self.compressor( + hidden_states, cos_compress_table, sin_compress_table, position_ids + ) + if self.indexer is not None: + compressed_idxs = self.indexer( + hidden_states, + q_lora, + cos_compress, + sin_compress, + cos_compress_table, + sin_compress_table, + position_ids, + seq_len, + ) + else: + compressed_idxs = _ref_compress_topk_idxs( + self.compress_ratio, + batch_size, + seq_len, + seq_len, + self.compressor.max_compressed_len, + hidden_states.device, + ) + kv_for_attention = torch.cat((kv, compressed_kv), dim=1) + topk_idxs = torch.cat((window_idxs, compressed_idxs), dim=-1).to(torch.int64) + attn_output = _ref_sparse_attention( + q, kv_for_attention, self.attn_sink, topk_idxs, self.softmax_scale + ) + else: + kv_heads = kv.unsqueeze(2).expand(batch_size, seq_len, self.n_heads, self.head_dim) + q_bh = q.transpose(1, 2).float() + kv_bh = kv_heads.transpose(1, 2).float() + scores = torch.matmul(q_bh, kv_bh.transpose(-1, -2)) + scores = scores * self.softmax_scale + query_pos = position_ids[:, :, None] + key_pos = position_ids[:, None, :] + causal = key_pos <= query_pos + window = (query_pos - key_pos) < self.window_size + scores = scores.masked_fill(~(causal.unsqueeze(1) & window.unsqueeze(1)), float("-inf")) + sink_logits = self.attn_sink.float().view(1, self.n_heads, 1, 1) + sink_logits = sink_logits.expand(batch_size, self.n_heads, seq_len, 1) + weights = torch.softmax(torch.cat([scores, sink_logits], dim=-1), dim=-1)[..., :-1] + attn_output = torch.matmul(weights, kv_bh).transpose(1, 2).to(hidden_states.dtype) + + out_nope, out_pe = torch.split( + attn_output, + [self.nope_head_dim, self.rope_head_dim], + dim=-1, + ) + out_pe = _interleaved_rope(out_pe, cos.unsqueeze(2), sin.unsqueeze(2), inverse=True) + attn_output = torch.cat((out_nope, out_pe), dim=-1) + attn_output = attn_output.view(batch_size, seq_len, self.n_groups, -1) + wo_a = self.wo_a.weight.float().view(self.n_groups, self.o_lora_rank, -1) + attn_output = torch.matmul( + attn_output.float().unsqueeze(-2), + wo_a.transpose(-1, -2), + ).squeeze(-2) + attn_output = attn_output.flatten(2) + return _linear_ref(attn_output.to(hidden_states.dtype), self.wo_b) + + +def _ref_dense_attention( + attn: nn.Module, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + config: DeepseekV4Config, +) -> torch.Tensor: + bsz, seq_len, _ = hidden_states.shape + num_heads = config.num_attention_heads + head_dim = config.head_dim + rope_dim = config.qk_rope_head_dim + nope_dim = head_dim - rope_dim + + q_lora = _ref_rmsnorm( + _linear_ref(hidden_states, attn.wq_a), + _get_linear(attn, "q_norm").weight, + config.rms_norm_eps, + ) + q = _linear_ref(q_lora, attn.wq_b).view(bsz, seq_len, num_heads, head_dim) + q = q * torch.rsqrt(q.float().square().mean(-1, keepdim=True) + config.rms_norm_eps).to(q.dtype) + + kv = _ref_rmsnorm( + _linear_ref(hidden_states, attn.wkv), + _get_linear(attn, "kv_norm").weight, + config.rms_norm_eps, + ).view(bsz, seq_len, 1, head_dim) + + cos, sin = _rope_cos_sin(position_ids, rope_dim, config.rope_theta) + cos = cos.unsqueeze(2) + sin = sin.unsqueeze(2) + q_nope, q_pe = torch.split(q, [nope_dim, rope_dim], dim=-1) + kv_nope, kv_pe = torch.split(kv, [nope_dim, rope_dim], dim=-1) + q = torch.cat([q_nope, _interleaved_rope(q_pe, cos, sin)], dim=-1) + kv = torch.cat([kv_nope, _interleaved_rope(kv_pe, cos, sin)], dim=-1) + + q_bh = q.transpose(1, 2).float() + kv_bh = kv.expand(bsz, seq_len, num_heads, head_dim).transpose(1, 2).float() + scores = torch.matmul(q_bh, kv_bh.transpose(-2, -1)) * (head_dim**-0.5) + query_pos = position_ids[:, :, None] + key_pos = position_ids[:, None, :] + causal = key_pos <= query_pos + window = (query_pos - key_pos) < config.sliding_window + scores = scores.masked_fill(~(causal.unsqueeze(1) & window.unsqueeze(1)), float("-inf")) + + sink_logits = attn.attn_sink.float().view(1, num_heads, 1, 1) + sink_logits = sink_logits.expand(bsz, num_heads, seq_len, 1) + weights = torch.softmax(torch.cat([scores, sink_logits], dim=-1), dim=-1)[..., :-1] + out = torch.matmul(weights, kv_bh).transpose(1, 2).to(hidden_states.dtype) + + out_nope, out_pe = torch.split(out, [nope_dim, rope_dim], dim=-1) + out = torch.cat([out_nope, _interleaved_rope(out_pe, cos, sin, inverse=True)], dim=-1) + + group_count = config.o_groups + group_width = (num_heads * head_dim) // group_count + out = out.reshape(bsz, seq_len, group_count, group_width) + wo_a_weight = attn.wo_a.weight.float().view(group_count, config.o_lora_rank, group_width) + out = torch.matmul(out.float().unsqueeze(-2), wo_a_weight.transpose(-1, -2)).squeeze(-2) + return _linear_ref(out.reshape(bsz, seq_len, group_count * config.o_lora_rank), attn.wo_b) + + +def _call_attention( + attn: nn.Module, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + config: DeepseekV4Config, +) -> torch.Tensor: + signature = inspect.signature(attn.forward) + if "position_embeddings" in signature.parameters: + return attn(hidden_states, _position_embeddings(config, position_ids), position_ids) + if "position_ids" in signature.parameters: + return attn(hidden_states, position_ids) + + cos, sin = _rope_cos_sin(position_ids, config.qk_rope_head_dim, config.rope_theta) + cos_comp, sin_comp = _rope_cos_sin( + position_ids, config.qk_rope_head_dim, config.compress_rope_theta + ) + return attn(hidden_states, cos, sin, cos_comp, sin_comp, cos_comp[0], sin_comp[0]) + + +def _decoder(model: DeepseekV4ForCausalLM) -> nn.Module: + return getattr(model, "model", model) + + +def test_small_config_exercises_flash_semantics() -> None: + config = _small_config() + model = DeepseekV4ForCausalLM(config).eval() + layers = _decoder(model).layers + + assert config.compress_ratios[0] == 0 + assert set(config.compress_ratios[1:]) == {4, 128} + assert config.num_hash_layers == 1 + assert config.n_shared_experts == 1 + assert config.hc_mult == 2 + assert config.swiglu_limit > 0 + assert all(hasattr(layer.attn, "attn_sink") for layer in layers) + assert hasattr(layers[0].ffn.gate, "tid2eid") + assert getattr(layers[1].ffn.gate, "bias", None) is not None + assert hasattr(layers[0].ffn.experts, "gate_up_proj_blocks") + assert layers[0].ffn.experts.gate_up_proj_blocks.dtype == torch.uint8 + + +def test_load_hook_preserves_converted_keys_and_drops_mtp_layer() -> None: + config = _small_config(num_hidden_layers=1, compress_ratios=(0,), num_hash_layers=0) + model = DeepseekV4ForCausalLM(config).eval() + embed = _sequential_tensor(model.embed.weight.shape) + + result = model.load_state_dict( + { + "embed.weight": embed, + "mtp.0.e_proj.weight": torch.ones(2, 2), + }, + strict=False, + ) + + assert result.unexpected_keys == [] + torch.testing.assert_close(model.embed.weight, embed, rtol=0, atol=0) + + +def test_load_hook_maps_wrapper_root_keys_and_shared_expert_aliases() -> None: + config = _small_config(num_hidden_layers=1, compress_ratios=(0,), num_hash_layers=0) + model = DeepseekV4ForCausalLM(config).eval() + embed = _sequential_tensor(model.embed.weight.shape) + head = _sequential_tensor(model.head.weight.shape, offset=10.0) + norm = _sequential_tensor(model.norm.weight.shape, offset=20.0) + shared_w1 = _sequential_tensor(model.layers[0].ffn.shared_experts.w1.weight.shape, offset=30.0) + shared_w3 = _sequential_tensor(model.layers[0].ffn.shared_experts.w3.weight.shape, offset=40.0) + shared_w2 = _sequential_tensor(model.layers[0].ffn.shared_experts.w2.weight.shape, offset=50.0) + + result = model.load_state_dict( + { + "model.embed_tokens.weight": embed, + "lm_head.weight": head, + "model.norm.weight": norm, + "model.layers.0.ffn.shared_experts.gate_proj.weight": shared_w1, + "model.layers.0.ffn.shared_experts.up_proj.weight": shared_w3, + "model.layers.0.ffn.shared_experts.down_proj.weight": shared_w2, + }, + strict=False, + ) + + assert result.unexpected_keys == [] + torch.testing.assert_close(model.embed.weight, embed, rtol=0, atol=0) + torch.testing.assert_close(model.head.weight, head, rtol=0, atol=0) + torch.testing.assert_close(model.norm.weight, norm, rtol=0, atol=0) + torch.testing.assert_close( + model.layers[0].ffn.shared_experts.w1.weight, shared_w1, rtol=0, atol=0 + ) + torch.testing.assert_close( + model.layers[0].ffn.shared_experts.w3.weight, shared_w3, rtol=0, atol=0 + ) + torch.testing.assert_close( + model.layers[0].ffn.shared_experts.w2.weight, shared_w2, rtol=0, atol=0 + ) + + +def test_rotary_embedding_returns_full_cached_tables() -> None: + config = _small_config(ad_rope_cache_len=16, max_position_embeddings=16) + rotary = dsv4.DeepseekV4RotaryEmbedding(config) + position_ids = _position_ids(2, 5, "cpu") + + cos_base, sin_base, cos_compress, sin_compress = rotary() + + assert cos_base.shape == (config.ad_rope_cache_len, config.qk_rope_head_dim // 2) + assert sin_base.shape == cos_base.shape + assert cos_compress.shape == cos_base.shape + assert sin_compress.shape == cos_base.shape + assert cos_base[position_ids].shape == (2, 5, config.qk_rope_head_dim // 2) + + +def test_rmsnorm_matches_reference() -> None: + norm = DeepseekV4RMSNorm(16, eps=1e-6) + with torch.no_grad(): + norm.weight.copy_(torch.linspace(0.5, 1.5, 16)) + x = torch.randn(2, 3, 16) + + actual = norm(x) + expected = _ref_rmsnorm(x, norm.weight, 1e-6) + + torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-6) + + +def test_mlp_swiglu_limit_matches_reference() -> None: + mlp = DeepseekV4MLP(16, 12, swiglu_limit=0.5).eval() + _set_mlp_weights(mlp) + x = torch.linspace(-2.0, 2.0, 2 * 3 * 16).view(2, 3, 16) + + actual = mlp(x) + expected = _ref_mlp(mlp, x) + + torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-6) + + +def test_deepseek_v4_chat_prompt_format_matches_reference() -> None: + assert ( + dsv4._format_deepseek_v4_chat_messages( + [{"role": "user", "content": "How big is the universe? "}] + ) + == "<｜begin▁of▁sentence｜><｜User｜>How big is the universe? <｜Assistant｜>" + ) + assert dsv4._format_deepseek_v4_chat_messages( + [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "u"}, + {"role": "assistant", "content": "a"}, + {"role": "user", "content": "u2"}, + ] + ) == ( + "<｜begin▁of▁sentence｜>sys<｜User｜>u<｜Assistant｜>" + "a<｜end▁of▁sentence｜><｜User｜>u2<｜Assistant｜>" + ) + + +def test_moe_shared_experts_use_configured_swiglu_limit() -> None: + config = _small_config(num_hidden_layers=1, compress_ratios=(0,)) + assert config.swiglu_limit > 0 + moe = DeepseekV4MoE(config, layer_idx=0).eval() + _set_mlp_weights(moe.shared_experts, offset=0.4) + x = torch.linspace(-2.0, 2.0, 2 * 3 * config.hidden_size).view(2, 3, config.hidden_size) + + assert moe.shared_experts.swiglu_limit == config.swiglu_limit + actual = moe.shared_experts(x) + expected = _ref_mlp(moe.shared_experts, x, swiglu_limit=config.swiglu_limit) + + torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-6) + + +def test_router_hash_and_score_routing_match_reference() -> None: + assert DeepseekV4Router is not None, "DeepSeek V4 router class must be exposed for tests" + config = _small_config() + hidden_states = torch.linspace(-1.0, 1.0, 5 * config.hidden_size).view(5, config.hidden_size) + input_ids = torch.tensor([0, 1, 2, 3, 4]) + + hash_router = DeepseekV4Router(config, layer_idx=0).eval() + _set_router_weights(hash_router) + expected_hash = _ref_router(hash_router, hidden_states, input_ids, config, hash_routing=True) + actual_hash = hash_router(hidden_states, input_ids) + torch.testing.assert_close(actual_hash[0], expected_hash[0]) + torch.testing.assert_close(actual_hash[1], expected_hash[1], rtol=1e-6, atol=1e-6) + + score_router = DeepseekV4Router(config, layer_idx=config.num_hash_layers).eval() + _set_router_weights(score_router) + expected_score = _ref_router(score_router, hidden_states, input_ids, config, hash_routing=False) + actual_score = score_router(hidden_states, input_ids) + torch.testing.assert_close(actual_score[0], expected_score[0]) + torch.testing.assert_close(actual_score[1], expected_score[1], rtol=1e-6, atol=1e-6) + + +@pytest.mark.parametrize("layer_idx", [0, 1], ids=["hash-routing", "score-routing"]) +def test_moe_uses_packed_mxfp4_routed_experts(layer_idx: int) -> None: + config = _small_config(num_hidden_layers=2, compress_ratios=(0, 0)) + moe = DeepseekV4MoE(config, layer_idx=layer_idx).eval() + _set_router_weights(moe.gate) + _set_mlp_weights(moe.shared_experts, offset=0.4) + + hidden_states = torch.linspace(-1.5, 1.5, 2 * 3 * config.hidden_size).view( + 2, 3, config.hidden_size + ) + input_ids = torch.tensor([[0, 1, 2], [3, 4, 5]]) + + actual = moe(hidden_states, input_ids) + + assert actual.shape == hidden_states.shape + assert torch.isfinite(actual).all() + assert moe.experts.gate_up_proj_blocks.shape == ( + config.n_routed_experts, + 2 * config.moe_intermediate_size, + config.hidden_size // 32, + 16, + ) + assert moe.experts.down_proj_blocks.shape == ( + config.n_routed_experts, + config.hidden_size, + config.moe_intermediate_size // 32, + 16, + ) + + +def test_dense_attention_with_sinks_matches_reference() -> None: + config = _small_config(num_hidden_layers=1, compress_ratios=(0,), num_hash_layers=0) + attn = DeepseekV4Attention(config, layer_idx=0).eval() + with torch.no_grad(): + attn.attn_sink.copy_(torch.tensor([-0.5, 0.0, 0.25, 0.75])) + hidden_states = torch.randn(2, 5, config.hidden_size) + position_ids = _position_ids(2, 5, hidden_states.device) + + actual = _call_attention(attn, hidden_states, position_ids, config) + expected = _ref_dense_attention(attn, hidden_states, position_ids, config) + + assert_rmse_close(actual, expected, rmse_ratio_tol=0.10, msg="Dense attention: ") + + +def test_torch_attention_with_sinks_accepts_bfloat16_inputs() -> None: + query = torch.randn(2, 5, 4, 8, dtype=torch.bfloat16) + key = torch.randn(2, 5, 4, 8, dtype=torch.bfloat16) + value = torch.randn(2, 5, 4, 8, dtype=torch.bfloat16) + sinks = torch.tensor([-0.5, 0.0, 0.25, 0.75], dtype=torch.bfloat16) + + actual = torch.ops.auto_deploy.torch_attention( + query, + key, + value, + dropout_p=0.0, + is_causal=True, + scale=0.5, + sinks=sinks, + sliding_window=4, + layout="bsnd", + ) + expected = torch.ops.auto_deploy.torch_attention( + query.float(), + key.float(), + value.float(), + dropout_p=0.0, + is_causal=True, + scale=0.5, + sinks=sinks.float(), + sliding_window=4, + layout="bsnd", + ) + + assert actual.dtype == torch.bfloat16 + torch.testing.assert_close(actual.float(), expected, rtol=0.02, atol=0.02) + + +def test_compressor_matches_independent_reference_with_fake_fp8_nope_path() -> None: + config = _small_config( + hidden_size=8, + num_attention_heads=1, + head_dim=128, + qk_rope_head_dim=64, + o_groups=1, + compress_ratios=(2,), + ad_rope_cache_len=16, + ad_compress_max_seq_len=6, + max_position_embeddings=16, + ) + compressor = DeepseekV4Compressor(config, compress_ratio=2, head_dim=config.head_dim).eval() + compressor = compressor.to(torch.bfloat16) + _copy_deterministic_linear(compressor.wkv, offset=0.0) + _copy_deterministic_linear(compressor.wgate, offset=0.17) + with torch.no_grad(): + compressor.ape.copy_( + torch.linspace(-0.2, 0.3, compressor.ape.numel()).view_as(compressor.ape) + ) + compressor.norm.weight.copy_(torch.linspace(0.6, 1.4, config.head_dim)) + + hidden_states = torch.linspace(-1.5, 1.75, 2 * 5 * config.hidden_size) + hidden_states = hidden_states.view(2, 5, config.hidden_size).to(torch.bfloat16) + position_ids = _position_ids(2, 5, hidden_states.device) + table_positions = torch.arange(config.ad_rope_cache_len, device=hidden_states.device) + cos_table, sin_table = _rope_cos_sin( + table_positions, + config.qk_rope_head_dim, + config.compress_rope_theta, + ) + + assert compressor.head_dim - compressor.rope_head_dim == 64 + actual = compressor(hidden_states, cos_table, sin_table, position_ids) + expected = _ref_compressor(compressor, hidden_states, cos_table, sin_table, position_ids) + + torch.testing.assert_close(actual.float(), expected.float(), rtol=0.01, atol=0.01) + + +def test_ratio4_indexer_matches_independent_reference_and_masks_invalid_prefix() -> None: + config = _small_config( + hidden_size=8, + q_lora_rank=8, + num_attention_heads=1, + head_dim=64, + qk_rope_head_dim=32, + index_n_heads=1, + index_head_dim=64, + index_topk=3, + compress_ratios=(4,), + ad_rope_cache_len=16, + ad_compress_max_seq_len=12, + max_position_embeddings=16, + ) + indexer = DeepseekV4Indexer(config, compress_ratio=4).eval() + _copy_deterministic_linear(indexer.wq_b, offset=0.11) + _copy_deterministic_linear(indexer.weights_proj, offset=-0.09) + _copy_deterministic_linear(indexer.compressor.wkv, offset=0.03) + _copy_deterministic_linear(indexer.compressor.wgate, offset=0.21) + with torch.no_grad(): + indexer.compressor.ape.copy_( + torch.linspace(-0.25, 0.35, indexer.compressor.ape.numel()).view_as( + indexer.compressor.ape + ) + ) + indexer.compressor.norm.weight.copy_(torch.linspace(0.7, 1.3, config.index_head_dim)) + + hidden_states = torch.linspace(-1.4, 1.6, 2 * 5 * config.hidden_size).view( + 2, 5, config.hidden_size + ) + q_lora = torch.linspace(-0.9, 1.1, 2 * 5 * config.q_lora_rank).view(2, 5, config.q_lora_rank) + position_ids = _position_ids(2, 5, hidden_states.device) + cos, sin = _rope_cos_sin(position_ids, config.qk_rope_head_dim, config.compress_rope_theta) + table_positions = torch.arange(config.ad_rope_cache_len, device=hidden_states.device) + cos_table, sin_table = _rope_cos_sin( + table_positions, + config.qk_rope_head_dim, + config.compress_rope_theta, + ) + offset = 10 + + assert config.index_head_dim % 32 == 0 + assert config.index_head_dim & (config.index_head_dim - 1) == 0 + actual = indexer( + hidden_states, + q_lora, + cos, + sin, + cos_table, + sin_table, + position_ids, + offset, + ) + expected = _ref_indexer( + indexer, + hidden_states, + q_lora, + cos, + sin, + cos_table, + sin_table, + position_ids, + offset, + ) + + torch.testing.assert_close(actual, expected) + assert torch.equal(actual[:, :3], torch.full_like(actual[:, :3], -1)) + assert (actual[:, 3:] >= 0).sum(dim=-1).eq(1).all() + + +def test_export_dynamic_shapes_finite_logits_and_expected_ops() -> None: + config = _small_config(num_hidden_layers=2, compress_ratios=(0, 4), num_hash_layers=1) + model = DeepseekV4ForCausalLM(config).eval() + input_ids = torch.randint(0, config.vocab_size, (2, 6)) + position_ids = _position_ids(2, 6, input_ids.device) + + gm = torch_export_to_gm( + model, + args=(input_ids,), + kwargs={"position_ids": position_ids}, + dynamic_shapes={ + "input_ids": {0: Dim.DYNAMIC, 1: Dim.DYNAMIC}, + "position_ids": {0: Dim.DYNAMIC, 1: Dim.DYNAMIC}, + }, + num_moe_experts_for_export=2, + ) + + with torch.no_grad(): + eager = model(input_ids=input_ids, position_ids=position_ids).logits + exported = _output_logits(gm(input_ids, position_ids=position_ids)) + assert torch.isfinite(exported).all() + assert_rmse_close(exported, eager, rmse_ratio_tol=0.05, msg="Export first shape: ") + + input_ids_2 = torch.randint(0, config.vocab_size, (1, 4)) + position_ids_2 = _position_ids(1, 4, input_ids_2.device) + with torch.no_grad(): + out_2 = _output_logits(gm(input_ids_2, position_ids=position_ids_2)) + assert out_2.shape == (1, 4, config.vocab_size) + assert torch.isfinite(out_2).all() + + target_names = [str(node.target) for node in gm.graph.nodes if node.op == "call_function"] + assert target_names.count("auto_deploy.torch_deepseek_v4_sparse_attention.default") == 2 + assert "auto_deploy.torch_linear_simple.default" in target_names + assert "auto_deploy.torch_mxfp4_moe_from_routing.default" in target_names + assert "auto_deploy.torch_moe.default" not in target_names + assert "auto_deploy.torch_attention.default" not in target_names + non_torch_ad_ops = sorted( + name + for name in set(target_names) + if name.startswith("auto_deploy.") and not name.startswith("auto_deploy.torch_") + ) + assert non_torch_ad_ops == [ + "auto_deploy.all_reduce.default", + "auto_deploy.view.default", + ] + + +def test_export_mxfp4_experts_apply_sharding_hints_rank1_ep_graph() -> None: + config = _small_config( + vocab_size=64, + hidden_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=1, + head_dim=8, + q_lora_rank=8, + qk_rope_head_dim=4, + o_groups=2, + o_lora_rank=8, + compress_ratios=(0,), + moe_intermediate_size=64, + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=2, + num_hash_layers=0, + ad_rope_cache_len=16, + ad_compress_max_seq_len=16, + max_position_embeddings=16, + ) + model = DeepseekV4ForCausalLM(config).eval() + input_ids = torch.randint(0, config.vocab_size, (2, 4)) + position_ids = _position_ids(2, 4, input_ids.device) + + gm = torch_export_to_gm( + model, + args=(input_ids,), + kwargs={"position_ids": position_ids}, + strict=False, + ) + + base_op = torch.ops.auto_deploy.torch_mxfp4_moe_from_routing.default + ep_op = torch.ops.auto_deploy.torch_mxfp4_moe_from_routing_ep.default + assert len(_call_nodes(gm, base_op)) == 1 + + gm_out = _make_apply_sharding_hints_optimizer( + world_size=2, + rank=1, + dist_backend="torch", + )(None, gm) + + assert len(_call_nodes(gm_out, base_op)) == 0 + ep_nodes = _call_nodes(gm_out, ep_op) + assert len(ep_nodes) == 1 + ep_node = ep_nodes[0] + + [expert_start, gate_up_blocks] = extract_op_args( + ep_node, + "expert_start", + "gate_up_blocks", + ) + assert expert_start == 2 + assert getattr(gate_up_blocks, "op", None) == "get_attr" + assert gm_out.get_buffer(gate_up_blocks.target).shape[0] == 2 + + dist_all_reduce_op = torch.ops.auto_deploy.torch_dist_all_reduce.default + all_reduce_nodes = _call_nodes(gm_out, dist_all_reduce_op) + assert len(all_reduce_nodes) == 3 + assert any(node.args[0] is ep_node for node in all_reduce_nodes) + assert len(_call_nodes(gm_out, torch.ops.auto_deploy.all_reduce.default)) == 0 + + layer = gm_out.get_submodule("layers.0") + assert layer.attn.attn_sink.shape == (config.num_attention_heads // 2,) + assert layer.attn.wq_b.weight.shape == ( + config.num_attention_heads * config.head_dim // 2, + config.q_lora_rank, + ) + assert layer.attn.wo_a.weight.shape == ( + config.o_lora_rank, + config.num_attention_heads * config.head_dim // config.o_groups, + ) + assert layer.attn.wo_b.weight.shape == ( + config.hidden_size, + config.o_groups * config.o_lora_rank // 2, + ) + + sparse_nodes = _call_nodes( + gm_out, + torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention.default, + ) + assert len(sparse_nodes) == 1 + attn_sink = sparse_nodes[0].args[2] + assert getattr(attn_sink, "op", None) == "get_attr" + assert gm_out.get_parameter(attn_sink.target).shape == (config.num_attention_heads // 2,) + + group_width = config.num_attention_heads * config.head_dim // config.o_groups + view_shapes = [ + extract_op_args(node, "shape")[0] + for node in _call_nodes(gm_out, torch.ops.auto_deploy.view.default) + ] + assert [2, 4, -1, config.head_dim] in view_shapes + assert [2, 4, -1, group_width] in view_shapes + assert [-1, config.o_lora_rank, group_width] in view_shapes + + +def test_export_mxfp4_experts_apply_sharding_hints_rank6_ep8_tp8_graph() -> None: + config = _small_config( + vocab_size=64, + hidden_size=64, + num_hidden_layers=1, + num_attention_heads=8, + num_key_value_heads=1, + head_dim=8, + q_lora_rank=8, + qk_rope_head_dim=4, + o_groups=8, + o_lora_rank=8, + compress_ratios=(0,), + moe_intermediate_size=64, + n_routed_experts=256, + n_shared_experts=1, + num_experts_per_tok=2, + num_hash_layers=0, + ad_rope_cache_len=16, + ad_compress_max_seq_len=16, + max_position_embeddings=16, + ) + model = DeepseekV4ForCausalLM(config).eval() + input_ids = torch.randint(0, config.vocab_size, (2, 4)) + position_ids = _position_ids(2, 4, input_ids.device) + + gm = torch_export_to_gm( + model, + args=(input_ids,), + kwargs={"position_ids": position_ids}, + strict=False, + ) + + base_op = torch.ops.auto_deploy.torch_mxfp4_moe_from_routing.default + ep_op = torch.ops.auto_deploy.torch_mxfp4_moe_from_routing_ep.default + assert len(_call_nodes(gm, base_op)) == 1 + + gm_out = _make_apply_sharding_hints_optimizer( + world_size=8, + rank=6, + dist_backend="torch", + )(None, gm) + + assert len(_call_nodes(gm_out, base_op)) == 0 + ep_nodes = _call_nodes(gm_out, ep_op) + assert len(ep_nodes) == 1 + ep_node = ep_nodes[0] + + [expert_start, gate_up_blocks, down_blocks] = extract_op_args( + ep_node, + "expert_start", + "gate_up_blocks", + "down_blocks", + ) + local_expert_count = config.n_routed_experts // 8 + assert expert_start == 192 + assert local_expert_count == 32 + assert getattr(gate_up_blocks, "op", None) == "get_attr" + assert getattr(down_blocks, "op", None) == "get_attr" + assert gm_out.get_buffer(gate_up_blocks.target).shape[0] == local_expert_count + assert gm_out.get_buffer(down_blocks.target).shape[0] == local_expert_count + + dist_all_reduce_op = torch.ops.auto_deploy.torch_dist_all_reduce.default + all_reduce_nodes = _call_nodes(gm_out, dist_all_reduce_op) + assert any(node.args[0] is ep_node for node in all_reduce_nodes) + assert len(_call_nodes(gm_out, torch.ops.auto_deploy.all_reduce.default)) == 0 + + layer = gm_out.get_submodule("layers.0") + local_head_count = config.num_attention_heads // 8 + assert local_head_count == 1 + assert layer.attn.attn_sink.shape == (local_head_count,) + assert layer.attn.wq_b.weight.shape == ( + local_head_count * config.head_dim, + config.q_lora_rank, + ) + assert layer.attn.wo_a.weight.shape == ( + config.o_lora_rank, + config.num_attention_heads * config.head_dim // config.o_groups, + ) + assert layer.attn.wo_b.weight.shape == ( + config.hidden_size, + config.o_groups * config.o_lora_rank // 8, + ) + + sparse_nodes = _call_nodes( + gm_out, + torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention.default, + ) + assert len(sparse_nodes) == 1 + attn_sink = sparse_nodes[0].args[2] + assert getattr(attn_sink, "op", None) == "get_attr" + assert gm_out.get_parameter(attn_sink.target).shape == (local_head_count,) + + +def test_ratio4_sparse_attention_sharding_only_splits_sink() -> None: + config = _small_config( + vocab_size=64, + hidden_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=1, + head_dim=8, + q_lora_rank=8, + qk_rope_head_dim=4, + o_groups=2, + o_lora_rank=8, + compress_ratios=(4,), + index_n_heads=2, + index_head_dim=8, + index_topk=2, + moe_intermediate_size=64, + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=2, + num_hash_layers=0, + ad_rope_cache_len=16, + ad_compress_max_seq_len=16, + max_position_embeddings=16, + ) + model = DeepseekV4ForCausalLM(config).eval() + input_ids = torch.randint(0, config.vocab_size, (2, 4)) + position_ids = _position_ids(2, 4, input_ids.device) + + gm = torch_export_to_gm( + model, + args=(input_ids,), + kwargs={"position_ids": position_ids}, + strict=False, + ) + gm_out = _make_apply_sharding_hints_optimizer( + world_size=2, + rank=1, + dist_backend="torch", + )(None, gm) + + layer = gm_out.get_submodule("layers.0") + assert layer.attn.attn_sink.shape == (config.num_attention_heads // 2,) + assert layer.attn.compressor.norm.weight.shape == (config.head_dim,) + assert layer.attn.indexer.compressor.norm.weight.shape == (config.index_head_dim,) + + sparse_nodes = _call_nodes( + gm_out, + torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention.default, + ) + assert len(sparse_nodes) == 1 + attn_sink = sparse_nodes[0].args[2] + compressor_norm = sparse_nodes[0].args[7] + indexer_norm = sparse_nodes[0].args[16] + assert getattr(attn_sink, "op", None) == "get_attr" + assert getattr(compressor_norm, "op", None) == "get_attr" + assert getattr(indexer_norm, "op", None) == "get_attr" + assert gm_out.get_parameter(attn_sink.target).shape == (config.num_attention_heads // 2,) + assert gm_out.get_parameter(compressor_norm.target).shape == (config.head_dim,) + assert gm_out.get_parameter(indexer_norm.target).shape == (config.index_head_dim,) + + +def test_factory_registration() -> None: + assert custom_models.DeepseekV4ForCausalLM is DeepseekV4ForCausalLM + assert "DeepseekV4ForCausalLM" in custom_models.__all__ + assert ( + AutoModelForCausalLMFactory._custom_model_mapping["DeepseekV4Config"] + is DeepseekV4ForCausalLM + ) + factory_cls = getattr(dsv4, "DeepseekV4AutoModelForCausalLMFactory", None) + assert factory_cls is not None, "DeepseekV4AutoModelForCausalLMFactory must be exposed" + assert factory_cls._custom_model_mapping["DeepseekV4Config"] is DeepseekV4ForCausalLM + assert ModelFactoryRegistry.get("DeepseekV4AutoModelForCausalLM") is factory_cls + + +def test_factory_example_inputs_cover_forward_contract(tmp_path) -> None: + config = _small_config( + num_hidden_layers=2, + compress_ratios=(0, 4), + ad_rope_cache_len=16, + ad_compress_max_seq_len=12, + max_position_embeddings=32, + ) + config.save_pretrained(tmp_path) + factory_cls = ModelFactoryRegistry.get("DeepseekV4AutoModelForCausalLM") + factory = factory_cls(model=str(tmp_path), max_seq_len=32) + + example_inputs = factory.get_example_inputs() + + assert set(example_inputs) == {"input_ids", "position_ids"} + input_ids = example_inputs["input_ids"] + position_ids = example_inputs["position_ids"] + assert input_ids.dtype == torch.long + assert position_ids.dtype == torch.long + assert input_ids.shape == (2, 8) + assert position_ids.shape == input_ids.shape + assert torch.equal(position_ids, _position_ids(2, 8, input_ids.device)) + + model = DeepseekV4ForCausalLM(config).eval() + with torch.no_grad(): + logits = model(**example_inputs).logits + assert logits.shape == (2, 8, config.vocab_size) diff --git a/tests/unittest/auto_deploy/conftest.py b/tests/unittest/auto_deploy/conftest.py new file mode 100644 index 000000000000..463ac0e513e0 --- /dev/null +++ b/tests/unittest/auto_deploy/conftest.py @@ -0,0 +1,45 @@ +# 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. + +"""Shared pytest configuration for AutoDeploy unit tests. + +AutoDeploy tests share lightweight helper modules under ``_utils_test``. Keeping +that directory on ``sys.path`` here lets tests import helpers by module name, +which mirrors the standalone llmc package's generated test conftest. + +This avoids repeating path manipulation in individual test files and keeps the +same import style valid in both places: + +* source tree: ``tests/unittest/auto_deploy/_utils_test`` +* standalone package: ``tests/_utils_test`` +""" + +import sys +from pathlib import Path + + +def _add_utils_test_dir_to_path() -> None: + """Make AutoDeploy test helpers importable by module name. + + The standalone package generator creates an equivalent conftest for copied + tests. Keeping this source-tree behavior aligned prevents tests from + depending on TensorRT-LLM-specific package names for test-only helpers. + """ + utils_test_dir = str(Path(__file__).resolve().parent / "_utils_test") + if utils_test_dir not in sys.path: + sys.path.insert(0, utils_test_dir) + + +_add_utils_test_dir_to_path() diff --git a/tests/unittest/auto_deploy/multigpu/transformations/library/test_apply_sharding_hints.py b/tests/unittest/auto_deploy/multigpu/transformations/library/test_apply_sharding_hints.py index c15b4bbea53d..cfa98977ebd5 100644 --- a/tests/unittest/auto_deploy/multigpu/transformations/library/test_apply_sharding_hints.py +++ b/tests/unittest/auto_deploy/multigpu/transformations/library/test_apply_sharding_hints.py @@ -27,20 +27,36 @@ import torch.nn as nn from _dist_test_utils import get_device_counts from _graph_test_helpers import run_test_transformed_gm +from torch.fx import Graph, GraphModule import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 import tensorrt_llm._torch.auto_deploy.distributed.common as dist_common import tensorrt_llm._torch.auto_deploy.transform.library # noqa: F401 from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm -from tensorrt_llm._torch.auto_deploy.transform.interface import SharedConfig +from tensorrt_llm._torch.auto_deploy.models.quant_checkpoint_layout import QuantizedCheckpointLayout +from tensorrt_llm._torch.auto_deploy.transform.interface import SharedConfig, Stages +from tensorrt_llm._torch.auto_deploy.transform.library.fused_moe_mxfp4 import ( + QuantizeMXFP4MOE as InsertMXFP4MLP, +) +from tensorrt_llm._torch.auto_deploy.transform.library.fused_moe_mxfp4 import ( + QuantizeMXFP4MOEConfig as MXFP4MLPConfig, +) +from tensorrt_llm._torch.auto_deploy.transform.library.sharding import _get_dist_ops from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer from tensorrt_llm._torch.auto_deploy.utils.dist_config import DistConfig -from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op +from tensorrt_llm._torch.auto_deploy.utils.node_utils import extract_op_args, is_op, shape pytestmark = pytest.mark.threadleak(enabled=False) FEATURES, HIDDEN = 32, 64 +DSV4_HIDDEN = 8 +DSV4_NUM_HEADS = 4 +DSV4_HEAD_DIM = 2 +DSV4_GROUPS = 2 +DSV4_GROUP_WIDTH = DSV4_NUM_HEADS * DSV4_HEAD_DIM // DSV4_GROUPS +DSV4_O_RANK = 3 + class HintedMLP(nn.Module): def __init__(self, features=FEATURES, hidden=HIDDEN): @@ -56,6 +72,110 @@ def forward(self, x): return h +class DeepSeekV4IRContractBlock(nn.Module): + """Small DeepSeek V4-shaped graph for sharding contract checks.""" + + def __init__(self): + super().__init__() + self.q_proj = nn.Linear(DSV4_HIDDEN, DSV4_NUM_HEADS * DSV4_HEAD_DIM, bias=False) + self.kv_proj = nn.Linear(DSV4_HIDDEN, DSV4_HEAD_DIM, bias=False) + self.wo_a = nn.Linear(DSV4_GROUP_WIDTH, DSV4_GROUPS * DSV4_O_RANK, bias=False) + self.wo_b = nn.Linear(DSV4_GROUPS * DSV4_O_RANK, DSV4_HIDDEN, bias=False) + self.attn_sink = nn.Parameter(torch.zeros(DSV4_NUM_HEADS, dtype=torch.float32)) + + def forward(self, x): + bsz, seq_len, _ = x.shape + q = torch.ops.auto_deploy.torch_linear_simple( + x, + self.q_proj.weight, + None, + tp_mode="colwise", + tp_min_local_shape=DSV4_HEAD_DIM, + layer_type="mla", + ) + q_left, q_right = torch.ops.auto_deploy.split_with_sizes( + q, + [DSV4_GROUP_WIDTH, DSV4_GROUP_WIDTH], + dim=-1, + enable_sharding=True, + layer_type="mla", + ) + q = torch.cat((q_left, q_right), dim=-1) + q = torch.ops.auto_deploy.view( + q, + [bsz, seq_len, DSV4_NUM_HEADS, DSV4_HEAD_DIM], + tp_scaled_dim=2, + layer_type="mla", + ) + + kv = torch.ops.auto_deploy.torch_linear_simple(x, self.kv_proj.weight, None) + topk_idxs = torch.zeros((bsz, seq_len, 1), dtype=torch.int64, device=x.device) + compressor_kv = q.new_empty(bsz, seq_len, 0) + compressor_gate = q.new_empty(bsz, seq_len, 0) + compressor_ape = q.new_empty(0, 0) + compressor_norm_weight = q.new_empty(0) + cos_table = q.new_empty(0, 0) + sin_table = q.new_empty(0, 0) + position_ids = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(bsz, -1) + indexer_q = q.new_empty(bsz, seq_len, 0, 0) + indexer_weights = q.new_empty(bsz, seq_len, 0) + indexer_compressor_kv = q.new_empty(bsz, seq_len, 0) + indexer_compressor_gate = q.new_empty(bsz, seq_len, 0) + indexer_compressor_ape = q.new_empty(0, 0) + indexer_compressor_norm_weight = q.new_empty(0) + attn_output = torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention( + q, + kv, + self.attn_sink, + topk_idxs, + compressor_kv, + compressor_gate, + compressor_ape, + compressor_norm_weight, + cos_table, + sin_table, + position_ids, + indexer_q, + indexer_weights, + indexer_compressor_kv, + indexer_compressor_gate, + indexer_compressor_ape, + indexer_compressor_norm_weight, + 1.0, + enable_sharding=True, + layer_type="mla", + ) + attn_output = torch.ops.auto_deploy.view( + attn_output, + [bsz, seq_len, DSV4_GROUPS, DSV4_GROUP_WIDTH], + tp_scaled_dim=2, + layer_type="mla", + ) + wo_a = torch.ops.auto_deploy.view( + self.wo_a.weight, + [DSV4_GROUPS, DSV4_O_RANK, DSV4_GROUP_WIDTH], + tp_scaled_dim=0, + layer_type="mla", + tp_min_local_shape=DSV4_O_RANK, + ) + attn_output = torch.ops.auto_deploy.torch_grouped_linear( + attn_output, + wo_a, + None, + tp_mode="colwise", + layer_type="mla", + tp_min_local_shape=DSV4_O_RANK, + ) + attn_output = torch.ops.auto_deploy.torch_linear_simple( + attn_output, + self.wo_b.weight, + None, + tp_mode="rowwise", + layer_type="mla", + ) + return torch.ops.auto_deploy.all_reduce(attn_output, layer_type="mla") + + # --------------------------------------------------------------------------- # test_sharding — multi-GPU end-to-end (follows test_tp_sharding.py pattern) # --------------------------------------------------------------------------- @@ -97,10 +217,22 @@ def test_sharding(world_size: int): # --------------------------------------------------------------------------- -def _make_optimizer(world_size: int, rank: int = 0): +def _make_optimizer( + world_size: int, + rank: int = 0, + dist_backend: str | None = None, + *, + simple_shard_only: bool = False, +): + apply_config = {"stage": "sharding"} + if dist_backend is not None: + apply_config["dist_backend"] = dist_backend + if simple_shard_only: + apply_config["simple_shard_only"] = True + opt = InferenceOptimizer( factory=None, - config={"apply_sharding_hints": {"stage": "sharding"}}, + config={"apply_sharding_hints": apply_config}, ) opt.shared_config = SharedConfig( local_rank=rank, @@ -141,3 +273,916 @@ def test_apply_hints(world_size, expect_skipped, expect_up_shape, expect_down_sh is_op(n, torch.ops.auto_deploy.torch_dist_all_reduce.default) for n in gm_out.graph.nodes ) assert has_dist_ar == (not expect_skipped) + + +def _export_deepseek_v4_contract_block(): + model = DeepSeekV4IRContractBlock() + x = torch.randn(2, 3, DSV4_HIDDEN) + return torch_export_to_gm(model, args=(x,), clone=True) + + +def _call_nodes(gm, op): + return [node for node in gm.graph.nodes if is_op(node, op)] + + +def _make_grouped_fp8_quantized_graph(*, num_groups=4, with_group_view=True): + batch_size, seq_len, rank, group_width = 2, 3, 8, 16 + + class Shell(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter( + torch.empty(num_groups * rank, group_width, dtype=torch.float8_e4m3fn) + ) + self.register_buffer("weight_scale_inv", torch.ones(num_groups, 1)) + + root = Shell() + graph = Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty(batch_size, seq_len, num_groups, group_width) + input_node = x + if with_group_view: + input_node = graph.call_function( + torch.ops.auto_deploy.view.default, + args=(x, [batch_size, seq_len, num_groups, group_width]), + kwargs={"tp_scaled_dim": 2, "layer_type": "mla"}, + ) + input_node.meta["val"] = torch.empty(batch_size, seq_len, num_groups, group_width) + weight = graph.get_attr("weight") + scale = graph.get_attr("weight_scale_inv") + out = graph.call_function( + torch.ops.auto_deploy.torch_fake_quant_grouped_finegrained_fp8_linear.default, + args=(input_node, weight, None, [], [scale], [], []), + kwargs={ + "tp_mode": "colwise", + "tp_min_local_shape": rank, + "layer_type": "mla", + "input_scale_fmt": "ue8m0", + }, + ) + out.meta["val"] = torch.empty(batch_size, seq_len, num_groups * rank) + graph.output(out) + return GraphModule(root, graph), num_groups, rank + + +def test_apply_hints_grouped_fp8_linear_trusts_group_sharded_view_input(): + gm, num_groups, rank = _make_grouped_fp8_quantized_graph(num_groups=8, with_group_view=True) + + gm_out = _make_optimizer(world_size=8, rank=7)(None, gm) + + grouped_nodes = _call_nodes( + gm_out, + torch.ops.auto_deploy.torch_fake_quant_grouped_finegrained_fp8_linear.default, + ) + assert len(grouped_nodes) == 1 + grouped_node = grouped_nodes[0] + grouped_input = grouped_node.args[0] + local_groups = num_groups // 8 + + assert is_op(grouped_input, torch.ops.auto_deploy.view.default) + assert not is_op(grouped_input, torch.ops.aten.slice.Tensor) + assert shape(grouped_input)[2] == local_groups + assert grouped_input.args[1][2] == -1 + assert gm_out.weight.shape == (local_groups * rank, 16) + assert gm_out.weight_scale_inv.shape == (local_groups, 1) + assert grouped_node.meta["val"].shape == (2, 3, local_groups * rank) + + +def test_apply_hints_grouped_fp8_linear_slices_plain_global_input_groups(): + gm, num_groups, rank = _make_grouped_fp8_quantized_graph(with_group_view=False) + + gm_out = _make_optimizer(world_size=2, rank=1)(None, gm) + + grouped_nodes = _call_nodes( + gm_out, + torch.ops.auto_deploy.torch_fake_quant_grouped_finegrained_fp8_linear.default, + ) + assert len(grouped_nodes) == 1 + grouped_node = grouped_nodes[0] + grouped_input = grouped_node.args[0] + local_groups = num_groups // 2 + + assert is_op(grouped_input, torch.ops.aten.slice.Tensor) + assert grouped_input.args[1:5] == (2, local_groups, num_groups, 1) + assert gm_out.weight.shape == (local_groups * rank, 16) + assert gm_out.weight_scale_inv.shape == (local_groups, 1) + assert grouped_node.meta["val"].shape == (2, 3, local_groups * rank) + + +def test_simple_shard_only_does_not_ordinary_shard_grouped_fp8_linear(): + gm, _, _ = _make_grouped_fp8_quantized_graph() + + gm_out = _make_optimizer(world_size=2, rank=1, simple_shard_only=True)(None, gm) + + info = gm_out.meta["_autodeploy"]["transform_history"]["apply_sharding_hints"] + assert info.num_matches == 0 + assert gm_out.weight.shape == (32, 16) + assert gm_out.weight_scale_inv.shape == (4, 1) + assert not _call_nodes(gm_out, torch.ops.auto_deploy.torch_dist_all_gather.default) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_apply_hints_default_dist_backend_uses_auto_selection(): + """Default IR sharding backend remains the existing auto-selected collective.""" + gm, _, _ = _export_hinted_mlp() + gm_out = _make_optimizer(world_size=2)(None, gm) + + _, expected_all_reduce = _get_dist_ops("auto") + assert len(_call_nodes(gm_out, expected_all_reduce)) == 1 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_apply_hints_torch_dist_backend_forces_torch_all_reduce(): + """Forcing dist_backend='torch' lowers all_reduce placeholders to torch collectives.""" + gm, _, _ = _export_hinted_mlp() + gm_out = _make_optimizer(world_size=2, dist_backend="torch")(None, gm) + + assert len(_call_nodes(gm_out, torch.ops.auto_deploy.torch_dist_all_reduce.default)) == 1 + assert len(_call_nodes(gm_out, torch.ops.auto_deploy.trtllm_dist_all_reduce.default)) == 0 + + +def test_deepseek_v4_ir_contract_linear_view_sparse_attention(): + """DeepSeek V4-shaped graph keeps TP, view, split, sink, and collective contracts.""" + gm = _export_deepseek_v4_contract_block() + gm_out = _make_optimizer(world_size=2)(None, gm) + + assert gm_out.q_proj.weight.shape == (DSV4_NUM_HEADS * DSV4_HEAD_DIM // 2, DSV4_HIDDEN) + assert gm_out.wo_a.weight.shape == (DSV4_O_RANK, DSV4_GROUP_WIDTH) + assert gm_out.wo_b.weight.shape == (DSV4_HIDDEN, DSV4_GROUPS * DSV4_O_RANK // 2) + assert gm_out.attn_sink.shape == (DSV4_NUM_HEADS // 2,) + + split_nodes = _call_nodes(gm_out, torch.ops.auto_deploy.split_with_sizes) + assert len(split_nodes) == 1 + [split_sizes] = extract_op_args(split_nodes[0], "split_sizes") + assert split_sizes == [DSV4_GROUP_WIDTH // 2, DSV4_GROUP_WIDTH // 2] + + view_shapes = [ + extract_op_args(node, "shape")[0] + for node in _call_nodes(gm_out, torch.ops.auto_deploy.view) + ] + assert [2, 3, -1, DSV4_HEAD_DIM] in view_shapes + assert [2, 3, -1, DSV4_GROUP_WIDTH] in view_shapes + assert [-1, DSV4_O_RANK, DSV4_GROUP_WIDTH] in view_shapes + + sparse_nodes = _call_nodes(gm_out, torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention) + assert len(sparse_nodes) == 1 + assert len(_call_nodes(gm_out, torch.ops.auto_deploy.torch_dist_all_reduce)) == 1 + + +def _annotate(node, value): + node.meta["val"] = value + return node + + +def _make_attr(graph, name, value): + node = graph.get_attr(name) + node.meta["val"] = value + return node + + +def _make_list_moe_graph(): + hidden_size = 4 + intermediate_size = 6 + num_experts = 4 + root = nn.Module() + for expert_idx in range(num_experts): + root.register_parameter( + f"w1_{expert_idx}", + nn.Parameter(torch.randn(intermediate_size, hidden_size)), + ) + root.register_parameter( + f"w2_{expert_idx}", + nn.Parameter(torch.randn(hidden_size, intermediate_size)), + ) + root.register_parameter( + f"w3_{expert_idx}", + nn.Parameter(torch.randn(intermediate_size, hidden_size)), + ) + + graph = torch.fx.Graph() + x = _annotate(graph.placeholder("x"), torch.empty(2, hidden_size)) + selected_experts = _annotate( + graph.placeholder("selected_experts"), torch.empty(2, 1, dtype=torch.int64) + ) + routing_weights = _annotate(graph.placeholder("routing_weights"), torch.empty(2, 1)) + + w1 = [_make_attr(graph, f"w1_{i}", getattr(root, f"w1_{i}")) for i in range(num_experts)] + w2 = [_make_attr(graph, f"w2_{i}", getattr(root, f"w2_{i}")) for i in range(num_experts)] + w3 = [_make_attr(graph, f"w3_{i}", getattr(root, f"w3_{i}")) for i in range(num_experts)] + moe = graph.call_function( + torch.ops.auto_deploy.torch_moe.default, + args=(x, selected_experts, routing_weights, w1, w2, w3), + kwargs={"layer_type": "moe"}, + ) + moe.meta["val"] = torch.empty(2, hidden_size) + graph.output(moe) + gm = torch.fx.GraphModule(root, graph) + gm.graph.lint() + gm.recompile() + return gm + + +def test_list_moe_ir_contract_leaves_ep_reduction_to_modeling(): + """List-based MoE EP sharding localizes experts without choosing a reduction site.""" + gm = _make_list_moe_graph() + gm_out = _make_optimizer(world_size=2)(None, gm) + moe_nodes = _call_nodes(gm_out, torch.ops.auto_deploy.torch_moe) + assert len(moe_nodes) == 1 + [w1_weight, w2_weight, w3_weight] = extract_op_args( + moe_nodes[0], "w1_weight", "w2_weight", "w3_weight" + ) + assert len(w1_weight) == 2 + assert len(w2_weight) == 2 + assert len(w3_weight) == 2 + assert len(_call_nodes(gm_out, torch.ops.auto_deploy.torch_dist_all_reduce)) == 0 + + +def _optional_auto_deploy_default(name): + try: + return getattr(torch.ops.auto_deploy, name).default + except AttributeError: + return None + + +_STACKED_MOE_EXPERT_ARG_NAMES = ( + "gate_up_blocks", + "gate_up_bias", + "gate_up_scales", + "down_blocks", + "down_bias", + "down_scales", +) + + +def _make_stacked_test_tensor(shape, dtype=torch.float32): + numel = 1 + for dim in shape: + numel *= dim + values = torch.arange(numel, dtype=torch.float32).reshape(shape) + if dtype == torch.uint8: + return values.remainder(251).to(torch.uint8) + return values.to(dtype) + + +def _make_stacked_mxfp4_graph( + base_op, routing_driven=False, include_optionals=False, expert_args_as_attrs=True +): + num_experts = 4 + hidden_size = 4 + intermediate_size = 3 + root = nn.Module() + tensors = { + "selected_experts": torch.tensor([[0], [3]], dtype=torch.int64), + "routing_weights": torch.ones(2, 1), + "router_weight": _make_stacked_test_tensor((num_experts, hidden_size)), + "router_bias": _make_stacked_test_tensor((num_experts,)), + "gate_up_blocks": _make_stacked_test_tensor( + (num_experts, 2 * intermediate_size, 1, 16), dtype=torch.uint8 + ), + "gate_up_bias": _make_stacked_test_tensor((num_experts, 2 * intermediate_size)), + "gate_up_scales": _make_stacked_test_tensor( + (num_experts, 2 * intermediate_size, 1), dtype=torch.uint8 + ), + "down_blocks": _make_stacked_test_tensor( + (num_experts, hidden_size, 1, 16), dtype=torch.uint8 + ), + "down_bias": _make_stacked_test_tensor((num_experts, hidden_size)), + "down_scales": _make_stacked_test_tensor((num_experts, hidden_size, 1), dtype=torch.uint8), + } + for name, tensor in tensors.items(): + if name in _STACKED_MOE_EXPERT_ARG_NAMES and not expert_args_as_attrs: + continue + root.register_buffer(name, tensor) + + graph = torch.fx.Graph() + x = _annotate(graph.placeholder("x"), torch.empty(2, hidden_size)) + attrs = {} + for name, tensor in tensors.items(): + if name in _STACKED_MOE_EXPERT_ARG_NAMES and not expert_args_as_attrs: + attrs[name] = _annotate(graph.placeholder(name), tensor) + else: + attrs[name] = _make_attr(graph, name, tensor) + if routing_driven: + args = ( + x, + attrs["selected_experts"], + attrs["routing_weights"], + attrs["gate_up_blocks"], + attrs["gate_up_bias"], + attrs["gate_up_scales"], + 1.0, + 10.0, + attrs["down_blocks"], + attrs["down_bias"], + attrs["down_scales"], + ) + if include_optionals: + args = args + ("gate_up", "gpt_oss", "moe") + else: + args = ( + x, + attrs["router_weight"], + attrs["router_bias"], + 1, + attrs["gate_up_blocks"], + attrs["gate_up_bias"], + attrs["gate_up_scales"], + 1.0, + 10.0, + attrs["down_blocks"], + attrs["down_bias"], + attrs["down_scales"], + ) + if include_optionals: + args = args + ("moe",) + moe = graph.call_function( + base_op, + args=args, + ) + moe.meta["val"] = torch.empty(2, hidden_size) + graph.output(moe) + gm = torch.fx.GraphModule(root, graph) + gm.graph.lint() + gm.recompile() + return gm + + +def _ensure_submodule(root, path): + submod = root + for name in path.split("."): + child = getattr(submod, name, None) + if child is None: + child = nn.Module() + submod.add_module(name, child) + submod = child + return submod + + +def _register_nested_buffer(root, target, tensor): + mod_name, _, attr_name = target.rpartition(".") + submod = _ensure_submodule(root, mod_name) + submod.register_buffer(attr_name, tensor) + + +def _make_deepseek_graph_mxfp4_graph(base_op, layer=3, num_experts=4): + hidden_size = 32 + intermediate_size = 32 + expert_targets = { + "gate_up_blocks": f"layers.{layer}.ffn.experts.gate_up_proj_blocks", + "gate_up_scales": f"layers.{layer}.ffn.experts.gate_up_proj_scales", + "down_blocks": f"layers.{layer}.ffn.experts.down_proj_blocks", + "down_scales": f"layers.{layer}.ffn.experts.down_proj_scales", + } + tensors = { + "router_weight": _make_stacked_test_tensor((num_experts, hidden_size)), + "router_bias": _make_stacked_test_tensor((num_experts,)), + "gate_up_blocks": _make_stacked_test_tensor( + (num_experts, 2 * intermediate_size, hidden_size // 32, 16), dtype=torch.uint8 + ), + "gate_up_bias": _make_stacked_test_tensor((num_experts, 2 * intermediate_size)), + "gate_up_scales": _make_stacked_test_tensor( + (num_experts, 2 * intermediate_size, hidden_size // 32), dtype=torch.uint8 + ), + "down_blocks": _make_stacked_test_tensor( + (num_experts, hidden_size, intermediate_size // 32, 16), dtype=torch.uint8 + ), + "down_bias": _make_stacked_test_tensor((num_experts, hidden_size)), + "down_scales": _make_stacked_test_tensor( + (num_experts, hidden_size, intermediate_size // 32), dtype=torch.uint8 + ), + } + root = nn.Module() + root.register_buffer("router_weight", tensors["router_weight"]) + root.register_buffer("router_bias", tensors["router_bias"]) + for name, target in expert_targets.items(): + _register_nested_buffer(root, target, tensors[name]) + root.register_buffer("gate_up_bias", tensors["gate_up_bias"]) + root.register_buffer("down_bias", tensors["down_bias"]) + + graph = torch.fx.Graph() + x = _annotate(graph.placeholder("x"), torch.empty(2, hidden_size)) + attrs = { + "router_weight": _make_attr(graph, "router_weight", tensors["router_weight"]), + "router_bias": _make_attr(graph, "router_bias", tensors["router_bias"]), + "gate_up_bias": _make_attr(graph, "gate_up_bias", tensors["gate_up_bias"]), + "down_bias": _make_attr(graph, "down_bias", tensors["down_bias"]), + } + for name, target in expert_targets.items(): + attrs[name] = _make_attr(graph, target, tensors[name]) + moe = graph.call_function( + base_op, + args=( + x, + attrs["router_weight"], + attrs["router_bias"], + 1, + attrs["gate_up_blocks"], + attrs["gate_up_bias"], + attrs["gate_up_scales"], + 1.0, + 10.0, + attrs["down_blocks"], + attrs["down_bias"], + attrs["down_scales"], + ), + ) + moe.meta["val"] = torch.empty(2, hidden_size) + graph.output(moe) + gm = torch.fx.GraphModule(root, graph) + gm.graph.lint() + gm.recompile() + return gm, expert_targets, hidden_size, intermediate_size + + +def _make_deepseek_flat_mxfp4_routing_graph(base_op, layer=3, num_experts=4): + hidden_size = 32 + intermediate_size = 32 + expert_targets = { + "gate_up_blocks": f"layers_{layer}_ffn_experts_gate_up_proj_blocks", + "gate_up_scales": f"layers_{layer}_ffn_experts_gate_up_proj_scales", + "down_blocks": f"layers_{layer}_ffn_experts_down_proj_blocks", + "down_scales": f"layers_{layer}_ffn_experts_down_proj_scales", + } + tensors = { + "selected_experts": torch.tensor([[0], [3]], dtype=torch.int64), + "routing_weights": torch.ones(2, 1), + "gate_up_blocks": _make_stacked_test_tensor( + (num_experts, 2 * intermediate_size, hidden_size // 32, 16), dtype=torch.uint8 + ), + "gate_up_bias": _make_stacked_test_tensor((num_experts, 2 * intermediate_size)), + "gate_up_scales": _make_stacked_test_tensor( + (num_experts, 2 * intermediate_size, hidden_size // 32), dtype=torch.uint8 + ), + "down_blocks": _make_stacked_test_tensor( + (num_experts, hidden_size, intermediate_size // 32, 16), dtype=torch.uint8 + ), + "down_bias": _make_stacked_test_tensor((num_experts, hidden_size)), + "down_scales": _make_stacked_test_tensor( + (num_experts, hidden_size, intermediate_size // 32), dtype=torch.uint8 + ), + } + root = nn.Module() + root.register_buffer("selected_experts", tensors["selected_experts"]) + root.register_buffer("routing_weights", tensors["routing_weights"]) + root.register_buffer("gate_up_bias", tensors["gate_up_bias"]) + root.register_buffer("down_bias", tensors["down_bias"]) + for name, target in expert_targets.items(): + root.register_buffer(target, tensors[name]) + + graph = torch.fx.Graph() + x = _annotate(graph.placeholder("x"), torch.empty(2, hidden_size)) + attrs = { + "selected_experts": _make_attr(graph, "selected_experts", tensors["selected_experts"]), + "routing_weights": _make_attr(graph, "routing_weights", tensors["routing_weights"]), + "gate_up_bias": _make_attr(graph, "gate_up_bias", tensors["gate_up_bias"]), + "down_bias": _make_attr(graph, "down_bias", tensors["down_bias"]), + } + for name, target in expert_targets.items(): + attrs[name] = _make_attr(graph, target, tensors[name]) + moe = graph.call_function( + base_op, + args=( + x, + attrs["selected_experts"], + attrs["routing_weights"], + attrs["gate_up_blocks"], + attrs["gate_up_bias"], + attrs["gate_up_scales"], + 1.0, + 10.0, + attrs["down_blocks"], + attrs["down_bias"], + attrs["down_scales"], + "up_gate", + "deepseek", + "moe", + num_experts, + ), + ) + moe.meta["val"] = torch.empty(2, hidden_size) + graph.output(moe) + gm = torch.fx.GraphModule(root, graph) + gm.graph.lint() + gm.recompile() + return gm, expert_targets, hidden_size, intermediate_size + + +class _MXFP4CheckpointLayoutFactory: + def __init__(self, checkpoint_layout): + self._checkpoint_layout = checkpoint_layout + + def get_quant_config(self): + return { + "checkpoint_layout": QuantizedCheckpointLayout( + checkpoint_consumers=(self._checkpoint_layout,) + ) + } + + +def _register_mxfp4_checkpoint_layout_hooks(gm, checkpoint_layout): + transform = InsertMXFP4MLP(MXFP4MLPConfig(stage=Stages.PATTERN_MATCHER)) + gm, info = transform._apply( + gm, + None, + _MXFP4CheckpointLayoutFactory(checkpoint_layout), + None, + ) + assert info.num_matches == 1 + return gm + + +def _make_deepseek_mxfp4_raw_state(layer, num_experts, hidden_size, intermediate_size): + def key(expert, projection, kind): + return f"layers.{layer}.ffn.experts.{expert}.{projection}.{kind}" + + def tensor(shape, offset): + numel = 1 + for dim in shape: + numel *= dim + values = (torch.arange(numel, dtype=torch.int64) + offset) % 251 + return values.to(torch.uint8).reshape(shape) + + state = {} + for expert in range(num_experts): + for projection, weight_shape, scale_shape, offset in ( + ( + "w1", + (intermediate_size, hidden_size // 2), + (intermediate_size, hidden_size // 32), + 17, + ), + ( + "w2", + (hidden_size, intermediate_size // 2), + (hidden_size, intermediate_size // 32), + 89, + ), + ( + "w3", + (intermediate_size, hidden_size // 2), + (intermediate_size, hidden_size // 32), + 151, + ), + ): + state[key(expert, projection, "weight")] = tensor(weight_shape, offset + expert) + state[key(expert, projection, "scale")] = tensor(scale_shape, offset + 31 + expert) + return state + + +def test_stacked_mxfp4_ir_contract_rewrites_to_matching_ep_variant(): + """Stacked MXFP4 MoE sharding is schema-based and rewrites to the matching EP op.""" + base_op = _optional_auto_deploy_default("triton_mxfp4_moe") + ep_op = _optional_auto_deploy_default("triton_mxfp4_moe_ep") + if base_op is None or ep_op is None: + pytest.skip("MXFP4 MoE custom ops are not registered in this environment") + + gm = _make_stacked_mxfp4_graph(base_op) + gm_out = _make_optimizer(world_size=2)(None, gm) + ep_nodes = _call_nodes(gm_out, ep_op) + assert len(ep_nodes) == 1 + [ep_size, ep_rank] = extract_op_args(ep_nodes[0], "ep_size", "ep_rank") + assert ep_size == 2 + assert ep_rank == 0 + + slice_nodes = _call_nodes(gm_out, torch.ops.aten.slice.Tensor) + expert_slices = [node for node in slice_nodes if node.args[1:5] == (0, 0, 2, 1)] + assert len(expert_slices) == 0 + expert_args = extract_op_args(ep_nodes[0], *_STACKED_MOE_EXPERT_ARG_NAMES) + assert all(getattr(arg, "op", None) == "get_attr" for arg in expert_args) + for name in _STACKED_MOE_EXPERT_ARG_NAMES: + assert getattr(gm_out, name).shape[0] == 2 + assert len(_call_nodes(gm_out, torch.ops.auto_deploy.torch_dist_all_reduce)) == 1 + + +def test_stacked_mxfp4_get_attr_expert_buffers_load_full_checkpoint_slices(): + """Full checkpoint buffers are split by load hooks after physical EP sharding.""" + base_op = _optional_auto_deploy_default("triton_mxfp4_moe") + if base_op is None or _optional_auto_deploy_default("triton_mxfp4_moe_ep") is None: + pytest.skip("MXFP4 MoE custom ops are not registered in this environment") + + rank = 1 + world_size = 2 + gm = _make_stacked_mxfp4_graph(base_op) + full_state = {name: tensor.clone() for name, tensor in gm.state_dict().items()} + gm_out = _make_optimizer(world_size=world_size, rank=rank)(None, gm) + + for name in _STACKED_MOE_EXPERT_ARG_NAMES: + getattr(gm_out, name).zero_() + + load_result = gm_out.load_state_dict( + {name: tensor.clone() for name, tensor in full_state.items()} + ) + assert load_result.missing_keys == [] + assert load_result.unexpected_keys == [] + + start = 2 + end = 4 + for name in _STACKED_MOE_EXPERT_ARG_NAMES: + assert torch.equal(getattr(gm_out, name), full_state[name][start:end]) + + +def test_stacked_mxfp4_non_attr_expert_args_keep_runtime_slices(): + """Dynamic expert tensors keep the old aten.slice runtime behavior.""" + base_op = _optional_auto_deploy_default("triton_mxfp4_moe") + ep_op = _optional_auto_deploy_default("triton_mxfp4_moe_ep") + if base_op is None or ep_op is None: + pytest.skip("MXFP4 MoE custom ops are not registered in this environment") + + gm = _make_stacked_mxfp4_graph(base_op, expert_args_as_attrs=False) + gm_out = _make_optimizer(world_size=2)(None, gm) + ep_nodes = _call_nodes(gm_out, ep_op) + assert len(ep_nodes) == 1 + + slice_nodes = _call_nodes(gm_out, torch.ops.aten.slice.Tensor) + expert_slices = [node for node in slice_nodes if node.args[1:5] == (0, 0, 2, 1)] + assert len(expert_slices) == 6 + expert_args = extract_op_args(ep_nodes[0], *_STACKED_MOE_EXPERT_ARG_NAMES) + assert all(getattr(arg, "target", None) == torch.ops.aten.slice.Tensor for arg in expert_args) + + +def test_stacked_mxfp4_rank1_deepseek_graph_pack_loads_high_expert_slice(): + """DeepSeek graph packing emits full buffers before sharding hooks split rank 1.""" + from tensorrt_llm._torch.auto_deploy.models.custom.modeling_deepseek_v4 import ( + build_deepseek_v4_packed_mxfp4_experts_layout, + ) + + base_op = _optional_auto_deploy_default("triton_mxfp4_moe") + if base_op is None or _optional_auto_deploy_default("triton_mxfp4_moe_ep") is None: + pytest.skip("MXFP4 MoE custom ops are not registered in this environment") + + layer = 3 + rank = 1 + world_size = 2 + num_experts = 4 + gm, expert_targets, hidden_size, intermediate_size = _make_deepseek_graph_mxfp4_graph( + base_op, layer=layer, num_experts=num_experts + ) + raw_expert_state = _make_deepseek_mxfp4_raw_state( + layer, num_experts, hidden_size, intermediate_size + ) + checkpoint = { + "router_weight": gm.router_weight.clone(), + "router_bias": gm.router_bias.clone(), + "gate_up_bias": gm.gate_up_bias.clone(), + "down_bias": gm.down_bias.clone(), + **raw_expert_state, + } + + layout = build_deepseek_v4_packed_mxfp4_experts_layout() + expected_full = layout.pack_experts( + raw_expert_state, + layer=layer, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_experts=num_experts, + ) + + gm = _register_mxfp4_checkpoint_layout_hooks(gm, layout) + gm_out = _make_optimizer(world_size=world_size, rank=rank)(None, gm) + experts = gm_out.get_submodule(f"layers.{layer}.ffn.experts") + assert experts.gate_up_proj_blocks.shape[0] == num_experts // world_size + + result = gm_out.load_state_dict(checkpoint) + assert result.missing_keys == [] + assert result.unexpected_keys == [] + + lo = num_experts // world_size + hi = num_experts + assert torch.equal( + gm_out.get_buffer(expert_targets["gate_up_blocks"]), + expected_full.gate_up_blocks[lo:hi], + ) + assert torch.equal( + gm_out.get_buffer(expert_targets["gate_up_scales"]), + expected_full.gate_up_scales[lo:hi], + ) + assert torch.equal( + gm_out.get_buffer(expert_targets["down_blocks"]), + expected_full.down_blocks[lo:hi], + ) + assert torch.equal( + gm_out.get_buffer(expert_targets["down_scales"]), + expected_full.down_scales[lo:hi], + ) + + +def test_stacked_mxfp4_rank1_deepseek_flat_graph_pack_loads_high_expert_slice(): + """Root FX buffers are packed before sharding hooks split rank 1.""" + from tensorrt_llm._torch.auto_deploy.models.custom.modeling_deepseek_v4 import ( + build_deepseek_v4_packed_mxfp4_experts_layout, + ) + + base_op = _optional_auto_deploy_default("torch_mxfp4_moe_from_routing") + ep_op = _optional_auto_deploy_default("torch_mxfp4_moe_from_routing_ep") + if base_op is None or ep_op is None: + pytest.skip("routing-driven MXFP4 MoE custom ops are not registered in this environment") + + layer = 3 + rank = 1 + world_size = 2 + num_experts = 4 + gm, expert_targets, hidden_size, intermediate_size = _make_deepseek_flat_mxfp4_routing_graph( + base_op, layer=layer, num_experts=num_experts + ) + raw_expert_state = _make_deepseek_mxfp4_raw_state( + layer, num_experts, hidden_size, intermediate_size + ) + layout = build_deepseek_v4_packed_mxfp4_experts_layout() + expected_full = layout.pack_experts( + raw_expert_state, + layer=layer, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_experts=num_experts, + ) + + gm = _register_mxfp4_checkpoint_layout_hooks(gm, layout) + checkpoint = { + "selected_experts": gm.selected_experts.clone(), + "routing_weights": gm.routing_weights.clone(), + "gate_up_bias": gm.gate_up_bias.clone(), + "down_bias": gm.down_bias.clone(), + **raw_expert_state, + } + + gm_out = _make_optimizer(world_size=world_size, rank=rank)(None, gm) + assert gm_out.get_buffer(expert_targets["gate_up_blocks"]).shape[0] == ( + num_experts // world_size + ) + + result = gm_out.load_state_dict(checkpoint) + assert result.missing_keys == [] + assert result.unexpected_keys == [] + + lo = num_experts // world_size + hi = num_experts + assert torch.equal( + gm_out.get_buffer(expert_targets["gate_up_blocks"]), + expected_full.gate_up_blocks[lo:hi], + ) + assert torch.equal( + gm_out.get_buffer(expert_targets["gate_up_scales"]), + expected_full.gate_up_scales[lo:hi], + ) + assert torch.equal( + gm_out.get_buffer(expert_targets["down_blocks"]), + expected_full.down_blocks[lo:hi], + ) + assert torch.equal( + gm_out.get_buffer(expert_targets["down_scales"]), + expected_full.down_scales[lo:hi], + ) + + +def test_stacked_mxfp4_deepseek_flat_graph_loads_production_shaped_rank7_slice(): + """Root FX buffers load the final EP slice for 256-expert routing-driven graphs.""" + from tensorrt_llm._torch.auto_deploy.models.custom.modeling_deepseek_v4 import ( + build_deepseek_v4_packed_mxfp4_experts_layout, + ) + + base_op = _optional_auto_deploy_default("torch_mxfp4_moe_from_routing") + ep_op = _optional_auto_deploy_default("torch_mxfp4_moe_from_routing_ep") + if base_op is None or ep_op is None: + pytest.skip("routing-driven MXFP4 MoE custom ops are not registered in this environment") + + layer = 3 + rank = 7 + world_size = 8 + num_experts = 256 + gm, expert_targets, hidden_size, intermediate_size = _make_deepseek_flat_mxfp4_routing_graph( + base_op, layer=layer, num_experts=num_experts + ) + raw_expert_state = _make_deepseek_mxfp4_raw_state( + layer, num_experts, hidden_size, intermediate_size + ) + layout = build_deepseek_v4_packed_mxfp4_experts_layout() + expected_full = layout.pack_experts( + raw_expert_state, + layer=layer, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_experts=num_experts, + ) + + gm = _register_mxfp4_checkpoint_layout_hooks(gm, layout) + checkpoint = { + "selected_experts": gm.selected_experts.clone(), + "routing_weights": gm.routing_weights.clone(), + "gate_up_bias": gm.gate_up_bias.clone(), + "down_bias": gm.down_bias.clone(), + **raw_expert_state, + } + + gm_out = _make_optimizer(world_size=world_size, rank=rank)(None, gm) + ep_nodes = _call_nodes(gm_out, ep_op) + assert len(ep_nodes) == 1 + expert_start, num_experts_total = extract_op_args( + ep_nodes[0], + "expert_start", + "num_experts_total", + ) + assert expert_start == 224 + assert num_experts_total == num_experts + + local_experts = num_experts // world_size + for target in expert_targets.values(): + local_buffer = gm_out.get_buffer(target) + assert local_buffer.shape[0] == local_experts + local_buffer.zero_() + + result = gm_out.load_state_dict(checkpoint) + assert result.missing_keys == [] + assert result.unexpected_keys == [] + + lo = expert_start + hi = num_experts + assert torch.equal( + gm_out.get_buffer(expert_targets["gate_up_blocks"]), + expected_full.gate_up_blocks[lo:hi], + ) + assert torch.equal( + gm_out.get_buffer(expert_targets["gate_up_scales"]), + expected_full.gate_up_scales[lo:hi], + ) + assert torch.equal( + gm_out.get_buffer(expert_targets["down_blocks"]), + expected_full.down_blocks[lo:hi], + ) + assert torch.equal( + gm_out.get_buffer(expert_targets["down_scales"]), + expected_full.down_scales[lo:hi], + ) + + +def test_stacked_mxfp4_routing_driven_ir_contract_rewrites_to_matching_ep_variant(): + """Routing-driven stacked MXFP4 EP op takes expert_start instead of ep_size/ep_rank.""" + base_op = _optional_auto_deploy_default("torch_mxfp4_moe_from_routing") + ep_op = _optional_auto_deploy_default("torch_mxfp4_moe_from_routing_ep") + if base_op is None or ep_op is None: + pytest.skip("routing-driven MXFP4 MoE custom ops are not registered in this environment") + + gm = _make_stacked_mxfp4_graph(base_op, routing_driven=True, include_optionals=True) + gm_out = _make_optimizer(world_size=2)(None, gm) + ep_nodes = _call_nodes(gm_out, ep_op) + assert len(ep_nodes) == 1 + [ + selected_experts, + routing_weights, + expert_start, + num_experts_total, + ep_size, + ep_rank, + gate_up_order, + swiglu_mode, + layer_type, + ] = extract_op_args( + ep_nodes[0], + "selected_experts", + "routing_weights", + "expert_start", + "num_experts_total", + "ep_size", + "ep_rank", + "gate_up_order", + "swiglu_mode", + "layer_type", + ) + assert selected_experts is not None + assert routing_weights is not None + assert expert_start == 0 + assert num_experts_total == 4 + assert ep_size is None + assert ep_rank is None + assert gate_up_order == "gate_up" + assert swiglu_mode == "gpt_oss" + assert layer_type == "moe" + + slice_nodes = _call_nodes(gm_out, torch.ops.aten.slice.Tensor) + expert_slices = [node for node in slice_nodes if node.args[1:5] == (0, 0, 2, 1)] + assert len(expert_slices) == 0 + expert_args = extract_op_args(ep_nodes[0], *_STACKED_MOE_EXPERT_ARG_NAMES) + assert all(getattr(arg, "op", None) == "get_attr" for arg in expert_args) + assert len(_call_nodes(gm_out, torch.ops.auto_deploy.torch_dist_all_reduce)) == 1 + + +def test_stacked_mxfp4_routing_driven_rank1_preserves_expert_start_with_torch_backend(): + """Rank 1 routing-driven MXFP4 EP rewrite keeps expert_start and torch all_reduce.""" + base_op = _optional_auto_deploy_default("torch_mxfp4_moe_from_routing") + ep_op = _optional_auto_deploy_default("torch_mxfp4_moe_from_routing_ep") + if base_op is None or ep_op is None: + pytest.skip("routing-driven MXFP4 MoE custom ops are not registered in this environment") + + gm = _make_stacked_mxfp4_graph(base_op, routing_driven=True, include_optionals=True) + gm_out = _make_optimizer(world_size=2, rank=1, dist_backend="torch")(None, gm) + + ep_nodes = _call_nodes(gm_out, ep_op) + assert len(ep_nodes) == 1 + expert_start, num_experts_total = extract_op_args( + ep_nodes[0], + "expert_start", + "num_experts_total", + ) + assert expert_start == 2 + assert num_experts_total == 4 + assert len(_call_nodes(gm_out, torch.ops.auto_deploy.torch_dist_all_reduce.default)) == 1 + assert len(_call_nodes(gm_out, torch.ops.auto_deploy.trtllm_dist_all_reduce.default)) == 0 diff --git a/tests/unittest/auto_deploy/singlegpu/compile/test_captured_graph.py b/tests/unittest/auto_deploy/singlegpu/compile/test_captured_graph.py index 9993b6cb8dea..d77225d41d4b 100644 --- a/tests/unittest/auto_deploy/singlegpu/compile/test_captured_graph.py +++ b/tests/unittest/auto_deploy/singlegpu/compile/test_captured_graph.py @@ -384,6 +384,58 @@ def get_args_kwargs(batch_size): assert model.forward_calls == calls_after_capture +def test_monolithic_decode_capture_falls_back_to_eager_for_prefill_shape(monkeypatch): + """Monolithic CUDA graph captures decode buckets and leaves non-matching shapes eager.""" + + class CountingModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.forward_calls = 0 + + def forward(self, input_ids): + self.forward_calls += 1 + return input_ids + 10 + + class ReplayGraph: + def __init__(self, compiled_model): + self.compiled_model = compiled_model + self.replay_calls = 0 + + def replay(self): + self.replay_calls += 1 + self.compiled_model._out_buffer_flat[0].copy_( + self.compiled_model._input_buffers[0] + 10 + ) + + model = CountingModel() + compiled_model = CapturedGraph(model, num_batched_inputs=1) + graphs = [] + + def fake_capture_one_graph(self, args, kwargs, refresh_args_static=None): + del args, kwargs, refresh_args_static + graph = ReplayGraph(self) + graphs.append(graph) + return graph, (2,) + + monkeypatch.setattr(CapturedGraph, "_capture_one_graph", fake_capture_one_graph) + + def get_args_kwargs(batch_size): + return (torch.ones(batch_size, 1),), {} + + compiled_model.capture_graph(get_args_kwargs, [2]) + calls_after_capture = model.forward_calls + + decode_output = compiled_model(torch.full((2, 1), 3.0)) + assert model.forward_calls == calls_after_capture + assert graphs[0].replay_calls == 1 + torch.testing.assert_close(decode_output, torch.full((2, 1), 13.0)) + + prefill_output = compiled_model(torch.full((1, 4), 5.0)) + assert model.forward_calls == calls_after_capture + 1 + assert graphs[0].replay_calls == 1 + torch.testing.assert_close(prefill_output, torch.full((1, 4), 15.0)) + + # ============================================================================ # Tests for CapturedGraph capture-time truncation # ============================================================================ @@ -847,7 +899,6 @@ def _make_dual_mode(self, piecewise_num_tokens=None): monolithic = MagicMock(spec=nn.Module) monolithic.return_value = torch.tensor([1.0]) - piecewise = MagicMock(spec=PiecewiseCapturedGraph) piecewise.piecewise_num_tokens = piecewise_num_tokens piecewise.original_model = MagicMock(return_value=torch.tensor([2.0])) diff --git a/tests/unittest/auto_deploy/singlegpu/compile/test_piecewise_utils.py b/tests/unittest/auto_deploy/singlegpu/compile/test_piecewise_utils.py index 96ccecc4d2e4..87a5fdb8388e 100644 --- a/tests/unittest/auto_deploy/singlegpu/compile/test_piecewise_utils.py +++ b/tests/unittest/auto_deploy/singlegpu/compile/test_piecewise_utils.py @@ -163,6 +163,11 @@ def test_known_attention_op_returns_true(self): assert is_dynamic_cached_op(node) is True assert _get_dynamic_op_policy(node) == DynamicOpPolicy.OUT_BUFFER + def test_known_sparse_attention_op_returns_true(self): + target = _FakeOpOverload("auto_deploy::torch_deepseek_v4_sparse_attention_with_cache") + node = _make_mock_node("call_function", target=target) + assert is_dynamic_cached_op(node) is True + def test_known_ssm_op_returns_true(self): target = _FakeOpOverload("auto_deploy::triton_cached_ssm") node = _make_mock_node("call_function", target=target) @@ -392,6 +397,16 @@ def test_needs_out_buffer_for_attention_still_true(self): gm = GraphModule(nn.Module(), graph) assert needs_out_buffer(gm) is True + def test_needs_out_buffer_for_sparse_attention(self): + """Sparse cached attention follows the dynamic op output-buffer ABI.""" + graph = Graph() + x = graph.placeholder("x") + attn_target = _FakeOpOverload("auto_deploy::torch_deepseek_v4_sparse_attention_with_cache") + attn = graph.create_node("call_function", attn_target, args=(x,), name="attn") + graph.output(attn) + gm = GraphModule(nn.Module(), graph) + assert needs_out_buffer(gm) is True + def test_split_does_not_reclassify_stream_switch(self): """The splitter must not reclassify stream-switch partitions as dynamic.""" gm = _build_graphmodule_with_stream_switch_ops( diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_deepseek_v4_sparse_attention.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_deepseek_v4_sparse_attention.py new file mode 100644 index 000000000000..f9732673db3f --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_deepseek_v4_sparse_attention.py @@ -0,0 +1,3173 @@ +# 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. + +"""Semantic tests for the DeepSeek V4 sparse attention source op.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch +from _model_test_utils import assert_rmse_close +from omegaconf import OmegaConf +from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode +from torch.export import Dim +from torch.fx import Graph + +import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: E402, F401 +from tensorrt_llm._torch.auto_deploy._compat import KvCacheConfig # noqa: E402 +from tensorrt_llm._torch.auto_deploy.custom_ops.attention import ( # noqa: E402 + deepseek_v4_sparse_attention as dsv4_sparse, +) +from tensorrt_llm._torch.auto_deploy.custom_ops.attention.deepseek_v4_sparse_attention import ( # noqa: E402 + DeepSeekV4SparseAttention, + FlashMLADeepSeekV4SparseAttention, +) +from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import ( # noqa: E402 + BatchInfo, + PagedResourceHandler, +) +from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm # noqa: E402 +from tensorrt_llm._torch.auto_deploy.models.custom.modeling_deepseek_v4 import ( # noqa: E402 + DeepseekV4Compressor, + DeepseekV4Config, + DeepseekV4Indexer, +) +from tensorrt_llm._torch.auto_deploy.shim.interface import CachedSequenceInterface # noqa: E402 +from tensorrt_llm._torch.auto_deploy.transform.interface import SharedConfig, Stages # noqa: E402 +from tensorrt_llm._torch.auto_deploy.transform.library.kvcache import ( # noqa: E402 + InsertCachedAttentionConfig, + InsertCachedDeepSeekV4SparseAttention, + _InsertCachedOperator, +) + + +def _repo_root() -> Path: + for parent in Path(__file__).resolve().parents: + if (parent / "tensorrt_llm/_torch/auto_deploy/config/default.yaml").is_file(): + return parent + raise AssertionError("Could not locate TensorRT-LLM repository root") + + +def _page_meta( + seq_lens: list[int], + input_positions: list[int], + slot_indices: list[int], + tokens_per_block: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + cu_num_pages = [0] + cache_loc = [] + last_page_len = [] + next_page_id = 0 + for seq_len, input_pos, slot_idx in zip(seq_lens, input_positions, slot_indices, strict=True): + total_len = input_pos + seq_len + if tokens_per_block is None: + seq_pages = [slot_idx] + lpl = max(total_len, 1) + else: + num_pages = max((max(total_len, 1) + tokens_per_block - 1) // tokens_per_block, 1) + seq_pages = list(range(next_page_id, next_page_id + num_pages)) + next_page_id += num_pages + lpl = ((max(total_len, 1) - 1) % tokens_per_block) + 1 + cache_loc.extend(seq_pages) + cu_num_pages.append(cu_num_pages[-1] + len(seq_pages)) + last_page_len.append(lpl) + + return ( + torch.tensor(cu_num_pages, dtype=torch.int32), + torch.tensor(cache_loc, dtype=torch.int32), + torch.tensor(last_page_len, dtype=torch.int32), + ) + + +def _context_meta(seq_len: int, tokens_per_block: int | None = None): + batch_info_host = BatchInfo() + batch_info_host.update([1, seq_len, 0, 0, 0, 0]) + cu_num_pages, cache_loc, last_page_len = _page_meta([seq_len], [0], [0], tokens_per_block) + return ( + batch_info_host.serialize(), + torch.tensor([seq_len], dtype=torch.int32), + torch.tensor([0], dtype=torch.int32), + torch.tensor([0], dtype=torch.int64), + torch.tensor([0, seq_len], dtype=torch.int32), + cu_num_pages, + cache_loc, + last_page_len, + ) + + +def _multi_context_meta(seq_lens: list[int], tokens_per_block: int | None = None): + total_tokens = sum(seq_lens) + cu_seqlen = [0] + for seq_len in seq_lens: + cu_seqlen.append(cu_seqlen[-1] + seq_len) + + batch_info_host = BatchInfo() + batch_info_host.update([len(seq_lens), total_tokens, 0, 0, 0, 0]) + slot_indices = list(range(len(seq_lens))) + cu_num_pages, cache_loc, last_page_len = _page_meta( + seq_lens, [0] * len(seq_lens), slot_indices, tokens_per_block + ) + return ( + batch_info_host.serialize(), + torch.tensor(seq_lens, dtype=torch.int32), + torch.zeros(len(seq_lens), dtype=torch.int32), + torch.tensor(slot_indices, dtype=torch.int64), + torch.tensor(cu_seqlen, dtype=torch.int32), + cu_num_pages, + cache_loc, + last_page_len, + ) + + +def _decode_meta(input_pos: int, tokens_per_block: int | None = None): + batch_info_host = BatchInfo() + batch_info_host.update([0, 0, 0, 0, 1, 1]) + cu_num_pages, cache_loc, last_page_len = _page_meta([1], [input_pos], [0], tokens_per_block) + return ( + batch_info_host.serialize(), + torch.tensor([1], dtype=torch.int32), + torch.tensor([input_pos], dtype=torch.int32), + torch.tensor([0], dtype=torch.int64), + torch.tensor([0, 1], dtype=torch.int32), + cu_num_pages, + cache_loc, + last_page_len, + ) + + +def _multi_decode_meta(input_positions: list[int], tokens_per_block: int | None = None): + seq_lens = [1] * len(input_positions) + cu_seqlen = list(range(len(input_positions) + 1)) + batch_info_host = BatchInfo() + batch_info_host.update([0, 0, 0, 0, len(input_positions), len(input_positions)]) + slot_indices = list(range(len(input_positions))) + cu_num_pages, cache_loc, last_page_len = _page_meta( + seq_lens, input_positions, slot_indices, tokens_per_block + ) + return ( + batch_info_host.serialize(), + torch.tensor(seq_lens, dtype=torch.int32), + torch.tensor(input_positions, dtype=torch.int32), + torch.tensor(slot_indices, dtype=torch.int64), + torch.tensor(cu_seqlen, dtype=torch.int32), + cu_num_pages, + cache_loc, + last_page_len, + ) + + +def _cuda_decode_meta(input_pos: int, slot_idx: int = 0, tokens_per_block: int | None = None): + ( + batch_info_host, + seq_len, + input_pos_tensor, + slot_idx_tensor, + cu_seqlen, + cu_num_pages, + cache_loc, + last_page_len, + ) = _decode_meta(input_pos, tokens_per_block=tokens_per_block) + input_pos_tensor.fill_(input_pos) + slot_idx_tensor.fill_(slot_idx) + if tokens_per_block is None: + cache_loc.fill_(slot_idx) + return ( + batch_info_host, + seq_len.cuda(), + input_pos_tensor.cuda(), + slot_idx_tensor.cuda(), + cu_seqlen.cuda(), + cu_num_pages.cuda(), + cache_loc.cuda(), + last_page_len.cuda(), + ) + + +def _sparse_attention_reference( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + batch_size, seq_len, num_heads, _ = q.shape + batch_idx = torch.arange(batch_size, device=q.device).view(batch_size, 1, 1) + batch_idx = batch_idx.expand(batch_size, seq_len, topk_idxs.shape[-1]) + + compute_dtype = torch.float32 if q.dtype in (torch.float16, torch.bfloat16) else q.dtype + gather_idxs = topk_idxs.to(torch.long).clamp(min=0) + selected_kv = kv[batch_idx, gather_idxs].to(compute_dtype) + logits = torch.matmul(q.to(compute_dtype), selected_kv.transpose(-1, -2)) + logits = logits * softmax_scale + logits = logits.masked_fill((topk_idxs < 0).unsqueeze(2), float("-inf")) + + sink_logits = attn_sink.to(dtype=compute_dtype).view(1, 1, num_heads, 1) + sink_logits = sink_logits.expand(batch_size, seq_len, num_heads, 1) + weights = torch.softmax(torch.cat([logits, sink_logits], dim=-1), dim=-1) + output = torch.matmul(weights[..., :-1], selected_kv) + return output.to(q.dtype) + + +def _empty_sparse_attention_tensors(q: torch.Tensor, kv: torch.Tensor) -> tuple[torch.Tensor, ...]: + del kv + return ( + q.new_empty(q.shape[0], q.shape[1], 0), + q.new_empty(q.shape[0], q.shape[1], 0), + q.new_empty(0, 0), + q.new_empty(0), + q.new_empty(0, 0), + q.new_empty(0, 0), + q.new_empty(q.shape[0], q.shape[1]), + q.new_empty(q.shape[0], q.shape[1], 0, 0), + q.new_empty(q.shape[0], q.shape[1], 0), + q.new_empty(q.shape[0], q.shape[1], 0), + q.new_empty(q.shape[0], q.shape[1], 0), + q.new_empty(0, 0), + q.new_empty(0), + ) + + +def _run_sparse_attention( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + softmax_scale: float = 1.0, +) -> torch.Tensor: + return torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention( + q, + kv, + attn_sink, + topk_idxs, + *_empty_sparse_attention_tensors(q, kv), + softmax_scale, + compress_ratio=0, + ) + + +def _run_cached_sparse_attention( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + metadata: tuple[torch.Tensor, ...], + swa_cache: torch.Tensor, + softmax_scale: float = 1.0, + window_size: int | None = None, + compress_ratio: int = 0, + out: torch.Tensor | None = None, +) -> torch.Tensor: + mhc_cache = swa_cache.new_empty(swa_cache.shape) + compressor_kv_cache = q.new_empty(swa_cache.shape[0], swa_cache.shape[1], 0) + compressor_gate_cache = q.new_empty(swa_cache.shape[0], swa_cache.shape[1], 0) + indexer_compressor_kv_cache = q.new_empty(swa_cache.shape[0], swa_cache.shape[1], 0) + indexer_compressor_gate_cache = q.new_empty(swa_cache.shape[0], swa_cache.shape[1], 0) + return torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention_with_cache( + q, + kv, + attn_sink, + topk_idxs, + *_empty_sparse_attention_tensors(q, kv), + *metadata, + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + softmax_scale, + window_size, + compress_ratio, + None, + 1e-6, + None, + out=out, + ) + + +def _run_flashmla_cached_sparse_attention( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + metadata: tuple[torch.Tensor, ...], + swa_cache: torch.Tensor, + softmax_scale: float = 1.0, + window_size: int | None = None, + compress_ratio: int = 0, + out: torch.Tensor | None = None, +) -> torch.Tensor: + mhc_cache = swa_cache.new_empty(swa_cache.shape) + compressor_kv_cache = q.new_empty(swa_cache.shape[0], swa_cache.shape[1], 0) + compressor_gate_cache = q.new_empty(swa_cache.shape[0], swa_cache.shape[1], 0) + indexer_compressor_kv_cache = q.new_empty(swa_cache.shape[0], swa_cache.shape[1], 0) + indexer_compressor_gate_cache = q.new_empty(swa_cache.shape[0], swa_cache.shape[1], 0) + return torch.ops.auto_deploy.flashmla_deepseek_v4_sparse_attention_with_cache( + q, + kv, + attn_sink, + topk_idxs, + *_empty_sparse_attention_tensors(q, kv), + *metadata, + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + softmax_scale, + window_size, + compress_ratio, + None, + 1e-6, + None, + out=out, + ) + + +def _run_sparse_attention_with_compressor( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + compressor_kv: torch.Tensor, + compressor_gate: torch.Tensor, + compressor: DeepseekV4Compressor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + position_ids: torch.Tensor, + softmax_scale: float = 1.0, + window_size: int | None = None, + compress_ratio: int = 0, + indexer_q: torch.Tensor | None = None, + indexer_weights: torch.Tensor | None = None, + indexer_compressor_kv: torch.Tensor | None = None, + indexer_compressor_gate: torch.Tensor | None = None, + indexer_compressor_ape: torch.Tensor | None = None, + indexer_compressor_norm_weight: torch.Tensor | None = None, +) -> torch.Tensor: + indexer_q = indexer_q if indexer_q is not None else q.new_empty(q.shape[0], q.shape[1], 0, 0) + indexer_weights = ( + indexer_weights if indexer_weights is not None else q.new_empty(q.shape[0], q.shape[1], 0) + ) + indexer_compressor_kv = ( + indexer_compressor_kv + if indexer_compressor_kv is not None + else q.new_empty(q.shape[0], q.shape[1], 0) + ) + indexer_compressor_gate = ( + indexer_compressor_gate + if indexer_compressor_gate is not None + else q.new_empty(q.shape[0], q.shape[1], 0) + ) + indexer_compressor_ape = ( + indexer_compressor_ape if indexer_compressor_ape is not None else q.new_empty(0, 0) + ) + indexer_compressor_norm_weight = ( + indexer_compressor_norm_weight + if indexer_compressor_norm_weight is not None + else q.new_empty(0) + ) + return torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention( + q, + kv, + attn_sink, + topk_idxs, + compressor_kv, + compressor_gate, + compressor.ape, + compressor.norm.weight, + cos_table, + sin_table, + position_ids, + indexer_q, + indexer_weights, + indexer_compressor_kv, + indexer_compressor_gate, + indexer_compressor_ape, + indexer_compressor_norm_weight, + softmax_scale, + False, + "mha_sparse", + 0, + window_size, + compress_ratio, + compressor.max_compressed_len, + kv.shape[-1], + compressor.rope_head_dim, + compressor.norm.eps, + ) + + +def _run_cached_sparse_attention_with_compressor( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + compressor_kv: torch.Tensor, + compressor_gate: torch.Tensor, + compressor: DeepseekV4Compressor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + position_ids: torch.Tensor, + metadata: tuple[torch.Tensor, ...], + swa_cache: torch.Tensor, + mhc_cache: torch.Tensor, + compressor_kv_cache: torch.Tensor, + compressor_gate_cache: torch.Tensor, + softmax_scale: float = 1.0, + window_size: int = 4, + compress_ratio: int = 4, + indexer_q: torch.Tensor | None = None, + indexer_weights: torch.Tensor | None = None, + indexer_compressor_kv: torch.Tensor | None = None, + indexer_compressor_gate: torch.Tensor | None = None, + indexer_compressor_ape: torch.Tensor | None = None, + indexer_compressor_norm_weight: torch.Tensor | None = None, + indexer_compressor_kv_cache: torch.Tensor | None = None, + indexer_compressor_gate_cache: torch.Tensor | None = None, +) -> torch.Tensor: + indexer_q = indexer_q if indexer_q is not None else q.new_empty(q.shape[0], q.shape[1], 0, 0) + indexer_weights = ( + indexer_weights if indexer_weights is not None else q.new_empty(q.shape[0], q.shape[1], 0) + ) + indexer_compressor_kv = ( + indexer_compressor_kv + if indexer_compressor_kv is not None + else q.new_empty(q.shape[0], q.shape[1], 0) + ) + indexer_compressor_gate = ( + indexer_compressor_gate + if indexer_compressor_gate is not None + else q.new_empty(q.shape[0], q.shape[1], 0) + ) + indexer_compressor_ape = ( + indexer_compressor_ape if indexer_compressor_ape is not None else q.new_empty(0, 0) + ) + indexer_compressor_norm_weight = ( + indexer_compressor_norm_weight + if indexer_compressor_norm_weight is not None + else q.new_empty(0) + ) + indexer_compressor_kv_cache = ( + indexer_compressor_kv_cache + if indexer_compressor_kv_cache is not None + else q.new_empty(swa_cache.shape[0], swa_cache.shape[1], 0) + ) + indexer_compressor_gate_cache = ( + indexer_compressor_gate_cache + if indexer_compressor_gate_cache is not None + else q.new_empty(swa_cache.shape[0], swa_cache.shape[1], 0) + ) + return torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention_with_cache( + q, + kv, + attn_sink, + topk_idxs, + compressor_kv, + compressor_gate, + compressor.ape, + compressor.norm.weight, + cos_table, + sin_table, + position_ids, + indexer_q, + indexer_weights, + indexer_compressor_kv, + indexer_compressor_gate, + indexer_compressor_ape, + indexer_compressor_norm_weight, + *metadata, + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + softmax_scale, + window_size, + compress_ratio, + compressor.max_compressed_len, + compressor.norm.eps, + compressor.rope_head_dim, + ) + + +def _rope_tables(max_seq_len: int, rope_dim: int) -> tuple[torch.Tensor, torch.Tensor]: + if rope_dim == 0: + return torch.empty(max_seq_len, 0), torch.empty(max_seq_len, 0) + positions = torch.arange(max_seq_len, dtype=torch.float32).unsqueeze(1) + freqs = torch.linspace(0.05, 0.25, rope_dim // 2, dtype=torch.float32).unsqueeze(0) + angles = positions * freqs + return angles.cos(), angles.sin() + + +def _compressor_case( + compress_ratio: int, + seq_len: int, + *, + compressed_capacity_tokens: int | None = None, + batch_size: int = 1, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + DeepseekV4Compressor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + hidden_size = 16 + head_dim = 8 + rope_dim = 4 + capacity = compressed_capacity_tokens or seq_len + config = DeepseekV4Config( + hidden_size=hidden_size, + num_hidden_layers=1, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=head_dim, + qk_rope_head_dim=rope_dim, + compress_ratios=(compress_ratio,), + ad_compress_max_seq_len=capacity, + ad_rope_cache_len=max(capacity, seq_len, 1), + ) + compressor = DeepseekV4Compressor(config, compress_ratio, head_dim).eval() + hidden_states = torch.randn(batch_size, seq_len, hidden_size) + compressor_kv, compressor_gate = compressor.project(hidden_states) + cos_table, sin_table = _rope_tables(max(capacity, seq_len, 1), rope_dim) + position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1).contiguous() + compressed_kv = compressor(hidden_states, cos_table, sin_table, position_ids) + return ( + compressor_kv, + compressor_gate, + compressed_kv, + compressor, + cos_table, + sin_table, + position_ids, + ) + + +def _visible_source_topk( + query_len: int, + input_pos: int, + kv_rows: int, + window_size: int, + compress_ratio: int, + max_compressed_len: int, + device: torch.device, +) -> torch.Tensor: + rows = [] + max_select = window_size + max_compressed_len + for token_offset in range(query_len): + query_pos = input_pos + token_offset + local_start = max(0, query_pos - window_size + 1) + selected = list(range(local_start, query_pos + 1)) + visible_compressed = min((query_pos + 1) // compress_ratio, max_compressed_len) + selected.extend(kv_rows + row_idx for row_idx in range(visible_compressed)) + selected.extend([-1] * (max_select - len(selected))) + rows.append(selected) + return torch.tensor([rows], dtype=torch.int64, device=device) + + +def _make_sparse_attention_caches( + max_seq_len: int, + head_dim: int, + compressor_state_dim: int, + fill_value: float = 0.0, + *, + num_slots: int = 1, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return ( + torch.full((num_slots, max_seq_len, head_dim), fill_value), + torch.full((num_slots, max_seq_len, head_dim), fill_value), + torch.full((num_slots, max_seq_len, compressor_state_dim), fill_value), + torch.full((num_slots, max_seq_len, compressor_state_dim), fill_value), + ) + + +def _make_paged_sparse_attention_caches( + max_seq_len: int, + tokens_per_block: int, + head_dim: int, + compressor_state_dim: int, + fill_value: float = 0.0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + num_pages = max((max_seq_len + tokens_per_block - 1) // tokens_per_block, 1) + return ( + torch.full((num_pages, tokens_per_block, head_dim), fill_value), + torch.full((num_pages, tokens_per_block, head_dim), fill_value), + torch.full((num_pages, tokens_per_block, compressor_state_dim), fill_value), + torch.full((num_pages, tokens_per_block, compressor_state_dim), fill_value), + ) + + +def _paged_cache_row( + cache: torch.Tensor, + logical_pos: int, + tokens_per_block: int, +) -> torch.Tensor: + return cache[logical_pos // tokens_per_block, logical_pos % tokens_per_block] + + +def _has_resource_with_suffix(resource_names: list[str], suffix: str) -> bool: + return any(name.endswith(suffix) for name in resource_names) + + +def test_flash_mla_sparse_prefill_contract_requires_dsv4_kwargs() -> None: + def sparse_prefill_without_dsv4_kwargs(q, kv, indices, sm_scale, d_v=512): + return q, kv, indices, sm_scale, d_v + + def dsv4_flash_mla_sparse_fwd( + q, + kv, + indices, + sm_scale, + d_v=512, + attn_sink=None, + topk_length=None, + out=None, + ): + return q, kv, indices, sm_scale, d_v, attn_sink, topk_length, out + + with pytest.raises(RuntimeError, match="DeepSeek V4 contract"): + dsv4_sparse._validate_flash_mla_sparse_prefill_contract(sparse_prefill_without_dsv4_kwargs) + + dsv4_sparse._validate_flash_mla_sparse_prefill_contract(dsv4_flash_mla_sparse_fwd) + + +def test_flash_mla_sparse_prefill_input_preparation_offsets_batches() -> None: + q = torch.empty(2, 2, 3, 4) + kv = torch.empty(2, 5, 4) + topk_idxs = torch.tensor( + [ + [[0, 4, -1, 5], [2, 99, 1, -2]], + [[0, 3, 4, -1], [1, 5, 2, 0]], + ], + dtype=torch.int64, + ) + + q_flat, kv_flat, indices, topk_length = dsv4_sparse._prepare_flash_mla_sparse_prefill_inputs( + q, kv, topk_idxs + ) + + assert q_flat.shape == (4, 3, 4) + assert kv_flat.shape == (10, 1, 4) + assert indices.shape == (4, 1, dsv4_sparse._FLASH_MLA_PREFILL_TOPK_MULTIPLE) + assert topk_length.tolist() == [4, 4, 4, 4] + expected_prefix = torch.tensor( + [[[0, 4, -1, -1]], [[2, -1, 1, -1]], [[5, 8, 9, -1]], [[6, -1, 7, 5]]], + dtype=torch.int32, + ) + torch.testing.assert_close(indices[:, :, :4], expected_prefix) + assert torch.all(indices[:, :, 4:] == -1) + + +def test_flash_mla_model1_cache_pack_writes_block_layout() -> None: + block_size = 4 + cache = torch.zeros( + 1, + block_size, + dsv4_sparse._FLASH_MLA_MODEL1_BYTES_PER_TOKEN, + dtype=torch.uint8, + ) + values = torch.zeros(2, dsv4_sparse._FLASH_MLA_MODEL1_HEAD_DIM, dtype=torch.bfloat16) + values[0, 0] = 1.0 + values[0, dsv4_sparse._FLASH_MLA_MODEL1_NOPE_DIM] = 3.0 + values[1, 0] = 2.0 + values[1, dsv4_sparse._FLASH_MLA_MODEL1_NOPE_DIM] = 4.0 + cu_num_pages, cache_loc, _ = _page_meta([2], [0], [0], tokens_per_block=block_size) + + dsv4_sparse._write_paged_cache_rows( + values, + cache, + seq_idx=0, + input_pos=0, + cu_num_pages_host=cu_num_pages, + cache_loc_host=cache_loc, + ) + + flat = cache.view(torch.uint8).view(1, -1)[0] + token_data_bytes = dsv4_sparse._FLASH_MLA_MODEL1_TOKEN_DATA_BYTES + scale_base = block_size * token_data_bytes + first_rope_start = dsv4_sparse._FLASH_MLA_MODEL1_NOPE_DIM + second_rope_start = token_data_bytes + dsv4_sparse._FLASH_MLA_MODEL1_NOPE_DIM + + assert flat[0].item() != 0 + assert flat[token_data_bytes].item() != 0 + first_rope = flat[first_rope_start : first_rope_start + 2].view(torch.bfloat16) + second_rope = flat[second_rope_start : second_rope_start + 2].view(torch.bfloat16) + torch.testing.assert_close(first_rope, torch.tensor([3.0], dtype=torch.bfloat16)) + torch.testing.assert_close(second_rope, torch.tensor([4.0], dtype=torch.bfloat16)) + assert flat[scale_base : scale_base + dsv4_sparse._FLASH_MLA_MODEL1_NUM_TILES].any() + second_scale = scale_base + dsv4_sparse._FLASH_MLA_MODEL1_SCALE_BYTES + assert flat[second_scale : second_scale + dsv4_sparse._FLASH_MLA_MODEL1_NUM_TILES].any() + + +def test_flash_mla_model1_cache_write_uses_vectorized_page_translation(monkeypatch) -> None: + block_size = 2 + cache = torch.zeros( + 3, + block_size, + dsv4_sparse._FLASH_MLA_MODEL1_BYTES_PER_TOKEN, + dtype=torch.uint8, + ) + values = torch.zeros(3, dsv4_sparse._FLASH_MLA_MODEL1_HEAD_DIM, dtype=torch.bfloat16) + values[:, dsv4_sparse._FLASH_MLA_MODEL1_NOPE_DIM] = torch.tensor( + [11.0, 12.0, 13.0], + dtype=torch.bfloat16, + ) + cu_num_pages = torch.tensor([0, 3], dtype=torch.int32) + cache_loc = torch.tensor([2, 0, 1], dtype=torch.int32) + + def scalar_page_lookup_should_not_run(*args, **kwargs): + raise AssertionError("scalar page lookup should not run") + + monkeypatch.setattr( + dsv4_sparse, + "_host_page_id_and_offset", + scalar_page_lookup_should_not_run, + ) + + dsv4_sparse._write_paged_cache_rows( + values, + cache, + seq_idx=0, + input_pos=1, + cu_num_pages_host=cu_num_pages, + cache_loc_host=cache_loc, + ) + + token_data_bytes = dsv4_sparse._FLASH_MLA_MODEL1_TOKEN_DATA_BYTES + rope_offset = dsv4_sparse._FLASH_MLA_MODEL1_NOPE_DIM + + def rope_value(page_id: int, page_offset: int) -> torch.Tensor: + flat_page = cache[page_id].view(-1) + start = page_offset * token_data_bytes + rope_offset + return flat_page[start : start + 2].view(torch.bfloat16) + + torch.testing.assert_close(rope_value(2, 1), torch.tensor([11.0], dtype=torch.bfloat16)) + torch.testing.assert_close(rope_value(0, 0), torch.tensor([12.0], dtype=torch.bfloat16)) + torch.testing.assert_close(rope_value(0, 1), torch.tensor([13.0], dtype=torch.bfloat16)) + assert not cache[1].any() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_flash_mla_model1_cache_pack_is_cuda_graph_capturable() -> None: + block_size = 4 + cache = torch.zeros( + 1, + block_size, + dsv4_sparse._FLASH_MLA_MODEL1_BYTES_PER_TOKEN, + dtype=torch.uint8, + device="cuda", + ) + values = torch.randn( + 2, + dsv4_sparse._FLASH_MLA_MODEL1_HEAD_DIM, + dtype=torch.bfloat16, + device="cuda", + ) + page_ids = torch.zeros(2, dtype=torch.int64, device="cuda") + page_offsets = torch.tensor([0, 1], dtype=torch.int64, device="cuda") + valid = torch.ones(2, dtype=torch.bool, device="cuda") + + dsv4_sparse._write_flash_mla_model1_rows_to_cache( + cache, + values, + page_ids, + page_offsets, + valid, + ) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + dsv4_sparse._write_flash_mla_model1_rows_to_cache( + cache, + values, + page_ids, + page_offsets, + valid, + ) + graph.replay() + torch.cuda.synchronize() + + flat = cache.view(torch.uint8).view(1, -1)[0] + token_data_bytes = dsv4_sparse._FLASH_MLA_MODEL1_TOKEN_DATA_BYTES + scale_base = block_size * token_data_bytes + assert flat[: token_data_bytes * 2].any() + assert flat[scale_base : scale_base + 2 * dsv4_sparse._FLASH_MLA_MODEL1_SCALE_BYTES].any() + + +def test_flash_mla_sparse_decode_backend_builds_swa_indices(monkeypatch) -> None: + num_decode = 2 + num_heads = 64 + head_dim = dsv4_sparse._FLASH_MLA_MODEL1_HEAD_DIM + tokens_per_block = 2 + metadata = _multi_decode_meta([2, 4], tokens_per_block=tokens_per_block) + q = torch.randn(num_decode, 1, num_heads, head_dim, dtype=torch.bfloat16) + kv = torch.randn(num_decode, 1, head_dim, dtype=torch.bfloat16) + attn_sink = torch.zeros(num_heads, dtype=torch.float32) + topk_idxs = torch.full((num_decode, 1, 1), -1, dtype=torch.int32) + swa_cache = torch.zeros( + 5, + tokens_per_block, + dsv4_sparse._FLASH_MLA_MODEL1_BYTES_PER_TOKEN, + dtype=torch.uint8, + ) + calls = [] + + def fake_flash_mla_sparse_decode( + q_decode, + swa_cache_arg, + attn_sink_arg, + indices, + topk_length, + softmax_scale, + extra_k_cache=None, + extra_indices=None, + extra_topk_length=None, + ): + calls.append( + { + "q_shape": tuple(q_decode.shape), + "cache_shape": tuple(swa_cache_arg.shape), + "attn_sink_shape": tuple(attn_sink_arg.shape), + "indices": indices.clone(), + "topk_length": topk_length.clone(), + "softmax_scale": softmax_scale, + "extra_k_cache": extra_k_cache, + "extra_indices": extra_indices, + "extra_topk_length": extra_topk_length, + } + ) + return q_decode + + monkeypatch.setattr(dsv4_sparse, "_flash_mla_sparse_decode", fake_flash_mla_sparse_decode) + + output = _run_flashmla_cached_sparse_attention( + q, + kv, + attn_sink, + topk_idxs, + metadata, + swa_cache, + softmax_scale=0.25, + window_size=3, + ) + + assert len(calls) == 1 + assert calls[0]["q_shape"] == (num_decode, num_heads, head_dim) + assert calls[0]["cache_shape"] == tuple(swa_cache.shape) + assert calls[0]["attn_sink_shape"] == (num_heads,) + assert calls[0]["softmax_scale"] == 0.25 + expected_prefix = torch.tensor([[[0, 1, 2]], [[6, 7, 8]]], dtype=torch.int32) + assert calls[0]["indices"].shape == ( + num_decode, + 1, + dsv4_sparse._FLASH_MLA_DECODE_TOPK_MULTIPLE, + ) + torch.testing.assert_close(calls[0]["indices"][:, :, :3], expected_prefix) + assert torch.all(calls[0]["indices"][:, :, 3:] == -1) + torch.testing.assert_close(calls[0]["topk_length"], torch.tensor([3, 3], dtype=torch.int32)) + assert calls[0]["extra_k_cache"] is None + torch.testing.assert_close( + output.reshape(num_decode, num_heads, head_dim), q.reshape(num_decode, num_heads, head_dim) + ) + + +def test_flashmla_sparse_backend_handles_mixed_prefill_decode(monkeypatch) -> None: + num_heads = 2 + head_dim = dsv4_sparse._FLASH_MLA_MODEL1_HEAD_DIM + tokens_per_block = 2 + batch_info_host = BatchInfo() + batch_info_host.update([1, 2, 0, 0, 1, 1]) + cu_num_pages, cache_loc, last_page_len = _page_meta( + [2, 1], + [0, 3], + [0, 1], + tokens_per_block=tokens_per_block, + ) + metadata = ( + batch_info_host.serialize(), + torch.tensor([2, 1], dtype=torch.int32), + torch.tensor([0, 3], dtype=torch.int32), + torch.tensor([0, 1], dtype=torch.int64), + torch.tensor([0, 2, 3], dtype=torch.int32), + cu_num_pages, + cache_loc, + last_page_len, + ) + q = torch.randn(1, 3, num_heads, head_dim, dtype=torch.bfloat16) + kv = torch.randn(1, 3, head_dim, dtype=torch.bfloat16) + attn_sink = torch.zeros(num_heads, dtype=torch.float32) + topk_idxs = torch.full((1, 3, 1), -1, dtype=torch.int32) + swa_cache = torch.zeros( + 3, + tokens_per_block, + dsv4_sparse._FLASH_MLA_MODEL1_BYTES_PER_TOKEN, + dtype=torch.uint8, + ) + calls = {"prefill": 0, "decode": []} + + def fake_flash_mla_sparse_prefill( + q_prefill, + kv_source, + attn_sink_arg, + topk_idxs_arg, + softmax_scale, + ): + del kv_source, attn_sink_arg, topk_idxs_arg, softmax_scale + calls["prefill"] += 1 + return q_prefill + 10 + + def fake_flash_mla_sparse_decode( + q_decode, + swa_cache_arg, + attn_sink_arg, + indices, + topk_length, + softmax_scale, + extra_k_cache=None, + extra_indices=None, + extra_topk_length=None, + ): + del swa_cache_arg, attn_sink_arg, softmax_scale, extra_k_cache, extra_indices + del extra_topk_length + calls["decode"].append( + { + "q": q_decode.clone(), + "indices": indices.clone(), + "topk_length": topk_length.clone(), + } + ) + return q_decode + 20 + + monkeypatch.setattr(dsv4_sparse, "_flash_mla_sparse_prefill", fake_flash_mla_sparse_prefill) + monkeypatch.setattr(dsv4_sparse, "_flash_mla_sparse_decode", fake_flash_mla_sparse_decode) + + output = _run_flashmla_cached_sparse_attention( + q, + kv, + attn_sink, + topk_idxs, + metadata, + swa_cache, + softmax_scale=0.25, + window_size=2, + ) + + assert calls["prefill"] == 1 + assert len(calls["decode"]) == 1 + assert calls["decode"][0]["q"].shape == (1, num_heads, head_dim) + assert calls["decode"][0]["topk_length"].tolist() == [2] + assert calls["decode"][0]["indices"].shape == (1, 1, 128) + torch.testing.assert_close(output[:, :2], q[:, :2] + 10) + torch.testing.assert_close(output[:, 2:3], q[:, 2:3] + 20) + + +def test_sparse_source_op_does_not_call_flashmla(monkeypatch) -> None: + def fail_if_called(*args, **kwargs): + del args, kwargs + raise AssertionError("semantic sparse attention must not call FlashMLA") + + monkeypatch.setattr(dsv4_sparse, "_flash_mla_sparse_prefill", fail_if_called) + + q = torch.randn(1, 2, 2, 4) + kv = torch.randn(1, 4, 4) + attn_sink = torch.randn(2) + topk_idxs = torch.tensor([[[0, 1], [2, 3]]], dtype=torch.int64) + + output = _run_sparse_attention(q, kv, attn_sink, topk_idxs, softmax_scale=0.5) + expected = dsv4_sparse._deepseek_v4_sparse_attention(q, kv, attn_sink, topk_idxs, 0.5) + torch.testing.assert_close(output, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_flashmla_sparse_backend_op_invokes_flashmla_kernel(monkeypatch) -> None: + capability = torch.cuda.get_device_capability() + if capability[0] < 9: + pytest.skip("FlashMLA sparse prefill requires Hopper or newer") + + flash_mla_sparse_fwd = dsv4_sparse._get_flash_mla_sparse_prefill() + if flash_mla_sparse_fwd is None: + pytest.skip("FlashMLA sparse prefill interface is not available") + + torch.manual_seed(123) + batch_size = 1 + seq_len = 2 + num_heads = 64 + head_dim = 512 + kv_rows = 256 + topk = 128 + softmax_scale = head_dim**-0.5 + device = torch.device("cuda") + + q = (torch.randn(batch_size, seq_len, num_heads, head_dim, device=device) / 8).to( + torch.bfloat16 + ) + kv = (torch.randn(batch_size, kv_rows, head_dim, device=device) / 8).to(torch.bfloat16) + attn_sink = torch.linspace(-0.2, 0.2, num_heads, device=device, dtype=torch.float32) + topk_idxs = torch.arange(topk, device=device, dtype=torch.int32).view(1, 1, topk) + topk_idxs = topk_idxs.expand(batch_size, seq_len, topk).contiguous() + metadata = _context_meta(seq_len) + swa_cache = torch.empty( + 1, + seq_len, + dsv4_sparse._FLASH_MLA_MODEL1_BYTES_PER_TOKEN, + device=device, + dtype=torch.uint8, + ) + + calls = [] + + def counted_flash_mla_sparse_fwd(*args, **kwargs): + q_flat, kv_flat, indices = args[:3] + calls.append((tuple(q_flat.shape), tuple(kv_flat.shape), tuple(indices.shape))) + return flash_mla_sparse_fwd(*args, **kwargs) + + monkeypatch.setattr( + dsv4_sparse, + "_get_flash_mla_sparse_prefill", + lambda: counted_flash_mla_sparse_fwd, + ) + + output = _run_flashmla_cached_sparse_attention( + q, + kv, + attn_sink, + topk_idxs, + metadata, + swa_cache, + softmax_scale, + ) + reference = dsv4_sparse._deepseek_v4_sparse_attention( + q, kv, attn_sink, topk_idxs, softmax_scale + ) + + assert calls == [((2, 64, 512), (256, 1, 512), (2, 1, 128))] + assert torch.allclose(output.float(), reference.float(), atol=5e-2, rtol=5e-2) + + +def test_sink_only_all_negative_topk_yields_finite_zero_output() -> None: + q = torch.randn(1, 2, 2, 4) + kv = torch.full((1, 5, 4), 1_000.0) + attn_sink = torch.tensor([-3.0, 2.0]) + topk_idxs = torch.full((1, 2, 4), -1, dtype=torch.int64) + + output = _run_sparse_attention(q, kv, attn_sink, topk_idxs, softmax_scale=0.5) + + assert torch.isfinite(output).all() + torch.testing.assert_close(output, torch.zeros_like(q), rtol=0, atol=0) + + +def test_duplicate_topk_indices_preserve_independent_probability_mass() -> None: + q = torch.tensor([[[[1.0, 0.0]]]]) + kv = torch.tensor([[[2.0, 0.0], [0.0, 2.0]]]) + attn_sink = torch.tensor([-20.0]) + topk_idxs = torch.tensor([[[0, 0, 1]]], dtype=torch.int64) + + output = _run_sparse_attention(q, kv, attn_sink, topk_idxs) + expected = _sparse_attention_reference(q, kv, attn_sink, topk_idxs, softmax_scale=1.0) + + assert_rmse_close(output, expected, rmse_ratio_tol=1e-6, msg="duplicate top-k: ") + assert output[0, 0, 0, 0] > output[0, 0, 0, 1] + + +def test_negative_indices_are_masked_before_softmax() -> None: + q = torch.tensor([[[[1.0, 0.0]]]]) + kv = torch.tensor([[[1000.0, 1000.0], [1.0, 0.0]]]) + attn_sink = torch.tensor([0.0]) + topk_idxs = torch.tensor([[[1, -1]]], dtype=torch.int32) + + output = _run_sparse_attention(q, kv, attn_sink, topk_idxs) + expected = _sparse_attention_reference(q, kv, attn_sink, topk_idxs, softmax_scale=1.0) + + assert_rmse_close(output, expected, rmse_ratio_tol=1e-6, msg="negative top-k mask: ") + assert output.abs().max() < 1.0 + + +def test_out_of_range_indices_are_masked_before_softmax() -> None: + q = torch.tensor([[[[1.0, 0.0]]]]) + kv = torch.tensor([[[1000.0, 1000.0], [1.0, 0.0]]]) + attn_sink = torch.tensor([0.0]) + topk_idxs = torch.tensor([[[1, 99]]], dtype=torch.int64) + masked_topk_idxs = torch.tensor([[[1, -1]]], dtype=torch.int64) + + output = _run_sparse_attention(q, kv, attn_sink, topk_idxs) + expected = _sparse_attention_reference(q, kv, attn_sink, masked_topk_idxs, softmax_scale=1.0) + + assert_rmse_close(output, expected, rmse_ratio_tol=1e-6, msg="high top-k mask: ") + assert output.abs().max() < 1.0 + + +def test_source_ratio0_matches_reference_for_mixed_patterns() -> None: + torch.manual_seed(11) + q = torch.randn(2, 4, 3, 6) + kv = torch.randn(2, 8, 6) + attn_sink = torch.tensor([-0.5, 0.25, 1.0]) + topk_idxs = torch.tensor( + [ + [[0, 1, -1, 1], [2, 2, 3, -1], [4, -1, -1, 5], [6, 0, 6, 1]], + [[7, 6, 5, 4], [3, -1, 3, 0], [-1, -1, -1, -1], [1, 2, 2, 7]], + ], + dtype=torch.int64, + ) + + output = _run_sparse_attention(q, kv, attn_sink, topk_idxs, softmax_scale=0.375) + expected = _sparse_attention_reference(q, kv, attn_sink, topk_idxs, softmax_scale=0.375) + + assert_rmse_close(output, expected, rmse_ratio_tol=1e-6, msg="mixed sparse attention: ") + + +def test_cached_ratio0_local_window_reads_past_kv_from_swa_cache() -> None: + q_prefill = torch.tensor([[[[1.0, 0.0]], [[0.0, 1.0]], [[1.0, 1.0]]]]) + kv_prefill = torch.tensor([[[1.0, 0.0], [0.0, 1.0], [2.0, 2.0]]]) + attn_sink = torch.tensor([-20.0]) + topk_prefill = torch.zeros(1, 3, 1, dtype=torch.int64) + swa_cache = torch.empty(1, 8, 2) + + _run_cached_sparse_attention( + q_prefill, + kv_prefill, + attn_sink, + topk_prefill, + _context_meta(seq_len=3), + swa_cache, + window_size=4, + ) + + q_decode = torch.tensor([[[[1.0, 0.5]]]]) + kv_decode = torch.tensor([[[3.0, -1.0]]]) + output = _run_cached_sparse_attention( + q_decode, + kv_decode, + attn_sink, + torch.zeros(1, 1, 1, dtype=torch.int64), + _decode_meta(input_pos=3), + swa_cache, + window_size=4, + ) + + expected_kv = torch.cat([kv_prefill, kv_decode], dim=1) + expected_topk = torch.tensor([[[0, 1, 2, 3]]], dtype=torch.int64) + expected = _sparse_attention_reference(q_decode, expected_kv, attn_sink, expected_topk, 1.0) + + torch.testing.assert_close(swa_cache[0, :4], expected_kv[0]) + assert_rmse_close(output, expected, rmse_ratio_tol=1e-6, msg="cached local window: ") + + +def test_cached_ratio0_decode_reads_and_writes_across_paged_boundary() -> None: + tokens_per_block = 2 + q_prefill = torch.tensor([[[[1.0, 0.0]], [[0.0, 1.0]], [[1.0, 1.0]]]]) + kv_prefill = torch.tensor([[[1.0, 0.0], [0.0, 1.0], [2.0, 2.0]]]) + attn_sink = torch.tensor([-20.0]) + topk_prefill = torch.zeros(1, 3, 1, dtype=torch.int64) + swa_cache, _, _, _ = _make_paged_sparse_attention_caches(4, tokens_per_block, 2, 0) + + _run_cached_sparse_attention( + q_prefill, + kv_prefill, + attn_sink, + topk_prefill, + _context_meta(seq_len=3, tokens_per_block=tokens_per_block), + swa_cache, + window_size=4, + ) + + q_decode = torch.tensor([[[[1.0, 0.5]]]]) + kv_decode = torch.tensor([[[3.0, -1.0]]]) + output = _run_cached_sparse_attention( + q_decode, + kv_decode, + attn_sink, + torch.zeros(1, 1, 1, dtype=torch.int64), + _decode_meta(input_pos=3, tokens_per_block=tokens_per_block), + swa_cache, + window_size=4, + ) + + expected_kv = torch.cat([kv_prefill, kv_decode], dim=1) + expected_topk = torch.tensor([[[0, 1, 2, 3]]], dtype=torch.int64) + expected = _sparse_attention_reference(q_decode, expected_kv, attn_sink, expected_topk, 1.0) + + torch.testing.assert_close(swa_cache[0], expected_kv[0, :2]) + torch.testing.assert_close(swa_cache[1], expected_kv[0, 2:4]) + assert_rmse_close(output, expected, rmse_ratio_tol=1e-6, msg="paged cached local window: ") + + +def test_cached_ratio0_flattened_prefill_uses_per_sequence_kv_slice() -> None: + q = torch.tensor([[[[1.0, 0.0]], [[0.0, 1.0]], [[1.0, 0.0]], [[0.0, 1.0]]]]) + kv = torch.tensor([[[1.0, 0.0], [0.0, 1.0], [10.0, 0.0], [0.0, 20.0]]]) + attn_sink = torch.tensor([-20.0]) + topk_idxs = torch.tensor([[[0], [1], [0], [1]]], dtype=torch.int64) + swa_cache = torch.empty(2, 8, 2) + + output = _run_cached_sparse_attention( + q, + kv, + attn_sink, + topk_idxs, + _multi_context_meta([2, 2]), + swa_cache, + window_size=2, + ) + + expected_seq0 = _run_sparse_attention(q[:, :2], kv[:, :2], attn_sink, topk_idxs[:, :2]) + expected_seq1 = _run_sparse_attention(q[:, 2:], kv[:, 2:], attn_sink, topk_idxs[:, 2:]) + expected = torch.cat((expected_seq0, expected_seq1), dim=1) + + torch.testing.assert_close(swa_cache[0, :2], kv[0, :2]) + torch.testing.assert_close(swa_cache[1, :2], kv[0, 2:]) + assert_rmse_close(output, expected, rmse_ratio_tol=1e-6, msg="flattened prefill: ") + + +def test_cached_ratio0_prefill_with_window_size_still_honors_topk_idxs() -> None: + q = torch.tensor([[[[1.0, 0.0]], [[0.0, 1.0]]]]) + kv = torch.tensor([[[1.0, 0.0], [0.0, 2.0]]]) + attn_sink = torch.tensor([-20.0]) + topk_idxs = torch.tensor([[[0], [1]]], dtype=torch.int64) + swa_cache = torch.empty(1, 8, 2) + + output = _run_cached_sparse_attention( + q, + kv, + attn_sink, + topk_idxs, + _context_meta(seq_len=2), + swa_cache, + window_size=2, + ) + expected = _run_sparse_attention(q, kv, attn_sink, topk_idxs) + + assert_rmse_close(output, expected, rmse_ratio_tol=1e-6, msg="cached prefill topk: ") + + +def test_cached_ratio0_topk_mode_preserves_duplicates_and_negative_mask() -> None: + q = torch.tensor([[[[1.0, 0.0]]]]) + kv = torch.tensor([[[2.0, 0.0], [100.0, 100.0], [1.0, 1.0]]]) + attn_sink = torch.tensor([-20.0]) + topk_idxs = torch.tensor([[[0, 0, -1, 2]]], dtype=torch.int64) + swa_cache = torch.empty(1, 8, 2) + + output = _run_cached_sparse_attention( + q, + kv, + attn_sink, + topk_idxs, + _context_meta(seq_len=1), + swa_cache, + ) + expected = _run_sparse_attention(q, kv, attn_sink, topk_idxs) + + assert_rmse_close(output, expected, rmse_ratio_tol=1e-6, msg="cached duplicate/mask: ") + assert output[0, 0, 0, 0] > output[0, 0, 0, 1] + + +def test_cached_ratio0_topk_decode_without_window_uses_cache_positions() -> None: + q_prefill = torch.zeros(1, 1, 1, 2) + kv_prefill = torch.tensor([[[2.0, 0.0]]]) + attn_sink = torch.tensor([-20.0]) + swa_cache = torch.empty(1, 8, 2) + + _run_cached_sparse_attention( + q_prefill, + kv_prefill, + attn_sink, + torch.zeros(1, 1, 1, dtype=torch.int64), + _context_meta(seq_len=1), + swa_cache, + ) + + q_decode = torch.tensor([[[[1.0, 0.0]]]]) + kv_decode = torch.tensor([[[3.0, 0.0]]]) + topk_decode = torch.tensor([[[0, 1]]], dtype=torch.int64) + + output = _run_cached_sparse_attention( + q_decode, + kv_decode, + attn_sink, + topk_decode, + _decode_meta(input_pos=1), + swa_cache, + ) + + expected_kv = torch.cat((kv_prefill, kv_decode), dim=1) + expected = _sparse_attention_reference(q_decode, expected_kv, attn_sink, topk_decode, 1.0) + + torch.testing.assert_close(swa_cache[0, :2], expected_kv[0]) + assert_rmse_close(output, expected, rmse_ratio_tol=1e-6, msg="cached top-k decode: ") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA") +def test_cached_ratio0_decode_cuda_graph_replay_uses_runtime_slot_and_input_pos() -> None: + q = torch.zeros(1, 1, 1, 2, device="cuda") + kv = torch.tensor([[[2.0, 0.0]]], device="cuda") + attn_sink = torch.tensor([-20.0], device="cuda") + topk_idxs = torch.tensor([[[1]]], dtype=torch.int64, device="cuda") + metadata = _cuda_decode_meta(input_pos=1, slot_idx=0) + swa_cache = torch.full((2, 6, 2), -99.0, device="cuda") + mhc_cache = torch.empty_like(swa_cache) + compressor_kv_cache = q.new_empty(2, 6, 0) + compressor_gate_cache = q.new_empty(2, 6, 0) + indexer_compressor_kv_cache = q.new_empty(2, 6, 0) + indexer_compressor_gate_cache = q.new_empty(2, 6, 0) + out = torch.empty_like(q) + empty_sparse_args = _empty_sparse_attention_tensors(q, kv) + + def run_op() -> None: + torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention_with_cache( + q, + kv, + attn_sink, + topk_idxs, + *empty_sparse_args, + *metadata, + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + 1.0, + None, + 0, + None, + 1e-6, + None, + out=out, + ) + + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + run_op() + stream.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_op() + + kv.copy_(torch.tensor([[[7.0, 3.0]]], device="cuda")) + topk_idxs.copy_(torch.tensor([[[3]]], dtype=torch.int64, device="cuda")) + metadata[2].copy_(torch.tensor([3], dtype=torch.int32, device="cuda")) + metadata[3].copy_(torch.tensor([1], dtype=torch.int64, device="cuda")) + metadata[6].copy_(torch.tensor([1], dtype=torch.int32, device="cuda")) + swa_cache[1, 3].fill_(-11.0) + + graph.replay() + torch.cuda.synchronize() + + expected = _sparse_attention_reference( + q, + kv, + attn_sink, + torch.zeros_like(topk_idxs), + softmax_scale=1.0, + ) + torch.testing.assert_close(swa_cache[1, 3], kv[0, 0]) + torch.testing.assert_close(out, expected) + + +def test_cached_ratio0_sink_only_negative_topk_yields_zero_output() -> None: + q = torch.tensor([[[[1.0, 0.0]]]]) + kv = torch.tensor([[[5.0, 5.0]]]) + attn_sink = torch.tensor([3.0]) + topk_idxs = torch.full((1, 1, 3), -1, dtype=torch.int64) + swa_cache = torch.empty(1, 4, 2) + + output = _run_cached_sparse_attention( + q, + kv, + attn_sink, + topk_idxs, + _decode_meta(input_pos=0), + swa_cache, + ) + + assert torch.isfinite(output).all() + torch.testing.assert_close(output, torch.zeros_like(q), rtol=0, atol=0) + + +def test_cached_ratio0_out_buffer_returns_dummy_and_fills_output() -> None: + q = torch.tensor([[[[1.0, 0.0]], [[0.0, 1.0]], [[7.0, 7.0]]]]) + kv = torch.tensor([[[2.0, 0.0], [0.0, 2.0], [100.0, 100.0]]]) + attn_sink = torch.tensor([-20.0]) + topk_idxs = torch.tensor([[[0], [1], [2]]], dtype=torch.int64) + + expected = _run_cached_sparse_attention( + q, + kv, + attn_sink, + topk_idxs, + _context_meta(seq_len=2), + torch.empty(1, 8, 2), + window_size=2, + ) + + out = torch.full_like(q, 123.0) + result = _run_cached_sparse_attention( + q, + kv, + attn_sink, + topk_idxs, + _context_meta(seq_len=2), + torch.empty(1, 8, 2), + window_size=2, + out=out, + ) + + assert result.numel() == 0 + torch.testing.assert_close(out, expected) + torch.testing.assert_close(out[:, 2:], torch.zeros_like(out[:, 2:]), rtol=0, atol=0) + + +@pytest.mark.parametrize("compress_ratio", [4, 128]) +def test_source_matches_expanded_sparse_construction(compress_ratio: int) -> None: + torch.manual_seed(123 + compress_ratio) + seq_len = compress_ratio + q = torch.randn(1, seq_len, 1, 8) + kv = torch.randn(1, seq_len, 8) + attn_sink = torch.tensor([-0.25]) + ( + compressor_kv, + compressor_gate, + compressed_kv, + compressor, + cos_table, + sin_table, + position_ids, + ) = _compressor_case(compress_ratio, seq_len) + topk_idxs = torch.arange(seq_len + compressor.max_compressed_len).view(1, 1, -1) + topk_idxs = topk_idxs.expand(1, seq_len, -1).to(torch.int64) + + output = _run_sparse_attention_with_compressor( + q, + kv, + attn_sink, + topk_idxs, + compressor_kv, + compressor_gate, + compressor, + cos_table, + sin_table, + position_ids, + compress_ratio=compress_ratio, + ) + expected = _run_sparse_attention( + q, + torch.cat((kv, compressed_kv), dim=1), + attn_sink, + topk_idxs, + ) + + assert_rmse_close( + output, + expected, + rmse_ratio_tol=1e-6, + msg=f"source ratio-{compress_ratio}: ", + ) + + +def test_cached_ratio0_prefill_and_decode_match_source() -> None: + torch.manual_seed(37) + total_len = 3 + prefill_len = 2 + q = torch.randn(1, total_len, 1, 8) + kv = torch.randn(1, total_len, 8) + attn_sink = torch.tensor([-0.25]) + ( + compressor_kv, + compressor_gate, + _, + compressor, + cos_table, + sin_table, + position_ids, + ) = _compressor_case(4, total_len) + swa_cache, mhc_cache, compressor_kv_cache, compressor_gate_cache = ( + _make_sparse_attention_caches( + total_len, + kv.shape[-1], + compressor_kv.shape[-1], + fill_value=777.0, + ) + ) + + topk_prefill = torch.tensor([[[0, 1], [1, 0]]], dtype=torch.int64) + output_prefill = _run_cached_sparse_attention_with_compressor( + q[:, :prefill_len], + kv[:, :prefill_len], + attn_sink, + topk_prefill, + compressor_kv[:, :prefill_len], + compressor_gate[:, :prefill_len], + compressor, + cos_table, + sin_table, + position_ids[:, :prefill_len], + _context_meta(seq_len=prefill_len), + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + window_size=None, + compress_ratio=0, + ) + expected_prefill = _run_sparse_attention_with_compressor( + q[:, :prefill_len], + kv[:, :prefill_len], + attn_sink, + topk_prefill, + compressor_kv[:, :prefill_len], + compressor_gate[:, :prefill_len], + compressor, + cos_table, + sin_table, + position_ids[:, :prefill_len], + window_size=None, + compress_ratio=0, + ) + + assert_rmse_close( + output_prefill, + expected_prefill, + rmse_ratio_tol=1e-6, + msg="cached ratio-0 prefill: ", + ) + + topk_decode = torch.tensor([[[0, 2]]], dtype=torch.int64) + output_decode = _run_cached_sparse_attention_with_compressor( + q[:, prefill_len:], + kv[:, prefill_len:], + attn_sink, + topk_decode, + compressor_kv[:, prefill_len:], + compressor_gate[:, prefill_len:], + compressor, + cos_table, + sin_table, + position_ids[:, prefill_len:], + _decode_meta(input_pos=prefill_len), + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + window_size=None, + compress_ratio=0, + ) + expected_decode = _run_sparse_attention_with_compressor( + q[:, prefill_len:], + kv, + attn_sink, + topk_decode, + compressor_kv, + compressor_gate, + compressor, + cos_table, + sin_table, + position_ids, + window_size=None, + compress_ratio=0, + ) + + torch.testing.assert_close(swa_cache[0, :total_len], kv[0]) + assert_rmse_close( + output_decode, + expected_decode, + rmse_ratio_tol=1e-6, + msg="cached ratio-0 decode: ", + ) + + +@pytest.mark.parametrize("compress_ratio", [4, 128]) +def test_cached_compressed_prefill_matches_source(compress_ratio: int) -> None: + torch.manual_seed(311 + compress_ratio) + seq_len = compress_ratio + q = torch.randn(1, seq_len, 1, 8) + kv = torch.randn(1, seq_len, 8) + attn_sink = torch.tensor([-0.25]) + ( + compressor_kv, + compressor_gate, + _, + compressor, + cos_table, + sin_table, + position_ids, + ) = _compressor_case(compress_ratio, seq_len) + swa_cache, mhc_cache, compressor_kv_cache, compressor_gate_cache = ( + _make_sparse_attention_caches( + seq_len, + kv.shape[-1], + compressor_kv.shape[-1], + fill_value=777.0, + ) + ) + topk_idxs = torch.arange(seq_len + compressor.max_compressed_len).view(1, 1, -1) + topk_idxs = topk_idxs.expand(1, seq_len, -1).to(torch.int64) + + output = _run_cached_sparse_attention_with_compressor( + q, + kv, + attn_sink, + topk_idxs, + compressor_kv, + compressor_gate, + compressor, + cos_table, + sin_table, + position_ids, + _context_meta(seq_len=seq_len), + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + compress_ratio=compress_ratio, + ) + expected = _run_sparse_attention_with_compressor( + q, + kv, + attn_sink, + topk_idxs, + compressor_kv, + compressor_gate, + compressor, + cos_table, + sin_table, + position_ids, + compress_ratio=compress_ratio, + ) + + assert_rmse_close( + output, + expected, + rmse_ratio_tol=1e-6, + msg=f"cached ratio-{compress_ratio} prefill: ", + ) + + +def test_cached_ratio128_token_input_pos_matches_full_source() -> None: + torch.manual_seed(183) + compress_ratio = 128 + total_len = 128 + compressed_capacity_tokens = 256 + prefill_len = total_len - 1 + window_size = 4 + q = torch.randn(1, total_len, 1, 8) + kv = torch.randn(1, total_len, 8) + attn_sink = torch.tensor([-0.5]) + ( + compressor_kv, + compressor_gate, + _, + compressor, + cos_table, + sin_table, + position_ids, + ) = _compressor_case( + compress_ratio, + total_len, + compressed_capacity_tokens=compressed_capacity_tokens, + ) + state_dim = compressor_kv.shape[-1] + swa_cache, mhc_cache, compressor_kv_cache, compressor_gate_cache = ( + _make_sparse_attention_caches( + compressed_capacity_tokens, + kv.shape[-1], + state_dim, + fill_value=777.0, + ) + ) + + topk_prefill = _visible_source_topk( + prefill_len, + 0, + prefill_len, + window_size, + compress_ratio, + compressor.max_compressed_len, + q.device, + ) + _run_cached_sparse_attention_with_compressor( + q[:, :prefill_len], + kv[:, :prefill_len], + attn_sink, + topk_prefill, + compressor_kv[:, :prefill_len], + compressor_gate[:, :prefill_len], + compressor, + cos_table, + sin_table, + position_ids[:, :prefill_len], + _context_meta(seq_len=prefill_len), + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + window_size=window_size, + compress_ratio=compress_ratio, + ) + + output = _run_cached_sparse_attention_with_compressor( + q[:, prefill_len:], + kv[:, prefill_len:], + attn_sink, + torch.zeros(1, 1, 1, dtype=torch.int64), + compressor_kv[:, prefill_len:], + compressor_gate[:, prefill_len:], + compressor, + cos_table, + sin_table, + position_ids[:, prefill_len:], + _decode_meta(input_pos=prefill_len), + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + window_size=window_size, + compress_ratio=compress_ratio, + ) + expected_topk = _visible_source_topk( + 1, + prefill_len, + total_len, + window_size, + compress_ratio, + compressor.max_compressed_len, + q.device, + ) + expected = _run_sparse_attention_with_compressor( + q[:, prefill_len:], + kv, + attn_sink, + expected_topk, + compressor_kv, + compressor_gate, + compressor, + cos_table, + sin_table, + position_ids, + window_size=window_size, + compress_ratio=compress_ratio, + ) + + torch.testing.assert_close(swa_cache[0, :total_len], kv[0]) + assert_rmse_close( + output, + expected, + rmse_ratio_tol=1e-6, + msg="cached token input_pos ratio-128: ", + ) + + +def test_cached_ratio128_decode_uses_offset_position_ids_for_compressed_row() -> None: + torch.manual_seed(229) + compress_ratio = 128 + total_len = 128 + compressed_capacity_tokens = 256 + prefill_len = total_len - 1 + position_offset = 17 + window_size = 4 + q = torch.randn(1, total_len, 1, 8) + kv = torch.randn(1, total_len, 8) + attn_sink = torch.tensor([-0.5]) + ( + compressor_kv, + compressor_gate, + _, + compressor, + _, + _, + position_ids, + ) = _compressor_case( + compress_ratio, + total_len, + compressed_capacity_tokens=compressed_capacity_tokens, + ) + cos_table, sin_table = _rope_tables( + position_offset + compressed_capacity_tokens, + compressor.rope_head_dim, + ) + position_ids = position_ids + position_offset + state_dim = compressor_kv.shape[-1] + swa_cache, mhc_cache, compressor_kv_cache, compressor_gate_cache = ( + _make_sparse_attention_caches( + compressed_capacity_tokens, + kv.shape[-1], + state_dim, + fill_value=777.0, + ) + ) + + topk_prefill = _visible_source_topk( + prefill_len, + 0, + prefill_len, + window_size, + compress_ratio, + compressor.max_compressed_len, + q.device, + ) + _run_cached_sparse_attention_with_compressor( + q[:, :prefill_len], + kv[:, :prefill_len], + attn_sink, + topk_prefill, + compressor_kv[:, :prefill_len], + compressor_gate[:, :prefill_len], + compressor, + cos_table, + sin_table, + position_ids[:, :prefill_len], + _context_meta(seq_len=prefill_len), + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + window_size=window_size, + compress_ratio=compress_ratio, + ) + + output = _run_cached_sparse_attention_with_compressor( + q[:, prefill_len:], + kv[:, prefill_len:], + attn_sink, + torch.zeros(1, 1, 1, dtype=torch.int64), + compressor_kv[:, prefill_len:], + compressor_gate[:, prefill_len:], + compressor, + cos_table, + sin_table, + position_ids[:, prefill_len:], + _decode_meta(input_pos=prefill_len), + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + window_size=window_size, + compress_ratio=compress_ratio, + ) + expected_topk = _visible_source_topk( + 1, + prefill_len, + total_len, + window_size, + compress_ratio, + compressor.max_compressed_len, + q.device, + ) + expected = _run_sparse_attention_with_compressor( + q[:, prefill_len:], + kv, + attn_sink, + expected_topk, + compressor_kv, + compressor_gate, + compressor, + cos_table, + sin_table, + position_ids, + window_size=window_size, + compress_ratio=compress_ratio, + ) + + assert_rmse_close( + output, + expected, + rmse_ratio_tol=1e-6, + msg="cached ratio-128 offset position_ids decode: ", + ) + + +def test_cached_ratio128_multi_decode_metadata_matches_source_and_writes_slots() -> None: + torch.manual_seed(217) + batch_size = 2 + compress_ratio = 128 + total_len = 128 + prefill_len = total_len - 1 + window_size = 4 + compressed_capacity_tokens = 256 + q = torch.randn(batch_size, total_len, 1, 8) + kv = torch.randn(batch_size, total_len, 8) + attn_sink = torch.tensor([-0.5]) + ( + compressor_kv, + compressor_gate, + compressed_kv, + compressor, + cos_table, + sin_table, + position_ids, + ) = _compressor_case( + compress_ratio, + total_len, + compressed_capacity_tokens=compressed_capacity_tokens, + batch_size=batch_size, + ) + swa_cache, mhc_cache, compressor_kv_cache, compressor_gate_cache = ( + _make_sparse_attention_caches( + compressed_capacity_tokens, + kv.shape[-1], + compressor_kv.shape[-1], + fill_value=777.0, + num_slots=batch_size, + ) + ) + + topk_prefill = _visible_source_topk( + prefill_len, + 0, + prefill_len, + window_size, + compress_ratio, + compressor.max_compressed_len, + q.device, + ).expand(batch_size, -1, -1) + _run_cached_sparse_attention_with_compressor( + q[:, :prefill_len], + kv[:, :prefill_len], + attn_sink, + topk_prefill, + compressor_kv[:, :prefill_len], + compressor_gate[:, :prefill_len], + compressor, + cos_table, + sin_table, + position_ids[:, :prefill_len], + _multi_context_meta([prefill_len, prefill_len]), + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + window_size=window_size, + compress_ratio=compress_ratio, + ) + torch.testing.assert_close(mhc_cache[:, 0], torch.full_like(mhc_cache[:, 0], 777.0)) + + output = _run_cached_sparse_attention_with_compressor( + q[:, prefill_len:], + kv[:, prefill_len:], + attn_sink, + torch.zeros(batch_size, 1, 1, dtype=torch.int64), + compressor_kv[:, prefill_len:], + compressor_gate[:, prefill_len:], + compressor, + cos_table, + sin_table, + position_ids[:, prefill_len:], + _multi_decode_meta([prefill_len, prefill_len]), + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + window_size=window_size, + compress_ratio=compress_ratio, + ) + expected_topk = _visible_source_topk( + 1, + prefill_len, + total_len, + window_size, + compress_ratio, + compressor.max_compressed_len, + q.device, + ).expand(batch_size, -1, -1) + expected = _run_sparse_attention_with_compressor( + q[:, prefill_len:], + kv, + attn_sink, + expected_topk, + compressor_kv, + compressor_gate, + compressor, + cos_table, + sin_table, + position_ids, + window_size=window_size, + compress_ratio=compress_ratio, + ) + + for slot_idx in range(batch_size): + torch.testing.assert_close(swa_cache[slot_idx, :total_len], kv[slot_idx]) + torch.testing.assert_close(mhc_cache[slot_idx, 0], compressed_kv[slot_idx, 0]) + torch.testing.assert_close( + mhc_cache[slot_idx, compress_ratio], + torch.full_like(mhc_cache[slot_idx, compress_ratio], 777.0), + ) + assert_rmse_close( + output, + expected, + rmse_ratio_tol=1e-6, + msg="cached ratio-128 multi decode: ", + ) + + +def test_cached_ratio4_decode_matches_source_with_learned_indexer_topk() -> None: + torch.manual_seed(44) + compress_ratio = 4 + prefill_len = 7 + total_len = 8 + window_size = 4 + q = torch.randn(1, total_len, 1, 8) + kv = torch.randn(1, total_len, 8) + attn_sink = torch.tensor([-0.5]) + ( + compressor_kv, + compressor_gate, + compressed_kv, + compressor, + cos_table, + sin_table, + position_ids, + ) = _compressor_case(compress_ratio, total_len) + indexer_config = DeepseekV4Config( + hidden_size=16, + num_hidden_layers=1, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=8, + q_lora_rank=8, + qk_rope_head_dim=4, + index_n_heads=1, + index_head_dim=8, + index_topk=1, + compress_ratios=(compress_ratio,), + ad_compress_max_seq_len=total_len, + ad_rope_cache_len=total_len, + ) + indexer = DeepseekV4Indexer(indexer_config, compress_ratio).eval() + hidden_states = torch.randn(1, total_len, indexer_config.hidden_size) + q_lora = torch.randn(1, total_len, indexer_config.q_lora_rank) + cos = cos_table[position_ids] + sin = sin_table[position_ids] + indexer_q, indexer_weights, indexer_compressor_kv, indexer_compressor_gate = indexer.project( + hidden_states, + q_lora, + cos, + sin, + ) + compressed_idxs_prefill = indexer( + hidden_states[:, :prefill_len], + q_lora[:, :prefill_len], + cos[:, :prefill_len], + sin[:, :prefill_len], + cos_table, + sin_table, + position_ids[:, :prefill_len], + prefill_len, + ) + compressed_idxs_decode = indexer( + hidden_states, + q_lora, + cos, + sin, + cos_table, + sin_table, + position_ids, + total_len, + )[:, prefill_len:] + swa_cache, mhc_cache, compressor_kv_cache, compressor_gate_cache = ( + _make_sparse_attention_caches( + 8, + kv.shape[-1], + compressor_kv.shape[-1], + fill_value=777.0, + ) + ) + indexer_compressor_kv_cache = torch.full( + (1, 8, indexer_compressor_kv.shape[-1]), + 777.0, + dtype=indexer_compressor_kv.dtype, + ) + indexer_compressor_gate_cache = torch.full_like(indexer_compressor_kv_cache, 777.0) + + local_prefill = _visible_source_topk( + prefill_len, + 0, + prefill_len, + window_size, + compress_ratio, + 0, + q.device, + ) + _run_cached_sparse_attention_with_compressor( + q[:, :prefill_len], + kv[:, :prefill_len], + attn_sink, + torch.cat((local_prefill, compressed_idxs_prefill), dim=-1), + compressor_kv[:, :prefill_len], + compressor_gate[:, :prefill_len], + compressor, + cos_table, + sin_table, + position_ids[:, :prefill_len], + _context_meta(seq_len=prefill_len), + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + window_size=window_size, + compress_ratio=compress_ratio, + indexer_q=indexer_q[:, :prefill_len], + indexer_weights=indexer_weights[:, :prefill_len], + indexer_compressor_kv=indexer_compressor_kv[:, :prefill_len], + indexer_compressor_gate=indexer_compressor_gate[:, :prefill_len], + indexer_compressor_ape=indexer.compressor.ape, + indexer_compressor_norm_weight=indexer.compressor.norm.weight, + indexer_compressor_kv_cache=indexer_compressor_kv_cache, + indexer_compressor_gate_cache=indexer_compressor_gate_cache, + ) + torch.testing.assert_close(mhc_cache[0, 0], compressed_kv[0, 0]) + + local_decode = _visible_source_topk( + 1, + prefill_len, + total_len, + window_size, + compress_ratio, + 0, + q.device, + ) + expected_topk = torch.cat((local_decode, compressed_idxs_decode), dim=-1) + output = _run_cached_sparse_attention_with_compressor( + q[:, prefill_len:], + kv[:, prefill_len:], + attn_sink, + torch.zeros(1, 1, window_size + indexer.index_topk, dtype=torch.int64), + compressor_kv[:, prefill_len:], + compressor_gate[:, prefill_len:], + compressor, + cos_table, + sin_table, + position_ids[:, prefill_len:], + _decode_meta(input_pos=prefill_len), + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + window_size=window_size, + compress_ratio=compress_ratio, + indexer_q=indexer_q[:, prefill_len:], + indexer_weights=indexer_weights[:, prefill_len:], + indexer_compressor_kv=indexer_compressor_kv[:, prefill_len:], + indexer_compressor_gate=indexer_compressor_gate[:, prefill_len:], + indexer_compressor_ape=indexer.compressor.ape, + indexer_compressor_norm_weight=indexer.compressor.norm.weight, + indexer_compressor_kv_cache=indexer_compressor_kv_cache, + indexer_compressor_gate_cache=indexer_compressor_gate_cache, + ) + expected = _run_sparse_attention_with_compressor( + q[:, prefill_len:], + kv, + attn_sink, + expected_topk, + compressor_kv, + compressor_gate, + compressor, + cos_table, + sin_table, + position_ids, + window_size=window_size, + compress_ratio=compress_ratio, + ) + + all_visible_topk = _visible_source_topk( + 1, + prefill_len, + total_len, + window_size, + compress_ratio, + compressor.max_compressed_len, + q.device, + ) + all_visible = _run_sparse_attention_with_compressor( + q[:, prefill_len:], + kv, + attn_sink, + all_visible_topk, + compressor_kv, + compressor_gate, + compressor, + cos_table, + sin_table, + position_ids, + window_size=window_size, + compress_ratio=compress_ratio, + ) + + torch.testing.assert_close(mhc_cache[0, compress_ratio], compressed_kv[0, 1]) + assert_rmse_close(output, expected, rmse_ratio_tol=1e-6, msg="cached ratio-4 indexer: ") + assert not torch.allclose(output, all_visible, rtol=1e-6, atol=1e-6) + + +def test_cached_ratio4_indexer_runtime_all_reduce_flips_selected_row( + monkeypatch: pytest.MonkeyPatch, +) -> None: + compress_ratio = 4 + window_size = 4 + max_seq_len = 12 + head_dim = 8 + index_head_dim = 2 + indexer_state_dim = 2 * index_head_dim + input_pos = 8 + q = torch.zeros(1, 1, 1, head_dim) + q[0, 0, 0, 0] = 1.0 + kv = torch.zeros(1, 1, head_dim) + attn_sink = torch.tensor([-20.0]) + compressor_config = DeepseekV4Config( + hidden_size=16, + num_hidden_layers=1, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=head_dim, + q_lora_rank=8, + qk_rope_head_dim=0, + compress_ratios=(compress_ratio,), + ad_compress_max_seq_len=max_seq_len, + ad_rope_cache_len=max_seq_len, + ) + compressor = DeepseekV4Compressor(compressor_config, compress_ratio, head_dim).eval() + compressor_kv = torch.zeros(1, 1, 2 * head_dim) + compressor_gate = torch.zeros_like(compressor_kv) + cos_table, sin_table = _rope_tables(max_seq_len, compressor.rope_head_dim) + position_ids = torch.tensor([[input_pos]], dtype=torch.int64) + swa_cache, mhc_cache, compressor_kv_cache, compressor_gate_cache = ( + _make_sparse_attention_caches( + max_seq_len, + head_dim, + compressor_kv.shape[-1], + ) + ) + mhc_cache[0, 0, 0] = 4.0 + mhc_cache[0, compress_ratio, 1] = 4.0 + + indexer_q = torch.tensor([[[[0.0, 1.0]]]]) + indexer_weights = torch.ones(1, 1, 1) + indexer_compressor_kv = torch.zeros(1, 1, indexer_state_dim) + indexer_compressor_gate = torch.zeros_like(indexer_compressor_kv) + indexer_compressor_kv_cache = torch.zeros(1, max_seq_len, indexer_state_dim) + indexer_compressor_gate_cache = torch.zeros_like(indexer_compressor_kv_cache) + indexer_compressor_kv_cache[0, :4, :index_head_dim] = torch.tensor([0.0, 1.0]) + indexer_compressor_kv_cache[0, :4, index_head_dim:] = torch.tensor([1.0, 0.0]) + indexer_compressor_kv_cache[0, 4:8, index_head_dim:] = torch.tensor([0.0, 1.0]) + indexer_compressor_ape = torch.zeros(compress_ratio, indexer_state_dim) + indexer_compressor_norm_weight = torch.ones(index_head_dim) + all_reduce_inputs: list[torch.Tensor] = [] + + def fake_all_reduce( + tensor: torch.Tensor, + op: object = None, + group: object = None, + async_op: bool = False, + ) -> None: + assert op == dsv4_sparse.dist_common.ReduceOp.SUM + assert group is None + assert not async_op + all_reduce_inputs.append(tensor.clone()) + tensor.copy_(torch.tensor([[0.0, 100.0, 1000.0]], dtype=tensor.dtype, device=tensor.device)) + + monkeypatch.setattr(dsv4_sparse.dist_common, "is_initialized", lambda: True) + monkeypatch.setattr(dsv4_sparse.dist_common, "get_world_size", lambda: 2) + monkeypatch.setattr(dsv4_sparse.dist_common, "all_reduce", fake_all_reduce) + + output = _run_cached_sparse_attention_with_compressor( + q, + kv, + attn_sink, + torch.zeros(1, 1, window_size + 1, dtype=torch.int64), + compressor_kv, + compressor_gate, + compressor, + cos_table, + sin_table, + position_ids, + _decode_meta(input_pos=input_pos), + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + window_size=window_size, + compress_ratio=compress_ratio, + indexer_q=indexer_q, + indexer_weights=indexer_weights, + indexer_compressor_kv=indexer_compressor_kv, + indexer_compressor_gate=indexer_compressor_gate, + indexer_compressor_ape=indexer_compressor_ape, + indexer_compressor_norm_weight=indexer_compressor_norm_weight, + indexer_compressor_kv_cache=indexer_compressor_kv_cache, + indexer_compressor_gate_cache=indexer_compressor_gate_cache, + ) + + assert len(all_reduce_inputs) == 1 + assert all_reduce_inputs[0][0, 0] > all_reduce_inputs[0][0, 1] + expected_topk = torch.arange(window_size + 1, dtype=torch.int64).view(1, 1, -1) + local_window = swa_cache[:, input_pos - window_size + 1 : input_pos + 1] + expected = _run_sparse_attention( + q, + torch.cat((local_window, mhc_cache[:, compress_ratio : compress_ratio + 1]), dim=1), + attn_sink, + expected_topk, + ) + without_reduce = _run_sparse_attention( + q, + torch.cat((local_window, mhc_cache[:, 0:1]), dim=1), + attn_sink, + expected_topk, + ) + torch.testing.assert_close(output, expected) + assert not torch.allclose(output, without_reduce, rtol=1e-6, atol=1e-6) + + +def test_cached_ratio128_emits_boundary_row_and_hides_future_rows() -> None: + torch.manual_seed(128) + compress_ratio = 128 + total_len = 128 + prefill_len = total_len - 1 + window_size = 4 + q = torch.randn(1, total_len, 1, 8) + kv = torch.randn(1, total_len, 8) + attn_sink = torch.tensor([-0.5]) + ( + compressor_kv, + compressor_gate, + compressed_kv, + compressor, + cos_table, + sin_table, + position_ids, + ) = _compressor_case( + compress_ratio, + total_len, + compressed_capacity_tokens=256, + ) + swa_cache, mhc_cache, compressor_kv_cache, compressor_gate_cache = ( + _make_sparse_attention_caches( + 256, + kv.shape[-1], + compressor_kv.shape[-1], + fill_value=777.0, + ) + ) + + _run_cached_sparse_attention_with_compressor( + q[:, :prefill_len], + kv[:, :prefill_len], + attn_sink, + _visible_source_topk( + prefill_len, + 0, + prefill_len, + window_size, + compress_ratio, + compressor.max_compressed_len, + q.device, + ), + compressor_kv[:, :prefill_len], + compressor_gate[:, :prefill_len], + compressor, + cos_table, + sin_table, + position_ids[:, :prefill_len], + _context_meta(seq_len=prefill_len), + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + window_size=window_size, + compress_ratio=compress_ratio, + ) + torch.testing.assert_close(mhc_cache[0, 0], torch.full_like(mhc_cache[0, 0], 777.0)) + + output = _run_cached_sparse_attention_with_compressor( + q[:, prefill_len:], + kv[:, prefill_len:], + attn_sink, + torch.zeros(1, 1, 1, dtype=torch.int64), + compressor_kv[:, prefill_len:], + compressor_gate[:, prefill_len:], + compressor, + cos_table, + sin_table, + position_ids[:, prefill_len:], + _decode_meta(input_pos=prefill_len), + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + window_size=window_size, + compress_ratio=compress_ratio, + ) + expected_topk = _visible_source_topk( + 1, + prefill_len, + total_len, + window_size, + compress_ratio, + compressor.max_compressed_len, + q.device, + ) + expected = _run_sparse_attention_with_compressor( + q[:, prefill_len:], + kv, + attn_sink, + expected_topk, + compressor_kv, + compressor_gate, + compressor, + cos_table, + sin_table, + position_ids, + window_size=window_size, + compress_ratio=compress_ratio, + ) + + torch.testing.assert_close(mhc_cache[0, 0], compressed_kv[0, 0]) + torch.testing.assert_close( + mhc_cache[0, compress_ratio], + torch.full_like(mhc_cache[0, compress_ratio], 777.0), + ) + assert_rmse_close(output, expected, rmse_ratio_tol=1e-6, msg="cached ratio-128 boundary: ") + + +def test_cached_ratio128_mhc_cache_uses_token_domain_paged_positions() -> None: + torch.manual_seed(129) + compress_ratio = 128 + tokens_per_block = 64 + total_len = 256 + prefill_len = total_len - 1 + window_size = 4 + q = torch.randn(1, total_len, 1, 8) + kv = torch.randn(1, total_len, 8) + attn_sink = torch.tensor([-0.25]) + ( + compressor_kv, + compressor_gate, + compressed_kv, + compressor, + cos_table, + sin_table, + position_ids, + ) = _compressor_case( + compress_ratio, + total_len, + compressed_capacity_tokens=total_len, + ) + swa_cache, mhc_cache, compressor_kv_cache, compressor_gate_cache = ( + _make_paged_sparse_attention_caches( + total_len, + tokens_per_block, + kv.shape[-1], + compressor_kv.shape[-1], + fill_value=777.0, + ) + ) + + topk_prefill = _visible_source_topk( + prefill_len, + 0, + prefill_len, + window_size, + compress_ratio, + compressor.max_compressed_len, + q.device, + ) + _run_cached_sparse_attention_with_compressor( + q[:, :prefill_len], + kv[:, :prefill_len], + attn_sink, + topk_prefill, + compressor_kv[:, :prefill_len], + compressor_gate[:, :prefill_len], + compressor, + cos_table, + sin_table, + position_ids[:, :prefill_len], + _context_meta(seq_len=prefill_len, tokens_per_block=tokens_per_block), + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + window_size=window_size, + compress_ratio=compress_ratio, + ) + torch.testing.assert_close( + _paged_cache_row(mhc_cache, 0, tokens_per_block), compressed_kv[0, 0] + ) + + output = _run_cached_sparse_attention_with_compressor( + q[:, prefill_len:], + kv[:, prefill_len:], + attn_sink, + torch.zeros(1, 1, 1, dtype=torch.int64), + compressor_kv[:, prefill_len:], + compressor_gate[:, prefill_len:], + compressor, + cos_table, + sin_table, + position_ids[:, prefill_len:], + _decode_meta(input_pos=prefill_len, tokens_per_block=tokens_per_block), + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + window_size=window_size, + compress_ratio=compress_ratio, + ) + expected_topk = _visible_source_topk( + 1, + prefill_len, + total_len, + window_size, + compress_ratio, + compressor.max_compressed_len, + q.device, + ) + expected = _run_sparse_attention_with_compressor( + q[:, prefill_len:], + kv, + attn_sink, + expected_topk, + compressor_kv, + compressor_gate, + compressor, + cos_table, + sin_table, + position_ids, + window_size=window_size, + compress_ratio=compress_ratio, + ) + + row1_pos = compress_ratio + torch.testing.assert_close( + _paged_cache_row(mhc_cache, row1_pos, tokens_per_block), + compressed_kv[0, 1], + ) + torch.testing.assert_close( + _paged_cache_row(swa_cache, prefill_len, tokens_per_block), + kv[0, prefill_len], + ) + assert_rmse_close(output, expected, rmse_ratio_tol=1e-6, msg="paged ratio-128 mhc: ") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA") +@pytest.mark.parametrize( + ("compress_ratio", "capture_pos", "replay_pos", "state_dim", "topk_width"), + [ + (4, 3, 7, 16, 5), + (128, 127, 255, 8, 6), + ], +) +def test_cached_compressed_decode_cuda_graph_replay_updates_runtime_compressed_row( + compress_ratio: int, + capture_pos: int, + replay_pos: int, + state_dim: int, + topk_width: int, +) -> None: + torch.manual_seed(1900 + compress_ratio) + max_seq_len = replay_pos + 1 + max_compressed_len = 2 + head_dim = 8 + window_size = 4 + q = torch.zeros(1, 1, 1, head_dim, device="cuda") + kv = torch.zeros(1, 1, head_dim, device="cuda") + attn_sink = torch.tensor([-20.0], device="cuda") + topk_idxs = torch.zeros(1, 1, topk_width, dtype=torch.int64, device="cuda") + metadata = _cuda_decode_meta(input_pos=capture_pos, slot_idx=0) + + compressor_kv_values = torch.randn(2, max_seq_len, state_dim, device="cuda") + compressor_gate_values = torch.randn(2, max_seq_len, state_dim, device="cuda") + compressor_kv = compressor_kv_values[0:1, capture_pos : capture_pos + 1].clone() + compressor_gate = compressor_gate_values[0:1, capture_pos : capture_pos + 1].clone() + compressor_ape = torch.zeros(compress_ratio, state_dim, device="cuda") + compressor_norm_weight = torch.ones(head_dim, device="cuda") + cos_table = torch.empty(max_seq_len, 0, device="cuda") + sin_table = torch.empty(max_seq_len, 0, device="cuda") + position_ids = torch.tensor([[capture_pos]], dtype=torch.int64, device="cuda") + + swa_cache = torch.zeros(2, max_seq_len, head_dim, device="cuda") + mhc_cache = torch.full((2, max_seq_len, head_dim), -77.0, device="cuda") + compressor_kv_cache = torch.zeros(2, max_seq_len, state_dim, device="cuda") + compressor_gate_cache = torch.zeros_like(compressor_kv_cache) + compressor_kv_cache[0, :capture_pos] = compressor_kv_values[0, :capture_pos] + compressor_gate_cache[0, :capture_pos] = compressor_gate_values[0, :capture_pos] + compressor_kv_cache[1, :replay_pos] = compressor_kv_values[1, :replay_pos] + compressor_gate_cache[1, :replay_pos] = compressor_gate_values[1, :replay_pos] + + if compress_ratio == 4: + indexer_q = torch.ones(1, 1, 1, 2, device="cuda") + indexer_weights = torch.ones(1, 1, 1, device="cuda") + indexer_compressor_kv = torch.randn(1, 1, 4, device="cuda") + indexer_compressor_gate = torch.randn(1, 1, 4, device="cuda") + indexer_compressor_ape = torch.zeros(compress_ratio, 4, device="cuda") + indexer_compressor_norm_weight = torch.ones(2, device="cuda") + indexer_compressor_kv_cache = torch.zeros(2, max_seq_len, 4, device="cuda") + indexer_compressor_gate_cache = torch.zeros_like(indexer_compressor_kv_cache) + else: + indexer_q = q.new_empty(1, 1, 0, 0) + indexer_weights = q.new_empty(1, 1, 0) + indexer_compressor_kv = q.new_empty(1, 1, 0) + indexer_compressor_gate = q.new_empty(1, 1, 0) + indexer_compressor_ape = q.new_empty(0, 0) + indexer_compressor_norm_weight = q.new_empty(0) + indexer_compressor_kv_cache = q.new_empty(2, max_seq_len, 0) + indexer_compressor_gate_cache = q.new_empty(2, max_seq_len, 0) + + out = torch.empty_like(q) + + def run_op() -> None: + torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention_with_cache( + q, + kv, + attn_sink, + topk_idxs, + compressor_kv, + compressor_gate, + compressor_ape, + compressor_norm_weight, + cos_table, + sin_table, + position_ids, + indexer_q, + indexer_weights, + indexer_compressor_kv, + indexer_compressor_gate, + indexer_compressor_ape, + indexer_compressor_norm_weight, + *metadata, + swa_cache, + mhc_cache, + compressor_kv_cache, + compressor_gate_cache, + indexer_compressor_kv_cache, + indexer_compressor_gate_cache, + 1.0, + window_size, + compress_ratio, + max_compressed_len, + 1e-6, + 0, + out=out, + ) + + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + run_op() + stream.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_op() + + compressor_kv.copy_(compressor_kv_values[1:2, replay_pos : replay_pos + 1]) + compressor_gate.copy_(compressor_gate_values[1:2, replay_pos : replay_pos + 1]) + metadata[2].copy_(torch.tensor([replay_pos], dtype=torch.int32, device="cuda")) + metadata[3].copy_(torch.tensor([1], dtype=torch.int64, device="cuda")) + metadata[6].copy_(torch.tensor([1], dtype=torch.int32, device="cuda")) + position_ids.copy_(torch.tensor([[replay_pos]], dtype=torch.int64, device="cuda")) + row_logical_pos = compress_ratio + mhc_cache[1, row_logical_pos].fill_(-33.0) + + graph.replay() + torch.cuda.synchronize() + + expected_row = dsv4_sparse._compressed_row_from_paged_state( + compressor_kv_cache, + compressor_gate_cache, + 0, + 1, + compress_ratio, + torch.tensor([0, 1], dtype=torch.int64), + torch.tensor([1], dtype=torch.int64), + compressor_ape, + compressor_norm_weight, + cos_table, + sin_table, + 1e-6, + 0, + compress_ratio, + head_dim, + state_dim, + compressor_kv.dtype, + ) + torch.testing.assert_close(mhc_cache[1, row_logical_pos], expected_row, rtol=1e-5, atol=1e-5) + + +def test_cached_sparse_attention_rejects_unsupported_compress_ratio() -> None: + q = torch.randn(1, 1, 1, 2) + kv = torch.randn(1, 1, 2) + attn_sink = torch.randn(1) + topk_idxs = torch.zeros(1, 1, 1, dtype=torch.int64) + swa_cache = torch.empty(1, 4, 2) + + with pytest.raises(ValueError, match="compress_ratio"): + _run_cached_sparse_attention( + q, + kv, + attn_sink, + topk_idxs, + _decode_meta(input_pos=0), + swa_cache, + compress_ratio=2, + ) + + +def test_cached_sparse_attention_rejects_short_metadata() -> None: + q = torch.randn(1, 1, 1, 2) + kv = torch.randn(1, 1, 2) + attn_sink = torch.randn(1) + topk_idxs = torch.zeros(1, 1, 1, dtype=torch.int64) + batch_info_host = BatchInfo() + batch_info_host.update([1, 1, 0, 0, 0, 0]) + cu_num_pages, cache_loc, last_page_len = _page_meta([1], [0], [0]) + metadata = ( + batch_info_host.serialize(), + torch.tensor([1], dtype=torch.int32), + torch.tensor([0], dtype=torch.int32), + torch.tensor([0], dtype=torch.int64), + torch.tensor([0], dtype=torch.int32), + cu_num_pages, + cache_loc, + last_page_len, + ) + + with pytest.raises(ValueError, match="cu_seqlen must have at least 2 elements"): + _run_cached_sparse_attention( + q, + kv, + attn_sink, + topk_idxs, + metadata, + torch.empty(1, 4, 2), + ) + + +def test_fake_tensor_shape_behavior() -> None: + q = torch.randn(2, 3, 2, 4) + kv = torch.randn(2, 6, 4) + attn_sink = torch.randn(2) + topk_idxs = torch.tensor( + [ + [[0, 1], [1, 2], [2, 3]], + [[3, 4], [4, 5], [5, -1]], + ], + dtype=torch.int64, + ) + + with FakeTensorMode(allow_non_fake_inputs=True) as fake_mode: + q_fake = fake_mode.from_tensor(q) + kv_fake = fake_mode.from_tensor(kv) + sink_fake = fake_mode.from_tensor(attn_sink) + topk_fake = fake_mode.from_tensor(topk_idxs) + output = _run_sparse_attention(q_fake, kv_fake, sink_fake, topk_fake, softmax_scale=0.5) + + assert isinstance(output, FakeTensor) + assert output.shape == q.shape + assert output.dtype == q.dtype + + +@pytest.mark.parametrize("compress_ratio", [0, 4, 128]) +def test_cached_fake_tensor_rank_behavior(compress_ratio: int) -> None: + q = torch.randn(1, 2, 1, 4) + kv = torch.randn(1, 4 if compress_ratio else 2, 4) + attn_sink = torch.randn(1) + topk_idxs = torch.zeros(1, 2, 1, dtype=torch.int64) + metadata = _context_meta(seq_len=2) + swa_cache = torch.empty(1, 8, 4) + + with FakeTensorMode(allow_non_fake_inputs=True) as fake_mode: + q_fake = fake_mode.from_tensor(q) + kv_fake = fake_mode.from_tensor(kv) + sink_fake = fake_mode.from_tensor(attn_sink) + topk_fake = fake_mode.from_tensor(topk_idxs) + metadata_fake = tuple(fake_mode.from_tensor(tensor) for tensor in metadata) + cache_fake = fake_mode.from_tensor(swa_cache) + output = _run_cached_sparse_attention( + q_fake, + kv_fake, + sink_fake, + topk_fake, + metadata_fake, + cache_fake, + window_size=2, + compress_ratio=compress_ratio, + ) + + assert isinstance(output, FakeTensor) + assert output.shape == q.shape + assert output.dtype == q.dtype + + with FakeTensorMode(allow_non_fake_inputs=True) as fake_mode: + with pytest.raises(ValueError, match="swa_cache must have rank 3"): + _run_cached_sparse_attention( + fake_mode.from_tensor(q), + fake_mode.from_tensor(kv), + fake_mode.from_tensor(attn_sink), + fake_mode.from_tensor(topk_idxs), + tuple(fake_mode.from_tensor(tensor) for tensor in metadata), + fake_mode.from_tensor(torch.empty(1, 8, 1, 4)), + window_size=2, + compress_ratio=compress_ratio, + ) + + +def test_export_with_dynamic_batch_sequence_and_topk() -> None: + class SparseAttentionModule(torch.nn.Module): + def forward( + self, + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + ) -> torch.Tensor: + return _run_sparse_attention(q, kv, attn_sink, topk_idxs, softmax_scale=0.5) + + batch = Dim("batch", min=1, max=4) + seq = Dim("seq", min=1, max=8) + kv_rows = Dim("kv_rows", min=4, max=12) + k_select = Dim("k_select", min=1, max=4) + + q = torch.randn(2, 3, 2, 4) + kv = torch.randn(2, 6, 4) + attn_sink = torch.randn(2) + topk_idxs = torch.tensor( + [ + [[0, 1], [2, 3], [4, 5]], + [[5, 4], [3, 2], [1, 0]], + ], + dtype=torch.int64, + ) + + exported = torch.export.export( + SparseAttentionModule(), + (q, kv, attn_sink, topk_idxs), + dynamic_shapes={ + "q": {0: batch, 1: seq}, + "kv": {0: batch, 1: kv_rows}, + "attn_sink": {}, + "topk_idxs": {0: batch, 1: seq, 2: k_select}, + }, + ) + + target_names = {str(node.target) for node in exported.graph.nodes if node.op == "call_function"} + assert "auto_deploy.torch_deepseek_v4_sparse_attention.default" in target_names + + q_alt = torch.randn(1, 4, 2, 4) + kv_alt = torch.randn(1, 7, 4) + sink_alt = torch.randn(2) + topk_alt = torch.tensor([[[0, 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 5]]]) + output = exported.module()(q_alt, kv_alt, sink_alt, topk_alt) + assert output.shape == q_alt.shape + + +class _TinyDeepSeekSparseModule(torch.nn.Module): + def __init__(self, compress_ratio: int = 0) -> None: + super().__init__() + self.compress_ratio = compress_ratio + + def forward( + self, + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + ) -> torch.Tensor: + state_dim = kv.shape[-1] * (2 if self.compress_ratio == 4 else 1) + compressor_kv = q.new_empty(q.shape[0], q.shape[1], state_dim) + compressor_gate = q.new_empty(q.shape[0], q.shape[1], state_dim) + compressor_ape = q.new_empty(self.compress_ratio, state_dim) + compressor_norm_weight = q.new_empty(kv.shape[-1]) + cos_table = q.new_empty(8, 0) + sin_table = q.new_empty(8, 0) + position_ids = torch.arange(q.shape[1], device=q.device).unsqueeze(0) + indexer_q = q.new_empty(q.shape[0], q.shape[1], 0, 0) + indexer_weights = q.new_empty(q.shape[0], q.shape[1], 0) + indexer_compressor_kv = q.new_empty(q.shape[0], q.shape[1], 0) + indexer_compressor_gate = q.new_empty(q.shape[0], q.shape[1], 0) + indexer_compressor_ape = q.new_empty(0, 0) + indexer_compressor_norm_weight = q.new_empty(0) + max_compressed_len = 2 if self.compress_ratio else None + rope_dim = 0 if self.compress_ratio else None + return torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention( + q, + kv, + attn_sink, + topk_idxs, + compressor_kv, + compressor_gate, + compressor_ape, + compressor_norm_weight, + cos_table, + sin_table, + position_ids, + indexer_q, + indexer_weights, + indexer_compressor_kv, + indexer_compressor_gate, + indexer_compressor_ape, + indexer_compressor_norm_weight, + 1.0, + False, + "mha_sparse", + 0, + 4, + self.compress_ratio, + max_compressed_len, + kv.shape[-1], + rope_dim, + 1e-6, + ) + + +class _TinyDenseAttentionModule(torch.nn.Module): + def forward(self, qkv: torch.Tensor) -> torch.Tensor: + return torch.ops.auto_deploy.torch_attention( + qkv, + qkv, + qkv, + None, + 0.0, + True, + 1.0, + None, + None, + None, + "bsnd", + ) + + +def test_deepseek_sparse_cache_initializers_use_schema_names_for_source_args() -> None: + graph = Graph() + q_node = graph.placeholder("q") + kv_node = graph.placeholder("kv") + compressor_kv_node = graph.placeholder("compressor_kv") + indexer_compressor_kv_node = graph.placeholder("indexer_compressor_kv") + + kv_node.meta["val"] = torch.empty(1, 2, 4, dtype=torch.float16) + compressor_kv_node.meta["val"] = torch.empty(1, 2, 8, dtype=torch.float32) + indexer_compressor_kv_node.meta["val"] = torch.empty(1, 2, 3, dtype=torch.float32) + + source_node = graph.call_function( + torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention.default, + args=(q_node,), + kwargs={ + "kv": kv_node, + "compressor_kv": compressor_kv_node, + "indexer_compressor_kv": indexer_compressor_kv_node, + }, + ) + + handlers = DeepSeekV4SparseAttention.get_cache_initializers(source_node, KvCacheConfig()) + + assert DeepSeekV4SparseAttention.get_num_qkv_args() == len(dsv4_sparse._SOURCE_TENSOR_ARG_NAMES) + assert isinstance(handlers["swa_cache"], PagedResourceHandler) + assert handlers["swa_cache"].token_shape == (4,) + assert handlers["swa_cache"].dtype == torch.float16 + assert handlers["compressor_kv_cache"].token_shape == (8,) + assert handlers["indexer_compressor_kv_cache"].token_shape == (3,) + + +def test_flashmla_deepseek_sparse_cache_initializers_use_packed_cache_layout() -> None: + graph = Graph() + q_node = graph.placeholder("q") + kv_node = graph.placeholder("kv") + compressor_kv_node = graph.placeholder("compressor_kv") + indexer_compressor_kv_node = graph.placeholder("indexer_compressor_kv") + + kv_node.meta["val"] = torch.empty(1, 2, 512, dtype=torch.bfloat16) + compressor_kv_node.meta["val"] = torch.empty(1, 2, 1024, dtype=torch.float32) + indexer_compressor_kv_node.meta["val"] = torch.empty(1, 2, 256, dtype=torch.float32) + + source_node = graph.call_function( + torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention.default, + args=(q_node,), + kwargs={ + "kv": kv_node, + "compressor_kv": compressor_kv_node, + "indexer_compressor_kv": indexer_compressor_kv_node, + }, + ) + + handlers = FlashMLADeepSeekV4SparseAttention.get_cache_initializers( + source_node, + KvCacheConfig(), + ) + + assert isinstance(handlers["swa_cache"], PagedResourceHandler) + assert handlers["swa_cache"].token_shape == (dsv4_sparse._FLASH_MLA_MODEL1_BYTES_PER_TOKEN,) + assert handlers["swa_cache"].dtype == torch.uint8 + assert handlers["mhc_cache"].token_shape == (dsv4_sparse._FLASH_MLA_MODEL1_BYTES_PER_TOKEN,) + assert handlers["mhc_cache"].dtype == torch.uint8 + assert handlers["compressor_kv_cache"].token_shape == (1024,) + assert handlers["indexer_compressor_kv_cache"].token_shape == (256,) + + +@pytest.mark.parametrize("compress_ratio", [0, 4, 128]) +def test_deepseek_sparse_cache_transform_rewrites_source_op_and_adds_resource( + compress_ratio: int, +) -> None: + q = torch.randn(1, 2, 1, 4) + kv = torch.randn(1, 4 if compress_ratio else 2, 4) + attn_sink = torch.randn(1) + topk_idxs = torch.tensor([[[0], [1]]], dtype=torch.int64) + gm = torch_export_to_gm( + _TinyDeepSeekSparseModule(compress_ratio=compress_ratio), + (q, kv, attn_sink, topk_idxs), + ) + cm = CachedSequenceInterface( + max_seq_len=8, + max_batch_size=2, + max_num_tokens=8, + device="cpu", + ) + + transform = InsertCachedDeepSeekV4SparseAttention( + InsertCachedAttentionConfig(stage=Stages.CACHE_INIT) + ) + gm, info = transform._apply(gm, cm, factory=None, shared_config=SharedConfig()) + + assert info.num_matches == 1 + targets = [node.target for node in gm.graph.nodes if node.op == "call_function"] + assert torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention.default not in targets + assert torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention_with_cache.default in targets + + placeholder_names = [node.target for node in gm.graph.nodes if node.op == "placeholder"] + resource_names = list(cm._resource_lookup) + for suffix in ( + "_swa_cache", + "_mhc_cache", + "_compressor_kv_cache", + "_compressor_gate_cache", + ): + assert _has_resource_with_suffix(placeholder_names, suffix) + assert _has_resource_with_suffix(resource_names, suffix) + + +def test_deepseek_sparse_cache_transform_can_select_flashmla_backend() -> None: + q = torch.randn(1, 2, 1, dsv4_sparse._FLASH_MLA_MODEL1_HEAD_DIM) + kv = torch.randn(1, 2, dsv4_sparse._FLASH_MLA_MODEL1_HEAD_DIM) + attn_sink = torch.randn(1) + topk_idxs = torch.tensor([[[0], [1]]], dtype=torch.int64) + gm = torch_export_to_gm( + _TinyDeepSeekSparseModule(compress_ratio=0), + (q, kv, attn_sink, topk_idxs), + ) + cm = CachedSequenceInterface( + max_seq_len=8, + max_batch_size=2, + max_num_tokens=8, + device="cpu", + ) + + transform = InsertCachedDeepSeekV4SparseAttention( + InsertCachedAttentionConfig( + stage=Stages.CACHE_INIT, + backend="flashmla_deepseek_v4_sparse", + ) + ) + gm, info = transform._apply(gm, cm, factory=None, shared_config=SharedConfig()) + + assert info.num_matches == 1 + targets = [node.target for node in gm.graph.nodes if node.op == "call_function"] + assert torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention.default not in targets + assert ( + torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention_with_cache.default not in targets + ) + assert torch.ops.auto_deploy.flashmla_deepseek_v4_sparse_attention_with_cache.default in targets + for name, handler in cm._resource_lookup.items(): + if name.endswith("_swa_cache") or name.endswith("_mhc_cache"): + assert handler.token_shape == (dsv4_sparse._FLASH_MLA_MODEL1_BYTES_PER_TOKEN,) + assert handler.dtype == torch.uint8 + + +def test_deepseek_v4_flash_registry_orders_sparse_cache_before_initialization() -> None: + root = _repo_root() + config_files = [ + root / "tensorrt_llm/_torch/auto_deploy/config/default.yaml", + root / "examples/auto_deploy/model_registry/configs/dashboard_default.yaml", + root / "examples/auto_deploy/model_registry/configs/world_size_8.yaml", + root / "examples/auto_deploy/model_registry/configs/deepseek_v4_flash.yaml", + ] + cfg = OmegaConf.merge(*(OmegaConf.load(path) for path in config_files)) + transform_names = list(cfg.transforms.keys()) + + assert cfg.transforms.insert_cached_deepseek_v4_sparse_attention.enabled is True + assert ( + cfg.transforms.insert_cached_deepseek_v4_sparse_attention.backend + == "flashmla_deepseek_v4_sparse" + ) + assert transform_names.index("insert_cached_deepseek_v4_sparse_attention") < ( + transform_names.index("initialize_cache") + ) + + +def test_deepseek_sparse_cache_transform_rejects_unsupported_compress_ratio() -> None: + q = torch.randn(1, 2, 1, 4) + kv = torch.randn(1, 2, 4) + attn_sink = torch.randn(1) + topk_idxs = torch.tensor([[[0], [1]]], dtype=torch.int64) + gm = torch_export_to_gm( + _TinyDeepSeekSparseModule(compress_ratio=2), + (q, kv, attn_sink, topk_idxs), + ) + cm = CachedSequenceInterface( + max_seq_len=8, + max_batch_size=2, + max_num_tokens=8, + device="cpu", + ) + transform = InsertCachedDeepSeekV4SparseAttention( + InsertCachedAttentionConfig(stage=Stages.CACHE_INIT) + ) + + with pytest.raises(RuntimeError, match="supports compress_ratio"): + transform._apply(gm, cm, factory=None, shared_config=SharedConfig()) + + +def test_dense_torch_attention_cache_insertion_remains_separate() -> None: + qkv = torch.randn(1, 2, 1, 4) + gm = torch_export_to_gm(_TinyDenseAttentionModule(), (qkv,)) + cm = CachedSequenceInterface( + max_seq_len=8, + max_batch_size=2, + max_num_tokens=8, + device="cpu", + ) + + deepseek_transform = InsertCachedDeepSeekV4SparseAttention( + InsertCachedAttentionConfig(stage=Stages.CACHE_INIT) + ) + gm_after_deepseek, deepseek_info = deepseek_transform._apply( + gm, cm, factory=None, shared_config=SharedConfig() + ) + assert deepseek_info.num_matches == 0 + + dense_transform = _InsertCachedOperator( + InsertCachedAttentionConfig(stage=Stages.CACHE_INIT, backend="torch") + ) + gm_after_dense, dense_info = dense_transform._apply( + gm_after_deepseek, cm, factory=None, shared_config=SharedConfig() + ) + + assert dense_info.num_matches == 1 + targets = [node.target for node in gm_after_dense.graph.nodes if node.op == "call_function"] + assert torch.ops.auto_deploy.torch_cached_attention_with_cache.default in targets + assert ( + torch.ops.auto_deploy.torch_deepseek_v4_sparse_attention_with_cache.default not in targets + ) diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.py index a689db818680..168e9d3605a6 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.py @@ -2,13 +2,15 @@ # SPDX-License-Identifier: Apache-2.0 import torch +from triton_kernels import target_info +from triton_kernels.tensor_details import layout from triton_kernels.tensor_details.layout import HopperMXValueLayout, StridedLayout from tensorrt_llm._torch.auto_deploy.custom_ops.fused_moe import mxfp4_moe def test_mxfp4_value_layout_uses_strided_layout_on_blackwell(monkeypatch): - monkeypatch.setattr(mxfp4_moe, "cuda_capability_geq", lambda major, minor=0: major >= 10) + monkeypatch.setattr(target_info, "cuda_capability_geq", lambda major, minor=0: major >= 10) value_layout, value_layout_opts = mxfp4_moe._mxfp4_value_layout(mx_axis=1) @@ -17,9 +19,9 @@ def test_mxfp4_value_layout_uses_strided_layout_on_blackwell(monkeypatch): def test_mxfp4_value_layout_keeps_default_layout_pre_blackwell(monkeypatch): - monkeypatch.setattr(mxfp4_moe, "cuda_capability_geq", lambda major, minor=0: False) + monkeypatch.setattr(target_info, "cuda_capability_geq", lambda major, minor=0: False) monkeypatch.setattr( - mxfp4_moe.layout, + layout, "make_default_matmul_mxfp4_w_layout", lambda mx_axis: (HopperMXValueLayout, {"mx_axis": mx_axis}), ) diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_torch_moe_swiglu_limit.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_torch_moe_swiglu_limit.py new file mode 100644 index 000000000000..5ed874564ba8 --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_torch_moe_swiglu_limit.py @@ -0,0 +1,158 @@ +# 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. + +import torch +import torch.nn.functional as F + +import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 +from tensorrt_llm._torch.auto_deploy._compat import ActivationType + + +def _inputs() -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + list[torch.Tensor], + list[torch.Tensor], + list[torch.Tensor], +]: + x = torch.tensor( + [ + [2.0, -1.0], + [-3.0, 0.5], + [1.5, 2.0], + ], + dtype=torch.float32, + ) + selected_experts = torch.tensor( + [ + [0, 1], + [1, 0], + [0, 1], + ], + dtype=torch.int64, + ) + routing_weights = torch.tensor( + [ + [0.75, 0.25], + [0.5, 0.5], + [0.25, 0.75], + ], + dtype=torch.float32, + ) + w1_weight = [ + torch.tensor([[2.0, -1.0], [0.5, 1.0]], dtype=torch.float32), + torch.tensor([[-1.0, 1.5], [2.0, 0.25]], dtype=torch.float32), + ] + w2_weight = [ + torch.tensor([[1.0, -0.5], [0.25, 1.5]], dtype=torch.float32), + torch.tensor([[-0.75, 1.25], [1.0, 0.5]], dtype=torch.float32), + ] + w3_weight = [ + torch.tensor([[-2.0, 0.5], [1.5, -1.0]], dtype=torch.float32), + torch.tensor([[1.0, -2.0], [-1.5, 0.75]], dtype=torch.float32), + ] + return x, selected_experts, routing_weights, w1_weight, w2_weight, w3_weight + + +def _reference_torch_moe( + x: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + w1_weight: list[torch.Tensor], + w2_weight: list[torch.Tensor], + w3_weight: list[torch.Tensor], + swiglu_limit: float, +) -> torch.Tensor: + x_flat = x.view(-1, x.shape[-1]) + output = torch.zeros_like(x_flat) + + for token_idx, token in enumerate(x_flat): + token_input = token.unsqueeze(0) + for route_idx, expert_idx_tensor in enumerate(selected_experts[token_idx]): + expert_idx = int(expert_idx_tensor.item()) + if expert_idx < 0 or expert_idx >= len(w1_weight): + continue + + gate_out = F.linear(token_input.to(w1_weight[expert_idx].dtype), w1_weight[expert_idx]) + up_out = F.linear(token_input.to(w3_weight[expert_idx].dtype), w3_weight[expert_idx]) + if swiglu_limit > 0.0: + gate_out = gate_out.clamp(max=swiglu_limit) + up_out = up_out.clamp(min=-swiglu_limit, max=swiglu_limit) + + expert_out = F.linear(F.silu(gate_out) * up_out, w2_weight[expert_idx]) + output[token_idx] += ( + expert_out.squeeze(0).to(output.dtype) * routing_weights[token_idx, route_idx] + ) + + return output.view_as(x) + + +def _torch_moe(swiglu_limit: float | None = None) -> torch.Tensor: + x, selected_experts, routing_weights, w1_weight, w2_weight, w3_weight = _inputs() + kwargs = { + "is_gated_mlp": True, + "act_fn": int(ActivationType.Silu), + } + if swiglu_limit is not None: + kwargs["swiglu_limit"] = swiglu_limit + + return torch.ops.auto_deploy.torch_moe( + x, + selected_experts, + routing_weights, + w1_weight, + w2_weight, + w3_weight, + **kwargs, + ) + + +def test_torch_moe_swiglu_non_positive_limit_preserves_default() -> None: + x, selected_experts, routing_weights, w1_weight, w2_weight, w3_weight = _inputs() + default_output = _torch_moe() + + for swiglu_limit in (0.0, -1.0): + actual = _torch_moe(swiglu_limit) + expected = _reference_torch_moe( + x, + selected_experts, + routing_weights, + w1_weight, + w2_weight, + w3_weight, + swiglu_limit, + ) + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(actual, default_output) + + +def test_torch_moe_swiglu_positive_limit_matches_reference_and_changes_output() -> None: + x, selected_experts, routing_weights, w1_weight, w2_weight, w3_weight = _inputs() + unlimited_output = _torch_moe(0.0) + limited_output = _torch_moe(1.0) + + expected_limited = _reference_torch_moe( + x, + selected_experts, + routing_weights, + w1_weight, + w2_weight, + w3_weight, + 1.0, + ) + + torch.testing.assert_close(limited_output, expected_limited) + assert not torch.allclose(limited_output, unlimited_output) diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_torch_mxfp4_moe.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_torch_mxfp4_moe.py new file mode 100644 index 000000000000..598d8c6f18c9 --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_torch_mxfp4_moe.py @@ -0,0 +1,974 @@ +# 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. + +import pytest +import torch +import torch.nn.functional as F + +import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: E402, F401 +from tensorrt_llm._torch.auto_deploy.models.custom.modeling_deepseek_v4 import ( # noqa: E402 + DeepseekV4PackedMxfp4ExpertsCheckpointLayout, + build_deepseek_v4_packed_mxfp4_experts_layout, +) +from tensorrt_llm._torch.auto_deploy.transform.interface import Stages # noqa: E402 +from tensorrt_llm._torch.auto_deploy.transform.library.fused_moe_mxfp4 import ( # noqa: E402 + QuantizeMXFP4MOE as InsertMXFP4MLP, +) +from tensorrt_llm._torch.auto_deploy.transform.library.fused_moe_mxfp4 import ( + QuantizeMXFP4MOEConfig as MXFP4MLPConfig, +) +from tensorrt_llm._torch.auto_deploy.transform.library.fused_moe_mxfp4 import ( + _mxfp4_target_names_from_node, + _resolve_mxfp4_expert_block_size, +) + +_E2M1_VALUES = torch.tensor( + [ + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ], + dtype=torch.float32, +) +_E8M0_EXPONENT_BIAS = 127 +_E8M0_DTYPE = getattr(torch, "float8_e8m0fnu", None) +_MXFP4_BLOCK_SIZE = 32 + + +def _scale_from_byte(scale_byte: int) -> float: + return float(2.0 ** (scale_byte - _E8M0_EXPONENT_BIAS)) + + +def _packed_mxfp4_bytes(tensor: torch.Tensor, name: str) -> torch.Tensor: + if tensor.dtype == torch.uint8: + return tensor.contiguous() + if tensor.dtype == torch.int8: + return tensor.contiguous().view(torch.uint8) + raise AssertionError(f"{name} should contain packed MXFP4 bytes, got {tensor.dtype}.") + + +def _e8m0_scale_bytes(tensor: torch.Tensor, name: str) -> torch.Tensor: + if tensor.dtype == torch.uint8: + return tensor.contiguous() + if _E8M0_DTYPE is not None and tensor.dtype == _E8M0_DTYPE: + return tensor.contiguous().view(torch.uint8) + raise AssertionError(f"{name} should contain raw E8M0 scale bytes, got {tensor.dtype}.") + + +def _scale_from_bytes(scale_bytes: torch.Tensor) -> torch.Tensor: + exponents = scale_bytes.to(torch.int32) - _E8M0_EXPONENT_BIAS + return torch.ldexp(torch.ones_like(exponents, dtype=torch.float32), exponents) + + +def _mxfp4_checkpoint_weight_dense_dequant( + weight: torch.Tensor, + scale: torch.Tensor, + *, + weight_name: str, + scale_name: str, +) -> torch.Tensor: + weight_bytes = _packed_mxfp4_bytes(weight, weight_name) + scale_bytes = _e8m0_scale_bytes(scale, scale_name) + assert weight_bytes.ndim == 2 + assert scale_bytes.ndim == 2 + assert weight_bytes.shape[0] == scale_bytes.shape[0] + assert weight_bytes.shape[1] == scale_bytes.shape[1] * (_MXFP4_BLOCK_SIZE // 2) + + blocks = weight_bytes.view( + weight_bytes.shape[0], + scale_bytes.shape[1], + _MXFP4_BLOCK_SIZE // 2, + ) + packed = blocks.to(torch.int64) + low = packed & 0x0F + high = (packed >> 4) & 0x0F + codes = torch.stack((low, high), dim=-1).reshape( + weight_bytes.shape[0], + scale_bytes.shape[1], + _MXFP4_BLOCK_SIZE, + ) + values = _E2M1_VALUES[codes] + decoded = values * _scale_from_bytes(scale_bytes).unsqueeze(-1) + return decoded.reshape(weight_bytes.shape[0], scale_bytes.shape[1] * _MXFP4_BLOCK_SIZE) + + +def _make_mxfp4_weight( + shape: tuple[int, ...], + *, + scale_byte: int = 127, + code_offset: int = 0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + cols = shape[-1] + assert cols % _MXFP4_BLOCK_SIZE == 0 + codes = (torch.arange(torch.tensor(shape).prod().item()) + code_offset) % len(_E2M1_VALUES) + codes = codes.reshape(shape).to(torch.uint8) + dense = _E2M1_VALUES[codes.long()] * _scale_from_byte(scale_byte) + + block_codes = codes.reshape(*shape[:-1], cols // _MXFP4_BLOCK_SIZE, _MXFP4_BLOCK_SIZE) + low = block_codes[..., 0::2] + high = block_codes[..., 1::2] + blocks = (low | (high << 4)).to(torch.uint8) + scales = torch.full(blocks.shape[:-1], scale_byte, dtype=torch.uint8) + return dense, blocks, scales + + +def _make_mxfp4_checkpoint_weight( + shape: tuple[int, int], + *, + scale_byte: int = 127, + code_offset: int = 0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + dense, blocks, scales = _make_mxfp4_weight( + shape, + scale_byte=scale_byte, + code_offset=code_offset, + ) + return dense, blocks.reshape(shape[0], shape[1] // 2), scales + + +def _deepseek_layout() -> DeepseekV4PackedMxfp4ExpertsCheckpointLayout: + return build_deepseek_v4_packed_mxfp4_experts_layout() + + +class _UnsupportedBlockSizeLayout: + expert_block_size = 64 + + +class _QuantConfigFactory: + def __init__(self, qcfg: dict[str, object]) -> None: + self._qcfg = qcfg + + def get_quant_config(self) -> dict[str, object]: + return self._qcfg + + +class _DenseMoeTransformModule(torch.nn.Module): + def __init__(self, hidden_size: int, intermediate_size: int, num_experts: int = 2) -> None: + super().__init__() + self.router_weight = torch.nn.Parameter(torch.randn(num_experts, hidden_size)) + self.router_bias = torch.nn.Parameter(torch.randn(num_experts)) + self.gate_up_w = torch.nn.Parameter( + torch.randn(num_experts, hidden_size, 2 * intermediate_size) + ) + self.gate_up_b = torch.nn.Parameter(torch.randn(num_experts, 2 * intermediate_size)) + self.down_w = torch.nn.Parameter(torch.randn(num_experts, intermediate_size, hidden_size)) + self.down_b = torch.nn.Parameter(torch.randn(num_experts, hidden_size)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + routing_weights = torch.ops.auto_deploy.torch_moe_router( + hidden_states, + self.router_weight, + self.router_bias, + 1, + ) + return torch.ops.auto_deploy.torch_moe_dense_mlp( + hidden_states, + routing_weights, + self.gate_up_w, + self.gate_up_b, + self.down_w, + self.down_b, + 1.0, + 10.0, + ) + + +class _Mxfp4KwargRuntimeModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("gate_up_blocks", torch.zeros((2, 64, 1, 16), dtype=torch.uint8)) + self.register_buffer("gate_up_scales", torch.zeros((2, 64, 1), dtype=torch.uint8)) + self.register_buffer("down_blocks", torch.zeros((2, 32, 1, 16), dtype=torch.uint8)) + self.register_buffer("down_scales", torch.zeros((2, 32, 1), dtype=torch.uint8)) + self.register_buffer("gate_up_bias", torch.zeros((2, 64), dtype=torch.float32)) + self.register_buffer("down_bias", torch.zeros((2, 32), dtype=torch.float32)) + + def forward( + self, + hidden_states: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + ) -> torch.Tensor: + return torch.ops.auto_deploy.torch_mxfp4_moe_from_routing( + hidden_states=hidden_states, + selected_experts=selected_experts, + routing_weights=routing_weights, + gate_up_blocks=self.gate_up_blocks, + gate_up_bias=self.gate_up_bias, + gate_up_scales=self.gate_up_scales, + alpha=1.0, + limit=10.0, + down_blocks=self.down_blocks, + down_bias=self.down_bias, + down_scales=self.down_scales, + gate_up_order="up_gate", + swiglu_mode="deepseek", + ) + + +def _deepseek_expert_key(layer: int, expert: int, projection: str, tensor_kind: str) -> str: + return f"layers.{layer}.ffn.experts.{expert}.{projection}.{tensor_kind}" + + +def _deepseek_packed_params_from_layout( + num_experts: int, + *, + layer: int = 0, + hidden_size: int = 32, + intermediate_size: int = 32, +) -> tuple[ + object, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + state = {} + w1_dense = [] + w2_dense = [] + w3_dense = [] + + for expert in range(num_experts): + w1, w1_packed, w1_scales = _make_mxfp4_checkpoint_weight( + (intermediate_size, hidden_size), + scale_byte=127, + code_offset=3 + 17 * expert, + ) + w2, w2_packed, w2_scales = _make_mxfp4_checkpoint_weight( + (hidden_size, intermediate_size), + scale_byte=127, + code_offset=7 + 19 * expert, + ) + w3, w3_packed, w3_scales = _make_mxfp4_checkpoint_weight( + (intermediate_size, hidden_size), + scale_byte=127, + code_offset=11 + 23 * expert, + ) + w1_dense.append(w1) + w2_dense.append(w2) + w3_dense.append(w3) + for projection, weight, scales in ( + ("w1", w1_packed, w1_scales), + ("w2", w2_packed, w2_scales), + ("w3", w3_packed, w3_scales), + ): + state[_deepseek_expert_key(layer, expert, projection, "weight")] = weight + state[_deepseek_expert_key(layer, expert, projection, "scale")] = scales + + packed = _deepseek_layout().pack_experts( + state, + layer=layer, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_experts=num_experts, + ) + return packed, torch.stack(w1_dense), torch.stack(w2_dense), torch.stack(w3_dense) + + +def _inputs(num_experts: int = 4) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + hidden_size = 32 + x = torch.linspace(-0.25, 0.25, steps=5 * hidden_size, dtype=torch.float32).reshape( + 5, hidden_size + ) + x[:, :4] = torch.tensor( + [ + [2.0, 0.0, 1.0, -1.0], + [0.0, 2.0, 1.0, 0.0], + [-1.0, 0.0, 3.0, 2.0], + [1.0, 1.0, 1.0, 1.0], + [0.0, -1.0, 2.0, 3.0], + ], + dtype=torch.float32, + ) + + router_weight = torch.zeros((num_experts, hidden_size), dtype=torch.float32) + for expert_idx in range(num_experts): + router_weight[expert_idx, expert_idx % 4] = 1.0 + 0.25 * expert_idx + router_bias = torch.linspace(0.3, -0.2, steps=num_experts, dtype=torch.float32) + return x, router_weight, router_bias + + +def _dense_reference( + hidden_states: torch.Tensor, + router_weight: torch.Tensor, + router_bias: torch.Tensor, + top_k: int, + gate_up_weight: torch.Tensor, + gate_up_bias: torch.Tensor, + alpha: float, + limit: float, + down_weight: torch.Tensor, + down_bias: torch.Tensor, + *, + expert_start: int = 0, +) -> torch.Tensor: + leading_shape = hidden_states.shape[:-1] + hidden_size = hidden_states.shape[-1] + x = hidden_states.reshape(-1, hidden_size) + router_logits = F.linear(x, router_weight, router_bias) + router_top_value, selected_experts = torch.topk(router_logits, top_k, dim=-1) + routing_weights = torch.softmax(router_top_value, dim=-1, dtype=torch.float32) + + output = torch.zeros((x.shape[0], hidden_size), dtype=torch.float32) + for local_expert_idx in range(gate_up_weight.shape[0]): + global_expert_idx = expert_start + local_expert_idx + token_idx, route_idx = torch.where(selected_experts == global_expert_idx) + if token_idx.numel() == 0: + continue + + gate_up = F.linear( + x[token_idx], gate_up_weight[local_expert_idx], gate_up_bias[local_expert_idx] + ) + gate, up = gate_up[..., ::2], gate_up[..., 1::2] + gate = gate.clamp(max=limit) + up = up.clamp(min=-limit, max=limit) + inter = (gate * torch.sigmoid(gate * alpha)) * (up + 1.0) + expert_out = F.linear(inter, down_weight[local_expert_idx], down_bias[local_expert_idx]) + output.index_add_(0, token_idx, expert_out * routing_weights[token_idx, route_idx, None]) + return output.reshape(*leading_shape, hidden_size) + + +def _dense_deepseek_routing_reference( + hidden_states: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + w1_weight: torch.Tensor, + w2_weight: torch.Tensor, + w3_weight: torch.Tensor, + *, + alpha: float, + limit: float, + expert_start: int = 0, +) -> torch.Tensor: + leading_shape = hidden_states.shape[:-1] + hidden_size = hidden_states.shape[-1] + x = hidden_states.reshape(-1, hidden_size) + selected_experts = selected_experts.reshape(x.shape[0], -1) + routing_weights = routing_weights.reshape_as(selected_experts).to(torch.float32) + + output = torch.zeros((x.shape[0], hidden_size), dtype=torch.float32) + for local_expert_idx in range(w1_weight.shape[0]): + global_expert_idx = expert_start + local_expert_idx + token_idx, route_idx = torch.where(selected_experts == global_expert_idx) + if token_idx.numel() == 0: + continue + + gate = F.linear(x[token_idx], w1_weight[local_expert_idx]) + up = F.linear(x[token_idx], w3_weight[local_expert_idx]) + if limit > 0.0: + gate = gate.clamp(max=limit) + up = up.clamp(min=-limit, max=limit) + inter = (gate * torch.sigmoid(gate * alpha)) * up + expert_out = F.linear(inter, w2_weight[local_expert_idx]) + output.index_add_(0, token_idx, expert_out * routing_weights[token_idx, route_idx, None]) + return output.reshape(*leading_shape, hidden_size) + + +def _mxfp4_params( + num_experts: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + hidden_size = 32 + intermediate_size = 32 + gate_up_weight, gate_up_blocks, gate_up_scales = _make_mxfp4_weight( + (num_experts, 2 * intermediate_size, hidden_size), + scale_byte=127, + code_offset=1, + ) + down_weight, down_blocks, down_scales = _make_mxfp4_weight( + (num_experts, hidden_size, intermediate_size), + scale_byte=128, + code_offset=5, + ) + return gate_up_weight, gate_up_blocks, gate_up_scales, down_weight, down_blocks, down_scales + + +def _expert_slice(num_experts: int, ep_size: int, ep_rank: int) -> tuple[int, int]: + base = num_experts // ep_size + lo = base * ep_rank + hi = num_experts if ep_rank == ep_size - 1 else base * (ep_rank + 1) + return lo, hi + + +def test_torch_mxfp4_moe_matches_dense_reference_with_bias_and_swiglu_limit() -> None: + num_experts = 4 + top_k = 2 + alpha = 1.7 + limit = 0.75 + x, router_weight, router_bias = _inputs(num_experts) + selected_experts = torch.topk(F.linear(x, router_weight, router_bias), top_k, dim=-1).indices + assert selected_experts.unique().numel() < selected_experts.numel() + gate_up_weight, gate_up_blocks, gate_up_scales, down_weight, down_blocks, down_scales = ( + _mxfp4_params(num_experts) + ) + gate_up_bias = torch.linspace(-0.2, 0.3, steps=num_experts * 64, dtype=torch.float32).reshape( + num_experts, 64 + ) + down_bias = torch.linspace(0.15, -0.1, steps=num_experts * 32, dtype=torch.float32).reshape( + num_experts, 32 + ) + + actual = torch.ops.auto_deploy.torch_mxfp4_moe( + x, + router_weight, + router_bias, + top_k, + gate_up_blocks, + gate_up_bias, + gate_up_scales, + alpha, + limit, + down_blocks, + down_bias, + down_scales, + ) + expected = _dense_reference( + x, + router_weight, + router_bias, + top_k, + gate_up_weight, + gate_up_bias, + alpha, + limit, + down_weight, + down_bias, + ) + unlimited = torch.ops.auto_deploy.torch_mxfp4_moe( + x, + router_weight, + router_bias, + top_k, + gate_up_blocks, + gate_up_bias, + gate_up_scales, + alpha, + 10.0, + down_blocks, + down_bias, + down_scales, + ) + + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + assert not torch.allclose(actual, unlimited) + + +def test_torch_mxfp4_moe_ep_partitions_experts_and_sums_to_full_output() -> None: + num_experts = 5 + ep_size = 3 + top_k = 3 + alpha = 1.2 + limit = 1.25 + x, router_weight, router_bias = _inputs(num_experts) + gate_up_weight, gate_up_blocks, gate_up_scales, down_weight, down_blocks, down_scales = ( + _mxfp4_params(num_experts) + ) + gate_up_bias = torch.linspace(0.2, -0.25, steps=num_experts * 64, dtype=torch.float32).reshape( + num_experts, 64 + ) + down_bias = torch.linspace(-0.1, 0.2, steps=num_experts * 32, dtype=torch.float32).reshape( + num_experts, 32 + ) + + full = torch.ops.auto_deploy.torch_mxfp4_moe( + x, + router_weight, + router_bias, + top_k, + gate_up_blocks, + gate_up_bias, + gate_up_scales, + alpha, + limit, + down_blocks, + down_bias, + down_scales, + ) + partial_sum = torch.zeros_like(full) + expected_sum = torch.zeros_like(full) + for ep_rank in range(ep_size): + lo, hi = _expert_slice(num_experts, ep_size, ep_rank) + partial_sum += torch.ops.auto_deploy.torch_mxfp4_moe_ep( + x, + router_weight, + router_bias, + top_k, + gate_up_blocks[lo:hi], + gate_up_bias[lo:hi], + gate_up_scales[lo:hi], + alpha, + limit, + down_blocks[lo:hi], + down_bias[lo:hi], + down_scales[lo:hi], + ep_size, + ep_rank, + ) + expected_sum += _dense_reference( + x, + router_weight, + router_bias, + top_k, + gate_up_weight[lo:hi], + gate_up_bias[lo:hi], + alpha, + limit, + down_weight[lo:hi], + down_bias[lo:hi], + expert_start=lo, + ) + + torch.testing.assert_close(partial_sum, full, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(partial_sum, expected_sum, rtol=1e-5, atol=1e-5) + + +def test_torch_mxfp4_moe_from_routing_matches_deepseek_layout_reference() -> None: + num_experts = 3 + hidden_size = 32 + intermediate_size = 32 + alpha = 1.0 + limit = 0.75 + x = torch.linspace(-0.3, 0.35, steps=4 * hidden_size, dtype=torch.float32).reshape( + 4, hidden_size + ) + selected_experts = torch.tensor( + [[0, 0], [2, 1], [1, 2], [2, 2]], + dtype=torch.int64, + ) + routing_weights = torch.tensor( + [[0.2, 0.35], [0.6, 0.1], [0.25, 0.45], [0.4, 0.15]], + dtype=torch.float32, + ) + packed, w1_weight, w2_weight, w3_weight = _deepseek_packed_params_from_layout(num_experts) + gate_up_bias = torch.zeros((num_experts, 2 * intermediate_size), dtype=torch.float32) + down_bias = torch.zeros((num_experts, hidden_size), dtype=torch.float32) + + actual = torch.ops.auto_deploy.torch_mxfp4_moe_from_routing( + x, + selected_experts, + routing_weights, + packed.gate_up_blocks, + gate_up_bias, + packed.gate_up_scales, + alpha, + limit, + packed.down_blocks, + down_bias, + packed.down_scales, + "up_gate", + "deepseek", + ) + expected = _dense_deepseek_routing_reference( + x, + selected_experts, + routing_weights, + w1_weight, + w2_weight, + w3_weight, + alpha=alpha, + limit=limit, + ) + + assert expected.abs().max() > 0 + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + + +def test_torch_mxfp4_moe_from_routing_matches_reference_across_token_chunks() -> None: + num_experts = 3 + hidden_size = 32 + intermediate_size = 32 + alpha = 1.0 + limit = 0.75 + num_tokens = 20 + x = torch.linspace(-0.3, 0.35, steps=num_tokens * hidden_size, dtype=torch.float32).reshape( + num_tokens, hidden_size + ) + token_ids = torch.arange(num_tokens, dtype=torch.int64) + selected_experts = torch.stack( + (token_ids % num_experts, (token_ids + 1) % num_experts), + dim=1, + ) + routing_weights = torch.stack( + ( + torch.linspace(0.15, 0.45, steps=num_tokens), + torch.linspace(0.4, 0.1, steps=num_tokens), + ), + dim=1, + ) + packed, w1_weight, w2_weight, w3_weight = _deepseek_packed_params_from_layout(num_experts) + gate_up_bias = torch.zeros((num_experts, 2 * intermediate_size), dtype=torch.float32) + down_bias = torch.zeros((num_experts, hidden_size), dtype=torch.float32) + + actual = torch.ops.auto_deploy.torch_mxfp4_moe_from_routing( + x, + selected_experts, + routing_weights, + packed.gate_up_blocks, + gate_up_bias, + packed.gate_up_scales, + alpha, + limit, + packed.down_blocks, + down_bias, + packed.down_scales, + "up_gate", + "deepseek", + ) + expected = _dense_deepseek_routing_reference( + x, + selected_experts, + routing_weights, + w1_weight, + w2_weight, + w3_weight, + alpha=alpha, + limit=limit, + ) + + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + + +def test_torch_mxfp4_moe_from_routing_ep_partitions_deepseek_layout_experts() -> None: + num_experts = 5 + ep_size = 3 + hidden_size = 32 + intermediate_size = 32 + alpha = 1.0 + limit = 1.0 + x = torch.linspace(-0.4, 0.4, steps=4 * hidden_size, dtype=torch.float32).reshape( + 4, hidden_size + ) + selected_experts = torch.tensor( + [[0, 4, 2], [3, 1, 4], [2, 2, 0], [1, 3, 4]], + dtype=torch.int64, + ) + routing_weights = torch.tensor( + [[0.2, 0.5, 0.1], [0.7, 0.15, 0.05], [0.25, 0.2, 0.1], [0.3, 0.25, 0.35]], + dtype=torch.float32, + ) + packed, _, _, _ = _deepseek_packed_params_from_layout(num_experts) + gate_up_bias = torch.zeros((num_experts, 2 * intermediate_size), dtype=torch.float32) + down_bias = torch.zeros((num_experts, hidden_size), dtype=torch.float32) + + full = torch.ops.auto_deploy.torch_mxfp4_moe_from_routing( + x, + selected_experts, + routing_weights, + packed.gate_up_blocks, + gate_up_bias, + packed.gate_up_scales, + alpha, + limit, + packed.down_blocks, + down_bias, + packed.down_scales, + "up_gate", + "deepseek", + ) + partial_sum = torch.zeros_like(full) + for ep_rank in range(ep_size): + lo, hi = _expert_slice(num_experts, ep_size, ep_rank) + partial_sum += torch.ops.auto_deploy.torch_mxfp4_moe_from_routing_ep( + x, + selected_experts, + routing_weights, + packed.gate_up_blocks[lo:hi], + gate_up_bias[lo:hi], + packed.gate_up_scales[lo:hi], + alpha, + limit, + packed.down_blocks[lo:hi], + down_bias[lo:hi], + packed.down_scales[lo:hi], + "up_gate", + "deepseek", + expert_start=lo, + ) + + torch.testing.assert_close(partial_sum, full, rtol=1e-5, atol=1e-5) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability(0) != (9, 0), + reason="Triton MXFP4 routed MoE Hopper path requires an SM90 GPU", +) +def test_triton_mxfp4_moe_from_routing_topk6_matches_torch_reference_on_hopper() -> None: + num_experts = 8 + ep_size = 4 + hidden_size = 32 + intermediate_size = 32 + alpha = 1.0 + limit = 1.0 + top_k = 6 + x = torch.linspace( + -0.15, + 0.2, + steps=6 * hidden_size, + dtype=torch.bfloat16, + device="cuda", + ).reshape(6, hidden_size) + selected_experts = torch.tensor( + [ + [0, 1, 2, 3, 4, 5], + [2, 3, 4, 5, 6, 7], + [1, 2, 3, 4, 5, 6], + [3, 4, 5, 6, 7, 0], + [0, 2, 4, 6, 1, 3], + [1, 3, 5, 7, 0, 2], + ], + dtype=torch.int64, + device="cuda", + ) + assert selected_experts.shape[1] == top_k + routing_weights = torch.linspace( + 0.1, + 0.9, + steps=selected_experts.numel(), + dtype=torch.float32, + device="cuda", + ).reshape_as(selected_experts) + packed, _, _, _ = _deepseek_packed_params_from_layout(num_experts) + gate_up_bias = torch.zeros( + (num_experts, 2 * intermediate_size), + dtype=torch.bfloat16, + device="cuda", + ) + down_bias = torch.zeros((num_experts, hidden_size), dtype=torch.bfloat16, device="cuda") + + reference = torch.ops.auto_deploy.torch_mxfp4_moe_from_routing( + x, + selected_experts, + routing_weights, + packed.gate_up_blocks.cuda(), + gate_up_bias, + packed.gate_up_scales.cuda(), + alpha, + limit, + packed.down_blocks.cuda(), + down_bias, + packed.down_scales.cuda(), + "up_gate", + "deepseek", + ) + full = torch.ops.auto_deploy.triton_mxfp4_moe_from_routing( + x, + selected_experts, + routing_weights, + packed.gate_up_blocks.cuda(), + gate_up_bias, + packed.gate_up_scales.cuda(), + alpha, + limit, + packed.down_blocks.cuda(), + down_bias, + packed.down_scales.cuda(), + "up_gate", + "deepseek", + ) + + partial_sum = torch.zeros_like(full) + for ep_rank in range(ep_size): + lo, hi = _expert_slice(num_experts, ep_size, ep_rank) + partial_sum += torch.ops.auto_deploy.triton_mxfp4_moe_from_routing_ep( + x, + selected_experts, + routing_weights, + packed.gate_up_blocks[lo:hi].cuda(), + gate_up_bias[lo:hi], + packed.gate_up_scales[lo:hi].cuda(), + alpha, + limit, + packed.down_blocks[lo:hi].cuda(), + down_bias[lo:hi], + packed.down_scales[lo:hi].cuda(), + "up_gate", + "deepseek", + expert_start=lo, + num_experts_total=num_experts, + ) + + assert full.dtype == torch.bfloat16 + assert partial_sum.dtype == torch.bfloat16 + torch.testing.assert_close(full, reference, rtol=5e-2, atol=5e-2) + torch.testing.assert_close(partial_sum, full, rtol=5e-2, atol=5e-2) + + lo, hi = _expert_slice(num_experts, ep_size, 1) + ep_args = ( + x, + selected_experts, + routing_weights, + packed.gate_up_blocks[lo:hi].cuda(), + gate_up_bias[lo:hi], + packed.gate_up_scales[lo:hi].cuda(), + alpha, + limit, + packed.down_blocks[lo:hi].cuda(), + down_bias[lo:hi], + packed.down_scales[lo:hi].cuda(), + "up_gate", + "deepseek", + ) + expected_ep = torch.ops.auto_deploy.triton_mxfp4_moe_from_routing_ep( + *ep_args, + expert_start=lo, + num_experts_total=num_experts, + ) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_ep = torch.ops.auto_deploy.triton_mxfp4_moe_from_routing_ep( + *ep_args, + expert_start=lo, + num_experts_total=num_experts, + ) + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close(captured_ep, expected_ep, rtol=5e-2, atol=5e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for CUDA graph capture") +def test_torch_mxfp4_moe_from_routing_ep_allows_cuda_graph_capture() -> None: + num_experts = 2 + hidden_size = 32 + intermediate_size = 32 + alpha = 1.0 + limit = 1.0 + x = torch.linspace( + -0.4, 0.4, steps=2 * hidden_size, dtype=torch.float32, device="cuda" + ).reshape(2, hidden_size) + selected_experts = torch.tensor([[0, 1], [1, 0]], dtype=torch.int64, device="cuda") + routing_weights = torch.tensor([[0.7, 0.3], [0.4, 0.6]], dtype=torch.float32, device="cuda") + packed, _, _, _ = _deepseek_packed_params_from_layout(num_experts) + gate_up_bias = torch.zeros( + (num_experts, 2 * intermediate_size), dtype=torch.float32, device="cuda" + ) + down_bias = torch.zeros((num_experts, hidden_size), dtype=torch.float32, device="cuda") + + op_args = ( + x, + selected_experts, + routing_weights, + packed.gate_up_blocks.cuda(), + gate_up_bias, + packed.gate_up_scales.cuda(), + alpha, + limit, + packed.down_blocks.cuda(), + down_bias, + packed.down_scales.cuda(), + "up_gate", + "deepseek", + ) + + expected = torch.ops.auto_deploy.torch_mxfp4_moe_from_routing_ep(*op_args, expert_start=0) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = torch.ops.auto_deploy.torch_mxfp4_moe_from_routing_ep(*op_args, expert_start=0) + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + + +def test_mxfp4_transform_backend_selector_respects_explicit_triton() -> None: + config = MXFP4MLPConfig(stage=Stages.PATTERN_MATCHER, backend="triton") + transform = InsertMXFP4MLP(config) + + assert transform._resolve_backend() == "triton" + + +def test_mxfp4_target_names_from_node_uses_op_schema_names_for_kwargs() -> None: + gm = torch.fx.symbolic_trace(_Mxfp4KwargRuntimeModule()) + node = next( + n + for n in gm.graph.nodes + if n.op == "call_function" + and n.target == torch.ops.auto_deploy.torch_mxfp4_moe_from_routing + ) + + assert _mxfp4_target_names_from_node(node) == { + "gate_up_blocks": "gate_up_blocks", + "gate_up_scales": "gate_up_scales", + "down_blocks": "down_blocks", + "down_scales": "down_scales", + } + + +def test_mxfp4_transform_resolves_expert_block_size_and_rejects_unsupported() -> None: + layout = _deepseek_layout() + + assert _resolve_mxfp4_expert_block_size({}, None) == _MXFP4_BLOCK_SIZE + assert _resolve_mxfp4_expert_block_size({}, layout) == _MXFP4_BLOCK_SIZE + assert ( + _resolve_mxfp4_expert_block_size({"expert_block_size": _MXFP4_BLOCK_SIZE}, layout) + == _MXFP4_BLOCK_SIZE + ) + + with pytest.raises(ValueError, match="supports expert_block_size=32.*got 64"): + _resolve_mxfp4_expert_block_size({"expert_block_size": 64}, None) + + with pytest.raises(ValueError, match="supports expert_block_size=32.*got 64"): + _resolve_mxfp4_expert_block_size({}, _UnsupportedBlockSizeLayout()) + + +@pytest.mark.parametrize( + ("qcfg", "checkpoint_layout", "expected_message"), + ( + ({"expert_block_size": True}, None, "should be an integer"), + ({"expert_block_size": 0}, None, "should be positive"), + ({"expert_block_size": 32.0}, None, "should be an integer"), + ({"expert_block_size": 32}, _UnsupportedBlockSizeLayout(), "mismatch"), + ), +) +def test_mxfp4_transform_rejects_invalid_expert_block_size_config( + qcfg: dict[str, object], + checkpoint_layout: object | None, + expected_message: str, +) -> None: + with pytest.raises(ValueError, match=expected_message): + _resolve_mxfp4_expert_block_size(qcfg, checkpoint_layout) + + +@pytest.mark.parametrize( + ("hidden_size", "intermediate_size", "expected_name"), + ( + (31, 32, "hidden_size"), + (32, 31, "intermediate_size"), + ), +) +def test_mxfp4_transform_rejects_expert_dims_not_divisible_by_block_size( + hidden_size: int, + intermediate_size: int, + expected_name: str, +) -> None: + gm = torch.fx.symbolic_trace(_DenseMoeTransformModule(hidden_size, intermediate_size)) + transform = InsertMXFP4MLP(MXFP4MLPConfig(stage=Stages.PATTERN_MATCHER)) + factory = _QuantConfigFactory({"quant_method": "mxfp4", "expert_block_size": 32}) + + with pytest.raises(ValueError, match=f"{expected_name}.*divisible.*32"): + transform._apply(gm, cm=None, factory=factory, shared_config=None) diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_trtllm_quant_mxfp4_trtllm_gen_moe.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_trtllm_quant_mxfp4_trtllm_gen_moe.py index 5c63535aed15..ab2356083a2c 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_trtllm_quant_mxfp4_trtllm_gen_moe.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_trtllm_quant_mxfp4_trtllm_gen_moe.py @@ -136,6 +136,28 @@ def _call_op(prep, *, act_dtype, x, router_weight, router_bias, top_k=2): ) +def _call_routed_op(prep, *, act_dtype, x, selected_experts, routing_weights): + return torch.ops.auto_deploy.trtllm_quant_mxfp4_trtllm_gen_moe_from_routing_fused( + x, + selected_experts, + routing_weights, + prep.fc1_weights_mxfp4, + prep.fc2_weights_mxfp4, + prep.fc1_weights_scale_ue8m0, + prep.fc2_weights_scale_ue8m0, + prep.fc1_bias_f32, + prep.fc2_bias_f32, + # DeepSeek defaults: beta=0 because its SwiGLU is silu(gate) * up. + torch.ones((prep.fc1_bias_f32.shape[0],), dtype=torch.float32, device=x.device), + torch.zeros((prep.fc1_bias_f32.shape[0],), dtype=torch.float32, device=x.device), + torch.full((prep.fc1_bias_f32.shape[0],), 10.0, dtype=torch.float32, device=x.device), + prep.valid_hidden_size, + prep.valid_intermediate_size, + act_dtype, + prep.fc1_bias_f32.shape[0], + ) + + # --------------------------------------------------------------------------- # Section 1: weight-preparation invariants (CPU/CUDA, SM-agnostic shuffle ops) # --------------------------------------------------------------------------- @@ -378,6 +400,54 @@ def test_prep_against_pt_reference_loader_byte_identical(): torch.testing.assert_close(prep.fc2_bias_f32, fc2_bias_ref_t, atol=0, rtol=0) +def test_up_gate_layout_prep_matches_equivalent_interleaved_layout(): + """DeepSeek stores gate/up as ``[up | gate]``; GPT-OSS stores interleaved rows.""" + device = "cuda" + gu_blocks, gu_scales, gu_bias, dn_blocks, dn_scales, dn_bias = _build_synthetic_mxfp4_inputs( + e=2, + h=128, + i=128, + device=device, + ) + + up_gate_blocks = torch.cat((gu_blocks[:, 1::2], gu_blocks[:, 0::2]), dim=1).contiguous() + up_gate_scales = torch.cat((gu_scales[:, 1::2], gu_scales[:, 0::2]), dim=1).contiguous() + up_gate_bias = torch.cat((gu_bias[:, 1::2], gu_bias[:, 0::2]), dim=1).contiguous() + + interleaved = prepare_trtllm_gen_moe_mxfp4_weights( + gu_blocks, + gu_scales, + gu_bias, + dn_blocks, + dn_scales, + dn_bias, + hidden_size=128, + intermediate_size=128, + tp_size=1, + tp_rank=0, + ) + up_gate = prepare_trtllm_gen_moe_mxfp4_weights( + up_gate_blocks, + up_gate_scales, + up_gate_bias, + dn_blocks, + dn_scales, + dn_bias, + hidden_size=128, + intermediate_size=128, + tp_size=1, + tp_rank=0, + gate_up_order="up_gate", + ) + + assert torch.equal(up_gate.fc1_weights_mxfp4, interleaved.fc1_weights_mxfp4) + assert torch.equal(up_gate.fc1_weights_scale_ue8m0, interleaved.fc1_weights_scale_ue8m0) + torch.testing.assert_close(up_gate.fc1_bias_f32, interleaved.fc1_bias_f32, atol=0, rtol=0) + assert torch.equal(up_gate.fc2_weights_mxfp4, interleaved.fc2_weights_mxfp4) + assert torch.equal(up_gate.fc2_weights_scale_ue8m0, interleaved.fc2_weights_scale_ue8m0) + torch.testing.assert_close(up_gate.fc2_bias_f32, interleaved.fc2_bias_f32, atol=0, rtol=0) + + # --------------------------------------------------------------------------- # Section 2: unified op act_dtype dispatch (Blackwell+ only) # --------------------------------------------------------------------------- @@ -449,3 +519,45 @@ def test_trtllm_quant_mxfp4_trtllm_gen_moe_fused_invalid_act_dtype(): router_weight=router_weight, router_bias=router_bias, ) + + +@skip_pre_blackwell +@pytest.mark.parametrize("act_dtype", ["bf16", "mxfp8"]) +def test_trtllm_quant_mxfp4_trtllm_gen_moe_from_routing_fused_act_dtype_dispatch(act_dtype): + """Routed DeepSeek-style op dispatches with external top-k tensors.""" + torch.manual_seed(0) + device = "cuda" + E, H, I_, top_k = 4, 256, 64, 2 + B = 8 + + inputs = _make_random_mxfp4_inputs( + num_experts=E, hidden_size=H, intermediate_size=I_, device=device + ) + prep = prepare_trtllm_gen_moe_mxfp4_weights( + *inputs, + hidden_size=H, + intermediate_size=I_, + tp_size=1, + tp_rank=0, + gate_up_order="up_gate", + ) + + x = torch.randn(B, H, dtype=torch.bfloat16, device=device) + token_ids = torch.arange(B, device=device, dtype=torch.int64) + selected_experts = torch.stack((token_ids % E, (token_ids + 1) % E), dim=1) + assert selected_experts.shape == (B, top_k) + routing_weights = torch.tensor([[0.65, 0.35]], dtype=torch.bfloat16, device=device).expand( + B, top_k + ) + + y = _call_routed_op( + prep, + act_dtype=act_dtype, + x=x, + selected_experts=selected_experts, + routing_weights=routing_weights, + ) + + assert y.shape == (B, H), f"unexpected output shape {tuple(y.shape)}, want {(B, H)}" + assert y.dtype == torch.bfloat16, f"unexpected output dtype {y.dtype}" + assert torch.isfinite(y).all(), f"non-finite output for act_dtype={act_dtype!r}" diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/quantization/test_quant.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/quantization/test_quant.py index 240904397e39..e2abd75b0381 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/quantization/test_quant.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/quantization/test_quant.py @@ -19,9 +19,11 @@ import torch.nn.functional as F from _torch_test_utils import fp4_compatible, fp8_compatible, trtllm_ops_available -import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 -from tensorrt_llm._torch.auto_deploy.utils.fp8_dequant import dequant_fp8_weight_two_dim_block_grid -from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import ( +import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: E402, F401 +from tensorrt_llm._torch.auto_deploy.utils.fp8_dequant import ( # noqa: E402 + dequant_fp8_weight_two_dim_block_grid, +) +from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import ( # noqa: E402 fp4_global_scale, pack_int4_in_uint8, unpack_uint8_to_int4_weight_2d, @@ -398,6 +400,170 @@ def test_finegrained_fp8_linear(M, N, K, bias): assert cos > 0.95, f"Cosine similarity too low: {cos}" +def _finegrained_fp8_dense_dequant_ref( + input_tensor, + weight_fp8, + weight_scale_inv, + block_size, + input_scale_fmt="", +): + block_n, block_k = block_size + fp8_max = torch.finfo(torch.float8_e4m3fn).max + + x_blocks = input_tensor.contiguous().view(-1, input_tensor.shape[-1] // block_k, block_k) + input_amax = x_blocks.abs().float().amax(dim=-1) + if input_scale_fmt.lower() == "ue8m0": + input_scale = torch.clamp(input_amax, min=1e-4) / fp8_max + input_scale = torch.pow(2.0, torch.ceil(torch.log2(input_scale))) + else: + input_scale = torch.clamp(input_amax / fp8_max, min=1e-12) + qinput = (x_blocks.float() / input_scale.unsqueeze(-1)).to(torch.float8_e4m3fn) + input_dequant = (qinput.float() * input_scale.unsqueeze(-1)).view_as(input_tensor) + + weight_scale = weight_scale_inv.repeat_interleave(block_n, dim=0).repeat_interleave( + block_k, dim=1 + ) + weight_dequant = weight_fp8.float() * weight_scale + return torch.nn.functional.linear( + input_dequant.to(input_tensor.dtype), + weight_dequant.to(input_tensor.dtype), + ) + + +def _grouped_finegrained_fp8_dense_dequant_ref( + input_tensor, + weight_fp8, + bias, + weight_scale_inv, + block_size, + input_scale_fmt="", +): + block_n, block_k = block_size + fp8_max = torch.finfo(torch.float8_e4m3fn).max + num_groups = input_tensor.shape[-2] + rank = weight_fp8.shape[0] // num_groups + + x_blocks = input_tensor.contiguous().view(*input_tensor.shape[:-1], -1, block_k) + input_amax = x_blocks.abs().float().amax(dim=-1) + if input_scale_fmt.lower() == "ue8m0": + input_scale = torch.clamp(input_amax, min=1e-4) / fp8_max + input_scale = torch.pow(2.0, torch.ceil(torch.log2(input_scale))) + else: + input_scale = torch.clamp(input_amax / fp8_max, min=1e-12) + qinput = (x_blocks.float() / input_scale.unsqueeze(-1)).to(torch.float8_e4m3fn) + input_dequant = (qinput.float() * input_scale.unsqueeze(-1)).view_as(input_tensor) + + weight_scale = weight_scale_inv.repeat_interleave(block_n, dim=0).repeat_interleave( + block_k, dim=1 + ) + weight_scale = weight_scale[: weight_fp8.shape[0], : weight_fp8.shape[1]] + weight_dequant = (weight_fp8.float() * weight_scale).to(input_tensor.dtype) + weight_grouped = weight_dequant.view(num_groups, rank, input_tensor.shape[-1]) + output = torch.matmul( + input_dequant.to(input_tensor.dtype).unsqueeze(-2), + weight_grouped.transpose(-1, -2), + ).squeeze(-2) + output = output.flatten(-2) + if bias is not None: + output = output + bias.reshape(weight_fp8.shape[0]).to(output.dtype) + return output + + +@pytest.mark.skipif(not fp8_compatible(), reason="Requires fp8 support") +def test_finegrained_fp8_linear_ue8m0_input_scale_quantizes_with_rounded_scale(): + block_size = [128, 128] + N, K = 128, 128 + + base = torch.linspace(1.0, 97.0, K, device="cuda", dtype=torch.float32) + input_tensor = torch.stack( + [ + base, + base * 0.75 + 3.0, + base * 0.5 + 7.0, + base * 0.25 + 11.0, + ] + ).to(torch.float16) + weight_fp8 = torch.ones(N, K, device="cuda", dtype=torch.float16).to(torch.float8_e4m3fn) + weight_scale_inv = torch.ones(1, 1, device="cuda", dtype=torch.float32) + + output_raw = torch.ops.auto_deploy.torch_fake_quant_finegrained_fp8_linear( + input_tensor, + weight_fp8, + None, + [], + [weight_scale_inv], + [], + [], + ) + output_ue8m0 = torch.ops.auto_deploy.torch_fake_quant_finegrained_fp8_linear( + input_tensor, + weight_fp8, + None, + [], + [weight_scale_inv], + [], + [], + input_scale_fmt="ue8m0", + ) + + ref_raw = _finegrained_fp8_dense_dequant_ref( + input_tensor, weight_fp8, weight_scale_inv, block_size + ) + ref_ue8m0 = _finegrained_fp8_dense_dequant_ref( + input_tensor, weight_fp8, weight_scale_inv, block_size, input_scale_fmt="ue8m0" + ) + + assert not torch.allclose(output_raw, output_ue8m0, rtol=1e-3, atol=1e-2) + torch.testing.assert_close(output_raw, ref_raw, rtol=0.02, atol=1.0) + torch.testing.assert_close(output_ue8m0, ref_ue8m0, rtol=0.02, atol=1.0) + + +@pytest.mark.parametrize("bias", [True, False]) +@pytest.mark.skipif(not fp8_compatible(), reason="Requires fp8 support") +def test_grouped_finegrained_fp8_linear_matches_dense_dequant_ue8m0_reference(bias): + block_size = [128, 128] + batch_size, seq_len, num_groups, rank, group_width = 2, 3, 2, 128, 128 + N = num_groups * rank + K = group_width + + input_tensor = torch.randn( + batch_size, seq_len, num_groups, K, device="cuda", dtype=torch.float16 + ) + weight = torch.randn(N, K, device="cuda", dtype=torch.float16) + bias_tensor = torch.randn(N, device="cuda", dtype=torch.float16) if bias else None + + fp8_max = torch.finfo(torch.float8_e4m3fn).max + weight_blocks = weight.float().view(num_groups, rank, K // block_size[1], block_size[1]) + amax = weight_blocks.abs().amax(dim=(1, 3)).to(torch.float32) + weight_scale_inv = torch.clamp(amax / fp8_max, min=torch.finfo(torch.float32).tiny) + weight_scale = weight_scale_inv.repeat_interleave(rank, dim=0).repeat_interleave( + block_size[1], dim=1 + ) + weight_fp8 = (weight.float() / weight_scale).to(torch.float8_e4m3fn) + + output = torch.ops.auto_deploy.torch_fake_quant_grouped_finegrained_fp8_linear( + input_tensor, + weight_fp8, + bias_tensor, + [], + [weight_scale_inv], + [], + [], + input_scale_fmt="ue8m0", + ) + ref = _grouped_finegrained_fp8_dense_dequant_ref( + input_tensor, + weight_fp8, + bias_tensor, + weight_scale_inv, + block_size, + input_scale_fmt="ue8m0", + ) + + assert output.shape == (batch_size, seq_len, N) + torch.testing.assert_close(output, ref, rtol=0.02, atol=1.0) + + def _fused_relu2_quantize_available(): return hasattr(torch.ops, "trtllm") and hasattr(torch.ops.trtllm, "fused_relu2_quantize") diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_hf.py b/tests/unittest/auto_deploy/singlegpu/models/test_hf.py index 55ff51603bdb..ca94e4b8736d 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_hf.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_hf.py @@ -41,56 +41,57 @@ def restore_custom_model_mapping(): AutoModelForCausalLMFactory._custom_model_mapping = old_mapping +@pytest.fixture +def mock_factory(): + with ( + patch.object(AutoModelForCausalLMFactory, "prefetch_checkpoint"), + patch.object(AutoModelForCausalLMFactory, "_load_quantization_config"), + ): + # Create factory instance with mocked methods to avoid HTTP requests + factory = AutoModelForCausalLMFactory(model="dummy_model") + # Set model path directly to avoid prefetch + factory._prefetched_model_path = "/dummy/path" + yield factory + + def test_hf_load_state_dict_with_device(): - """Test that hf_load_state_dict_with_device correctly patches modeling.load_state_dict.""" - # Create mock for original load_state_dict original_load_state_dict = MagicMock() - # Test with CPU device with patch.object(modeling, "load_state_dict", original_load_state_dict): with hf_load_state_dict_with_device(device="cpu"): - # Call the patched function modeling.load_state_dict("dummy_checkpoint") - - # Check that device was set correctly original_load_state_dict.assert_called_once_with( "dummy_checkpoint", device_map={"": "cpu"} ) - - # Reset mock for next test original_load_state_dict.reset_mock() - # Check that original behavior is restored modeling.load_state_dict("dummy_checkpoint", device_map="original_device_map") original_load_state_dict.assert_called_once_with( "dummy_checkpoint", device_map="original_device_map" ) original_load_state_dict.reset_mock() - # Test with CUDA device (if available) if torch.cuda.is_available(): with patch.object(modeling, "load_state_dict", original_load_state_dict): with hf_load_state_dict_with_device(device="cuda"): - # Call the patched function modeling.load_state_dict("dummy_checkpoint") - - # Check that device was set correctly original_load_state_dict.assert_called_once_with( "dummy_checkpoint", device_map={"": "cuda"} ) -@pytest.fixture -def mock_factory(): +def test_disable_preload_uses_accelerate_loader(mock_factory): + model = SimpleModel() + ckpt_file = "/dummy/path/model.safetensors.index.json" + with ( - patch.object(AutoModelForCausalLMFactory, "prefetch_checkpoint"), - patch.object(AutoModelForCausalLMFactory, "_load_quantization_config"), + patch.object(mock_factory, "_get_checkpoint_file", return_value=ckpt_file) as get_mock, + patch("tensorrt_llm._torch.auto_deploy.models.hf.load_checkpoint_in_model") as load_mock, ): - # Create factory instance with mocked methods to avoid HTTP requests - factory = AutoModelForCausalLMFactory(model="dummy_model") - # Set model path directly to avoid prefetch - factory._prefetched_model_path = "/dummy/path" - yield factory + mock_factory._load_checkpoint(model, "cpu", disable_preload=True) + + get_mock.assert_called_once_with(mock_factory.model) + load_mock.assert_called_once_with(model, checkpoint=ckpt_file, full_state_dict=False) def test_recursive_update_config(mock_factory): diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_mxfp4_moe.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_mxfp4_moe.py index 4d513e390e26..860bbe2dfbbb 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_mxfp4_moe.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_mxfp4_moe.py @@ -22,12 +22,13 @@ import torch.nn as nn import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 (op registration) -import tensorrt_llm._torch.auto_deploy.transform.library.fused_moe_mxfp4 # noqa: F401 +import tensorrt_llm._torch.auto_deploy.transform.library.fused_moe_mxfp4 as fused_moe_mxfp4 from tensorrt_llm._torch.auto_deploy.custom_ops.fused_moe.prepare_trtllm_gen_moe_mxfp4_weights import ( prepare_trtllm_gen_moe_mxfp4_weights, ) from tensorrt_llm._torch.auto_deploy.transform.interface import SharedConfig, TransformRegistry from tensorrt_llm._torch.auto_deploy.utils.dist_config import DistConfig +from tensorrt_llm._torch.auto_deploy.utils.node_utils import extract_op_args # The transform calls ``prepare_trtllm_gen_moe_mxfp4_weights`` which itself # invokes ``torch.ops.trtllm.shuffle_matrix`` — registered CUDA-only. @@ -139,6 +140,67 @@ def _build_pre_fuse_gm(raw_tensors: Tuple[torch.Tensor, ...]) -> torch.fx.GraphM return torch.fx.GraphModule(root, graph) +def _build_pre_fuse_routed_gm( + raw_tensors: Tuple[torch.Tensor, ...], + *, + ep: bool = False, +) -> torch.fx.GraphModule: + """Build a DeepSeek-style graph that still calls the torch reference routed op.""" + gu_blocks, gu_scales, gu_bias, dn_blocks, dn_scales, dn_bias = raw_tensors + + root = nn.Module() + root.experts = nn.Module() + raw_specs = [ + ("gate_up_proj_blocks", gu_blocks), + ("gate_up_proj_scales", gu_scales), + ("gate_up_proj_bias", gu_bias), + ("down_proj_blocks", dn_blocks), + ("down_proj_scales", dn_scales), + ("down_proj_bias", dn_bias), + ] + for name, t in raw_specs: + root.experts.register_parameter(name, nn.Parameter(t.clone(), requires_grad=False)) + + graph = torch.fx.Graph() + hidden = graph.placeholder("hidden") + selected_experts = graph.placeholder("selected_experts") + routing_weights = graph.placeholder("routing_weights") + + gu_blocks_n = graph.get_attr("experts.gate_up_proj_blocks") + gu_scales_n = graph.get_attr("experts.gate_up_proj_scales") + gu_bias_n = graph.get_attr("experts.gate_up_proj_bias") + dn_blocks_n = graph.get_attr("experts.down_proj_blocks") + dn_scales_n = graph.get_attr("experts.down_proj_scales") + dn_bias_n = graph.get_attr("experts.down_proj_bias") + + args = ( + hidden, + selected_experts, + routing_weights, + gu_blocks_n, + gu_bias_n, + gu_scales_n, + 1.0, + 10.0, + dn_blocks_n, + dn_bias_n, + dn_scales_n, + "up_gate", + "deepseek", + "moe", + ) + if ep: + target = torch.ops.auto_deploy.torch_mxfp4_moe_from_routing_ep.default + args = (*args, 1, E) + else: + target = torch.ops.auto_deploy.torch_mxfp4_moe_from_routing.default + args = (*args, E) + + moe = graph.call_function(target, args=args) + graph.output(moe) + return torch.fx.GraphModule(root, graph) + + def _run_fuse(gm: torch.fx.GraphModule, dist_config: DistConfig): """Apply just ``FuseMXFP4Moe`` with the given ``dist_config``.""" shared_config = SharedConfig( @@ -159,6 +221,12 @@ def _moe_node(gm: torch.fx.GraphModule) -> torch.fx.Node: return nodes[0] +def _single_call_node(gm: torch.fx.GraphModule, target_op) -> torch.fx.Node: + nodes = [n for n in gm.graph.nodes if n.op == "call_function" and n.target is target_op] + assert len(nodes) == 1, f"expected exactly one {target_op} node, found {len(nodes)}" + return nodes[0] + + # --------------------------------------------------------------------------- # TP=1 — single-rank: raw → prepared swap, no bias /= moe_tp_size # --------------------------------------------------------------------------- @@ -290,3 +358,78 @@ def test_fuse_mxfp4_moe_idempotent_on_already_prepped_graph(): _, info2 = _run_fuse(gm, dc) assert info2.skipped is True, "second run should skip — no raw HF buffers left to prep" assert info2.num_matches == 0 + + +def test_fuse_mxfp4_moe_routed_target_policy(): + """Routed DeepSeek MXFP4 target policy is explicit and hardware-gated.""" + assert ( + fused_moe_mxfp4.FuseMXFP4Moe._routed_target_for_sm(100) + == torch.ops.auto_deploy.trtllm_quant_mxfp4_trtllm_gen_moe_from_routing_fused.default + ) + assert ( + fused_moe_mxfp4.FuseMXFP4Moe._routed_target_for_sm(103) + == torch.ops.auto_deploy.trtllm_quant_mxfp4_trtllm_gen_moe_from_routing_fused.default + ) + assert ( + fused_moe_mxfp4.FuseMXFP4Moe._routed_target_for_sm(90) + == torch.ops.auto_deploy.triton_mxfp4_moe_from_routing.default + ) + assert fused_moe_mxfp4.FuseMXFP4Moe._routed_target_for_sm(80) is None + + +def test_fuse_mxfp4_moe_routed_h100_rewrites_to_triton_from_routing(): + """On an actual SM90 GPU, the routed path selects and runs the Triton W4A16 op.""" + if fused_moe_mxfp4.get_sm_version() != 90: + pytest.skip("H100/SM90 runtime coverage only") + + raw = _make_raw_mxfp4_tensors(device="cuda") + gm = _build_pre_fuse_routed_gm(raw) + hidden = torch.linspace( + -0.15, + 0.2, + steps=6 * H, + dtype=torch.bfloat16, + device="cuda", + ).reshape(6, H) + selected_experts = torch.tensor( + [[0, 1], [2, 3], [1, 2], [3, 0], [0, 2], [1, 3]], + dtype=torch.int64, + device="cuda", + ) + routing_weights = torch.tensor( + [[0.60, 0.40], [0.55, 0.45], [0.35, 0.65], [0.25, 0.75], [0.70, 0.30], [0.20, 0.80]], + dtype=torch.float32, + device="cuda", + ) + gu_blocks, gu_scales, gu_bias, dn_blocks, dn_scales, dn_bias = raw + reference = torch.ops.auto_deploy.torch_mxfp4_moe_from_routing( + hidden, + selected_experts, + routing_weights, + gu_blocks, + gu_bias, + gu_scales, + 1.0, + 10.0, + dn_blocks, + dn_bias, + dn_scales, + "up_gate", + "deepseek", + ) + + dc = DistConfig(world_size=1, rank=0, tp_size=1, moe_tp_size=1, moe_ep_size=1) + _, info = _run_fuse(gm, dc) + + target = torch.ops.auto_deploy.triton_mxfp4_moe_from_routing.default + n = _single_call_node(gm, target) + assert info.skipped is False + assert info.num_matches == 1 + assert n.args[3].target == "experts.gate_up_proj_blocks" + (num_experts_total,) = extract_op_args(n, "num_experts_total") + assert num_experts_total == E + assert hasattr(gm.experts, "gate_up_proj_blocks") + assert not hasattr(gm.experts, "fc1_w_trtllm") + out = gm(hidden, selected_experts, routing_weights) + assert out.dtype == torch.bfloat16 + torch.testing.assert_close(out, reference, rtol=5e-2, atol=5e-2) diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_rmsnorm.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_rmsnorm.py index 5c3c52366132..391fb41b02c1 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_rmsnorm.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_rmsnorm.py @@ -68,6 +68,30 @@ def forward(self, x): return x +class DirectRMSNormFP32Weight(torch.nn.Module): + def __init__(self, hidden_size, eps=1e-6): + super().__init__() + self.weight = torch.nn.Parameter( + torch.ones(hidden_size, device="cuda", dtype=torch.float32) + ) + self.eps = eps + + def forward(self, hidden_states): + return torch.ops.auto_deploy.torch_rmsnorm(hidden_states, self.weight, self.eps) + + +def _node_dtype(node): + val = node.meta.get("val") + if hasattr(val, "dtype"): + return val.dtype + + tensor_meta = node.meta.get("tensor_meta") + if hasattr(tensor_meta, "dtype"): + return tensor_meta.dtype + + return None + + def _run_test(model, op, variant): def checker(gm): return any(is_op(n, op) for n in gm.graph.nodes) @@ -148,3 +172,34 @@ def test_fuse_rmsnorm_preserves_node_metadata(): ] assert len(rms_nodes) >= 1 assert all("val" in n.meta and hasattr(n.meta["val"], "dtype") for n in rms_nodes) + + +def test_flashinfer_rmsnorm_casts_fp32_weight_to_input_dtype(): + model = DirectRMSNormFP32Weight(1024) + x = torch.randn(2, 1024, device="cuda", dtype=torch.float16) + dynamic_shapes = {0: Dim.DYNAMIC} + gm = torch_export_to_gm(model, args=(x,), dynamic_shapes=(dynamic_shapes,), clone=True) + + gm_transformed = InferenceOptimizer( + None, + { + "fuse_rmsnorm": { + "stage": "post_load_fusion", + "gated_rmsnorm_backend": "triton", + "rmsnorm_backend": "flashinfer", + }, + }, + )(None, gm) + + rms_nodes = [ + n for n in gm_transformed.graph.nodes if is_op(n, torch.ops.auto_deploy.flashinfer_rms_norm) + ] + assert len(rms_nodes) == 1 + + rms_node = rms_nodes[0] + weight_arg = rms_node.args[1] + assert is_op(weight_arg, torch.ops.aten.to.dtype) + assert weight_arg.args[1] == torch.float16 + assert _node_dtype(rms_node.args[0]) == torch.float16 + assert _node_dtype(weight_arg) == torch.float16 + assert _node_dtype(rms_node.args[0]) == _node_dtype(weight_arg) diff --git a/tests/unittest/auto_deploy/singlegpu/utils/test_example_configs.py b/tests/unittest/auto_deploy/singlegpu/utils/test_example_configs.py index a375f755f2a8..16c23e10527d 100644 --- a/tests/unittest/auto_deploy/singlegpu/utils/test_example_configs.py +++ b/tests/unittest/auto_deploy/singlegpu/utils/test_example_configs.py @@ -35,6 +35,7 @@ _EXCLUDED_FILES = { "models.yaml", # model registry index, not an LlmArgs config "flux_transforms.yaml", # for build_and_run_flux.py, different schema + "nemotron_fp8_ir_test.yaml", # full ExperimentConfig with top-level model/args } # Dummy model name used during validation (model path is not resolved during construction) diff --git a/tests/unittest/auto_deploy/singlegpu/utils/test_node_utils_sharding.py b/tests/unittest/auto_deploy/singlegpu/utils/test_node_utils_sharding.py index 502d9ac6d872..2d7135df8a05 100644 --- a/tests/unittest/auto_deploy/singlegpu/utils/test_node_utils_sharding.py +++ b/tests/unittest/auto_deploy/singlegpu/utils/test_node_utils_sharding.py @@ -137,6 +137,31 @@ def test_enable_sharding_node_linear(): assert ShardableNode.from_node(lin_nodes[0]) is not None +def test_enable_sharding_node_grouped_finegrained_fp8_linear(): + class Shell(nn.Module): + def __init__(self): + super().__init__() + self.w = nn.Parameter(torch.empty(32, 16, dtype=torch.float8_e4m3fn)) + self.register_buffer("weight_scale_inv", torch.ones(1, 1)) + + root = Shell() + graph = fx.Graph() + x = graph.placeholder("x") + w = graph.get_attr("w") + scale = graph.get_attr("weight_scale_inv") + out = graph.call_function( + torch.ops.auto_deploy.torch_fake_quant_grouped_finegrained_fp8_linear.default, + args=(x, w, None, [], [scale], [], []), + kwargs={"tp_mode": "colwise", "tp_min_local_shape": 8, "layer_type": "mla"}, + ) + graph.output(out) + gm = GraphModule(root, graph) + + grouped_nodes = [n for n in _call_function_nodes(gm) if ShardableNode.from_node(n) is not None] + assert len(grouped_nodes) == 1 + assert ShardableNode.from_node(grouped_nodes[0]) is not None + + def test_enable_sharding_node_view(): graph = fx.Graph() x = graph.placeholder("x") diff --git a/tests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.py b/tests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.py index a8790288ce85..5ef80a69fe7d 100644 --- a/tests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.py +++ b/tests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.py @@ -22,7 +22,11 @@ ) from tensorrt_llm._torch.auto_deploy.transform.library.sharding import _shard_fp4_weight_scale from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import ( + ceil_pow2_scale, + fake_fp4_act_quant, + fake_fp8_act_quant, fp4_global_scale, + hadamard_rotate, modelopt_fp4_scale_to_cutlass_fp4_scale, ) @@ -130,3 +134,165 @@ def test_fp8_load_hook_maps_prequantized_scales_with_prefix(): ) assert prefix + "layer.proj.activation_scale" not in mock_state_dict assert prefix + "layer.proj.weight_scale_inv" not in mock_state_dict + + +def _ref_fake_fp8_act_quant(x: torch.Tensor, block_size: int = 64) -> torch.Tensor: + dim = x.shape[-1] + if dim == 0 or dim % block_size != 0: + return x + + dtype = x.dtype + x_float = x.float() + grouped = x_float.reshape(*x_float.shape[:-1], dim // block_size, block_size) + scale = ceil_pow2_scale(grouped.abs().amax(dim=-1, keepdim=True), 448.0, 1.0e-4) + quant = torch.clamp(grouped / scale, -448.0, 448.0).to(dtype).float() + return (quant * scale).reshape_as(x_float).to(dtype) + + +def _ref_fake_fp4_act_quant(x: torch.Tensor, block_size: int = 32) -> torch.Tensor: + dim = x.shape[-1] + if dim == 0 or dim % block_size != 0: + return x + + dtype = x.dtype + x_float = x.float() + grouped = x_float.reshape(*x_float.shape[:-1], dim // block_size, block_size) + scale = ceil_pow2_scale(grouped.abs().amax(dim=-1, keepdim=True), 6.0, 6.0 * 2.0**-126) + normalized = torch.clamp(grouped / scale, -6.0, 6.0) + levels = normalized.new_tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) + level_idx = (normalized.abs().unsqueeze(-1) - levels).abs().argmin(dim=-1) + quant = levels[level_idx] * normalized.sign() + return (quant * scale).reshape_as(x_float).to(dtype) + + +def _ref_hadamard_rotate(x: torch.Tensor) -> torch.Tensor: + dim = x.shape[-1] + if dim <= 1: + return x + if dim & (dim - 1): + raise ValueError(f"Hadamard rotation requires power-of-two dimension, got {dim}.") + + original_shape = x.shape + out = x.float() + width = 1 + while width < dim: + out = out.reshape(*out.shape[:-1], dim // (2 * width), 2, width) + left = out[..., 0, :] + right = out[..., 1, :] + out = torch.cat((left + right, left - right), dim=-1).reshape(original_shape) + width *= 2 + return (out * (dim**-0.5)).to(x.dtype) + + +def test_fake_fp4_act_quant_exports_with_block_aligned_dim() -> None: + class Wrapper(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return fake_fp4_act_quant(x, block_size=32) + + x = torch.tensor( + [ + 0.0, + 0.25, + 0.2501, + 0.75, + 0.7501, + 1.25, + 1.2501, + 1.75, + 1.7501, + 2.5, + 2.5001, + 3.5, + 3.5001, + 5.0, + 5.0001, + 6.0, + -0.25, + -0.2501, + -0.75, + -0.7501, + -1.25, + -1.2501, + -1.75, + -1.7501, + -2.5, + -2.5001, + -3.5, + -3.5001, + -5.0, + -5.0001, + -6.0, + 4.0, + ], + dtype=torch.float32, + ).reshape(1, 32) + wrapper = Wrapper().eval() + + expected = wrapper(x) + exported = torch.export.export(wrapper, (x,), strict=False).module() + actual = exported(x) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(expected, _ref_fake_fp4_act_quant(x), rtol=0, atol=0) + assert expected[0, 1] == 0.0 + assert expected[0, 3] == 0.5 + assert expected[0, 13] == 4.0 + assert expected[0, 17] == -0.5 + + +@pytest.mark.parametrize( + ("quant_fn", "block_size", "groups"), + [ + pytest.param(fake_fp8_act_quant, 64, 2, id="fp8"), + pytest.param(fake_fp4_act_quant, 32, 4, id="fp4"), + ], +) +def test_fake_act_quant_export_keeps_prefix_inferred_for_sharding( + quant_fn, block_size: int, groups: int +) -> None: + class Wrapper(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return quant_fn(x, block_size=block_size) + + x = torch.randn(2, 3, 64, 128) + gm = torch.export.export(Wrapper().eval(), (x,), strict=False).module() + grouping_shapes = [] + for node in gm.graph.nodes: + if node.target != torch.ops.aten.reshape.default: + continue + shape = node.args[1] + if isinstance(shape, (list, tuple)) and list(shape[-2:]) == [groups, block_size]: + grouping_shapes.append(shape) + + assert grouping_shapes + assert all(list(shape) == [-1, groups, block_size] for shape in grouping_shapes) + + +@pytest.mark.parametrize("head_count", [1, 8, 64]) +def test_hadamard_rotate_matches_reference_for_sharded_prefix_shapes(head_count: int) -> None: + x = torch.randn(2, 3, head_count, 1024, dtype=torch.bfloat16) + + actual = hadamard_rotate(x) + expected = _ref_hadamard_rotate(x) + + assert actual.shape == x.shape + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_hadamard_rotate_export_keeps_prefix_inferred_for_sharding() -> None: + class Wrapper(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return hadamard_rotate(x) + + x = torch.randn(2, 3, 64, 128) + gm = torch.export.export(Wrapper().eval(), (x,), strict=False).module() + hadamard_reshape_shapes = [] + for node in gm.graph.nodes: + if node.target != torch.ops.aten.reshape.default: + continue + shape = node.args[1] + if isinstance(shape, (list, tuple)) and len(shape) == 4: + hadamard_reshape_shapes.append(shape) + + assert hadamard_reshape_shapes + assert all(shape[0] == -1 for shape in hadamard_reshape_shapes)