[Bug]: illegal memory access when attempting to use FP8 KV cache with DeepSeek R1 W4A8 on Hopper

System Info

Host: 4xH200, Ubuntu 24.04, CUDA 12.8

Container image: nvcr.io/nvidia/tensorrt-llm/release:1.2.0rc0

Who can help?

Quantization: @Tracin

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

• Image: nvcr.io/nvidia/tensorrt-llm/release:1.2.0rc0
• Model: https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8

cat >/root/data/trtllm-config.yml<<EOF
kv_cache_config:
  dtype: fp8
  enable_block_reuse: true
EOF
trtllm-serve /root/data/Barrrrry/DeepSeek-R1-W4AFP8 --backend pytorch --max_seq_len 9216 --max_batch_size 4096 --max_num_tokens 16384 --tp_size 4 --trust_remote_code --extra_llm_api_options /root/data/trtllm-config.yml

Results in:

/usr/local/lib/python3.12/dist-packages/torch/cuda/__init__.py:63: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
  import pynvml  # type: ignore[import]
/usr/local/lib/python3.12/dist-packages/modelopt/torch/utils/import_utils.py:32: UserWarning: Failed to import huggingface plugin due to: AttributeError("module 'transformers.modeling_utils' has no attribute 'Conv1D'"). You may ignore this warning if you do not need this plugin.
  warnings.warn(
/usr/local/lib/python3.12/dist-packages/modelopt/torch/__init__.py:36: UserWarning: transformers version 4.56.0 is incompatible with nvidia-modelopt and may cause issues. Please install recommended version with `pip install nvidia-modelopt[hf]` if working with HF models.
  _warnings.warn(
[TensorRT-LLM] TensorRT LLM version: 1.2.0rc0
/usr/local/lib/python3.12/dist-packages/pydantic/_internal/_fields.py:198: UserWarning: Field name "schema" in "ResponseFormat" shadows an attribute in parent "OpenAIBaseModel"
  warnings.warn(
[10/02/2025-17:27:44] [TRT-LLM] [I] Using LLM with PyTorch backend
[10/02/2025-17:27:44] [TRT-LLM] [I] Set nccl_plugin to None.
[10/02/2025-17:27:44] [TRT-LLM] [I] neither checkpoint_format nor checkpoint_loader were provided, checkpoint_format will be set to HF.
[10/02/2025-17:27:44] [TRT-LLM] [I] start MpiSession with 4 workers
[10/02/2025-17:27:44] [TRT-LLM] [I] Found /root/data/Barrrrry/DeepSeek-R1-W4AFP8/hf_quant_config.json, pre-quantized checkpoint is used.
[10/02/2025-17:27:44] [TRT-LLM] [I] Setting quant_algo=MIXED_PRECISION form HF quant config.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
You are using a model of type deepseek_v3 to instantiate a model of type . This is not supported for all configurations of models and can yield errors.
/root/data/Barrrrry/DeepSeek-R1-W4AFP8
rank 0 using MpiPoolSession to spawn MPI processes
[10/02/2025-17:27:44] [TRT-LLM] [I] Generating a new HMAC key for server proxy_request_queue
[10/02/2025-17:27:44] [TRT-LLM] [I] Generating a new HMAC key for server worker_init_status_queue
[10/02/2025-17:27:44] [TRT-LLM] [I] Generating a new HMAC key for server proxy_result_queue
[10/02/2025-17:27:44] [TRT-LLM] [I] Generating a new HMAC key for server proxy_stats_queue
[10/02/2025-17:27:44] [TRT-LLM] [I] Generating a new HMAC key for server proxy_kv_cache_events_queue
/usr/local/lib/python3.12/dist-packages/torch/cuda/__init__.py:63: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
  import pynvml  # type: ignore[import]
/usr/local/lib/python3.12/dist-packages/torch/cuda/__init__.py:63: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
  import pynvml  # type: ignore[import]
/usr/local/lib/python3.12/dist-packages/torch/cuda/__init__.py:63: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
  import pynvml  # type: ignore[import]
/usr/local/lib/python3.12/dist-packages/torch/cuda/__init__.py:63: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
  import pynvml  # type: ignore[import]
Multiple distributions found for package optimum. Picked distribution: optimum
Multiple distributions found for package optimum. Picked distribution: optimum
Multiple distributions found for package optimum. Picked distribution: optimum
Multiple distributions found for package optimum. Picked distribution: optimum
/usr/local/lib/python3.12/dist-packages/modelopt/torch/utils/import_utils.py:32: UserWarning: Failed to import huggingface plugin due to: AttributeError("module 'transformers.modeling_utils' has no attribute 'Conv1D'"). You may ignore this warning if you do not need this plugin.
  warnings.warn(
/usr/local/lib/python3.12/dist-packages/modelopt/torch/__init__.py:36: UserWarning: transformers version 4.56.0 is incompatible with nvidia-modelopt and may cause issues. Please install recommended version with `pip install nvidia-modelopt[hf]` if working with HF models.
  _warnings.warn(
/usr/local/lib/python3.12/dist-packages/modelopt/torch/utils/import_utils.py:32: UserWarning: Failed to import huggingface plugin due to: AttributeError("module 'transformers.modeling_utils' has no attribute 'Conv1D'"). You may ignore this warning if you do not need this plugin.
  warnings.warn(
/usr/local/lib/python3.12/dist-packages/modelopt/torch/utils/import_utils.py:32: UserWarning: Failed to import huggingface plugin due to: AttributeError("module 'transformers.modeling_utils' has no attribute 'Conv1D'"). You may ignore this warning if you do not need this plugin.
  warnings.warn(
/usr/local/lib/python3.12/dist-packages/modelopt/torch/__init__.py:36: UserWarning: transformers version 4.56.0 is incompatible with nvidia-modelopt and may cause issues. Please install recommended version with `pip install nvidia-modelopt[hf]` if working with HF models.
  _warnings.warn(
/usr/local/lib/python3.12/dist-packages/modelopt/torch/__init__.py:36: UserWarning: transformers version 4.56.0 is incompatible with nvidia-modelopt and may cause issues. Please install recommended version with `pip install nvidia-modelopt[hf]` if working with HF models.
  _warnings.warn(
/usr/local/lib/python3.12/dist-packages/modelopt/torch/utils/import_utils.py:32: UserWarning: Failed to import huggingface plugin due to: AttributeError("module 'transformers.modeling_utils' has no attribute 'Conv1D'"). You may ignore this warning if you do not need this plugin.
  warnings.warn(
/usr/local/lib/python3.12/dist-packages/modelopt/torch/__init__.py:36: UserWarning: transformers version 4.56.0 is incompatible with nvidia-modelopt and may cause issues. Please install recommended version with `pip install nvidia-modelopt[hf]` if working with HF models.
  _warnings.warn(
[TensorRT-LLM] TensorRT LLM version: 1.2.0rc0
[TensorRT-LLM] TensorRT LLM version: 1.2.0rc0
[TensorRT-LLM] TensorRT LLM version: 1.2.0rc0
[TensorRT-LLM] TensorRT LLM version: 1.2.0rc0
/usr/local/lib/python3.12/dist-packages/pydantic/_internal/_fields.py:198: UserWarning: Field name "schema" in "ResponseFormat" shadows an attribute in parent "OpenAIBaseModel"
  warnings.warn(
/usr/local/lib/python3.12/dist-packages/pydantic/_internal/_fields.py:198: UserWarning: Field name "schema" in "ResponseFormat" shadows an attribute in parent "OpenAIBaseModel"
  warnings.warn(
/usr/local/lib/python3.12/dist-packages/pydantic/_internal/_fields.py:198: UserWarning: Field name "schema" in "ResponseFormat" shadows an attribute in parent "OpenAIBaseModel"
  warnings.warn(
/usr/local/lib/python3.12/dist-packages/pydantic/_internal/_fields.py:198: UserWarning: Field name "schema" in "ResponseFormat" shadows an attribute in parent "OpenAIBaseModel"
  warnings.warn(
[TensorRT-LLM][INFO] Refreshed the MPI local session
[TensorRT-LLM][INFO] Refreshed the MPI local session
[TensorRT-LLM][INFO] Refreshed the MPI local session
[TensorRT-LLM][INFO] Refreshed the MPI local session
[10/02/2025-17:27:57] [TRT-LLM] [RANK 0] [I] ATTENTION RUNTIME FEATURES:  AttentionRuntimeFeatures(chunked_prefill=False, cache_reuse=True, has_speculative_draft_tokens=False, chunk_size=16384, chunked_prefill_buffer_batch_size=4)
/root/data/Barrrrry/DeepSeek-R1-W4AFP8
`torch_dtype` is deprecated! Use `dtype` instead!
/root/data/Barrrrry/DeepSeek-R1-W4AFP8
`torch_dtype` is deprecated! Use `dtype` instead!
/root/data/Barrrrry/DeepSeek-R1-W4AFP8
`torch_dtype` is deprecated! Use `dtype` instead!
/root/data/Barrrrry/DeepSeek-R1-W4AFP8
[10/02/2025-17:27:57] [TRT-LLM] [RANK 0] [I] Validating KV Cache config against kv_cache_dtype="fp8"
`torch_dtype` is deprecated! Use `dtype` instead!
[10/02/2025-17:27:58] [TRT-LLM] [RANK 0] [I] Use 85.40 GB for model weights.
[10/02/2025-17:27:58] [TRT-LLM] [RANK 0] [I] Prefetching 350.76GB checkpoint files.
[10/02/2025-17:27:58] [TRT-LLM] [RANK 0] [I] Prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00059-of-00070.safetensors to memory...
[10/02/2025-17:27:58] [TRT-LLM] [RANK 0] [I] Prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00008-of-00070.safetensors to memory...
[10/02/2025-17:27:58] [TRT-LLM] [RANK 0] [I] Prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00002-of-00070.safetensors to memory...
[10/02/2025-17:27:58] [TRT-LLM] [RANK 0] [I] Prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00001-of-00070.safetensors to memory...
[10/02/2025-17:27:58] [TRT-LLM] [RANK 0] [I] Prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00023-of-00070.safetensors to memory...
[10/02/2025-17:27:58] [TRT-LLM] [RANK 0] [I] Prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00048-of-00070.safetensors to memory...
[10/02/2025-17:27:58] [TRT-LLM] [RANK 0] [I] Prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00065-of-00070.safetensors to memory...
[10/02/2025-17:27:58] [TRT-LLM] [RANK 0] [I] Prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00011-of-00070.safetensors to memory...
[10/02/2025-17:27:58] [TRT-LLM] [RANK 0] [I] Prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00046-of-00070.safetensors to memory...
[10/02/2025-17:27:58] [TRT-LLM] [RANK 0] [I] Prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00052-of-00070.safetensors to memory...
[10/02/2025-17:27:58] [TRT-LLM] [RANK 0] [I] Prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00066-of-00070.safetensors to memory...
[10/02/2025-17:27:58] [TRT-LLM] [RANK 0] [I] Prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00039-of-00070.safetensors to memory...
[10/02/2025-17:27:58] [TRT-LLM] [RANK 0] [I] Prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00022-of-00070.safetensors to memory...
[10/02/2025-17:27:58] [TRT-LLM] [RANK 0] [I] Prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00018-of-00070.safetensors to memory...
[10/02/2025-17:27:58] [TRT-LLM] [RANK 0] [I] Prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00045-of-00070.safetensors to memory...
[10/02/2025-17:27:58] [TRT-LLM] [RANK 0] [I] Prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00044-of-00070.safetensors to memory...
[10/02/2025-17:28:11] [TRT-LLM] [RANK 0] [I] Finished prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00008-of-00070.safetensors.
[10/02/2025-17:28:11] [TRT-LLM] [RANK 0] [I] Prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00049-of-00070.safetensors to memory...
[10/02/2025-17:28:15] [TRT-LLM] [RANK 0] [I] Finished prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00066-of-00070.safetensors.
[10/02/2025-17:28:15] [TRT-LLM] [RANK 0] [I] Prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00056-of-00070.safetensors to memory...
[10/02/2025-17:28:15] [TRT-LLM] [RANK 0] [I] Finished prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00011-of-00070.safetensors.
[10/02/2025-17:28:15] [TRT-LLM] [RANK 0] [I] Finished prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00001-of-00070.safetensors.
[10/02/2025-17:28:15] [TRT-LLM] [RANK 0] [I] Finished prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00065-of-00070.safetensors.
[10/02/2025-17:28:15] [TRT-LLM] [RANK 0] [I] Finished prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00002-of-00070.safetensors.
[10/02/2025-17:28:16] [TRT-LLM] [RANK 0] [I] Finished prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00052-of-00070.safetensors.
[10/02/2025-17:28:16] [TRT-LLM] [RANK 0] [I] Finished prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00039-of-00070.safetensors.
[10/02/2025-17:28:16] [TRT-LLM] [RANK 0] [I] Finished prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00059-of-00070.safetensors.
[10/02/2025-17:28:17] [TRT-LLM] [RANK 0] [I] Finished prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00022-of-00070.safetensors.
[10/02/2025-17:28:17] [TRT-LLM] [RANK 0] [I] Finished prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00018-of-00070.safetensors.
[10/02/2025-17:28:18] [TRT-LLM] [RANK 0] [I] Finished prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00023-of-00070.safetensors.
[10/02/2025-17:28:19] [TRT-LLM] [RANK 0] [I] Finished prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00046-of-00070.safetensors.
[10/02/2025-17:28:19] [TRT-LLM] [RANK 0] [I] Finished prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00045-of-00070.safetensors.
[10/02/2025-17:28:19] [TRT-LLM] [RANK 0] [I] Finished prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00044-of-00070.safetensors.
[10/02/2025-17:28:19] [TRT-LLM] [RANK 0] [I] Finished prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00048-of-00070.safetensors.
[10/02/2025-17:28:22] [TRT-LLM] [RANK 0] [I] Finished prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00056-of-00070.safetensors.
[10/02/2025-17:28:22] [TRT-LLM] [RANK 0] [I] Finished prefetching /root/data/Barrrrry/DeepSeek-R1-W4AFP8/model-00049-of-00070.safetensors.
Loading safetensors weights in parallel: 100%|██████████| 71/71 [00:11<00:00,  6.42it/s]
Loading safetensors weights in parallel: 100%|██████████| 71/71 [00:11<00:00,  6.41it/s]
Loading safetensors weights in parallel: 100%|██████████| 71/71 [00:11<00:00,  6.39it/s]
Loading safetensors weights in parallel: 100%|██████████| 71/71 [00:11<00:00,  6.37it/s]
Loading weights: 100%|██████████| 1644/1644 [00:32<00:00, 50.70it/s]
Post loading weights: 100%|██████████| 1640/1640 [00:00<00:00, 695657.22it/s]
Model init total -- 68.78s
Loading weights: 100%|██████████| 1644/1644 [00:32<00:00, 50.65it/s]
Loading weights: 100%|██████████| 1644/1644 [00:32<00:00, 50.74it/s]
Post loading weights: 100%|██████████| 1640/1640 [00:00<00:00, 712223.91it/s]
Loading weights: 100%|██████████| 1644/1644 [00:32<00:00, 50.62it/s]
Post loading weights: 100%|██████████| 1640/1640 [00:00<00:00, 747924.17it/s]
Model init total -- 68.65s
Post loading weights:   0%|          | 0/1640 [00:00<?, ?it/s]Model init total -- 68.69s
Post loading weights: 100%|██████████| 1640/1640 [00:00<00:00, 753908.22it/s]
Model init total -- 68.74s
[10/02/2025-17:29:10] [TRT-LLM] [RANK 0] [I] Change tokens_per_block to: 64 for using FlashMLA
[10/02/2025-17:29:10] [TRT-LLM] [RANK 0] [I] Using Sampler: TorchSampler
[10/02/2025-17:29:10] [TRT-LLM] [RANK 0] [I] max_tokens is set by kv_cache_config.max_tokens: 16512
[10/02/2025-17:29:10] [TRT-LLM] [RANK 0] [I] Adjusted attention window size to 9217 in blocks_per_window
[TensorRT-LLM][INFO] Max KV cache blocks per sequence: 145 [window size=9217], tokens per block=64, primary blocks=258, secondary blocks=0
[TensorRT-LLM][INFO] Max KV cache blocks per sequence: 145 [window size=9217], tokens per block=64, primary blocks=258, secondary blocks=0
[TensorRT-LLM][INFO] Max KV cache blocks per sequence: 145 [window size=9217], tokens per block=64, primary blocks=258, secondary blocks=0
[TensorRT-LLM][INFO] Max KV cache blocks per sequence: 145 [window size=9217], tokens per block=64, primary blocks=258, secondary blocks=0
[TensorRT-LLM][INFO] Number of tokens per block: 64.
[TensorRT-LLM][INFO] [MemUsageChange] Allocated 0.54 GiB for max tokens in paged KV cache (16512).
[TensorRT-LLM][INFO] Number of tokens per block: 64.
[TensorRT-LLM][INFO] [MemUsageChange] Allocated 0.54 GiB for max tokens in paged KV cache (16512).
[TensorRT-LLM][INFO] Number of tokens per block: 64.
[TensorRT-LLM][INFO] [MemUsageChange] Allocated 0.54 GiB for max tokens in paged KV cache (16512).
[TensorRT-LLM][INFO] Number of tokens per block: 64.
[TensorRT-LLM][INFO] [MemUsageChange] Allocated 0.54 GiB for max tokens in paged KV cache (16512).
[10/02/2025-17:29:10] [TRT-LLM] [RANK 0] [I] max_seq_len=9217, max_num_requests=4096, max_num_tokens=16384, max_batch_size=4096
[10/02/2025-17:29:10] [TRT-LLM] [RANK 0] [I] cache_transceiver is disabled
[10/02/2025-17:29:10] [TRT-LLM] [RANK 0] [I] [Autotuner] Autotuning process starts ...
[TensorRT-LLM][WARNING] Attention workspace size is not enough, increase the size from 0 bytes to 638239744 bytes
[TensorRT-LLM][WARNING] Attention workspace size is not enough, increase the size from 0 bytes to 638239744 bytes
[TensorRT-LLM][WARNING] Attention workspace size is not enough, increase the size from 0 bytes to 638239744 bytes
CUDA error (../tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmha_v2_cu/fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_tma_ws_sm90.cu:446): an illegal memory access was encountered
[TensorRT-LLM][WARNING] Attention workspace size is not enough, increase the size from 0 bytes to 638239744 bytes
CUDA error (../tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmha_v2_cu/fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_tma_ws_sm90.cu:446): an illegal memory access was encountered
CUDA error (../tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmha_v2_cu/fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_tma_ws_sm90.cu:446): an illegal memory access was encountered
CUDA error (../tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmha_v2_cu/fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_tma_ws_sm90.cu:446): an illegal memory access was encountered
--------------------------------------------------------------------------
Child job 2 terminated normally, but 1 process returned
a non-zero exit code. Per user-direction, the job has been aborted.
--------------------------------------------------------------------------

Start by reproducing the reported trtllm-serve command with the provided kv_cache_config.yml, DeepSeek-R1-W4AFP8 model, FP8 KV cache, and PyTorch backend. Compare behavior with FP8 KV caching disabled or with other supported configurations to isolate the failure. Done means the reported illegal memory access no longer occurs for this setup.