Merge branch 'feat/deepseek_v4' into user/mingyangh/dsv4-oproj-fp8-fuse

Author: mingyangHao

cpp/include/tensorrt_llm/executor/executor.h

@@ -745,6 +745,7 @@ class Request
     ~Request();
 
     [[nodiscard]] VecTokens getInputTokenIds() const;
+    [[nodiscard]] SizeType32 getNumInputTokens() const;
     [[nodiscard]] SizeType32 getMaxTokens() const;
     [[nodiscard]] bool getStreaming() const;
     [[nodiscard]] SamplingConfig getSamplingConfig() const;

cpp/tensorrt_llm/executor/request.cpp

@@ -79,6 +79,11 @@ VecTokens Request::getInputTokenIds() const
     return mImpl->getInputTokenIds();
 }
 
+SizeType32 Request::getNumInputTokens() const
+{
+    return mImpl->getNumInputTokens();
+}
+
 SizeType32 Request::getMaxTokens() const
 {
     return mImpl->getMaxNewTokens();

cpp/tensorrt_llm/executor/requestImpl.h

@@ -120,6 +120,11 @@ class Request::Impl
         return mInputTokenIds;
     }
 
+    [[nodiscard]] SizeType32 getNumInputTokens() const
+    {
+        return static_cast<SizeType32>(mInputTokenIds.size());
+    }
+
     [[nodiscard]] SizeType32 getMaxNewTokens() const
     {
         return mMaxNewTokens;

cpp/tensorrt_llm/nanobind/executor/request.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -26,6 +26,7 @@
 #include "tensorrt_llm/runtime/cudaStream.h"
 
 #include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
 #include <nanobind/stl/chrono.h>
 #include <nanobind/stl/function.h>
 #include <nanobind/stl/list.h>
@@ -645,47 +646,69 @@ void initRequestBindings(nb::module_& m)
             nb::cast<std::optional<tle::CacheSaltIDType>>(state[33]), nb::cast<std::optional<tle::IdType>>(state[34]));
     };
 
+    // Convert input_token_ids to VecTokens. Fast path: a 1-D contiguous int32
+    // ndarray is memcpy'd into the vector (no per-element PyLong cast, which is
+    // O(ISL) on the GIL-held submit path). Anything else (list[int], etc.) falls
+    // back to the default nanobind sequence cast, so this is fully back-compatible.
+    // This complements PR #15134 (which bytes-encodes Request *pickling*); here we
+    // target Request *construction* on the RpcWorker.submit / _enqueue_request path.
+    auto toVecTokens = [](nb::handle ids) -> tle::VecTokens
+    {
+        nb::ndarray<int32_t const, nb::ndim<1>, nb::c_contig> arr;
+        if (nb::try_cast(ids, arr, /*convert=*/false))
+        {
+            tle::VecTokens out(arr.shape(0));
+            if (arr.shape(0) > 0)
+            {
+                std::memcpy(out.data(), arr.data(), arr.shape(0) * sizeof(int32_t));
+            }
+            return out;
+        }
+        return nb::cast<tle::VecTokens>(ids);
+    };
+
     nb::class_<tle::Request> request(m, "Request", nb::dynamic_attr());
     request
