[None][fix] Runtime fixes and small kernel changes from serving Qwen3.8-Flash-Next-FP8

cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu

@@ -337,6 +337,8 @@ struct MnnvlAllReduceKernelParams
     uint32_t* bufferFlags;
     bool waitForResults;
     QuantizationSFLayout layout;
+    // One-shot only: false releases PDL dependents before the reduction (their prologue overlaps it).
+    bool triggerCompletionAtEnd = true;
 };
 
 template <typename PackedType, typename T>
@@ -591,6 +593,10 @@ __global__ void __launch_bounds__(1024) oneshotAllreduceFusionKernel(MnnvlAllRed
     int threadOffset = token * params.tokenDim + packedIdx * kELTS_PER_THREAD;
 
     cudaGridDependencySynchronize();
+    if (!params.triggerCompletionAtEnd)
+    {
+        cudaTriggerProgrammaticLaunchCompletion();
+    }
 #else
     int packedIdx = blockIdx.y * blockDim.x + threadIdx.x;
     int token = blockIdx.x;
@@ -673,7 +679,10 @@ __global__ void __launch_bounds__(1024) oneshotAllreduceFusionKernel(MnnvlAllRed
         }
     }
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-    cudaTriggerProgrammaticLaunchCompletion();
+    if (params.triggerCompletionAtEnd)
+    {
+        cudaTriggerProgrammaticLaunchCompletion();
+    }
 #endif
     if constexpr (ar_fusion::HasResidual<Pattern>)
     {
@@ -845,6 +854,7 @@ void oneshotAllreduceFusionOp(AllReduceFusionParams const& params)
             reinterpret_cast<T*>(params.bufferPtrLocal), mcPtr, params.quantOut, params.scaleOut, params.scaleFactor,
             numTokens, tokenDim, params.nRanks, params.rank, static_cast<float>(params.epsilon), params.bufferFlags,
             false, params.layout};
