[Bug]: crash on nano v3 trtllm-serve with concurrency 40 during runtime - due to moe multi-stream

System Info

H100

TRTLLM main: a792c23dcf8bdf4d05ce1eb8b01a709d894cc18d
git merge --squash origin/gk/fix_ssm_head_parallel

set in examples/auto_deploy/nano_v3.yaml:
world_size: 8

pip install aiperf

aiperf profile --model "/lustre/fs1/portfolios/coreai/users/chengzhang/models/NVIDIA-Nemotron-Nano-31B-A3-v3-Mcore-FP8-HF/" --url "0.0.0.0:8082" --endpoint-type "chat" --streaming --concurrency 40 --request-count 1120 --isl 1000 --osl 1000 --artifact-dir "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_modelopt/users/egeva/sweep_results/260105_0308_ad_tp1/isl_1000_osl_1000_conc_224" --num-warmup-requests 1 --extra-inputs "{\"ignore_eos\": true, \"min_tokens\": 1000}"

trtllm-serve /lustre/fs1/portfolios/coreai/users/chengzhang/models/NVIDIA-Nemotron-Nano-31B-A3-v3-Mcore-FP8-HF/   --host 0.0.0.0   --port 8082   --trust_remote_code   --backend _autodeploy   --extra_llm_api_options examples/auto_deploy/nano_v3.yaml 2>&1 | tee log.txt

error:

INFO: 127.0.0.1:41166 - "POST /v1/chat/completions HTTP/1.1" 200 OK
/opt/pytorch/pytorch/aten/src/ATen/native/cuda/TensorCompare.cu:112: _assert_async_cuda_kernel: block: [0,0,0], thread: [0,0,0] Assertion failed. /opt/pytorch/pytorch/aten/src/ATen/native/cuda/TensorCompare.cu:112: _assert_async_cuda_kernel: block: [0,0,0], thread: [0,0,0] Assertion failed.
/opt/pytorch/pytorch/aten/src/ATen/native/cuda/TensorCompare.cu:112: _assert_async_cuda_kernel: block: [0,0,0], thread: [0,0,0] Assertion `` failed.
Traceback (most recent call last):
File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_modelopt/users/egeva/TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/py_executor.py", line 2259, in _forward_step
outputs = forward(scheduled_requests, self.resource_manager,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_modelopt/users/egeva/TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/py_executor.py", line 2241, in forward
return self.model_engine.forward(
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 120, in decorate_context
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_modelopt/users/egeva/TensorRT-LLM/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py", line 356, in wrapper
return _call_func()
^^^^^^^^^^^^
File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_modelopt/users/egeva/TensorRT-LLM/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py", line 352, in _call_func
return func(self, scheduled_requests, resource_manager, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_modelopt/users/egeva/TensorRT-LLM/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py", line 814, in forward
"logits": self._compute_logits(),
^^^^^^^^^^^^^^^^^^^^^^
File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_modelopt/users/egeva/TensorRT-LLM/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py", line 784, in _compute_logits
logits: torch.Tensor = self.model(**self.cache_seq_interface.named_args)[0]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_modelopt/users/egeva/TensorRT-LLM/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py", line 155, in forward
return self.model(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph_module.py", line 837, in call_wrapped
return self._wrapped_call(self, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph_module.py", line 413, in call
raise e
File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph_module.py", line 400, in call
return super(self.cls, obj).call(*args, **kwargs) # type: ignore[misc]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1881, in _call_impl
return inner()
^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1829, in inner
result = forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "<eval_with_key>.3403", line 361, in forward
triton_ssm_prepare_metadata_default = torch.ops.auto_deploy.triton_ssm_prepare_metadata.default(position_ids, batch_info_host, seq_len, cu_seqlen, 128); seq_len = None
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 841, in call
return self._op(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_library/custom_ops.py", line 343, in backend_impl
result = self._backend_fns[device_type](*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_compile.py", line 53, in inner
return disable_fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 1044, in _fn
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_library/custom_ops.py", line 376, in wrapped_fn
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_modelopt/users/egeva/TensorRT-LLM/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_mamba.py", line 59, in _triton_ssm_prepare_metadata
chunk_indices, chunk_offsets = cu_seqlens_to_chunk_indices_offsets(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_modelopt/users/egeva/TensorRT-LLM/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py", line 63, in cu_seqlens_to_chunk_indices_offsets
N = math.ceil(total_seqlens / chunk_size) + (cu_seqlens[:-1] % chunk_size
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch.AcceleratorError: CUDA error: device-side assert triggered
Search for cudaErrorAssert' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information. CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect. For debugging consider passing CUDA_LAUNCH_BLOCKING=1 Compile with TORCH_USE_CUDA_DSA` to enable device-side assertions.

[01/05/2026-07:04:47] [TRT-LLM] [RANK 1] [E] Encountered an error in forward function: CUDA error: device-side assert triggered
Search for cudaErrorAssert' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information. CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect. For debugging consider passing CUDA_LAUNCH_BLOCKING=1 Compile with TORCH_USE_CUDA_DSA` to enable device-side assertions.

/opt/pytorch/pytorch/aten/src/ATen/native/cuda/TensorCompare.cu:112: _assert_async_cuda_kernel: block: [0,0,0], thread: [0,0,0Traceback (most recent call last):
] Assertion `` failed.
File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_modelopt/users/egeva/TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/py_executor.py", line 2387, in _update_requests
self.sampler.update_requests(sample_state, resource_manager)
File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 120, in decorate_context
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_modelopt/users/egeva/TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/sampler.py", line 1758, in update_requests
state.sampler_event.synchronize()
File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_modelopt/users/egeva/TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/sampler.py", line 108, in synchronize
self.cuda_event.synchronize()
File "/usr/local/lib/python3.12/dist-packages/torch/cuda/streams.py", line 231, in synchronize
super().synchronize()
torch.AcceleratorError: CUDA error: device-side assert triggered
Search for cudaErrorAssert' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information. CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect. For debugging consider passing CUDA_LAUNCH_BLOCKING=1 Compile with TORCH_USE_CUDA_DSA` to enable device-side assertions.

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

na

Expected behavior

na

actual behavior

na

additional notes

na

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.