[None][feat] AutoDeploy: Add DSv4 attention & MoE backends

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"
     },

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,33 +142,49 @@ 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)
 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}
-    $<$<COMPILE_LANGUAGE:CXX>:-std=c++17>
+    $<$<COMPILE_LANGUAGE:CXX>:-std=c++20>
     $<$<COMPILE_LANGUAGE:CXX>:-O3>
     $<$<COMPILE_LANGUAGE:CXX>:-fPIC>
     $<$<COMPILE_LANGUAGE:CXX>:-DNDEBUG>
     $<$<COMPILE_LANGUAGE:CXX>:-Wno-deprecated-declarations>
     $<$<COMPILE_LANGUAGE:CXX>:-Wno-c++11-narrowing>
     $<$<COMPILE_LANGUAGE:CXX>:-fno-lto>
+    # The FlashMLA API source only includes pybind11/pybind11.h, but the sparse
+    # interface exposes std::optional<at::Tensor>. Force Torch's pybind casters
+    # into that translation unit without patching the fetched source.
+    $<$<COMPILE_LANGUAGE:CXX>:-include>
+    $<$<COMPILE_LANGUAGE:CXX>:torch/csrc/utils/pybind.h>
     $<$<COMPILE_LANGUAGE:CUDA>:-O3>
-    $<$<COMPILE_LANGUAGE:CUDA>:-std=c++17>
+    $<$<COMPILE_LANGUAGE:CUDA>:-std=c++20>
     $<$<COMPILE_LANGUAGE:CUDA>:-DNDEBUG>
     $<$<COMPILE_LANGUAGE:CUDA>:-D_USE_MATH_DEFINES>
     $<$<COMPILE_LANGUAGE:CUDA>:-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)

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

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

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] = {}

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,

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

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 |

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",