-        .def(nb::init<tle::VecTokens,                           // inputTokenIds
-                 tle::SizeType32,                               // maxTokens
-                 bool,                                          // streaming
-                 tle::SamplingConfig const&,                    // samplingConfig
-                 tle::OutputConfig const&,                      // outputConfig
-                 std::optional<tle::SizeType32> const&,         // endId
-                 std::optional<tle::SizeType32> const&,         // padId
-                 std::optional<std::vector<SizeType32>>,        // positionIds
-                 std::optional<std::list<tle::VecTokens>>,      // badWords
-                 std::optional<std::list<tle::VecTokens>>,      // stopWords
-                 std::optional<tle::Tensor>,                    // embeddingBias
-                 std::optional<tle::ExternalDraftTokensConfig>, // externalDraftTokensConfig
-                 std::optional<tle::PromptTuningConfig>,        // pTuningConfig
-                 std::optional<tle::MultimodalInput>,           // multimodalInput
-                 std::optional<tle::Tensor>,                    // multimodalEmbedding
-                 std::optional<tle::MropeConfig>,               // mRopeConfig
-                 std::optional<tle::LoraConfig>,                // loraConfig
-                 std::optional<tle::LookaheadDecodingConfig>,   // lookaheadConfig
-                 std::optional<tle::KvCacheRetentionConfig>,    // kvCacheRetentionConfig
-                 std::optional<std::string>,                    // logitsPostProcessorName
-                 std::optional<tle::LogitsPostProcessor>,       // logitsPostProcessor
-                 std::optional<tle::VecTokens>,                 // encoderInputTokenIds
-                 std::optional<tle::IdType>,                    // clientId
-                 bool,                                          // returnAllGeneratedTokens
-                 tle::PriorityType,                             // priority
-                 tle::RequestType,                              // type
-                 std::optional<tle::ContextPhaseParams>,        // contextPhaseParams
-                 std::optional<tle::Tensor>,                    // encoderInputFeatures
-                 std::optional<tle::SizeType32>,                // encoderOutputLength
-                 std::optional<tle::Tensor>,                    // crossAttentionMask
-                 SizeType32,                                    // numReturnSequences
-                 std::optional<tle::EagleConfig>,               // eagleConfig
-                 std::optional<tle::Tensor>,                    // skipCrossAttnBlocks
-                 std::optional<tle::GuidedDecodingParams>,      // guidedDecodingParams
-                 std::optional<tle::SizeType32>,                // languageAdapterUid
-                 std::optional<tle::MillisecondsType>,          // allottedTimeMs
-                 std::optional<tle::CacheSaltIDType>,           // cacheSaltID
-                 std::optional<tle::IdType>                     // disaggRequestId
-                 >(),
+        .def(
+            "__init__",
+            [toVecTokens](tle::Request* self,
+                nb::handle input_token_ids, // list[int] or int32 ndarray
+                tle::SizeType32 max_tokens, bool streaming, tle::SamplingConfig const& sampling_config,
+                tle::OutputConfig const& output_config, std::optional<tle::SizeType32> const& end_id,
+                std::optional<tle::SizeType32> const& pad_id, std::optional<std::vector<SizeType32>> position_ids,
+                std::optional<std::list<tle::VecTokens>> bad_words, std::optional<std::list<tle::VecTokens>> stop_words,
+                std::optional<tle::Tensor> embedding_bias,
+                std::optional<tle::ExternalDraftTokensConfig> external_draft_tokens_config,
+                std::optional<tle::PromptTuningConfig> prompt_tuning_config,
+                std::optional<tle::MultimodalInput> multimodal_input, std::optional<tle::Tensor> multimodal_embedding,
+                std::optional<tle::MropeConfig> mrope_config, std::optional<tle::LoraConfig> lora_config,
+                std::optional<tle::LookaheadDecodingConfig> lookahead_config,
+                std::optional<tle::KvCacheRetentionConfig> kv_cache_retention_config,
+                std::optional<std::string> logits_post_processor_name,
+                std::optional<tle::LogitsPostProcessor> logits_post_processor,
+                std::optional<tle::VecTokens> encoder_input_token_ids, std::optional<tle::IdType> client_id,
+                bool return_all_generated_tokens, tle::PriorityType priority, tle::RequestType type,
+                std::optional<tle::ContextPhaseParams> context_phase_params,
+                std::optional<tle::Tensor> encoder_input_features, std::optional<tle::SizeType32> encoder_output_length,
+                std::optional<tle::Tensor> cross_attention_mask, SizeType32 num_return_sequences,
+                std::optional<tle::EagleConfig> eagle_config, std::optional<tle::Tensor> skip_cross_attn_blocks,
+                std::optional<tle::GuidedDecodingParams> guided_decoding_params,
+                std::optional<tle::SizeType32> language_adapter_uid,
+                std::optional<tle::MillisecondsType> allotted_time_ms,
+                std::optional<tle::CacheSaltIDType> cache_salt_id, std::optional<tle::IdType> disagg_request_id)
+            {
+                new (self) tle::Request(toVecTokens(input_token_ids), max_tokens, streaming, sampling_config,
+                    output_config, end_id, pad_id, std::move(position_ids), std::move(bad_words), std::move(stop_words),
+                    std::move(embedding_bias), std::move(external_draft_tokens_config), std::move(prompt_tuning_config),
+                    std::move(multimodal_input), std::move(multimodal_embedding), std::move(mrope_config),
+                    std::move(lora_config), std::move(lookahead_config), std::move(kv_cache_retention_config),
+                    std::move(logits_post_processor_name), std::move(logits_post_processor),
+                    std::move(encoder_input_token_ids), client_id, return_all_generated_tokens, priority, type,
+                    std::move(context_phase_params), std::move(encoder_input_features), encoder_output_length,
+                    std::move(cross_attention_mask), num_return_sequences, std::move(eagle_config),
+                    std::move(skip_cross_attn_blocks), std::move(guided_decoding_params), language_adapter_uid,
+                    allotted_time_ms, cache_salt_id, disagg_request_id);
+            },
             // clang-format off
         nb::arg("input_token_ids"),
         nb::arg("max_tokens"),
@@ -726,8 +749,9 @@ void initRequestBindings(nb::module_& m)
         nb::arg("allotted_time_ms") = nb::none(),
                 nb::arg("cache_salt_id") = nb::none(),
         nb::arg("disagg_request_id") = nb::none()
-    )         // clang-format on
+    ) // clang-format on
         .def_prop_ro("input_token_ids", &tle::Request::getInputTokenIds)
