[TRTLLM-16460][feat] modeling_v2: DeepSeek-V4.1-Flash-NVFP4 target (sm_103, unified parallelism)

.github/CODEOWNERS

@@ -394,6 +394,8 @@
 /tests/unittest/_torch/modeling_v2 @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju
 /tests/integration/defs/accuracy/test_modeling_v2_deepseek_v3.py @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju
 /tests/integration/defs/accuracy/test_modeling_v2_gpt_oss.py @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju
+/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v41.py @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju
+/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v41_replay.py @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju
 # The two catalog categories whose contracts state kernel behaviour the
 # attention and MoE owners are the authority on; co-owned rather than reassigned.
 /tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju @xxi-nv @yuxianq

cpp/tensorrt_llm/kernels/indexerKCacheGather.cu

@@ -108,7 +108,7 @@ __global__ void indexerKCacheGatherUnifiedKernel(uint8_t const* __restrict__ k_c
     // Convert flat index to memory offset using strides (k cache pool from cpp kv cache manager is non-contiguous)
     int64_t src_offset = flatIndexToMemoryOffset(flat_idx, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3,
         cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3);
-    int64_t dst_offset = token_idx * head_dim + head_dim_idx;
+    int64_t dst_offset = static_cast<int64_t>(token_idx) * head_dim + head_dim_idx;
 
     // 4 bytes read from non-contiguous cache, write to contiguous output
     *reinterpret_cast<uint32_t*>(&out_fp8[dst_offset]) = *reinterpret_cast<uint32_t const*>(&k_cache[src_offset]);
@@ -118,7 +118,7 @@ __global__ void indexerKCacheGatherUnifiedKernel(uint8_t const* __restrict__ k_c
     {
         int64_t src_offset_scale = flatIndexToMemoryOffset(flat_idx_scale_base, cache_dim_0, cache_dim_1, cache_dim_2,
             cache_dim_3, cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3);
-        int64_t dst_offset_scale = token_idx * scale_size; // scale_size = 4
+        int64_t dst_offset_scale = static_cast<int64_t>(token_idx) * scale_size; // scale_size = 4
 
         // 4 bytes read for scale
         *reinterpret_cast<uint32_t*>(&out_scale[dst_offset_scale])

cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu

@@ -103,7 +103,7 @@ __global__ void indexerKCacheScatterUnifiedKernel(uint8_t const* __restrict__ k_
     // Convert flat index to memory offset using strides (k cache pool from cpp kv cache manager is non-contiguous)
     int64_t dst_offset = flatIndexToMemoryOffset(flat_idx, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3,
         cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3);
-    int64_t src_offset = token_idx * head_dim + head_dim_idx;
+    int64_t src_offset = static_cast<int64_t>(token_idx) * head_dim + head_dim_idx;
 
     // 4 bytes write
     *reinterpret_cast<uint32_t*>(&k_cache[dst_offset]) = *reinterpret_cast<uint32_t const*>(&k_fp8_bytes[src_offset]);
@@ -113,7 +113,7 @@ __global__ void indexerKCacheScatterUnifiedKernel(uint8_t const* __restrict__ k_
     {
         int64_t dst_offset_scale = flatIndexToMemoryOffset(flat_idx_scale_base, cache_dim_0, cache_dim_1, cache_dim_2,
             cache_dim_3, cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3);
-        int64_t src_offset_scale = token_idx * scale_size; // scale_size = 4
+        int64_t src_offset_scale = static_cast<int64_t>(token_idx) * scale_size; // scale_size = 4
 
         // 4 bytes write for scale
         *reinterpret_cast<uint32_t*>(&k_cache[dst_offset_scale])

cpp/tensorrt_llm/kernels/moe/routing/customMoeRoutingKernels.cu

@@ -310,8 +310,6 @@ INSTANTIATE_RENORM_MOE_ROUTING(half, __nv_bfloat16, int32_t, true);
 INSTANTIATE_RENORM_MOE_ROUTING(__nv_bfloat16, __nv_bfloat16, int32_t, true);
 #endif
 
-static constexpr int kTOPK = 6;
-
 // CUDA kernel for gate forward
 // Input: pre-computed scores from linear(x, weight) done outside kernel
 // Template parameters:
@@ -425,7 +423,7 @@ __global__ void gate_forward_kernel(
 
 // C++ wrapper function (output tensors passed as parameters)
 // All tensors are float32
-template <int nExperts, bool hash>
+template <int nExperts, int topK, bool hash>
 void launch_gate_forward_kernel(float* scores_in, float* bias, int* input_ids, int* tid2eid, float* out_weights,
     int* out_indices, int batch_size, float route_scale, bool const* image_mask, float const* vision_bias,
     cudaStream_t stream)
@@ -434,18 +432,35 @@ void launch_gate_forward_kernel(float* scores_in, float* bias, int* input_ids, i
     constexpr int threads_per_block = warps_per_block * WARP_SIZE;
     int const blocks = (batch_size + warps_per_block - 1) / warps_per_block;
 
-    gate_forward_kernel<nExperts, kTOPK, hash><<<blocks, threads_per_block, 0, stream>>>(scores_in, bias, input_ids,
+    gate_forward_kernel<nExperts, topK, hash><<<blocks, threads_per_block, 0, stream>>>(scores_in, bias, input_ids,
         tid2eid, out_weights, out_indices, batch_size, route_scale, image_mask, vision_bias);
 }
 
+template <int nExperts, int topK>
+void dispatch_gate_forward_kernel(float* scores_in, float* bias, int* input_ids, int* tid2eid, float* out_weights,
+    int* out_indices, int batch_size, float route_scale, bool is_hash, bool const* image_mask, float const* vision_bias,
+    cudaStream_t stream)
+{
+    if (is_hash)
+    {
+        launch_gate_forward_kernel<nExperts, topK, true>(scores_in, nullptr, input_ids, tid2eid, out_weights,
+            out_indices, batch_size, route_scale, nullptr, nullptr, stream);
+    }
+    else
+    {
+        launch_gate_forward_kernel<nExperts, topK, false>(scores_in, bias, nullptr, nullptr, out_weights, out_indices,
+            batch_size, route_scale, image_mask, vision_bias, stream);
+    }
+}
+
 void gate_forward(void* scores_in, // [batch_size, nExperts] - pre-computed from linear(x, weight)
     void* bias,                    // nullptr if hash mode
     void* input_ids,               // nullptr if non-hash mode
     void* tid2eid,                 // nullptr if non-hash mode
     void* out_weights,             // [batch_size, topK] - pre-allocated
     void* out_indices,             // [batch_size, topK] - pre-allocated
-    int batch_size, int n_experts, float route_scale, bool is_hash, bool const* image_mask, float const* vision_bias,
-    cudaStream_t stream)
+    int batch_size, int n_experts, int topk, float route_scale, bool is_hash, bool const* image_mask,
+    float const* vision_bias, cudaStream_t stream)
 {
     auto* scores = static_cast<float*>(scores_in);
     auto* bias_ptr = static_cast<float*>(bias);
@@ -454,33 +469,27 @@ void gate_forward(void* scores_in, // [batch_size, nExperts] - pre-computed from
     auto* weights = static_cast<float*>(out_weights);
     auto* indices = static_cast<int*>(out_indices);
 
-    switch (n_experts)
+    // Supported (n_experts, topk) pairs: 256/6 and 384/6 (DeepSeek-V4 / V4.1 routed MoE), 128/3 (DeepSeek-V4.1 DSpark
+    // draft MoE).
+    if (n_experts == 256 && topk == 6)
     {
-    case 256:
-        if (is_hash)
-        {
-            launch_gate_forward_kernel<256, true>(scores, nullptr, input_ids_ptr, tid2eid_ptr, weights, indices,
-                batch_size, route_scale, nullptr, nullptr, stream);
-        }
-        else
-        {
-            launch_gate_forward_kernel<256, false>(scores, bias_ptr, nullptr, nullptr, weights, indices, batch_size,
-                route_scale, image_mask, vision_bias, stream);
-        }
-        break;
-    case 384:
-        if (is_hash)
-        {
-            launch_gate_forward_kernel<384, true>(scores, nullptr, input_ids_ptr, tid2eid_ptr, weights, indices,
-                batch_size, route_scale, nullptr, nullptr, stream);
-        }
-        else
-        {
-            launch_gate_forward_kernel<384, false>(scores, bias_ptr, nullptr, nullptr, weights, indices, batch_size,
-                route_scale, image_mask, vision_bias, stream);
-        }
-        break;
-    default: TLLM_CHECK_WITH_INFO(false, "gate_forward only supports n_experts 256 or 384");
+        dispatch_gate_forward_kernel<256, 6>(scores, bias_ptr, input_ids_ptr, tid2eid_ptr, weights, indices, batch_size,
+            route_scale, is_hash, image_mask, vision_bias, stream);
+    }
+    else if (n_experts == 384 && topk == 6)
+    {
+        dispatch_gate_forward_kernel<384, 6>(scores, bias_ptr, input_ids_ptr, tid2eid_ptr, weights, indices, batch_size,
+            route_scale, is_hash, image_mask, vision_bias, stream);
+    }
+    else if (n_experts == 128 && topk == 3)
+    {
+        dispatch_gate_forward_kernel<128, 3>(scores, bias_ptr, input_ids_ptr, tid2eid_ptr, weights, indices, batch_size,
+            route_scale, is_hash, image_mask, vision_bias, stream);
+    }
+    else
+    {
+        TLLM_CHECK_WITH_INFO(false,
+            "gate_forward supports (n_experts, topk) in {(256, 6), (384, 6), (128, 3)}, got (%d, %d)", n_experts, topk);
     }
     sync_check_cuda_error(stream);
 }

cpp/tensorrt_llm/kernels/moe/routing/customMoeRoutingKernels.h

@@ -38,8 +38,8 @@ void gate_forward(void* scores_in, // [batch_size, nExperts] - pre-computed from
     void* tid2eid,                 // nullptr if non-hash mode
     void* out_weights,             // [batch_size, topK] - pre-allocated
     void* out_indices,             // [batch_size, topK] - pre-allocated
-    int batch_size, int n_experts, float route_scale, bool is_hash, bool const* image_mask, float const* vision_bias,
-    cudaStream_t stream);
+    int batch_size, int n_experts, int topk, float route_scale, bool is_hash, bool const* image_mask,
+    float const* vision_bias, cudaStream_t stream);
 } // namespace kernels
 
 TRTLLM_NAMESPACE_END

cpp/tensorrt_llm/thop/moe/routing/moeGateOp.cpp

@@ -26,8 +26,6 @@ TRTLLM_NAMESPACE_BEGIN
 namespace torch_ext
 {
 
-static constexpr int kTOPK = 6;
-
 // PyTorch wrapper function that calls the kernel implementation
 void gate_forward(th::Tensor scores_in, // [batch_size, nExperts] - pre-computed from linear(x, weight)
     th::Tensor bias,                    // empty tensor if hash mode
@@ -38,9 +36,9 @@ void gate_forward(th::Tensor scores_in, // [batch_size, nExperts] - pre-computed
     int64_t topk, double route_scale, bool is_hash, std::optional<th::Tensor> image_mask,
     std::optional<th::Tensor> vision_bias)
 {
-    TORCH_CHECK(topk == kTOPK, "topk must be ", kTOPK);
     auto const n_experts = scores_in.size(1);
-    TORCH_CHECK(n_experts == 256 || n_experts == 384, "n_experts must be 256 or 384");
+    TORCH_CHECK((n_experts == 256 && topk == 6) || (n_experts == 384 && topk == 6) || (n_experts == 128 && topk == 3),
+        "(n_experts, topk) must be one of (256, 6), (384, 6), (128, 3); got (", n_experts, ", ", topk, ")");
     TORCH_CHECK(scores_in.scalar_type() == torch::kFloat32, "scores_in must be float32");
     TORCH_CHECK(out_weights.scalar_type() == torch::kFloat32, "out_weights must be float32");
     TORCH_CHECK(out_indices.scalar_type() == torch::kInt32, "out_indices must be int32");
@@ -71,7 +69,8 @@ void gate_forward(th::Tensor scores_in, // [batch_size, nExperts] - pre-computed
     kernels::gate_forward(scores_in.data_ptr<float>(), is_hash ? nullptr : bias.data_ptr<float>(),
         is_hash ? input_ids.data_ptr<int>() : nullptr, is_hash ? tid2eid.data_ptr<int>() : nullptr,
         out_weights.data_ptr<float>(), out_indices.data_ptr<int>(), batch_size, static_cast<int>(n_experts),
-        static_cast<float>(route_scale), is_hash, image_mask.has_value() ? image_mask->data_ptr<bool>() : nullptr,
+        static_cast<int>(topk), static_cast<float>(route_scale), is_hash,
+        image_mask.has_value() ? image_mask->data_ptr<bool>() : nullptr,
         vision_bias.has_value() ? vision_bias->data_ptr<float>() : nullptr, stream);
 }
 

examples/llm-api/quickstart_advanced.py

@@ -22,9 +22,9 @@
 from tensorrt_llm.llmapi import (AttentionDpConfig, AutoDecodingConfig,
                                  CudaGraphConfig, DFlashDecodingConfig,
                                  DraftTargetDecodingConfig,
-                                 Eagle3DecodingConfig, KvCacheConfig, MoeConfig,
-                                 MTPDecodingConfig, NGramDecodingConfig,
-                                 TorchCompileConfig)
+                                 DSparkDecodingConfig, Eagle3DecodingConfig,
+                                 KvCacheConfig, MoeConfig, MTPDecodingConfig,
+                                 NGramDecodingConfig, TorchCompileConfig)
 
 example_prompts = [
     "Hello, my name is",
@@ -320,6 +320,13 @@ def setup_llm(args, **kwargs):
         spec_config = DFlashDecodingConfig(
             max_draft_len=args.spec_decode_max_draft_len,
             speculative_model=args.draft_model_dir)
+    elif spec_decode_algo == "DSPARK":
+        # Without --draft_model_dir the drafter is the target checkpoint's own
+        # mtp.* stages; TorchLlmArgs resolves that, and requires
+        # max_draft_len to equal the drafter's trained block size.
+        spec_config = DSparkDecodingConfig(
+            max_draft_len=args.spec_decode_max_draft_len,
+            speculative_model=args.draft_model_dir)
     elif spec_decode_algo == "DRAFT_TARGET":
         spec_config = DraftTargetDecodingConfig(
             max_draft_len=args.spec_decode_max_draft_len,

tensorrt_llm/_torch/_experimental/modeling_v2/README.md

@@ -20,7 +20,7 @@ export TRTLLM_MODELING_V2=require        # off | auto | require
 | `TRTLLM_MODELING_V2` | Behaviour |
 |---|---|
 | unset or `off` (default) | The resolver returns immediately. Nothing in this package is imported and behaviour is byte-for-byte what it is today. |
-| `auto` | Uses a target when one matches this exact configuration; falls back to the built-in implementation when none does. |
+| `auto` | Uses a target when one matches this exact configuration; falls back to the built-in implementation when none does. A family whose targets bring their own input processor (DeepSeek-V4.1) raises instead of falling back, because the processor is chosen per checkpoint and the built-in model expects another one. |
 | `require` | Raises instead of falling back, naming the criterion that did not match. |
 
 **Use `require` for anything you will attribute to modeling_v2.** Under `auto`,
@@ -36,8 +36,10 @@ its workers resolve the built-in is the silent split this package exists to
 prevent.
 
 An environment variable rather than an LLM-API field is a deliberate trade: it
-keeps the entire concept inside this package, so the only change modeling_v2
-needs anywhere else is the `_resolve_class` hook. The cost is that the switch
+keeps the entire concept inside this package, so the only changes modeling_v2
+needs anywhere else are the `_resolve_class` hook and, for a target family that
+reads its own multimodal inputs, one call in `create_input_processor`
+(`modeling_v2_input_processor`). The cost is that the switch
 does not appear in a run's recorded `llm.args` and cannot be set through
 `--extra_llm_api_options`.
 
@@ -256,10 +258,11 @@ the names but not the submodule paths, so all five files importing them
 failed at import. They now use the canonical path.
 
 The lesson generalizes: with no pin, a target's contact with private engine
-surface is checked only by running it. Each target declares that surface as
-`REQUIRED_TRTLLM_OPS`, `test_modeling_v2_target_contract.py` asserts every name
-in it exists, and the first-forward metadata field check catches the rest --
-which turns a drifting engine from a wrong answer into a loud failure.
+surface is checked only by running it. The step contract check does that by
+calling `_build_step_args` on the metadata it was handed, which reads every
+field the target consumes -- so a rename upstream becomes a loud failure rather
+than a wrong answer. Declaring the same surface a second time, as a list of
+names beside the code that reads it, only adds something to forget to update.
 
 ## Two facts this migration surfaced about upstream
 

tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py

@@ -17,8 +17,9 @@
 ``off`` (unset, the default) is byte-for-byte today's behaviour: the resolver returns
 immediately and nothing in this package is imported. ``auto`` uses a target
 when one matches and falls back to the built-in implementation when none
-does. ``require`` raises instead of falling back -- see ``_router_index`` for
-why that mode is not optional.
+does, unless the checkpoint's targets bring their own input processor, which
+the built-in model must not run behind. ``require`` raises instead of falling
+back -- see ``_router_index`` for why that mode is not optional.
 
 This package sits *beside* ``_torch/models/`` rather than inside it, which is
 load-bearing twice over: ``is_builtin_zoo_module`` matches on the zoo's
@@ -31,11 +32,14 @@
 from ._router_index import (
     MODELING_V2_ENV,
     MODELING_V2_ROUTERS,
+    MODELING_V2_VALIDATE_ENV,
     ROUTING_BACKENDS,
     ModelingV2Context,
     ModelingV2Mode,
     assert_backend_can_route,
+    modeling_v2_input_processor,
     modeling_v2_resolve,
+    step_contract_enabled,
 )
 
 __all__ = [
@@ -44,6 +48,9 @@
     "MODELING_V2_ROUTERS",
     "ModelingV2Context",
     "ModelingV2Mode",
+    "MODELING_V2_VALIDATE_ENV",
     "assert_backend_can_route",
+    "modeling_v2_input_processor",
     "modeling_v2_resolve",
+    "step_contract_enabled",
 ]