+        kernelParams.triggerCompletionAtEnd = params.triggerCompletionAtEnd;
 
         switch (params.nRanks)
         {

cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h

@@ -52,6 +52,8 @@ struct AllReduceFusionParams
     ar_fusion::AllReduceFusionPattern pattern
         = ar_fusion::AllReduceFusionPattern::kAllReduce; //!< Fused epilogue pattern
     bool fp8Payload = false; //!< Two-shot residual RMSNorm calls exchange E4M3 (lossy; see twoshotAllreduceFusionOp)
+    //! One-shot only: false releases PDL dependents before the reduction instead of after it
+    bool triggerCompletionAtEnd = true;
 
     //! @}
 

cpp/tensorrt_llm/thop/allreduceOp.cpp

@@ -2226,7 +2226,7 @@ bool hasMnnvlNormOutput(AllReduceFusionOp fusionOp)
 std::vector<torch::Tensor> mnnvlFusionAllReduce(torch::Tensor& input, torch::optional<torch::Tensor> const& gamma,
     torch::optional<torch::Tensor> const& residual_in, torch::optional<double> epsilon, torch::Tensor& comm_buffer,
     torch::Tensor& buffer_flags, bool rmsnorm_fusion, torch::optional<torch::Tensor> const& scale, int64_t fusion_op_,
-    bool fp8_payload)
+    bool fp8_payload, bool trigger_completion_at_end)
 {
     auto* mcast_mem = tensorrt_llm::common::findMcastDevMemBuffer(comm_buffer.data_ptr());
     TORCH_CHECK(
@@ -2334,6 +2334,7 @@ std::vector<torch::Tensor> mnnvlFusionAllReduce(torch::Tensor& input, torch::opt
     allreduce_params.input = input.const_data_ptr();
     allreduce_params.pattern = pattern;
     allreduce_params.fp8Payload = fp8_payload;
+    allreduce_params.triggerCompletionAtEnd = trigger_completion_at_end;
     allreduce_params.layout = tensorrt_llm::QuantizationSFLayout::SWIZZLED;
 
     if (fusionOp == AllReduceFusionOp::NONE)
@@ -2513,7 +2514,8 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
         // CUDA graphs). Ordering comes from the input/output data deps.
         "mnnvl_fusion_allreduce(Tensor input, Tensor? gamma, Tensor? residual, "
         "float? epsilon, Tensor comm_buffer, Tensor buffer_flags, bool rmsnorm_fusion, "
-        "Tensor? scale=None, int fusion_op=0, bool fp8_payload=False) -> "
+        "Tensor? scale=None, int fusion_op=0, bool fp8_payload=False, "
+        "bool trigger_completion_at_end=True) -> "
         "Tensor[]");
     m.def(
         "allreduce("

tensorrt_llm/_torch/attention/backends/sparse/qsa/kernels.py

@@ -34,7 +34,7 @@ def _qsa_pdl_enabled(rows: int) -> bool:
     return rows <= 4 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" and get_sm_version() >= 90
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["num_rows"])
 def _qsa_unscale_block_table_kernel(
     scaled_block_table,
     block_table,
@@ -104,7 +104,19 @@ def triton_qsa_unscale_block_table(
     )
 
 
-@triton.jit
+# The counts follow the batch and the row-metadata views start at any row: keep both out of the
+# cache key (here and in the kernels below).
+@triton.jit(
+    do_not_specialize=["num_rows"],
+    do_not_specialize_on_alignment=[
+        "kv_lens",
+        "seq_lens",
+        "request_indices",
+        "logical_positions",
+        "sequence_lengths",
+        "visible_blocks",
+    ],
+)
 def _qsa_decode_token_mapping_kernel(
     kv_lens,
     seq_lens,
@@ -190,7 +202,17 @@ def triton_qsa_decode_token_mapping(
 QSA_GENERATION_MAPPING_MAX_SEQS = 256
 
 
-@triton.jit(do_not_specialize=["num_seqs", "num_rows"])
+@triton.jit(
+    do_not_specialize=["num_seqs", "num_rows"],
+    do_not_specialize_on_alignment=[
+        "kv_lens",
+        "seq_lens",
+        "request_indices",
+        "logical_positions",
+        "sequence_lengths",
+        "visible_blocks",
+    ],
+)
 def _qsa_generation_token_mapping_kernel(
     kv_lens,
     seq_lens,
@@ -836,7 +858,7 @@ def triton_qsa_decode_pre_indexer(
     return q
 
 
-@triton.jit
+@triton.jit(do_not_specialize_on_alignment=["logical_positions", "request_indices"])
 def _qsa_prefill_compress_kernel(
     logical_positions,
     request_indices,
@@ -1045,7 +1067,7 @@ def triton_qsa_prefill_compress(
     )
 
 
-@triton.jit
+@triton.jit(do_not_specialize_on_alignment=["query_positions", "sequence_lengths"])
 def _expand_qsa_block_indices_kernel(
     block_indices,
     query_positions,
@@ -1198,7 +1220,7 @@ def triton_expand_qsa_block_indices(
     return output
 
 
-@triton.jit
+@triton.jit(do_not_specialize_on_alignment=["query_positions", "request_indices"])
 def _qsa_paged_index_scores_kernel(
     q,
     index_cache,
@@ -1403,7 +1425,10 @@ def _qsa_paged_index_scores_prefetch_tile(
     )
 
 
-@triton.jit
+@triton.jit(
+    do_not_specialize=["num_rows"],
+    do_not_specialize_on_alignment=["query_positions", "request_indices"],
+)
 def _qsa_paged_index_scores_qtile_kernel(
     q,
     index_cache,
@@ -2409,7 +2434,7 @@ def triton_qsa_paged_sparse_gqa(
     return output
 
 
-@triton.jit
+@triton.jit(do_not_specialize_on_alignment=["request_ptr", "visible_ptr"])
 def _qsa_mtp_reuse_topk_kernel(
     shared_ptr,
     shared_visible_ptr,

tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py

@@ -192,7 +192,8 @@ def _(input,
           rmsnorm_fusion,
           scale=None,
           fusion_op: int = 0,
-          fp8_payload: bool = False):
+          fp8_payload: bool = False,
+          trigger_completion_at_end: bool = True):
         from tensorrt_llm.functional import AllReduceFusionOp
         op = AllReduceFusionOp(fusion_op)
         if op == AllReduceFusionOp.NONE and rmsnorm_fusion:

tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py

@@ -361,18 +361,16 @@ def make_routing_dummy_topk(
         assert topk_weights_for_tuner is None
         assert topk_ids_for_tuner is None
 
-    # Define hook to recreate dummy tensors when shape changes during profiling
+    # Define hook to recreate dummy tensors for every profiled bucket
     def recreate_dummy_topk_if_needed(
             inputs: List[torch.Tensor]) -> List[torch.Tensor]:
-        """Recreate dummy topk tensors if token count changed during profiling."""
+        """Give every profiled bucket its own dummy routing: the AutoTuner fills a
+        bucket's int top-k by repeating the tuned call's rows."""
         current_num_tokens = inputs[hidden_states_index].shape[0]
 
         # Only recreate if we originally created dummies
         if need_dummy_topk:
-            # Check if shape changed
-            if inputs[-1] is not None and inputs[-1].shape[
-                    0] != current_num_tokens:
-                # Recreate with new shape
+            if inputs[-1] is not None:
                 topk_weights_for_tuner, topk_ids_for_tuner = make_selected_dummy_topk(
                     current_num_tokens, inputs[hidden_states_index].device,
                     None)

tensorrt_llm/_torch/modules/mamba/causal_conv1d_triton.py

@@ -29,7 +29,7 @@
 from tensorrt_llm._utils import get_sm_version
 
 
-@triton.jit()
+@triton.jit(do_not_specialize=["num_cache_lines"])
 def _causal_conv1d_fwd_kernel(  # continuous batching
     # Pointers to matrices
     x_ptr,  # (dim, cu_seqlen) holding `batch` of actual sequences + padded sequences
@@ -43,7 +43,7 @@ def _causal_conv1d_fwd_kernel(  # continuous batching
     # Matrix dimensions
     dim: tl.constexpr,
     seqlen: tl.int32,  # cu_seqlen
-    num_cache_lines: tl.constexpr,  # added to support vLLM larger cache lines
+    num_cache_lines,  # runtime: the slot count follows startup memory
     # Strides
     stride_x_seq: tl.constexpr,  # stride to get to next sequence,
     stride_x_dim: tl.constexpr,  # stride to get to next feature-value,
@@ -563,7 +563,16 @@ def grid(META):
 # When calculating token 3's convolution, it should conv to token 1 (parent) and token 0 (grand-parent)
 # When calculating token 2's convolution, it should conv to token 0 (parent)
 # This kernel is a fused kernel which will also produce retrieve_parent_token based on retrieve_next_token & retrieve_next_sibling
-@triton.jit()
+# The request count and the index vectors' offsets change every step: keep them out of the cache key.
+@triton.jit(
+    do_not_specialize=["num_cache_lines", "batch"],
+    do_not_specialize_on_alignment=[
+        "cache_seqlens_ptr",
+        "conv_state_indices_ptr",
+        "num_accepted_tokens_ptr",
+        "intermediate_state_indices_ptr",
+    ],
+)
 def _causal_conv1d_update_kernel(
     # Pointers to matrices
     x_ptr,  # (batch, dim, seqlen)
@@ -584,7 +593,7 @@ def _causal_conv1d_update_kernel(
     dim: tl.constexpr,
     seqlen: tl.constexpr,
     state_len: tl.constexpr,
-    num_cache_lines: tl.constexpr,  # added to support vLLM larger cache lines
+    num_cache_lines,  # runtime: the slot count follows startup memory
     # Strides
     stride_x_seq: tl.constexpr,
     stride_x_dim: tl.constexpr,

tensorrt_llm/_torch/speculative/fused_accept.py

@@ -37,7 +37,11 @@
 MAX_DRAFT_LEN = 31  # columns of the accept tile (k + 1 <= 32)
 
 
-@triton.jit
+# The counts and the token views' offsets change every step: keep them out of the cache key.
+@triton.jit(
+    do_not_specialize=["num_contexts", "num_gens"],
+    do_not_specialize_on_alignment=["target_ptr", "draft_ptr"],
+)
 def _strict_accept_kernel(
     target_ptr,  # [num_contexts + num_gens * (K + 1)] sampled target tokens (any int dtype)
     draft_ptr,  # [num_gens, K] draft tokens of the previous step
@@ -139,7 +143,8 @@ def fused_strict_accept(
             pool = counter = num_accepted  # unused in mode 1
         pool_mask = pool_size - 1
     draft = draft_tokens.contiguous() if num_gens > 0 else num_accepted
-    block_b = max(16, min(64, triton.next_power_of_2(max(num_gens, num_contexts, 1))))
+    # One tile for every batch: the kernel loops over row blocks.
+    block_b = 64
     _strict_accept_kernel[(1,)](
         target_tokens,
         draft,
@@ -165,7 +170,8 @@ def fused_strict_accept(
     return accepted, num_accepted
 
 
-@triton.jit
+# The counts change every step: keep them out of the cache key.
+@triton.jit(do_not_specialize=["num_contexts", "num_ctx_tokens", "batch_size"])
 def _step0_epilogue_kernel(
     seq_lens_ptr,
     num_accepted_ptr,

tensorrt_llm/_torch/speculative/spec_tail_kernels.py

@@ -44,7 +44,8 @@
 
 
 @triton.jit(
-    do_not_specialize=["draft_ptr", "stride_draft_row", "stride_draft_col", "num_ctx", "num_gen"]
+    do_not_specialize=["draft_ptr", "stride_draft_row", "stride_draft_col", "num_ctx", "num_gen"],
+    do_not_specialize_on_alignment=["target_ptr"],
 )
 def _strict_accept_kernel(
     target_ptr,

tensorrt_llm/executor/result.py

@@ -6,6 +6,7 @@
 import json
 import math
 import time
+import traceback
 import weakref
 from dataclasses import dataclass, field
 from queue import Empty, Queue
@@ -867,8 +868,8 @@ def record_stats(self,
                         if math.isfinite(ppl):
                             metrics_stats[MetricNames.PREFILL_PERPLEXITY] = ppl
                 except (ValueError, TypeError):
-                    logger.debug("Failed to compute prefill perplexity",
-                                 exc_info=True)
+                    logger.debug("Failed to compute prefill perplexity:\n"
+                                 f"{traceback.format_exc()}")
 
         # Candidate-scoped metric: generation perplexity is computed per candidate
         # from this candidate's logprobs.
@@ -895,8 +896,8 @@ def record_stats(self,
                         mean_lp = sum(gen_lps) / len(gen_lps)
                         gen_ppl = math.exp(-mean_lp)
                 except (ValueError, TypeError):
-                    logger.debug("Failed to compute generation perplexity",
-                                 exc_info=True)
+                    logger.debug("Failed to compute generation perplexity:\n"
+                                 f"{traceback.format_exc()}")
             if gen_ppl is not None and math.isfinite(gen_ppl):
                 metrics_stats[MetricNames.GENERATION_PERPLEXITY] = gen_ppl
 

tensorrt_llm/executor/rpc/rpc_client.py

@@ -6,6 +6,7 @@
 import os
 import threading
 import time
+import traceback
 import uuid
 from datetime import datetime
 from typing import Any, AsyncIterator, Dict, Optional
@@ -402,8 +403,8 @@ async def _response_reader(self):
                             raise
                         except Exception as e:
                             # Log the error but continue reading unless it's a critical error
-                            logger.error(f"Error processing response: {e}",
-                                         exc_info=True)
+                            logger.error(f"Error processing response: {e}\n"
+                                         f"{traceback.format_exc()}")
 
                             # For critical errors, propagate and exit
                             if isinstance(e, (ConnectionError, zmq.ZMQError)):

tensorrt_llm/serve/openai_server.py

@@ -1085,10 +1085,10 @@ def _validate_internal_disagg_request(
             getattr(self, "_internal_disagg_auth_key", None), request, headers)
 
     def _get_scheduling_params(
-            self, request: ChatCompletionRequest,
+            self, request: Union[ChatCompletionRequest, CompletionRequest],
             raw_request: Optional[Request]) -> SchedulingParams:
         return SchedulingParams(
-            agent_hierarchy=request.agent_hierarchy,
+            agent_hierarchy=getattr(request, "agent_hierarchy", None),
             subagent_affinity_id=validate_subagent_affinity(
                 getattr(self, "_internal_disagg_auth_key", None), request,
                 getattr(self, "server_role", None),
@@ -2024,8 +2024,8 @@ async def _iteration_stats_collector_loop(self):
                         self._record_iteration_stat(llm_stat)
                 except Exception as e:
                     # Log errors but continue collecting stats
-                    logger.error(f"Error collecting iteration stats: {e}",
-                                 exc_info=True)
+                    logger.error(f"Error collecting iteration stats: {e}\n"
+                                 f"{traceback.format_exc()}")
                     # Brief sleep to avoid tight loop on persistent errors
                     await asyncio.sleep(0.1)
         except asyncio.CancelledError:
@@ -3112,6 +3112,8 @@ async def generator_wrapper(generator: AsyncIterator[Any]):
                 request, None if raw_request is None else raw_request.headers)
             conversation_params = to_llm_conversation_params(
                 request.conversation_params)
+            scheduling_params = self._get_scheduling_params(
+                request, raw_request)
             for idx, prompt in enumerate(prompts):
                 postproc_args = CompletionPostprocArgs.from_request(request)
                 self._apply_spec_decode_stats_opt_in(postproc_args)
@@ -3149,6 +3151,7 @@ async def generator_wrapper(generator: AsyncIterator[Any]):
                     lora_request=request.lora_request,
                     disaggregated_params=disaggregated_params,
                     conversation_params=conversation_params,
+                    scheduling_params=scheduling_params,
                     trace_headers=trace_headers,
                     kv_hints=request.kv_hints,
                     priority=request.priority if request.priority is not None

tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml

@@ -77,6 +77,7 @@ l0_gb200_multi_gpus:
   - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_checkpoint_preserves_cuda_graph_addresses
   - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_checkpoint_rejects_wrong_group_membership
   - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_workspace_growth_keeps_captured_graphs
+  - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_oneshot_early_pdl_release
   - unittest/_torch/modeling_v2/comm/test_modeling_v2_allgather_op_matrix.py
   - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_checkpoint_preserves_moe_graph_addresses
   - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_engine_checkpoint_coordination