+        .def_prop_ro("num_input_tokens", &tle::Request::getNumInputTokens)
         .def_prop_ro("max_tokens", &tle::Request::getMaxTokens)
         .def_prop_rw("streaming", &tle::Request::getStreaming, &tle::Request::setStreaming)
         .def_prop_rw("sampling_config", &tle::Request::getSamplingConfig, &tle::Request::setSamplingConfig)

jenkins/L0_Test.groovy

@@ -3358,7 +3358,7 @@ def launchTestJobs(pipeline, testFilter)
     fullSet += SBSATestConfigs.keySet()
 
     SBSASlurmTestConfigs = [
-        "GB200-4_GPUs-PyTorch-DS-1": ["auto:gb200-x4", "l0_gb200_multi_gpus_ds", 1, 1, 4],
+        "GB200-4_GPUs-PyTorch-DS-1": ["auto:gb200-x4-split", "l0_gb200_multi_gpus_ds", 1, 1, 4],
     ]
     fullSet += SBSASlurmTestConfigs.keySet()
 

tensorrt_llm/_torch/attention_backend/flashinfer.py

@@ -23,9 +23,12 @@
 except RuntimeError:
     # Override TORCH_CUDA_ARCH_LIST for JIT compilation of flashinfer kernels
     # since the existed TORCH_CUDA_ARCH_LIST may be too general and flashinfer requires sm75+.
-    capability = torch.cuda.get_device_capability()
-    arch_list = f"{capability[0]}.{capability[1]}"
-    os.environ["TORCH_CUDA_ARCH_LIST"] = arch_list
+    # Guard on a visible GPU: with CUDA_VISIBLE_DEVICES="" (pure client) the
+    # capability query would force a CUDA context at import time.
+    if torch.cuda.is_available():
+        capability = torch.cuda.get_device_capability()
+        arch_list = f"{capability[0]}.{capability[1]}"
+        os.environ["TORCH_CUDA_ARCH_LIST"] = arch_list
 
 from tensorrt_llm._utils import prefer_pinned
 

tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py

@@ -1174,6 +1174,14 @@ def _run_overlapped_indexer_prepare(
                 self.k_cache_update_event.record()
         else:
             weights, k_fp8, k_scale = pre_aux
+            # pre_aux tensors were allocated on aux_stream; record on the
+            # consuming stream so the caching allocator can't recycle them mid-use.
+            cur_stream = torch.cuda.current_stream()
+            weights.record_stream(cur_stream)
+            if k_fp8 is not None:
+                k_fp8.record_stream(cur_stream)
+            if k_scale is not None:
+                k_scale.record_stream(cur_stream)
             q = self._qk_projection_and_rope(qr, position_ids)
 
         q_fp8, q_scale = self._quantize_q(q)