diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 5fc507320ae1..4705db4e6e95 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -12,6 +12,7 @@ The following is a table of supported models for the PyTorch backend: | `DeepSeekV2ForCausalLM` [^5] | DeepSeek V2 | `deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct` | | `DeepseekV3ForCausalLM` | DeepSeek-V3, Kimi-K2 | `deepseek-ai/DeepSeek-V3` | | `DeepseekV32ForCausalLM` | DeepSeek-V3.2 | `deepseek-ai/DeepSeek-V3.2` | +| `DeepseekV4ForCausalLM` [^11] | DeepSeek-V4 | `deepseek-ai/DeepSeek-V4-Pro` | | `ExaoneForCausalLM` [^5] | EXAONE 3.5 | `LGAI-EXAONE/EXAONE-3.5-32B-Instruct` | | `Exaone4ForCausalLM` | EXAONE 4.0 | `LGAI-EXAONE/EXAONE-4.0-32B` | | `ExaoneMoEForCausalLM` | K-EXAONE | `LGAI-EXAONE/K-EXAONE-236B-A23B` | @@ -32,7 +33,7 @@ The following is a table of supported models for the PyTorch backend: | `LlamaForCausalLM` | Llama 3.1, Llama 3, Llama 2, LLaMA | `meta-llama/Meta-Llama-3.1-70B` | | `Llama4ForConditionalGeneration` | Llama 4 | `meta-llama/Llama-4-Scout-17B-16E-Instruct` | | `MiniMaxM2ForCausalLM` [^5] | MiniMax M2/M2.1/M2.7 | `MiniMaxAI/MiniMax-M2.7` | -| `MiniMaxM3SparseForConditionalGeneration` [^11]| MiniMax-M3 | `MiniMaxAI/MiniMax-M3` | +| `MiniMaxM3SparseForConditionalGeneration` [^12]| MiniMax-M3 | `MiniMaxAI/MiniMax-M3` | | `MistralForCausalLM` | Mistral | `mistralai/Mistral-7B-v0.1` | | `MixtralForCausalLM` | Mixtral | `mistralai/Mixtral-8x7B-v0.1` | | `MllamaForConditionalGeneration` | Llama 3.2 | `meta-llama/Llama-3.2-11B-Vision` | @@ -63,6 +64,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl | -------------------------------- | ----------------- | ---------- | -------------------------- | --------------------- | --------------- | --- | ---------------- | ----------------- | ------ | ------------- | ---------------- | -------------- | ------------------------ | --------------------- | --------------- | | `DeepseekV3ForCausalLM` | Yes | Yes | Yes | Yes | Yes [^1] | Yes | No | No | No | Yes | Yes | Yes [^2] | N/A | Yes | Yes | | `DeepseekV32ForCausalLM` | Yes | Yes | Yes | Yes | Yes | Yes | No | No | No | Yes | Yes | Yes | N/A | Yes | Yes | +| `DeepseekV4ForCausalLM` [^11] | Yes | Yes | Yes | Untested | Yes | Yes | No | No | No | Yes | Yes | Untested | Yes | Untested | Untested | | `Glm4MoeForCausalLM` | Yes | Yes | Yes | Untested | Yes | Yes | No | No | No | Yes | Yes | Untested | N/A | Yes | Yes | | `Qwen3MoeForCausalLM` | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | N/A | Yes | Yes | | `Qwen3NextForCausalLM` [^3] | Yes | Yes | Yes | Untested | Yes | No | No | No | No | Yes | Yes | No | No | Untested | Untested | @@ -73,7 +75,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl | `NemotronHForCausalLM` | Yes | Yes | Yes | Yes | Yes | Yes | No | No | No | Yes | Yes | Yes | N/A | Untested | Untested | | `Gemma4ForConditionalGeneration` | Untested | Yes | Untested | No | Yes | No | No | No | No | Yes | Untested | No | Yes | Untested | Untested | | `Step3p7ForConditionalGeneration`| Yes | Yes | Yes | Untested | Untested | Yes | No | No | No | Yes | Untested | Untested | Yes | Untested | Untested | -| `MiniMaxM3SparseForConditionalGeneration` [^11] | Yes | Yes | Yes | Untested | Untested | No | No | No | No | Yes | Untested | No | N/A | Untested | Untested | +| `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes | Yes | Yes | Untested | Untested | No | No | No | No | Yes | Untested | No | N/A | Untested | Untested | [^1]: Chunked Prefill for MLA can only be enabled on SM100/SM103. [^2]: KV cache reuse for MLA can only be enabled on SM90/SM100/SM103 and in BF16/FP8 KV cache dtype. @@ -84,7 +86,8 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl [^8]: Supports text and image inputs. The vision tower runs in BF16 even when the text decoder is quantized (FP8 block-scale or NVFP4). The text decoder is also usable standalone (text-only) via the `Step3p5ForCausalLM` architecture. [^9]: Audio modality only supported on E2B/E4B variants. [^10]: Audio requires a checkpoint with a `sound_config` and is supported only on the full (non-disaggregated) model path, not the EPD disaggregated path. -[^11]: Supports text, image, and video inputs over the block-sparse attention path. The model is served in BF16. The text decoder is also usable standalone (text-only) via the `MiniMaxM3SparseForCausalLM` architecture. KV cache reuse and MTP are not supported on the sparse-attention path in this release. +[^11]: DeepSeek-V4 is only supported on Blackwell GPUs (`SM100+`). See the [DeepSeek-V4 example README](../../../examples/models/core/deepseek_v4/README.md) for setup and parallelism. +[^12]: Supports text, image, and video inputs over the block-sparse attention path. The published MXFP8 checkpoint is dequantized on load so the runtime sees an effectively BF16 model. The text decoder is also usable standalone (text-only) via the `MiniMaxM3SparseForCausalLM` architecture. KV cache reuse and MTP are not supported on the sparse-attention path in this release. # Multimodal Feature Support Matrix (PyTorch Backend) @@ -105,7 +108,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl | `Qwen3VLForConditionalGeneration` | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | L + I + V | | `Qwen3VLMoeForConditionalGeneration` | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | L + I + V | | `Step3p7ForConditionalGeneration` | Yes | Yes | Untested | Yes | Untested | Untested | Untested | Untested | L + I | -| `MiniMaxM3SparseForConditionalGeneration` [^11] | Yes | Yes | Untested | Yes | Untested | No | Untested | Untested | L + I + V | +| `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes | Yes | Untested | Yes | Untested | No | Untested | Untested | L + I + V | Note: - L: Language diff --git a/examples/models/core/deepseek_v4/README.md b/examples/models/core/deepseek_v4/README.md new file mode 100644 index 000000000000..fe7d4976c114 --- /dev/null +++ b/examples/models/core/deepseek_v4/README.md @@ -0,0 +1,488 @@ +# DeepSeek-V4 + +This guide walks you through the examples to run DeepSeek-V4 models using NVIDIA TensorRT LLM with +the PyTorch backend. + +DeepSeek-V4 uses the `DeepseekV4ForCausalLM` architecture in TensorRT LLM. Compared with +DeepSeek-V3/R1/V3.2, it has a separate model implementation and sparse attention path. Use the +commands in this guide as starting points and tune the parallelism and memory settings for your +checkpoint and workload. + +Please refer to [this guide](https://nvidia.github.io/TensorRT-LLM/installation/build-from-source-linux.html) +for how to build TensorRT LLM from source and start a TRT-LLM Docker container. + +> [!NOTE] +> This guide assumes that you replace placeholder values such as `` with the +> appropriate paths. Commands in this guide target the PyTorch backend. + + +## Table of Contents + +- [DeepSeek-V4](#deepseek-v4) + - [Table of Contents](#table-of-contents) + - [Hardware Requirements](#hardware-requirements) + - [Downloading the Model Weights](#downloading-the-model-weights) + - [Quick Start](#quick-start) + - [Run a single inference](#run-a-single-inference) + - [Run chat-style prompts](#run-chat-style-prompts) + - [Multi-Token Prediction (MTP)](#multi-token-prediction-mtp) + - [Benchmarking](#benchmarking) + - [Evaluation](#evaluation) + - [Serving](#serving) + - [trtllm-serve](#trtllm-serve) + - [OpenAI-compatible request](#openai-compatible-request) + - [Advanced Configuration](#advanced-configuration) + - [Parallelism](#parallelism) + - [Sparse attention](#sparse-attention) + - [KV cache](#kv-cache) + - [Quantized checkpoints](#quantized-checkpoints) + - [Notes and Troubleshooting](#notes-and-troubleshooting) + + +## Hardware Requirements + +DeepSeek-V4 is only supported on Blackwell GPUs (`SM100+`) in the current PyTorch backend +implementation. Pre-Blackwell GPUs are not supported for this model path. + +DeepSeek-V4 has two model scales, and each scale provides Base and Instruct checkpoints. The table +below follows the model list published on the +[DeepSeek-V4 Hugging Face model card](https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro): + +| Checkpoint | Total Params | Activated Params | Context Length | Precision | +| --- | --- | --- | --- | --- | +| DeepSeek-V4-Flash-Base | 284B | 13B | 1M | FP8 Mixed | +| DeepSeek-V4-Flash | 284B | 13B | 1M | FP4 + FP8 Mixed | +| DeepSeek-V4-Pro-Base | 1.6T | 49B | 1M | FP8 Mixed | +| DeepSeek-V4-Pro | 1.6T | 49B | 1M | FP4 + FP8 Mixed | + +The minimum number of GPUs depends on the model scale, checkpoint precision, KV cache budget, +maximum sequence length, and runtime batch size. For initial bring-up, an 8xB200 node is enough for +Flash checkpoints and the FP4 + FP8 mixed DeepSeek-V4-Pro checkpoint. DeepSeek-V4-Pro-Base is larger +because it uses FP8 mixed precision; if you want to keep the deployment on a single node, use an +8xB300 node. Multi-node Blackwell deployments are still recommended for larger KV cache budgets, +longer context windows, or higher throughput targets. Tune `--tp_size`, `--ep_size`, +`--max_num_tokens`, and the KV cache memory fraction for your deployment target. + +DeepSeek-V4 requires KV cache block sizes of 128 or 256 tokens. TensorRT LLM defaults DeepSeek-V4 to +`tokens_per_block=128`, but scripts that set their own KV cache config should pass this explicitly. + + +## Downloading the Model Weights + +Choose one of the DeepSeek-V4 checkpoint IDs: + +| Checkpoint | Hugging Face model ID | Prompt format | +| --- | --- | --- | +| DeepSeek-V4-Flash-Base | `deepseek-ai/DeepSeek-V4-Flash-Base` | Raw completion | +| DeepSeek-V4-Flash | `deepseek-ai/DeepSeek-V4-Flash` | Chat/Instruct | +| DeepSeek-V4-Pro-Base | `deepseek-ai/DeepSeek-V4-Pro-Base` | Raw completion | +| DeepSeek-V4-Pro | `deepseek-ai/DeepSeek-V4-Pro` | Chat/Instruct | + +Then download the weights: + +```bash +git lfs install +MODEL_ID=deepseek-ai/DeepSeek-V4-Flash +git clone https://huggingface.co/${MODEL_ID} +``` + +At minimum, the checkpoint config should identify the architecture as DeepSeek-V4: + +```json +{ + "architectures": ["DeepseekV4ForCausalLM"], + "model_type": "deepseek_v4" +} +``` + +Do not replace the full checkpoint config with this minimal snippet. TensorRT LLM also reads +DeepSeek-V4-specific sparse attention fields such as `compress_ratios`, `window_size` or +`sliding_window`, and indexer settings from the checkpoint config unless you provide a complete +override through `sparse_attention_config`. + + +## Quick Start + +### Run a single inference + +To quickly run DeepSeek-V4, use [examples/llm-api/quickstart_advanced.py](../../../llm-api/quickstart_advanced.py): + +```bash +cd examples/llm-api +python quickstart_advanced.py \ + --model_dir \ + --tp_size 8 \ + --moe_ep_size 8 \ + --tokens_per_block 128 \ + --max_num_tokens 8192 \ + --max_seq_len 4096 \ + --kv_cache_fraction 0.5 +``` + +The command above assumes one 8-GPU node. If you use a different number of GPUs, adjust `--tp_size` +and `--moe_ep_size` so that the requested parallelism matches your available world size. DeepSeek-V4 +checkpoints advertise a 1M-token context window; for bring-up, set `--max_seq_len` and the KV cache +memory fraction explicitly, then increase them according to your memory budget. + +### Run chat-style prompts + +DeepSeek-V4 Instruct checkpoints (`DeepSeek-V4-Flash` and `DeepSeek-V4-Pro`) use the checkpoint +reference chat/message format. TensorRT LLM provides a `deepseek_v4` tokenizer wrapper for this +format. Use `custom_tokenizer="deepseek_v4"` only with Instruct checkpoints and chat-style prompts. + +Base checkpoints (`DeepSeek-V4-Flash-Base` and `DeepSeek-V4-Pro-Base`) are completion models. For +Base checkpoints, do not apply a chat template and do not pass `custom_tokenizer="deepseek_v4"`; +send raw text prompts instead. + +```python +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm.llmapi import KvCacheConfig + +def main(): + llm = LLM( + model="", + backend="pytorch", + tensor_parallel_size=8, + moe_expert_parallel_size=8, + custom_tokenizer="deepseek_v4", + kv_cache_config=KvCacheConfig( + tokens_per_block=128, + free_gpu_memory_fraction=0.5, + ), + max_seq_len=4096, + max_num_tokens=8192, + ) + + messages = [{"role": "user", "content": "Explain TensorRT LLM in one paragraph."}] + prompt = llm.tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + ) + outputs = llm.generate([prompt], SamplingParams(max_tokens=128)) + print(outputs[0].outputs[0].text) + + +if __name__ == "__main__": + main() +``` + +### Multi-Token Prediction (MTP) + +If the checkpoint contains MTP layers, run MTP speculative decoding with the one-model flow: + +```bash +cd examples/llm-api +python quickstart_advanced.py \ + --model_dir \ + --tp_size 8 \ + --moe_ep_size 8 \ + --tokens_per_block 128 \ + --max_num_tokens 8192 \ + --max_seq_len 4096 \ + --kv_cache_fraction 0.5 \ + --spec_decode_algo MTP \ + --spec_decode_max_draft_len N \ + --use_one_model +``` + +`N` is the number of draft tokens to predict. Start with `N=1` for bring-up, then increase it after +validating accuracy and latency for your workload. + + +## Benchmarking + +The following example prepares a synthetic dataset and runs `trtllm-bench` throughput on one 8-GPU +Blackwell node: + +```bash +trtllm-bench --model \ + --model_path \ + prepare-dataset \ + --output /tmp/deepseek_v4_1k1k.txt \ + token-norm-dist \ + --input-mean 1024 \ + --output-mean 1024 \ + --input-stdev 0 \ + --output-stdev 0 \ + --num-requests 256 + +cat > /tmp/deepseek_v4_config.yml < \ + --model_path \ + throughput \ + --tp 8 \ + --ep 8 \ + --dataset /tmp/deepseek_v4_1k1k.txt \ + --max_batch_size 256 \ + --max_num_tokens 8192 \ + --concurrency 2048 \ + --num_requests 6144 \ + --kv_cache_free_gpu_mem_fraction 0.9 \ + --config /tmp/deepseek_v4_config.yml +``` + +The example enables attention DP because it is typically beneficial for high-throughput, large-batch +workloads. It also uses FP8 KV cache (`kv_cache_config.dtype: fp8`), which is the recommended +starting point for benchmarking DeepSeek-V4 throughput. For checkpoints with MTP layers, enable MTP +for benchmarking as well: use `num_nextn_predict_layers: 1` for throughput-oriented runs, and use +`num_nextn_predict_layers: 3` for low-latency runs. When `enable_attention_dp` is enabled, +`--max_batch_size` is the maximum batch size per local rank; use `--concurrency` high enough to +saturate all ranks. Tune `--max_batch_size`, `--max_num_tokens`, `--concurrency`, MTP depth, and the +KV cache memory fraction for the target ISL/OSL distribution. + + +## Evaluation + +Evaluate model accuracy using `trtllm-eval`. The following commands are for Instruct checkpoints and +apply the DeepSeek-V4 chat template through `--custom_tokenizer deepseek_v4` and +`--apply_chat_template`. For Base checkpoints, remove both flags because Base models expect raw +completion prompts. `--custom_tokenizer` is a top-level `trtllm-eval` option, so keep it before the +dataset subcommand such as `mmlu`, `gsm8k`, or `gpqa_diamond`. + +1. Prepare a configuration file: + +```bash +cat > ./deepseek_v4_config.yml < \ + --tp_size 8 \ + --ep_size 8 \ + --max_batch_size 16 \ + --max_num_tokens 8192 \ + --max_seq_len 4096 \ + --custom_tokenizer deepseek_v4 \ + --config ./deepseek_v4_config.yml \ + mmlu \ + --apply_chat_template +``` + +3. Evaluate GSM8K with an Instruct checkpoint: + +```bash +trtllm-eval --model \ + --tp_size 8 \ + --ep_size 8 \ + --max_batch_size 16 \ + --max_num_tokens 8192 \ + --max_seq_len 4096 \ + --custom_tokenizer deepseek_v4 \ + --config ./deepseek_v4_config.yml \ + gsm8k \ + --apply_chat_template \ + --system_prompt "Solve the problem carefully. End your response with a final line exactly in the form #### , using the simplest numeric form without units or trailing zeros." +``` + +The `--system_prompt` constrains the answer format so that the lm-eval `strict-match` +regex (which expects a final `#### ` line) can pick up the model's answer. +Without it, DeepSeek-V4 Instruct checkpoints often return the correct value in a +free-form sentence, which `flexible-extract` recovers but `strict-match` does not. + +4. Evaluate GPQA Diamond with an Instruct checkpoint: + +```bash +trtllm-eval --model \ + --tp_size 8 \ + --ep_size 8 \ + --max_batch_size 16 \ + --max_num_tokens 8192 \ + --max_seq_len 4096 \ + --custom_tokenizer deepseek_v4 \ + --config ./deepseek_v4_config.yml \ + gpqa_diamond \ + --apply_chat_template +``` + + +## Serving + +### trtllm-serve + +Create a serving config: + +```bash +cat > ./deepseek_v4_serve.yml < \ + --backend pytorch \ + --host 0.0.0.0 \ + --port 8000 \ + --tp_size 8 \ + --ep_size 8 \ + --max_seq_len 4096 \ + --custom_tokenizer deepseek_v4 \ + --config ./deepseek_v4_serve.yml +``` + +The `/v1/chat/completions` API applies chat formatting on the server side, so clients should send +OpenAI-style `messages` rather than preformatted prompt strings. For Base checkpoints, use the same +command but remove `--custom_tokenizer deepseek_v4`. Increase `max_seq_len`, `max_batch_size`, and +the KV cache memory fraction after validating the memory budget for your target deployment. + +### OpenAI-compatible request + +For Instruct checkpoints, send a chat-completions request: + +```bash +curl http://localhost:8000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "", + "messages": [ + { + "role": "user", + "content": "Write a short summary of TensorRT LLM." + } + ], + "stream": true, + "max_tokens": 128 + }' +``` + +For Base checkpoints, use the text completions API with a raw prompt: + +```bash +curl http://localhost:8000/v1/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "", + "prompt": "TensorRT LLM is", + "stream": true, + "max_tokens": 128 + }' +``` + + +## Advanced Configuration + +### Parallelism + +DeepSeek-V4 supports the same main PyTorch backend parallelism knobs used by other large MoE models: + +- Tensor parallelism (`--tp_size` or `tensor_parallel_size`) shards attention and dense weights. +- Pipeline parallelism (`--pp_size` or `pipeline_parallel_size`) distributes model layers across + pipeline stages, which can help fit larger checkpoints or larger KV cache budgets across more + GPUs. +- Expert parallelism (`--ep_size` or `moe_expert_parallel_size`) distributes routed experts. +- Attention DP (`enable_attention_dp: true`) keeps attention data-parallel across ranks and is + commonly used for high-throughput, large-batch serving. + +For latency-oriented tests, start without attention DP. For throughput-oriented tests, enable +attention DP in YAML: + +```yaml +enable_attention_dp: true +attention_dp_config: + batching_wait_iters: 0 + enable_balance: true + timeout_iters: 60 +``` + +When attention DP is enabled, remember that `max_batch_size` is local-rank batch size. Increase +`concurrency` and `num_requests` accordingly when benchmarking. + +### Sparse attention + +If `sparse_attention_config` is not provided, TensorRT LLM configures DeepSeek-V4 sparse attention +from the model config. It reads fields such as `compress_ratios`, `window_size` or `sliding_window`, +and indexer settings, then constructs the corresponding `DeepSeekV4SparseAttentionConfig`. + +If `sparse_attention_config` is provided, user values override the corresponding sparse attention +settings, subject to the current implementation constraints: `window_size` must be `128`, and +`compress_ratios` must use supported ratios (`1`, `4`, or `128`). If checkpoint `compress_ratios` +are present and longer than the user-provided list, TensorRT LLM keeps the checkpoint list to avoid +silently changing the sparse attention layout. + +Example YAML override: + +```yaml +sparse_attention_config: + algorithm: deepseek_v4 + window_size: 128 + index_topk: 512 +``` + +### KV cache + +DeepSeek-V4 uses `DeepseekV4CacheManager`, a `KvCacheManagerV2` subclass. This manager can describe +different cache layer types per model layer, so DeepSeek-V4 can map sliding-window, compressed, +indexer, and compressor-state caches according to the sparse attention layout from the model config +or user-provided `sparse_attention_config`. + +DeepSeek-V4 KV cache requires: + +- `tokens_per_block` set to `128` or `256`. +- `max_beam_width=1`. +- Blackwell GPUs for the current implementation. + +Use a lower `free_gpu_memory_fraction`, `max_batch_size`, or `max_num_tokens` if the workload runs +out of memory during initialization or prefill. + +### Quantized checkpoints + +TensorRT LLM detects supported quantization metadata from the checkpoint directory, including +`hf_quant_config.json`, `quantization_config`, or `dtypes.json`. For DeepSeek-V4 checkpoints with +MXFP4 routed MoE expert weights, TensorRT LLM automatically applies the routed-expert quantization +configuration. + + +## Notes and Troubleshooting + +- `DeepseekV4CacheManager requires tokens_per_block in [128, 256]`: pass + `--tokens_per_block 128` in `quickstart_advanced.py` or set + `kv_cache_config.tokens_per_block: 128` in YAML. +- `DeepSeek-V4 is not supported on pre-blackwell GPUs`: run on Blackwell GPUs (`SM100+`). +- Out-of-memory during initialization or prefill: reduce `max_batch_size`, `max_num_tokens`, or + `kv_cache_config.free_gpu_memory_fraction`. For bring-up on 8xB200, set `max_seq_len` explicitly + instead of using the checkpoint's 1M-token context length. +- Chat formatting issues with `trtllm-serve` or `trtllm-eval` on Instruct checkpoints: pass + `--custom_tokenizer deepseek_v4`. Do not use this tokenizer wrapper for Base checkpoints. +- Tool-call chat formatting is not supported by the DeepSeek-V4 tokenizer wrapper yet. diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index f0081dc10c4a..adcd7d7ecc17 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -1885,20 +1885,53 @@ def prepare_one_prefill_chunk( ) @staticmethod - def prepare_for_update_k_cache(metadata: DSAtrtllmAttentionMetadata, - indexer_params: IndexerParams): - """ - Prepare indexer for the update_k_cache stage. + def build_indexer_params( + metadata: DSAtrtllmAttentionMetadata) -> Optional[IndexerParams]: + kv_cache_manager = metadata.kv_cache_manager + if kv_cache_manager is None or not hasattr(kv_cache_manager, + 'index_head_dim'): + return None - Compute slot_mapping for all requests (both context and generation) - This maps each token to its flat cache position for vectorized KV cache updates + head_dim = metadata.indexer_head_dim + data_bytes_per_token = head_dim // 2 if getattr( + kv_cache_manager, 'use_fp4', False) else head_dim + compress_ratio = _effective_compress_ratio_divisor( + _select_indexer_compress_ratio(metadata.compress_ratios)) + return IndexerParams( + num_contexts=metadata.num_contexts, + num_generations=metadata.num_generations, + num_ctx_tokens=metadata.num_ctx_tokens, + head_dim=head_dim, + quant_block_size=metadata.indexer_quant_block_size, + tokens_per_block=metadata._tokens_per_block, + compress_ratio=compress_ratio, + request_ids=metadata.request_ids, + num_past_tokens=metadata.kv_cache_params.num_cached_tokens_per_seq, + seq_lens=metadata.seq_lens, + data_bytes_per_token=data_bytes_per_token, + ) + + @staticmethod + def recompute_slot_mappings(metadata: DSAtrtllmAttentionMetadata, + indexer_params: Optional[IndexerParams] = None): + """Recompute slot_mapping_fp8/scale from current block offsets. + + This is the subset of prepare_for_update_k_cache() that maps each new + compressed KV token to its flat cache position. It is safe to call in + isolation after a caller has swapped the active KV cache manager and + block-offset buffers, such as during draft KV-cache replay. """ + if indexer_params is None: + indexer_params = Indexer.build_indexer_params(metadata) + if indexer_params is None: + return + batch_size = indexer_params.batch_size tokens_per_block = indexer_params.tokens_per_block head_dim = indexer_params.head_dim new_kv_tokens = indexer_params.new_kv_tokens total_new_kv_tokens = new_kv_tokens.sum().item() - data_bytes_per_token = head_dim // 2 if metadata.kv_cache_manager.use_fp4 else head_dim + data_bytes_per_token = indexer_params.data_bytes_per_token # Compute global positions for all kv tokens in the batch (fully vectorized) req_indices = torch.repeat_interleave( @@ -1941,6 +1974,17 @@ def prepare_for_update_k_cache(metadata: DSAtrtllmAttentionMetadata, metadata.host_slot_mapping_scale[:total_new_kv_tokens], non_blocking=True) + @staticmethod + def prepare_for_update_k_cache(metadata: DSAtrtllmAttentionMetadata, + indexer_params: IndexerParams): + """ + Prepare indexer for the update_k_cache stage. + + Compute slot_mapping for all requests (both context and generation) + This maps each token to its flat cache position for vectorized KV cache updates + """ + Indexer.recompute_slot_mappings(metadata, indexer_params) + @staticmethod def prepare_for_chunked_prefill(metadata: DSAtrtllmAttentionMetadata, indexer_params: IndexerParams): @@ -2088,39 +2132,16 @@ def prepare(metadata: DSAtrtllmAttentionMetadata): # This can happen when the metadata is being used with a draft KV cache manager # during MTP speculative decoding, which uses a regular KVCacheManager instead # of DSACacheManager. - kv_cache_manager = metadata.kv_cache_manager - if kv_cache_manager is None or not hasattr(kv_cache_manager, - 'index_head_dim'): + indexer_params = Indexer.build_indexer_params(metadata) + if indexer_params is None: return num_contexts = metadata.num_contexts num_generations = metadata.num_generations num_ctx_tokens = metadata.num_ctx_tokens - request_ids = metadata.request_ids seq_lens = metadata.seq_lens - head_dim = metadata.indexer_head_dim - quant_block_size = metadata.indexer_quant_block_size - num_past_tokens = metadata.kv_cache_params.num_cached_tokens_per_seq compress_ratio = _effective_compress_ratio_divisor( _select_indexer_compress_ratio(metadata.compress_ratios)) - # MXFP4 indexer K cache packs two E2M1 codes per byte so the data - # footprint is half head_dim; FP8 keeps one byte per element. - use_fp4 = getattr(kv_cache_manager, 'use_fp4', False) - data_bytes_per_token = head_dim // 2 if use_fp4 else head_dim - - indexer_params = IndexerParams( - num_contexts=num_contexts, - num_generations=num_generations, - num_ctx_tokens=num_ctx_tokens, - head_dim=head_dim, - quant_block_size=quant_block_size, - tokens_per_block=metadata._tokens_per_block, - compress_ratio=compress_ratio, - request_ids=request_ids, - num_past_tokens=num_past_tokens, - seq_lens=seq_lens, - data_bytes_per_token=data_bytes_per_token, - ) # Store compressed KV token count for context requests metadata.num_ctx_kv_tokens = indexer_params.new_kv_tokens[: num_contexts].sum( diff --git a/tensorrt_llm/_torch/configs/__init__.py b/tensorrt_llm/_torch/configs/__init__.py index c4d9fb25ecfd..ab5cf9e81650 100644 --- a/tensorrt_llm/_torch/configs/__init__.py +++ b/tensorrt_llm/_torch/configs/__init__.py @@ -1,4 +1,5 @@ from tensorrt_llm._torch.configs.deepseek_v3 import DeepseekV3Config +from tensorrt_llm._torch.configs.deepseekv4 import DeepseekV4Config from tensorrt_llm._torch.configs.laguna import LagunaConfig @@ -19,6 +20,7 @@ def _register_custom_configs_with_transformers() -> None: custom_configs = { "deepseek_v32": DeepseekV3Config, "kimi_k2": DeepseekV3Config, + "deepseek_v4": DeepseekV4Config, "laguna": LagunaConfig, } for model_type, config_class in custom_configs.items(): @@ -30,4 +32,4 @@ def _register_custom_configs_with_transformers() -> None: _register_custom_configs_with_transformers() del _register_custom_configs_with_transformers -__all__ = ["DeepseekV3Config", "LagunaConfig"] +__all__ = ["DeepseekV3Config", "DeepseekV4Config", "LagunaConfig"] diff --git a/tensorrt_llm/_torch/configs/deepseekv4.py b/tensorrt_llm/_torch/configs/deepseekv4.py new file mode 100644 index 000000000000..6ee24a86840a --- /dev/null +++ b/tensorrt_llm/_torch/configs/deepseekv4.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging + +logger = logging.get_logger(__name__) + + +# Default Engram embedding vocabulary size per n-gram level. +# Derived from the DeepSeek-V3 tokenizer vocab size (129280) scaled by 5. +_DEFAULT_ENGRAM_VOCAB_SIZE = 129280 * 5 + + +# This is a temporary workaround for DeepSeek-V4 model as HF does not support it yet +# TODO: Remove this once HF supports DeepSeek-V4 +class DeepseekV4Config(PretrainedConfig): + model_type = "deepseek_v4" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=129280, + hidden_size=4096, + intermediate_size=14336, + moe_intermediate_size=2048, + num_hidden_layers=43, + num_nextn_predict_layers=0, + num_attention_heads=64, + num_key_value_heads=64, + n_shared_experts=1, + n_routed_experts=256, + ep_size=1, + routed_scaling_factor=1.5, + kv_lora_rank=448, + q_lora_rank=1024, + qk_rope_head_dim=64, + v_head_dim=512, + qk_nope_head_dim=448, + topk_method="noaux_tc", + n_group=8, + topk_group=4, + num_experts_per_tok=6, + moe_layer_freq=1, + first_k_dense_replace=3, + norm_topk_prob=True, + scoring_func="sqrtsoftplus", + hidden_act="silu", + max_position_embeddings=65536, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=None, + bos_token_id=0, + eos_token_id=1, + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + attention_dropout=0.0, + index_head_dim=128, + index_n_heads=64, + index_topk=512, + o_groups=8, + o_lora_rank=1024, + n_hash_layers=3, + hc_mult=4, + hc_sinkhorn_iters=20, + hc_eps=1e-6, + window_size=128, + compress_rope_theta=160000, + compress_ratios=None, + swiglu_limit=10.0, + has_engram=False, + engram_vocab_size=None, + engram_max_ngram_size=3, + engram_n_embed_per_ngram=512, + engram_n_head_per_ngram=8, + engram_kernel_size=4, + engram_pad_id=2, + engram_layer_ids=None, + engram_seed=0, + **kwargs, + ): + # DeepSeek-V4 HF config uses `sliding_window`, `num_hash_layers`, + # and `head_dim` for these internal fields. + # Accept them as aliases for the internal names the rest of the code uses. + if "sliding_window" in kwargs: + window_size = kwargs.pop("sliding_window") + if "num_hash_layers" in kwargs: + n_hash_layers = kwargs.pop("num_hash_layers") + if "head_dim" in kwargs: + v_head_dim = kwargs.pop("head_dim") + if "score_func" in kwargs: + scoring_func = kwargs.pop("score_func") + + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_nextn_predict_layers = num_nextn_predict_layers + self.num_attention_heads = num_attention_heads + self.n_shared_experts = n_shared_experts + self.n_routed_experts = n_routed_experts + self.ep_size = ep_size + self.routed_scaling_factor = routed_scaling_factor + self.kv_lora_rank = kv_lora_rank + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.topk_method = topk_method + self.n_group = n_group + self.topk_group = topk_group + self.num_experts_per_tok = num_experts_per_tok + self.moe_layer_freq = moe_layer_freq + self.first_k_dense_replace = first_k_dense_replace + self.norm_topk_prob = norm_topk_prob + self.scoring_func = scoring_func + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.window_size = window_size + + # indexer + self.index_head_dim = index_head_dim + self.index_n_heads = index_n_heads + self.index_topk = index_topk + + # output projection + self.o_groups = o_groups + self.o_lora_rank = o_lora_rank + + # hash layer and mhc + self.n_hash_layers = n_hash_layers + self.hc_mult = hc_mult + self.hc_sinkhorn_iters = hc_sinkhorn_iters + self.hc_eps = hc_eps + + # kv compression + self.compress_rope_theta = compress_rope_theta + self.compress_ratios = compress_ratios + self.swiglu_limit = swiglu_limit + + # Engram configuration + self.has_engram = has_engram + self.engram_vocab_size = ( + engram_vocab_size + if engram_vocab_size is not None + else [_DEFAULT_ENGRAM_VOCAB_SIZE, _DEFAULT_ENGRAM_VOCAB_SIZE] + ) + self.engram_max_ngram_size = engram_max_ngram_size + self.engram_n_embed_per_ngram = engram_n_embed_per_ngram + self.engram_n_head_per_ngram = engram_n_head_per_ngram + self.engram_kernel_size = engram_kernel_size + self.engram_pad_id = engram_pad_id + self.engram_layer_ids = engram_layer_ids if engram_layer_ids is not None else [1, 15] + self.engram_seed = engram_seed + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/tensorrt_llm/_torch/memory_buffer_utils.py b/tensorrt_llm/_torch/memory_buffer_utils.py index 3be29e30a745..9dfcca890278 100644 --- a/tensorrt_llm/_torch/memory_buffer_utils.py +++ b/tensorrt_llm/_torch/memory_buffer_utils.py @@ -130,6 +130,27 @@ def get_memory_buffers(): return _buffer +def clear_memory_buffers(): + """Clear all non-reserved buffers from the global buffer pool. + + This should be called after warmup phases (e.g., autotuner warmup) that + may allocate large workspace buffers sized for the maximum token count. + Clearing these prevents them from inflating the memory baseline used by + the KV cache profiler, which would otherwise reduce the memory available + for activations during inference. + """ + global _buffer + for buffer_name, blocks in list(_buffer.buffers.items()): + remaining = [b for b in blocks if b.is_reserved] + freed = [b for b in blocks if not b.is_reserved] + for block in freed: + del block.buffer + if remaining: + _buffer.buffers[buffer_name] = remaining + else: + del _buffer.buffers[buffer_name] + + _shared_pool = None diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 1c5d4b9e5c0b..78fbd82da645 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -2,6 +2,7 @@ import errno import json import os +import struct import tempfile from dataclasses import dataclass, field from pathlib import Path @@ -20,6 +21,7 @@ from tensorrt_llm.bindings import LayerType as LayerTypeCpp from tensorrt_llm.functional import AllReduceStrategy from tensorrt_llm.llmapi.llm_args import (DeepSeekSparseAttentionConfig, + DeepSeekV4SparseAttentionConfig, KvCacheConfig, MoeLoadBalancerConfig, MultimodalConfig) from tensorrt_llm.logger import logger @@ -40,6 +42,9 @@ TConfig = TypeVar("TConfig", bound=transformers.PretrainedConfig) +_DEEPSEEK_V4_ARCHITECTURES = {"DeepseekV4ForCausalLM"} +_DEEPSEEK_V4_ROUTED_EXPERT_WEIGHT = "layers.0.ffn.experts.0.w1.weight" + def _unified_kv_pool_includes_mamba( is_disagg: bool, spec_config: Optional['SpeculativeConfig']) -> bool: @@ -326,6 +331,11 @@ def resolve_moe_backend(moe_backend: str, if moe_backend.upper() != "AUTO": return moe_backend + if architecture in _DEEPSEEK_V4_ARCHITECTURES: + sm_version = get_sm_version() + if 100 <= sm_version < 120: + return "TRTLLM" + if architecture == "GptOssForCausalLM": sm_version = get_sm_version() # Select the best performing backend based on SM version @@ -546,6 +556,124 @@ def load_hf_quant_config(hf_quant_config, moe_backend, checkpoint_dir=None): quant_config.exclude_modules = default_exclude return quant_config, layer_quant_config + @staticmethod + def _read_safetensors_header(path: Path) -> Dict[str, Any]: + with open(path, "rb") as f: + header_size = struct.unpack(" Optional[Dict]: + checkpoint_path = Path(checkpoint_dir) + candidates = [] + index_path = checkpoint_path / "model.safetensors.index.json" + if index_path.exists(): + with open(index_path) as f: + index = json.load(f) + shard_name = index.get("weight_map", {}).get(tensor_name) + if shard_name is not None: + candidates.append(checkpoint_path / shard_name) + + candidates.extend(sorted(checkpoint_path.glob("*.safetensors"))) + seen = set() + for candidate in candidates: + if candidate in seen or not candidate.exists(): + continue + seen.add(candidate) + header = ModelConfig._read_safetensors_header(candidate) + if tensor_name in header: + return header[tensor_name] + return None + + @staticmethod + def _detect_deepseek_v4_routed_moe_layout( + checkpoint_dir: str) -> Optional[str]: + tensor_info = ModelConfig._get_safetensors_header_for_tensor( + checkpoint_dir, _DEEPSEEK_V4_ROUTED_EXPERT_WEIGHT) + if tensor_info is None: + return None + + dtype = tensor_info.get("dtype") + shape = tensor_info.get("shape", []) + if dtype == "I8" and len(shape) == 2: + return "mxfp4" + if dtype == "U8": + return "nvfp4" + return None + + @staticmethod + def _is_deepseek_v4_base_checkpoint(checkpoint_dir: str) -> bool: + tensor_info = ModelConfig._get_safetensors_header_for_tensor( + checkpoint_dir, _DEEPSEEK_V4_ROUTED_EXPERT_WEIGHT) + if tensor_info is None: + return False + + return ModelConfig._detect_deepseek_v4_routed_moe_layout( + checkpoint_dir) not in ("mxfp4", "nvfp4") + + @staticmethod + def _has_deepseek_v4_layer_only_modelopt_quant_config( + quant_config_file: str) -> bool: + with open(quant_config_file) as f: + quant_config_dict = json.load(f) + + quantization_config = quant_config_dict.get('quantization', {}) + return (quantization_config.get('quant_algo', None) is None and + quantization_config.get('quantized_layers', None) is not None) + + @staticmethod + def _set_deepseek_v4_routed_moe_quant_config(pretrained_config, + checkpoint_dir: str, + moe_backend: str, + layer_quant_config, + spec_config=None, + require_layout: bool = False): + layout = ModelConfig._detect_deepseek_v4_routed_moe_layout( + checkpoint_dir) + if layout not in ("mxfp4", "nvfp4"): + if require_layout: + raise ValueError( + "DeepSeek-V4 checkpoint has layer-specific quantized_layers " + "in hf_quant_config.json, but the routed MoE layout could " + "not be detected from safetensors metadata. Expected " + f"{_DEEPSEEK_V4_ROUTED_EXPERT_WEIGHT} to use dtype I8 " + "for MXFP4 or U8 for NVFP4.") + return layer_quant_config + + experts_quant_config = QuantConfig() + if layout == "mxfp4": + experts_quant_config.quant_algo = ModelConfig.get_mxfp4_quant_algo( + moe_backend) + experts_quant_config.group_size = 32 + else: + experts_quant_config.quant_algo = QuantAlgo.NVFP4 + experts_quant_config.group_size = 16 + experts_quant_config.exclude_modules = [ + 'block.*.attn.out', 'block.*.mlp.gate', 'block.*.attn.qkv', + 'embedding', 'unembedding' + ] + + if layer_quant_config is None: + layer_quant_config = {} + else: + layer_quant_config = dict(layer_quant_config) + + num_moe_layers = pretrained_config.num_hidden_layers + if (spec_config is not None + and spec_config.spec_dec_mode.is_mtp_one_model()): + num_moe_layers += spec_config.num_nextn_predict_layers + + for layer_idx in range(num_moe_layers): + layer_quant_config[ + f"model.layers.{layer_idx}.mlp.experts"] = experts_quant_config + + logger.info( + "Detected DeepSeek-V4 routed MoE %s checkpoint layout; using " + "%s for routed experts.", layout.upper(), + experts_quant_config.quant_algo) + return layer_quant_config + @staticmethod def load_quant_config_from_dtypes_json(dtypes_json_file, moe_backend: str): quant_config = QuantConfig() @@ -603,6 +731,56 @@ def from_pretrained(cls, checkpoint_dir: str, trust_remote_code=False, **kwargs): + + def update_sparse_attention_indexer_config(pretrained_config, kwargs): + sparse_attention_config = kwargs.get('sparse_attention_config') + if sparse_attention_config: + index_n_heads = sparse_attention_config.index_n_heads or pretrained_config.index_n_heads + index_head_dim = sparse_attention_config.index_head_dim or pretrained_config.index_head_dim + index_topk = sparse_attention_config.index_topk or pretrained_config.index_topk + indexer_max_chunk_size = sparse_attention_config.indexer_max_chunk_size + skip_indexer_for_short_seqs = sparse_attention_config.skip_indexer_for_short_seqs + # Pass-through DSA tuning flags so user-set values survive the + # V4 sparse_attention_config rebuild below. The V3.2 path + # already threads these explicitly (lines 723-727); without + # this block the V4 rebuild silently drops any user override + # back to subclass defaults (e.g., enable_heuristic_topk=False + # even when the user set it to True in --extra_llm_api_options). + use_cute_dsl_topk = sparse_attention_config.use_cute_dsl_topk + use_cute_dsl_paged_mqa_logits = sparse_attention_config.use_cute_dsl_paged_mqa_logits + q_split_threshold = sparse_attention_config.q_split_threshold + indexer_rope_interleave = sparse_attention_config.indexer_rope_interleave + enable_heuristic_topk = sparse_attention_config.enable_heuristic_topk + indexer_k_dtype = sparse_attention_config.indexer_k_dtype + else: + index_n_heads = pretrained_config.index_n_heads + index_head_dim = pretrained_config.index_head_dim + index_topk = pretrained_config.index_topk + indexer_max_chunk_size = None + skip_indexer_for_short_seqs = True + # Defaults match DeepSeekSparseAttentionConfig field defaults. + use_cute_dsl_topk = False + use_cute_dsl_paged_mqa_logits = False + q_split_threshold = 8192 + indexer_rope_interleave = False + enable_heuristic_topk = False + indexer_k_dtype = "fp8" + indexer_config = {} + indexer_config['index_n_heads'] = index_n_heads + indexer_config['index_head_dim'] = index_head_dim + indexer_config['index_topk'] = index_topk + indexer_config['indexer_max_chunk_size'] = indexer_max_chunk_size + indexer_config[ + 'skip_indexer_for_short_seqs'] = skip_indexer_for_short_seqs + indexer_config['use_cute_dsl_topk'] = use_cute_dsl_topk + indexer_config[ + 'use_cute_dsl_paged_mqa_logits'] = use_cute_dsl_paged_mqa_logits + indexer_config['q_split_threshold'] = q_split_threshold + indexer_config['indexer_rope_interleave'] = indexer_rope_interleave + indexer_config['enable_heuristic_topk'] = enable_heuristic_topk + indexer_config['indexer_k_dtype'] = indexer_k_dtype + return indexer_config + # Use file lock to prevent race conditions when multiple processes # try to import/cache the same remote model config file with config_file_lock(): @@ -658,6 +836,78 @@ def from_pretrained(cls, indexer_rope_interleave=indexer_rope_interleave, enable_heuristic_topk=enable_heuristic_topk, indexer_k_dtype=indexer_k_dtype) + elif pretrained_config.architectures[ + 0] == "DeepseekV4ForCausalLM": + if cls._is_deepseek_v4_base_checkpoint(checkpoint_dir): + logger.warning( + "Support for DeepSeek-V4 Base checkpoints is " + "experimental. For better supported behavior, use " + "a DeepSeek-V4 Instruct checkpoint.") + indexer_config = update_sparse_attention_indexer_config( + pretrained_config, kwargs) + checkpoint_compress_ratios = getattr( + pretrained_config, 'compress_ratios', None) + num_base_layers = pretrained_config.num_hidden_layers + spec_config = kwargs.get('spec_config', None) + if (spec_config is not None + and getattr(spec_config, 'num_nextn_predict_layers', + None) is None): + spec_config.num_nextn_predict_layers = getattr( + pretrained_config, 'num_nextn_predict_layers', 1) + mtp_enabled = (spec_config is not None and + spec_config.spec_dec_mode.is_mtp_one_model()) + sparse_attention_config = kwargs.get( + 'sparse_attention_config') + checkpoint_window_size = getattr(pretrained_config, + 'window_size', None) + if checkpoint_window_size is None: + checkpoint_window_size = getattr( + pretrained_config, 'sliding_window', None) + if sparse_attention_config: + compress_ratios = sparse_attention_config.compress_ratios + window_size = sparse_attention_config.window_size + if 'window_size' not in sparse_attention_config.model_fields_set: + window_size = checkpoint_window_size + else: + compress_ratios = checkpoint_compress_ratios + window_size = checkpoint_window_size + + if (checkpoint_compress_ratios is not None + and (compress_ratios is None + or len(checkpoint_compress_ratios) + > len(compress_ratios))): + compress_ratios = checkpoint_compress_ratios + + if window_size is None: + window_size = checkpoint_window_size + if window_size is None: + window_size = pretrained_config.sliding_window + + # Normalize checkpoint-facing ratio 0 (SWA-only/uncompressed) + # to 1 internally so cache allocation math works. The + # external config keeps the original semantics. + compress_ratios = [ + ratio if ratio > 0 else 1 for ratio in compress_ratios + ] + + # Only synthesize ratios for extra MTP layers. The base + # model ratios must come from the checkpoint or an + # explicit user override; padding a short default list for + # non-MTP changes sparse attention semantics. + if mtp_enabled: + mtp_num_layers = spec_config.num_nextn_predict_layers + total_layers = num_base_layers + mtp_num_layers + if len(compress_ratios) < total_layers: + compress_ratios = list(compress_ratios) + [1] * ( + total_layers - len(compress_ratios)) + + indexer_k_dtype = indexer_config.pop('indexer_k_dtype') + kwargs[ + 'sparse_attention_config'] = DeepSeekV4SparseAttentionConfig( + compress_ratios=compress_ratios, + window_size=window_size, + indexer_k_dtype=indexer_k_dtype, + **indexer_config) else: raise ValueError( "checkpoint_dir is None. Cannot load model config without a valid checkpoint directory." @@ -733,6 +983,7 @@ def _recursive_update_config(config: transformers.PretrainedConfig, quant_config = QuantConfig() layer_quant_config = None + require_deepseek_v4_routed_moe_layout = False requested_moe_backend = kwargs.get('moe_backend', 'AUTO') architecture = pretrained_config.architectures[ 0] if pretrained_config.architectures else "" @@ -755,8 +1006,33 @@ def _recursive_update_config(config: transformers.PretrainedConfig, getattr(pretrained_config, "quantization_config", None), source_file="hf_quant_config.json", ) + if architecture in _DEEPSEEK_V4_ARCHITECTURES: + require_deepseek_v4_routed_moe_layout = ( + cls._has_deepseek_v4_layer_only_modelopt_quant_config( + quant_config_file)) quant_config, layer_quant_config = cls._build_modelopt_quant_config( normalized, checkpoint_dir, moe_backend_hint) + hf_quant_config = getattr(pretrained_config, "quantization_config", + None) + if quant_config.quant_algo is None and hf_quant_config is not None: + hf_quant_config, hf_layer_quant_config = cls.load_hf_quant_config( + hf_quant_config, + moe_backend_hint, + checkpoint_dir=checkpoint_dir) + if hf_quant_config.quant_algo is not None: + logger.info( + "Using quantization_config from config.json as global " + "quantization because hf_quant_config.json does not set " + "a global quant_algo.") + quant_config = hf_quant_config + if hf_layer_quant_config is not None: + if layer_quant_config is None: + layer_quant_config = hf_layer_quant_config + else: + layer_quant_config = { + **hf_layer_quant_config, + **layer_quant_config, + } # quantized ckpt in other formats elif getattr(pretrained_config, "quantization_config", None) is not None: @@ -775,6 +1051,15 @@ def _recursive_update_config(config: transformers.PretrainedConfig, quant_config=quant_config, ) + if architecture in _DEEPSEEK_V4_ARCHITECTURES: + layer_quant_config = cls._set_deepseek_v4_routed_moe_quant_config( + pretrained_config, + checkpoint_dir, + kwargs['moe_backend'], + layer_quant_config, + kwargs.get('spec_config', None), + require_layout=require_deepseek_v4_routed_moe_layout) + model_config = cls(pretrained_config=pretrained_config, quant_config=quant_config, quant_config_dict=layer_quant_config, diff --git a/tensorrt_llm/_torch/models/__init__.py b/tensorrt_llm/_torch/models/__init__.py index ec78bf51fe06..2c9a503424e8 100644 --- a/tensorrt_llm/_torch/models/__init__.py +++ b/tensorrt_llm/_torch/models/__init__.py @@ -13,6 +13,7 @@ from .modeling_clip import CLIPVisionModel from .modeling_cohere2 import Cohere2ForCausalLM from .modeling_deepseekv3 import DeepseekV3ForCausalLM +from .modeling_deepseekv4 import DeepseekV4ForCausalLM from .modeling_exaone4 import Exaone4ForCausalLM from .modeling_exaone4_5 import Exaone4_5_ForConditionalGeneration from .modeling_exaone_moe import ExaoneMoeForCausalLM @@ -81,6 +82,7 @@ "Mistral3VLM", "MistralForCausalLM", "MixtralForCausalLM", + "DeepseekV4ForCausalLM", "NemotronH_Nano_VL_V2", "NemotronForCausalLM", "NemotronHForCausalLM", diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py new file mode 100644 index 000000000000..48a46deaf2d6 --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -0,0 +1,2587 @@ +# -------------------------------------------------- +# Portions of this code were derived from DeepSeek‑V3: +# https://github.com/deepseek-ai/DeepSeek-V3 +# +# MIT License + +# Copyright (c) 2023 DeepSeek + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# -------------------------------------------------- + +import copy +import math +import os +from typing import TYPE_CHECKING, Dict, List, Optional + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + +import torch +import triton +import triton.language as tl +from torch import nn +from tqdm import tqdm +from transformers import PretrainedConfig + +import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils +from tensorrt_llm._ipc_utils import can_access_peer +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams +from ..attention_backend.sparse.deepseek_v4.deepseek_v4 import DeepseekV4TrtllmAttentionMetadata +from ..distributed import ( + AllReduce, + AllReduceFusionOp, + AllReduceParams, + MoEAllReduce, + MoEAllReduceParams, + allgather, +) +from ..model_config import ModelConfig +from ..modules.attention import MLA +from ..modules.decoder_layer import DecoderLayer +from ..modules.embedding import Embedding +from ..modules.engram import Engram, EngramConfig, EngramHashProvider +from ..modules.fused_moe import ( + CutlassFusedMoE, + DeepSeekV4MoeRoutingMethod, + MoE, + MoEWeightLoadingMode, + TritonFusedMoE, + TRTLLMGenFusedMoE, + create_moe, + get_moe_cls, +) +from ..modules.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE +from ..modules.fused_moe.fused_moe_wide_ep import WideEPMoE +from ..modules.gated_mlp import GatedMLP +from ..modules.linear import Linear, TensorParallelMode, WeightsLoadingConfig +from ..modules.mhc.hyper_connection import HCHead, HCState, mHC +from ..modules.multi_stream_utils import maybe_execute_in_parallel +from ..modules.rms_norm import RMSNorm +from ..peft.lora.layer import LoraLayer +from ..speculative import SpecMetadata, get_num_extra_kv_tokens +from ..utils import ( + AuxStreamType, + EventType, + Fp4QuantizedTensor, + create_lm_head_tp_mapping, + replace_parameter_and_save_metadata, +) +from .modeling_speculative import SpecDecOneEngineForCausalLM +from .modeling_utils import DecoderModel, EagerFusionConfig, filter_weights, register_auto_model + + +@triton.jit +def weight_dequant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr): + """ + Dequantizes weights using the provided scaling factors and stores the result. + + Args: + x_ptr (tl.pointer): Pointer to the quantized weights. + s_ptr (tl.pointer): Pointer to the scaling factors. + y_ptr (tl.pointer): Pointer to the output buffer for dequantized weights. + M (int): Number of rows in the weight matrix. + N (int): Number of columns in the weight matrix. + BLOCK_SIZE (tl.constexpr): Size of the block for tiling. + + Returns: + None + """ + pid_m = tl.program_id(axis=0) + pid_n = tl.program_id(axis=1) + n = tl.cdiv(N, BLOCK_SIZE) + offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + offs = offs_m[:, None] * N + offs_n[None, :] + mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + x = tl.load(x_ptr + offs, mask=mask).to(tl.float32) + s = tl.load(s_ptr + pid_m * n + pid_n) + y = x * s + tl.store(y_ptr + offs, y, mask=mask) + + +def weight_dequant(x: torch.Tensor, s: torch.Tensor, block_size: int = 128) -> torch.Tensor: + """ + Dequantizes the given weight tensor using the provided scale tensor. + + Args: + x (torch.Tensor): The quantized weight tensor of shape (M, N). + s (torch.Tensor): The scale tensor of shape (M, N). + block_size (int, optional): The block size to use for dequantization. Defaults to 128. + + Returns: + torch.Tensor: The dequantized weight tensor of the same shape as `x`. + + Raises: + AssertionError: If `x` or `s` are not contiguous or if their dimensions are not 2. + """ + assert x.is_contiguous() and s.is_contiguous(), "Input tensors must be contiguous" + assert x.dim() == 2 and s.dim() == 2, "Input tensors must have 2 dimensions" + if s.dtype != torch.float32: + s = s.float() + M, N = x.size() + y = torch.empty_like(x, dtype=torch.get_default_dtype()) + grid = lambda meta: (triton.cdiv(M, meta["BLOCK_SIZE"]), triton.cdiv(N, meta["BLOCK_SIZE"])) # noqa: E731 + weight_dequant_kernel[grid](x, s, y, M, N, BLOCK_SIZE=block_size) + return y + + +def _deepseek_v4_pos_embd_params( + config: PretrainedConfig, + model_config: ModelConfig, + layer_idx: Optional[int], + predicted_tokens_per_seq: int = 1, +) -> PositionalEmbeddingParams: + # Match the spec/overlap headroom that py_executor_creator applies + # to model_engine_max_seq_len before sizing the KV cache. The overlap + # branch is included unconditionally because disable_overlap_scheduler + # isn't reachable here; the overshoot when overlap is actually off is + # at most predicted_tokens_per_seq positions (= 1 + max_total_draft_tokens, + # which can exceed max_draft_len for tree-shaped drafts). + runtime_max_seq_len = model_config.max_seq_len + if runtime_max_seq_len is not None and model_config.spec_config is not None: + runtime_max_seq_len += 2 * max(0, predicted_tokens_per_seq - 1) + runtime_max_seq_len += get_num_extra_kv_tokens(model_config.spec_config) + rope_params = RopeParams.from_config(config) + if runtime_max_seq_len is not None: + rope_params.max_seq_len = runtime_max_seq_len + + compress_ratios = None + if model_config.sparse_attention_config is not None: + compress_ratios = getattr(model_config.sparse_attention_config, "compress_ratios", None) + if not compress_ratios: + compress_ratios = getattr(config, "compress_ratios", None) + + compress_ratio = 0 + if compress_ratios and layer_idx is not None: + compress_ratio = compress_ratios[min(layer_idx, len(compress_ratios) - 1)] + + if compress_ratio > 1: + rope_params.theta = getattr(config, "compress_rope_theta", rope_params.theta) + rope_params.scale_type = RotaryScalingType.yarn + # DeepSeek-V4 reference applies YaRN frequency interpolation but does + # not scale the cos/sin amplitudes. + rope_params.mscale = 0.0 + rope_params.mscale_all_dim = 0.0 + pos_type = PositionEmbeddingType.yarn + else: + # DeepSeek-V4 reference uses base RoPE, without YaRN, for SWA-only + # layers. Internal configs normalize checkpoint ratio 0 to 1, so both + # values are treated as uncompressed here. + rope_params.theta = getattr(config, "rope_theta", rope_params.theta) + rope_params.scale_type = RotaryScalingType.none + rope_params.scale = 1.0 + pos_type = PositionEmbeddingType.rope_gptj + + return PositionalEmbeddingParams( + type=pos_type, + rope=rope_params, + is_neox=False, + ) + + +@torch.compile(dynamic=True) +def moe_reduce_add_shared_output(routed_output, shared_output): + routed_output = torch.sum(routed_output, dim=1, keepdim=False) + return shared_output + routed_output + + +# Per-attention parameter renames inside `attn.` / `mtp.0.attn.`. +# Maps checkpoint name component → model name component. +_ATTN_PARAM_RENAME = { + "wq_a": "q_a_proj", + "wq_b": "q_b_proj", + "wkv": "kv_a_proj_with_mqa", + "wo_b": "o_b_proj", + "q_norm": "q_a_layernorm", + "kv_norm": "kv_a_layernorm", +} + +# Shared expert leaf rename: checkpoint w1/w3/w2 → model gate/up/down. The +# loader's `params_map` then fuses `gate_proj` + `up_proj` into `gate_up_proj`. +_SHARED_EXPERT_RENAME = { + "w1": "gate_proj", + "w3": "up_proj", + "w2": "down_proj", +} + + +def _resolve_enable_fused_hc(config: PretrainedConfig) -> bool: + """Resolve the DeepSeek-V4 fused HC boundary-fusion knob.""" + env = os.environ.get("TRTLLM_MHC_ENABLE_FUSED_HC") + if env is not None: + return env not in ("0", "false", "False") + return bool(getattr(config, "enable_fused_hc", True)) + + +def _copy_deepseek_v4_fused_a_weight_scale( + module: Linear, fused_a: torch.Tensor, fused_a_scale: torch.Tensor +) -> None: + if tuple(module.weight_scale.shape) == tuple(fused_a_scale.shape): + module.weight_scale.data.copy_(fused_a_scale) + return + + if ( + module.weight_scale.ndim == 2 + and fused_a_scale.ndim == 2 + and module.weight_scale.shape[0] >= fused_a_scale.shape[0] + and module.weight_scale.shape[1] == fused_a_scale.shape[1] + ): + module.weight_scale.data[0 : fused_a_scale.shape[0]].copy_(fused_a_scale) + return + + expected_fp8_scale_shape = ( + math.ceil(fused_a.shape[0] / 128), + math.ceil(fused_a.shape[1] / 128), + ) + can_rebuild_fp8_scale = ( + fused_a_scale.ndim == 2 + and tuple(fused_a_scale.shape) == expected_fp8_scale_shape + and tuple(module.weight.shape) == tuple(fused_a.shape) + ) + if not can_rebuild_fp8_scale: + raise RuntimeError( + "DeepSeek-V4 fused A projection weight_scale shape mismatch: " + f"module={tuple(module.weight_scale.shape)}, checkpoint={tuple(fused_a_scale.shape)}, " + f"weight={tuple(module.weight.shape)}, checkpoint_weight={tuple(fused_a.shape)}" + ) + + corrected_weight_scale = nn.Parameter( + torch.empty( + fused_a_scale.shape, + dtype=module.weight_scale.dtype, + device=module.weight_scale.device, + ), + requires_grad=False, + ) + replace_parameter_and_save_metadata( + module, + "weight_scale", + corrected_weight_scale, + module.rebuild_tensor_metadata, + ) + + module.weight_scale.data.copy_(fused_a_scale) + + +def _remap_deepseek_v4_checkpoint_keys( + weights: Dict, num_hidden_layers: int, kv_lora_rank: int = 448 +) -> Dict: + """Convert DeepSeek-V4 checkpoint keys to model named-parameter keys. + + Why: the upstream DS-V4 release uses keys like ``layers.X.attn.wkv.weight``, + ``mtp.0.ffn.experts.0.w1.scale``, ``embed.weight``, ``head.weight``. The + TRT-LLM model exposes them as ``model.layers.X.self_attn.kv_a_proj_with_mqa.weight``, + ``model.layers.{N}.mlp.experts.0.w1.weight_scale_inv``, ``model.embed_tokens.weight``, + ``lm_head.weight``. This pass renames keys, fuses split projections that the + model represents as one tensor, and synthesizes default values for params + the model has but the checkpoint does not (so the existing loader's strict + key lookup does not raise). + + Caveats — limitations carried as TODOs (kept here so future readers can + audit them in one place rather than chasing comments): + * ``self_attn.indexer.wk.weight`` and ``self_attn.indexer.k_norm.{weight,bias}`` + are zero-filled — V4 indexer's k path is served by the compressor, so + the base ``Indexer.wk`` / ``k_norm`` are unused at forward time. + * Routed experts can use either FP8 block scales or the packed MXFP4 + layout. The first routed expert weight determines the scale suffix used + for all routed expert tensors in the shard. + * ``self_attn.o_a_proj`` is loaded by the DeepSeek-V4 loader because it + is a direct MLA parameter, not a child Linear module. + * ``mtp.0.head.weight`` is dropped — DeepSeekV4MTP reuses the main + ``lm_head`` via ``shared_head``. Flash omits this key entirely; Flash-Base + carries it but matches the main head, so we let the main head win. + """ + mtp_layer_prefix = f"model.layers.{num_hidden_layers}" + routed_moe_scale_name = "weight_scale_inv" + for key, value in weights.items(): + if ( + key.startswith("layers.") + and ".ffn.experts." in key + and key.endswith(".weight") + and getattr(value, "ndim", 0) == 2 + and value.dtype in (torch.int8, torch.uint8) + ): + routed_moe_scale_name = "weight_scale" + break + + def _rename_attn_subkey(rest: str) -> Optional[str]: + # rest examples: "wq_a.weight", "wq_a.scale", "wo_a.weight", + # "attn_sink", "compressor.wkv.weight", "indexer.wq_b.scale", + # "kv_norm.weight" + # ``attn_sink`` is loaded by the ``mqa`` branch in the per-module + # loader, which reads it under the parent ``self_attn.attn_sink`` + # key. Pass through unchanged. + if rest == "attn_sink": + return "attn_sink" + # `wo_a` is an nn.Parameter on the model side (not a Linear), so + # `wo_a.weight` carries the value directly into `o_a_proj` without + # a trailing ``.weight``. Retain `.scale` so the loader can dequantize + # FP8 block-scaled checkpoints before assigning the bf16 parameter. + if rest == "wo_a.weight": + return "o_a_proj" + if rest == "wo_a.scale": + return "o_a_proj.weight_scale_inv" + # Compressor / indexer paths — pass through with .scale rename, plus + # wkv+wgate fusion handled separately below. + if rest.startswith("compressor.") or rest.startswith("indexer."): + return rest.replace(".scale", ".weight_scale_inv") + head, sep, tail = rest.partition(".") + new_head = _ATTN_PARAM_RENAME.get(head, head) + if tail == "scale": + tail = "weight_scale_inv" + return f"{new_head}.{tail}" if sep else new_head + + def _rename_ffn_subkey(rest: str) -> str: + # Examples: + # gate.weight / gate.tid2eid → gate.weight / gate.tid2eid + # gate.bias → gate.e_score_correction_bias + # experts... → experts... + # shared_experts.. → shared_experts._proj. + if rest == "gate.bias": + return "gate.e_score_correction_bias" + if rest.startswith("experts.") and rest.endswith(".scale"): + return f"{rest[: -len('.scale')]}.{routed_moe_scale_name}" + rest = rest.replace(".scale", ".weight_scale_inv") + # Non-hashed layers carry the routing logit bias as `gate.bias`; the + # model wires it through `DeepseekV4Gate.e_score_correction_bias`. + if rest == "gate.bias": + return "gate.e_score_correction_bias" + if rest.startswith("shared_experts."): + parts = rest.split(".") + if len(parts) >= 2 and parts[1] in _SHARED_EXPERT_RENAME: + parts[1] = _SHARED_EXPERT_RENAME[parts[1]] + rest = ".".join(parts) + return rest + + def _rename_layer_subkey(rest: str) -> Optional[str]: + # rest examples: "attn_norm.weight", "ffn_norm.weight", + # "hc_attn_fn", "attn.wkv.weight", "ffn.experts.0.w1.weight" + if rest == "attn_norm.weight": + return "input_layernorm.weight" + if rest == "ffn_norm.weight": + return "post_attention_layernorm.weight" + # ``hc_attn_*`` and ``hc_ffn_*`` are loaded by ``load_flat_hc_weights`` + # which builds the lookup key as ``f"{stem}_{a}"`` with the parent + # module path as the stem (e.g. ``model.layers.0.hc_attn_fn``). Pass + # the flat-underscore form through unchanged so that lookup succeeds. + if rest.startswith("hc_attn_") or rest.startswith("hc_ffn_"): + return rest + if rest.startswith("attn."): + new_sub = _rename_attn_subkey(rest[len("attn.") :]) + return None if new_sub is None else f"self_attn.{new_sub}" + if rest.startswith("ffn."): + return f"mlp.{_rename_ffn_subkey(rest[len('ffn.') :])}" + return rest + + out: Dict[str, torch.Tensor] = {} + # Pending fusions: collected first, materialized at the end so we don't + # depend on iteration order. + compressor_split: Dict[str, Dict[str, torch.Tensor]] = {} + + def _record_compressor_part(model_key: str, part: str, tensor: torch.Tensor): + # model_key looks like "...self_attn.compressor..weight" or with + # ".indexer.compressor.". Strip trailing "..weight" → "." + # is wgate or wkv; we want a stable "fusion bucket" key that ends at + # the parent compressor. + bucket = model_key.rsplit(f".{part}.", 1)[0] + compressor_split.setdefault(bucket, {})[part] = tensor + + def _emit_or_collect(model_key: str, tensor: torch.Tensor): + """Route a (model_key, tensor) pair and collect compressor fusions.""" + if ".compressor." in model_key and ( + model_key.endswith(".wkv.weight") or model_key.endswith(".wgate.weight") + ): + part = "wkv" if model_key.endswith(".wkv.weight") else "wgate" + _record_compressor_part(model_key, part, tensor) + return + if ( + routed_moe_scale_name == "weight_scale" + and ".mlp.experts." in model_key + and (model_key.endswith(".weight") or model_key.endswith(".weight_scale")) + and tensor.dtype != torch.uint8 + ): + tensor = tensor.view(torch.uint8) + out[model_key] = tensor + + for k, v in weights.items(): + # Top-level keys that don't go through the layer/mtp branches. + if k == "embed.weight": + out["model.embed_tokens.weight"] = v + continue + if k == "head.weight": + out["lm_head.weight"] = v + continue + if k == "norm.weight": + out["model.norm.weight"] = v + continue + if k.startswith("hc_head_"): + # ``load_flat_hc_weights`` looks up ``f"{stem}_{a}"`` with stem set + # to the module's full name path (``model.hc_head``). Emit the flat + # checkpoint key under the parent prefix so the lookup matches. + out[f"model.{k}"] = v + continue + + # mtp.0.head.weight is intentionally dropped (Flash-Base only); see + # docstring caveats. + if k == "mtp.0.head.weight": + continue + + # mtp.0.* — route to model.layers.{num_hidden_layers}.* + if k.startswith("mtp.0."): + rest = k[len("mtp.0.") :] + # MTP-only keys: enorm, hnorm map directly; norm maps to + # shared_head.norm; hc_head_* maps under shared_head; e_proj / + # h_proj are loaded as two separate Linear modules (no fused + # eh_proj on the MTP layer). + if rest in ("enorm.weight", "hnorm.weight"): + out[f"{mtp_layer_prefix}.{rest}"] = v + continue + if rest == "norm.weight": + out[f"{mtp_layer_prefix}.shared_head.norm.weight"] = v + continue + if rest.startswith("hc_head_"): + # The MTP HCHead is wired at ``...shared_head.hc_head``, so + # ``load_flat_hc_weights`` matches ``names[-1] == "hc_head"`` + # first and computes ``stem = ".".join(names)`` with the full + # module path. Emit at that full path with the flat suffix + # so the lookup matches. + out[f"{mtp_layer_prefix}.shared_head.{rest}"] = v + continue + for proj in ("e_proj", "h_proj"): + if rest.startswith(f"{proj}."): + suffix = rest[len(f"{proj}.") :] + if suffix == "scale": + suffix = "weight_scale_inv" + out[f"{mtp_layer_prefix}.{proj}.{suffix}"] = v + break + else: + # General per-layer transform reused for the MTP layer. + new_rest = _rename_layer_subkey(rest) + if new_rest is None: + continue + _emit_or_collect(f"{mtp_layer_prefix}.{new_rest}", v) + continue + + # layers..* — route to model.layers..* + if k.startswith("layers."): + parts = k.split(".", 2) + if len(parts) < 3: + continue + layer_idx, rest = parts[1], parts[2] + new_rest = _rename_layer_subkey(rest) + if new_rest is None: + continue + _emit_or_collect(f"model.layers.{layer_idx}.{new_rest}", v) + continue + + # Anything else: pass through. This catches future top-level keys + # added by the upstream release without silently dropping them. + out[k] = v + + # Materialize compressor wkv_gate fusion: model has a single Linear with + # out_features = state_dim * 2; checkpoint splits it as `wkv` (kv half) + # and `wgate` (gate half). Concatenating wkv first matches the order the + # compressor kernels expect (kv_score = [kv | gate]). + for bucket, parts in compressor_split.items(): + if "wkv" not in parts or "wgate" not in parts: + # Partial — emit what we have so the loader fails loudly with a + # specific missing key rather than a silent shape mismatch. + for name, tensor in parts.items(): + out[f"{bucket}.{name}.weight"] = tensor + continue + out[f"{bucket}.wkv_gate.weight"] = torch.cat([parts["wkv"], parts["wgate"]], dim=0) + + return out + + +class DeepseekV4WeightLoader: + def __init__(self, model, is_draft_model: bool = False): + self.model = model + self.config = model.config + self.model_config = model.model_config + self.is_draft_model = is_draft_model + + def load_weights(self, weights: Dict, skip_modules: List[str] = []): + # If the checkpoint uses raw DS-V4 keys (layers.X.attn.wkv.weight, + # mtp.0.*, embed.weight, head.weight), rewrite them to the model's + # named-parameter keys before iterating modules. The detection is by + # presence of any top-level "layers." key; HF-style checkpoints use + # "model.layers." and skip this branch. + if any(k == "embed.weight" or k.startswith("layers.") for k in weights): + weights = _remap_deepseek_v4_checkpoint_keys( + weights, + num_hidden_layers=self.config.num_hidden_layers, + kv_lora_rank=self.config.kv_lora_rank, + ) + # Synthesize defaults (with correct shape pulled from the model) + # for parameters the model has but the V4 checkpoint omits. We do + # this in one place vs scattering zero-fills through the per- + # module branches so the missing-key contract is auditable here. + # indexer.k_norm.weight → ones + # indexer.k_norm.bias → zeros + # indexer.wk.weight → zeros (V4 indexer's k path is + # served by compressor; wk unused) + _ones_suffixes = ("self_attn.indexer.k_norm.weight",) + _zeros_suffixes = ( + "self_attn.indexer.k_norm.bias", + "self_attn.indexer.wk.weight", + ) + model_params = dict(self.model.named_parameters()) + for pname, p in model_params.items(): + if pname in weights: + continue + if any(pname.endswith(s) for s in _ones_suffixes): + weights[pname] = torch.ones_like(p, device="cpu") + elif any(pname.endswith(s) for s in _zeros_suffixes): + weights[pname] = torch.zeros_like(p, device="cpu") + + def requantize_weight_with_new_scale( + weight, weight_scale, old_scale_2, new_scale_2, device + ): + """ + Dequantize FP4 weights and requantize with a new scale. + + Args: + weight: FP4 quantized weight tensor 2D [,] + weight_scale: FP8 per-block scaling factors + old_scale_2: original global scale (amax/(448*6)) + new_scale_2: new global scale (amax/(448*6)) + device: target device for computation + + Returns: + (requantized_weight, new_weight_scale) + """ + # Remember original dtype of weight_scale + original_scale_dtype = weight_scale.dtype + original_scale_shape = weight_scale.shape + + # Dequantize + dequant_shape = (weight.shape[0], weight.shape[1] * 2) + weight_dequant = ( + torch.ops.tensorrt_llm.e2m1_and_ufp8sf_scale_to_float_v2( + weight.contiguous(), + weight_scale.flatten().view(fp4_utils.float4_sf_dtype).contiguous(), + old_scale_2, + 16, + 1, + True, + ) + .to(dtype=torch.bfloat16) + .reshape(dequant_shape) + ) + + # Requantize using the new_scale_2 + weight_requant, weight_scale_requant = torch.ops.trtllm.fp4_quantize( + weight_dequant.to(device), + 1.0 / new_scale_2.to(device), + 16, # scaling_vector_size + False, + ) + + # Ensure the returned scale has the same dtype as the input scale + return weight_requant.cpu(), weight_scale_requant.reshape(original_scale_shape).view( + original_scale_dtype + ).cpu() + + def rename_moe_weight(weights: Dict, rename_rules: Dict): + result = {} + for key, value in weights.items(): + new_key = key + for old, new in rename_rules.items(): + new_key = new_key.replace(old, new) + result[new_key] = value + return result + + ## Prepare weights for TP + def split(v, tp_size, idx, dim=0): + if tp_size == 1: + return v + if len(v.shape) == 1: + return torch.chunk(v, tp_size)[idx].contiguous() + return torch.chunk(v, tp_size, dim=dim)[idx].contiguous() + + def split_matrix_tp(v, tensor_parallel, rank, dim): + return split(v, tensor_parallel, rank, dim=dim) + + def load_kv_b_proj_and_k_b_proj_trans(module_name: str, is_scale: bool) -> torch.Tensor: + weight_name = "weight" if not is_scale else "weight_scale_inv" + local_qk_nope_head_dim = qk_nope_head_dim if not is_scale else qk_nope_head_dim // 128 + local_v_head_dim = v_head_dim if not is_scale else v_head_dim // 128 + local_kv_lora_rank = kv_lora_rank if not is_scale else kv_lora_rank // 128 + + kv_b_proj = weights[f"{module_name}.{weight_name}"][:].unflatten( + 0, + [ + num_heads, + local_qk_nope_head_dim + local_v_head_dim, + ], + ) + + if not self.model_config.mapping.enable_attention_dp: + kv_b_proj = split_matrix_tp(kv_b_proj, tp_size, tp_rank, 0) + k_nope_weight, v_weight = kv_b_proj.split( + [local_qk_nope_head_dim, local_v_head_dim], + dim=1, + ) + weight_divisor = 1 if self.model_config.mapping.enable_attention_dp else tp_size + local_num_heads = num_heads // weight_divisor + + k_nope_weight_trans = k_nope_weight.transpose(2, 1).contiguous() + + kv_b_proj = torch.concat( + [ + k_nope_weight.reshape( + local_num_heads * local_qk_nope_head_dim, local_kv_lora_rank + ), + v_weight.reshape(local_num_heads * local_v_head_dim, local_kv_lora_rank), + ], + dim=0, + ) + + return kv_b_proj, k_nope_weight_trans + + def load_kv_b_proj_and_k_b_proj_trans_dequant(module_name: str) -> torch.Tensor: + weight_name = "weight" + local_qk_nope_head_dim = qk_nope_head_dim + local_v_head_dim = v_head_dim + local_kv_lora_rank = kv_lora_rank + + kv_b_proj = weights[f"{module_name}.{weight_name}"][:].cuda() + + weight_name = "weight_scale_inv" + kv_b_proj_scale = weights[f"{module_name}.{weight_name}"][:].cuda() + + kv_b_proj = weight_dequant(kv_b_proj, kv_b_proj_scale) + kv_b_proj = kv_b_proj.unflatten( + 0, + [ + num_heads, + local_qk_nope_head_dim + local_v_head_dim, + ], + ) + if not self.model_config.mapping.enable_attention_dp: + kv_b_proj = split_matrix_tp(kv_b_proj, tp_size, tp_rank, 0) + k_nope_weight, v_weight = kv_b_proj.split( + [local_qk_nope_head_dim, local_v_head_dim], + dim=1, + ) + weight_divisor = 1 if self.model_config.mapping.enable_attention_dp else tp_size + local_num_heads = num_heads // weight_divisor + + k_nope_weight_trans = k_nope_weight.transpose(2, 1).contiguous() + + kv_b_proj = torch.concat( + [ + k_nope_weight.reshape( + local_num_heads * local_qk_nope_head_dim, local_kv_lora_rank + ), + v_weight.reshape(local_num_heads * local_v_head_dim, local_kv_lora_rank), + ], + dim=0, + ) + + return kv_b_proj, k_nope_weight_trans + + def load_o_a_proj(module_name: str, module) -> None: + weight_name = f"{module_name}.o_a_proj" + o_a_proj = weights[weight_name][:].unflatten( + 0, + [ + module.num_groups, + module.o_lora_rank, + ], + ) + + scale_name = f"{weight_name}.weight_scale_inv" + o_a_proj_scale = weights.get(scale_name) + if o_a_proj_scale is not None: + o_a_proj_scale = o_a_proj_scale[:].unflatten( + 0, + [ + module.num_groups, + module.o_lora_rank // 128, + ], + ) + elif o_a_proj.dtype == torch.float8_e4m3fn: + raise KeyError(f"Missing FP8 block scale for {weight_name}") + + if not self.model_config.mapping.enable_attention_dp: + o_a_proj = split_matrix_tp(o_a_proj, tp_size, tp_rank, 0) + if o_a_proj_scale is not None: + o_a_proj_scale = split_matrix_tp(o_a_proj_scale, tp_size, tp_rank, 0) + + # Skip the BF16 dequant when the destination is FP8 (the cute_dsl + # FP8 BMM path on SM100 consumes the native FP8 weight directly). + if o_a_proj_scale is not None and module.o_a_proj.dtype != torch.float8_e4m3fn: + o_a_proj = weight_dequant( + o_a_proj.reshape(-1, o_a_proj.shape[-1]).contiguous().cuda(), + o_a_proj_scale.reshape(-1, o_a_proj_scale.shape[-1]).contiguous().cuda(), + ).view(o_a_proj.shape) + + module.o_a_proj.data.copy_( + o_a_proj.reshape(module.o_a_proj.shape).to(module.o_a_proj.dtype) + ) + + if getattr(module, "o_a_proj_scale", None) is not None and o_a_proj_scale is not None: + module.o_a_proj_scale.data.copy_( + o_a_proj_scale.reshape(module.o_a_proj_scale.shape) + ) + + if getattr(module, "o_a_proj_dequant", None) is not None and o_a_proj_scale is not None: + module.o_a_proj_dequant.data.copy_( + o_a_proj.reshape(module.o_a_proj_dequant.shape).to( + module.o_a_proj_dequant.dtype + ) + ) + + def split_kv_b_proj(kv_b_proj: torch.Tensor, is_scale: bool) -> torch.Tensor: + local_qk_nope_head_dim = qk_nope_head_dim if not is_scale else qk_nope_head_dim // 128 + local_v_head_dim = v_head_dim if not is_scale else v_head_dim // 128 + + weight_divisor = 1 if self.model_config.mapping.enable_attention_dp else tp_size + local_num_heads = num_heads // weight_divisor + + k_b_proj, v_b_proj = kv_b_proj.split( + [local_num_heads * local_qk_nope_head_dim, local_num_heads * local_v_head_dim], + dim=0, + ) + k_b_proj = k_b_proj.view([local_num_heads, local_qk_nope_head_dim, -1]) + v_b_proj = v_b_proj.view([local_num_heads, local_v_head_dim, -1]) + + if cp_size > 1: + local_cp_heads = local_num_heads // cp_size + k_b_proj = k_b_proj[cp_rank * local_cp_heads : (cp_rank + 1) * local_cp_heads] + v_b_proj = v_b_proj[cp_rank * local_cp_heads : (cp_rank + 1) * local_cp_heads] + + return k_b_proj, v_b_proj + + is_lite = self.config.q_lora_rank is None + num_heads = self.config.num_attention_heads + qk_nope_head_dim = getattr(self.config, "qk_nope_head_dim", None) + v_head_dim = getattr(self.config, "v_head_dim", getattr(self.config, "head_dim", None)) + kv_lora_rank = getattr(self.config, "kv_lora_rank", None) + + tp_rank = self.model_config.mapping.tp_rank + tp_size = self.model_config.mapping.tp_size + cp_rank = self.model_config.mapping.cp_rank + cp_size = self.model_config.mapping.cp_size + + params_map = {"gate_up_proj": ["gate_proj", "up_proj"]} + all_named_modules = dict(self.model.named_modules()) + + def load_flat_hc_weights(module, names: List[str]) -> bool: + """Load mHC / HCHead from flat ckpt keys: ``_{fn,base,scale}``. + + V4 / DeepSeek-V4 checkpoints store these as flat names (e.g. + ``hc_attn_fn``) rather than structured (``hc_attn.fn``). + ``shared_head.hc_head`` (MTP) is rewritten to flat ``hc_head_*`` + under the same parent. + """ + if names[-1] in ("hc_attn", "hc_ffn", "hc_head"): + stem = ".".join(names) + elif names[-2:] == ["shared_head", "hc_head"]: + stem = ".".join(names[:-2] + ["hc_head"]) + else: + return False + keys = {a: f"{stem}_{a}" for a in ("fn", "base", "scale")} + if not all(k in weights for k in keys.values()): + return False + for attr, key in keys.items(): + getattr(module, attr).data.copy_(weights[key][:]) + return True + + for name, module in tqdm(all_named_modules.items(), desc="Loading weights"): + if name.startswith("draft_model"): + continue + names = name.split(".") + parent_module_name = ".".join(names[:-1]) + + if names[-1] == "mqa": + parent_attn_name = ".".join(names[:-1]) + attn_sink_key = f"{parent_attn_name}.attn_sink" + if attn_sink_key in weights: + sink_full = weights[attn_sink_key][:] + if not self.model_config.mapping.enable_attention_dp: + sink_full = split(sink_full, tp_size, tp_rank) + sink_full = sink_full.to(torch.float32).cuda() + module.attn_sink = nn.Parameter(sink_full, requires_grad=False) + continue + + if len(module._parameters) <= 0: + continue + elif any(skip_module in name for skip_module in skip_modules): + continue + else: + if "model.layers" in name and int(names[2]) >= self.config.num_hidden_layers: + mtp_layer_idx = int(names[2]) - self.config.num_hidden_layers + names[2] = str( + mtp_layer_idx % self.config.num_nextn_predict_layers + + self.config.num_hidden_layers + ) + name = ".".join(names) + if names[-1] == "kv_b_proj": + # TODO: remove weight_dequant after enabling fp8_bmm + dequant_kv_b_proj = ( + self.model_config.quant_config.is_module_excluded_from_quantization( + names[-1] + ) + ) + if dequant_kv_b_proj: + kv_b_proj, k_b_proj_trans = load_kv_b_proj_and_k_b_proj_trans_dequant(name) + else: + kv_b_proj, k_b_proj_trans = load_kv_b_proj_and_k_b_proj_trans( + name, is_scale=False + ) + module.weight.data.copy_(kv_b_proj.reshape(module.weight.shape)) + + attn_module = all_named_modules[parent_module_name] + _, v_b_proj = split_kv_b_proj(module.weight.data, is_scale=False) + attn_module.v_b_proj = nn.Parameter(v_b_proj, requires_grad=False) + + attn_module.k_b_proj_trans.data.copy_( + k_b_proj_trans.reshape(attn_module.k_b_proj_trans.shape) + ) + + if getattr(module, "weight_scale", None) is not None and not dequant_kv_b_proj: + kv_b_proj_scale, k_b_proj_trans_scale = load_kv_b_proj_and_k_b_proj_trans( + name, is_scale=True + ) + module.weight_scale.copy_( + kv_b_proj_scale.reshape(module.weight_scale.shape) + ) + attn_module.k_b_proj_trans_scale.copy_( + k_b_proj_trans_scale.reshape(attn_module.k_b_proj_trans_scale.shape) + ) + + _, v_b_proj_scale = split_kv_b_proj(module.weight_scale.data, is_scale=True) + attn_module.v_b_proj_scale = nn.Parameter( + v_b_proj_scale, requires_grad=False + ) + + if attn_module.k_b_proj_trans_dequant is not None: + attn_module.k_b_proj_trans_dequant.data.copy_( + weight_dequant( + k_b_proj_trans.view(-1, k_b_proj_trans.shape[-1]).cuda(), + k_b_proj_trans_scale.view( + -1, k_b_proj_trans_scale.shape[-1] + ).cuda(), + ) + .view(*attn_module.k_b_proj_trans_dequant.shape) + .to(attn_module.k_b_proj_trans_dequant.dtype) + ) + if attn_module.v_b_proj_dequant is not None: + attn_module.v_b_proj_dequant.data.copy_( + weight_dequant( + v_b_proj.view(-1, v_b_proj.shape[-1]).cuda(), + v_b_proj_scale.view(-1, v_b_proj_scale.shape[-1]).cuda(), + ) + .view(*attn_module.v_b_proj_dequant.shape) + .to(attn_module.v_b_proj_dequant.dtype) + ) + elif names[-1] == "kv_a_proj_with_mqa": + nvfp4_fused_a = ( + self.model_config.get_quant_config().layer_quant_mode.has_nvfp4() + and weights[f"{'.'.join(names[:-1])}.kv_a_proj_with_mqa.weight"].dtype + == fp4_utils.float4_e2m1x2 + and weights[f"{'.'.join(names[:-1])}.q_a_proj.weight"].dtype + == fp4_utils.float4_e2m1x2 + ) + if nvfp4_fused_a: + ########### input_scale + kv_a_proj_with_mqa_input_scale = weights[ + f"{'.'.join(names[:-1])}.kv_a_proj_with_mqa.input_scale" + ] + if not is_lite: + q_a_proj_input_scale = weights[ + f"{'.'.join(names[:-1])}.q_a_proj.input_scale" + ] + assert kv_a_proj_with_mqa_input_scale == q_a_proj_input_scale, ( + "kv_a_proj_with_mqa.input_scale and q_a_proj.input_scale should be the same" + ) + # modelopt ckpt stores amax/(448*6), convert to (448*6)/amax + shared_input_scale = kv_a_proj_with_mqa_input_scale + module.input_scale.data.copy_(1.0 / shared_input_scale) + E2M1_MAX = 6.0 + module.inv_input_scale.data.copy_(module.input_scale / E2M1_MAX) + ########### weight_scale_2 + need_requant_kv_a_proj_with_mqa = False + need_requant_q_a_proj = False + kv_a_proj_with_mqa_scale_2 = weights[ + f"{'.'.join(names[:-1])}.kv_a_proj_with_mqa.weight_scale_2" + ] + shared_weight_scale_2 = kv_a_proj_with_mqa_scale_2 + if not is_lite: + q_a_proj_scale_2 = weights[ + f"{'.'.join(names[:-1])}.q_a_proj.weight_scale_2" + ] + if kv_a_proj_with_mqa_scale_2 < q_a_proj_scale_2: + shared_weight_scale_2 = q_a_proj_scale_2 + need_requant_kv_a_proj_with_mqa = True + elif q_a_proj_scale_2 < kv_a_proj_with_mqa_scale_2: + need_requant_q_a_proj = True + + ########### alpha + alpha = shared_input_scale.float() * shared_weight_scale_2.float() + module.alpha.data.copy_(alpha) + module.scalar_alpha = alpha.item() + + ########### weights + kv_a_proj_with_mqa = weights[ + f"{'.'.join(names[:-1])}.kv_a_proj_with_mqa.weight" + ][:] + + if not is_lite: + q_a_proj = weights[f"{'.'.join(names[:-1])}.q_a_proj.weight"][:] + + ########### weight_scale + kv_a_proj_with_mqa_scale = weights[ + f"{'.'.join(names[:-1])}.kv_a_proj_with_mqa.weight_scale" + ][:] + kv_a_proj_with_mqa_scale = torch.ops.trtllm.block_scale_interleave( + kv_a_proj_with_mqa_scale.view(fp4_utils.float4_sf_dtype) + ) + if not is_lite: + q_a_proj_scale = weights[ + f"{'.'.join(names[:-1])}.q_a_proj.weight_scale" + ][:] + q_a_proj_scale = torch.ops.trtllm.block_scale_interleave( + q_a_proj_scale.view(fp4_utils.float4_sf_dtype) + ) + + ########### requantize + if need_requant_kv_a_proj_with_mqa: + # requant kv_a_proj_with_mqa + kv_a_proj_with_mqa, kv_a_proj_with_mqa_scale = ( + requantize_weight_with_new_scale( + kv_a_proj_with_mqa, + kv_a_proj_with_mqa_scale, + kv_a_proj_with_mqa_scale_2, + shared_weight_scale_2, + device=module.weight.device, + ) + ) + if need_requant_q_a_proj: + # requant q_a_proj + q_a_proj, q_a_proj_scale = requantize_weight_with_new_scale( + q_a_proj, + q_a_proj_scale, + q_a_proj_scale_2, + shared_weight_scale_2, + device=module.weight.device, + ) + + ########### fuse and load weights + if not is_lite: + fused_a = torch.cat([q_a_proj, kv_a_proj_with_mqa], dim=0) + else: + fused_a = kv_a_proj_with_mqa + + # For DeepseekV32: kv_a_proj_with_mqa is oversized + # to include indexer k weights, which is filled in post_load_weights. + module.weight.data[0 : fused_a.shape[0]].copy_(fused_a) + + ########### fuse weight_scale + if not is_lite: + fused_a_scale = torch.cat( + [q_a_proj_scale, kv_a_proj_with_mqa_scale], dim=0 + ) + else: + fused_a_scale = kv_a_proj_with_mqa_scale + # For DeepseekV32: kv_a_proj_with_mqa is oversized + # to include indexer k weights, which is filled in post_load_weights. + _copy_deepseek_v4_fused_a_weight_scale(module, fused_a, fused_a_scale) + else: + fused_a = weights[f"{'.'.join(names[:-1])}.kv_a_proj_with_mqa.weight"][:] + if not is_lite: + q_a_proj = weights[f"{'.'.join(names[:-1])}.q_a_proj.weight"][:] + fused_a = torch.cat([q_a_proj, fused_a], dim=0) + + if f"{'.'.join(names[:-1])}.kv_a_proj_with_mqa.weight_scale_inv" in weights: + fused_a_scale = weights[ + f"{'.'.join(names[:-1])}.kv_a_proj_with_mqa.weight_scale_inv" + ] + if not is_lite: + q_a_proj_scale = weights[ + f"{'.'.join(names[:-1])}.q_a_proj.weight_scale_inv" + ][:] + fused_a_scale = torch.cat([q_a_proj_scale, fused_a_scale], dim=0) + + _copy_deepseek_v4_fused_a_weight_scale(module, fused_a, fused_a_scale) + # For DeepseekV32: kv_a_proj_with_mqa is oversized + # to include indexer k weights, which is filled in post_load_weights. + module.weight.data[0 : fused_a.shape[0]].copy_(fused_a) + elif names[-1] in params_map: + module_weights = [] + for new_name in params_map[names[-1]]: + module_weights.append( + filter_weights(".".join(names[:-1] + [new_name]), weights) + ) + module.load_weights(weights=module_weights) + elif names[-1] == "experts": + module_weights = filter_weights(name, weights) + module_weights = rename_moe_weight( + module_weights, + { + "down_proj": "w2", + "up_proj": "w3", + "gate_proj": "w1", + }, + ) + module.load_weights(weights=[module_weights]) + elif names[-1] == "backend" and isinstance(module, MoE): + # Special case: ConfigurableMoE.backend (TRTLLMGenFusedMoE) + # Currently saved MoE weights don't include 'backend' in their names. + # After MoE refactoring, ConfigurableMoE now has a backend submodule, + # and weights loading is done in the backend, so module name includes '.backend'. + # We need to use parent module name (without .backend) to match saved weight names. + # After MoE refactoring is fully complete, all paths will follow this branch. + parent_name = ".".join(names[:-1]) + module_weights = filter_weights(parent_name, weights) + module_weights = rename_moe_weight( + module_weights, + { + "down_proj": "w2", + "up_proj": "w3", + "gate_proj": "w1", + }, + ) + module.load_weights(weights=[module_weights]) + elif names[-1] == "self_attn": + if f"{name}.o_a_proj" in weights: + load_o_a_proj(name, module) + attn_sink_key = f"{name}.attn_sink" + if attn_sink_key in weights: + sink_full = weights[attn_sink_key][:] + if not self.model_config.mapping.enable_attention_dp: + sink_full = split(sink_full, tp_size, tp_rank) + sink_full = sink_full.to(torch.float32).cuda() + module.mqa.attn_sink = nn.Parameter(sink_full, requires_grad=False) + continue + elif names[-1] == "mqa": + # DeepseekV4TrtllmAttention owns the optional attn_sink + # (per-head fp32, already TP-sharded). The checkpoint key + # uses the parent attention module name, not the .mqa + # suffix. When the key is absent we leave module.attn_sink + # as None so DeepseekV4TrtllmAttention.forward does not pass + # attention_sinks to the kernel. + parent_attn_name = ".".join(names[:-1]) + attn_sink_key = f"{parent_attn_name}.attn_sink" + if attn_sink_key in weights: + sink_full = weights[attn_sink_key][:] + if not self.model_config.mapping.enable_attention_dp: + sink_full = split(sink_full, tp_size, tp_rank) + sink_full = sink_full.to(torch.float32).cuda() + module.attn_sink = nn.Parameter(sink_full, requires_grad=False) + continue + elif names[-1] == "next_layer_layernorm": + continue + elif isinstance(module, (mHC, HCHead)) and load_flat_hc_weights(module, names): + continue + elif names[-1] in ("engram",): + # Engram is a container module with no direct parameters; + # its leaf sub-modules (multi_head_embedding, kv_proj, + # short_conv) and direct parameters (key_norm_weight, + # query_norm_weight) are loaded via the generic path. + continue + else: + module_weights = filter_weights(name, weights) + if hasattr(module, "load_weights"): + module.load_weights(weights=[module_weights]) + else: + for n, p in module.named_parameters(): + p.data.copy_(module_weights[n][:]) + + +@torch.compile(options={"max-autotune": True}) +def _get_last_token_states(hidden_states, attn_metadata): + last_tokens = ( + torch.cumsum( + attn_metadata.seq_lens_cuda, + dim=0, + dtype=torch.long, + ) + - 1 + ) + return hidden_states[last_tokens] + + +class DeepseekV4MTPHead(nn.Module): + def __init__(self, model_config: ModelConfig[PretrainedConfig]): + super().__init__() + config = model_config.pretrained_config + self.model_config = model_config + + self.norm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + self.hc_head = HCHead(config.hc_mult, config.hidden_size) + self.hc_mult = config.hc_mult + self.hidden_dim = config.hidden_size + + self.mapping_lm_head_tp = None + + def forward( + self, + hidden_states: torch.Tensor, + lm_head: Linear, + attn_metadata: DeepseekV4TrtllmAttentionMetadata, + return_context_logits: bool = False, + ) -> torch.Tensor: + if not return_context_logits: + if attn_metadata is not None: + hidden_states = _get_last_token_states(hidden_states, attn_metadata) + else: + hidden_states = hidden_states[-1].unsqueeze(0) + + hidden_states = hidden_states.reshape(-1, self.hc_mult, self.hidden_dim) + hidden_states = self.hc_head(hidden_states) + hidden_states = self.norm(hidden_states) + + enable_attention_dp = self.model_config.mapping.enable_attention_dp + enable_lm_head_tp_in_adp = ( + enable_attention_dp and self.model_config.mapping.enable_lm_head_tp_in_adp + ) + + # Add pre-lm gather logic + if enable_lm_head_tp_in_adp: + # ADP + LM TP mode: perform All-Gather before LM_head + self.mapping_lm_head_tp = create_lm_head_tp_mapping( + self.model_config.mapping, hidden_states.shape[0] + ) + hidden_states = allgather(hidden_states, self.mapping_lm_head_tp, dim=0) + + # Temporarily disable gather_output when not in ADP mode or (in ADP mode and LM TP is enabled) + if not enable_attention_dp or enable_lm_head_tp_in_adp: + lm_head.gather_output = False + logits = lm_head( + hidden_states, mapping_lm_head_tp=self.mapping_lm_head_tp, is_spec_decoding_head=True + ) + if not enable_attention_dp or enable_lm_head_tp_in_adp: + lm_head.gather_output = True + return logits + + +class DeepseekV4LogitsProcessor(nn.Module): + def __init__( + self, + model_config: ModelConfig[PretrainedConfig], + hc_head: HCHead, + norm: RMSNorm, + ): + super().__init__() + config = model_config.pretrained_config + self.model_config = model_config + self.hc_mult = config.hc_mult + self.hidden_dim = config.hidden_size + # Keep HCHead and final norm owned by DeepseekV4Model. This processor only + # borrows them, so checkpoint loading and PP weight removal still happen + # through the model's normal module tree. + object.__setattr__(self, "_hc_head", hc_head) + object.__setattr__(self, "_norm", norm) + + def forward( + self, + hidden_states: torch.Tensor, + lm_head: Linear, + attn_metadata: DeepseekV4TrtllmAttentionMetadata, + return_context_logits: bool = False, + ) -> torch.Tensor: + if not self.model_config.mapping.is_last_pp_rank(): + return lm_head(hidden_states).float() + + if not return_context_logits: + if attn_metadata is not None: + hidden_states = _get_last_token_states(hidden_states, attn_metadata) + else: + hidden_states = hidden_states[-1] + + hidden_states = hidden_states.reshape(-1, self.hc_mult, self.hidden_dim) + hidden_states = self._hc_head(hidden_states) + hidden_states = self._norm(hidden_states) + return lm_head(hidden_states).float() + + +class DeepseekV4Linear(Linear): + """ + A wrapper around Linear because we may optionally use min-latency kernels depending on input shapes. + """ + + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = True, + dtype: torch.dtype = None, + mapping: Optional[Mapping] = None, + tensor_parallel_mode: Optional[TensorParallelMode] = None, + gather_output: bool = False, # COLUMN parallel only + quant_config: Optional[QuantConfig] = None, + weights_loading_config: Optional[WeightsLoadingConfig] = None, + reduce_output: bool = True, # ROW parallel only + skip_create_weights_in_init: bool = False, + use_custom_cublas_mm: bool = False, + lora: Optional[LoraLayer] = None, + ): + super().__init__( + in_features, + out_features, + bias, + dtype, + mapping, + tensor_parallel_mode, + gather_output, + quant_config, + weights_loading_config, + reduce_output, + skip_create_weights_in_init, + use_custom_cublas_mm, + lora, + ) + + def apply_linear( + self, + input, + bias, + lora_params: Optional[dict] | None = None, + layer_idx: Optional[int] | None = None, + ): + num_tokens = input.shape[0] + if not self.has_any_quant and 1 <= num_tokens <= 16 and get_sm_version() not in [120, 121]: + output = torch.ops.trtllm.dsv3_fused_a_gemm_op(input, self.weight.t(), bias, None) + else: + output = super().apply_linear(input, bias, lora_params, layer_idx) + return output + + +class DeepseekV4Attention(MLA): + def __init__( + self, + model_config: ModelConfig[PretrainedConfig], + layer_idx: Optional[int] = None, + aux_stream: Optional[torch.cuda.Stream] = None, + mapping_with_cp: Optional[Mapping] = None, + reduce_output: bool = True, + ): + config = model_config.pretrained_config + assert config.qk_rope_head_dim == 64, ( + "DeepseekV4Attention only supports qk_rope_head_dim=64" + ) + assert config.kv_lora_rank == 448, "DeepseekV4Attention only supports kv_lora_rank=448" + predicted_tokens_per_seq = ( + model_config.spec_config.tokens_per_gen_step + if model_config.spec_config is not None + else 1 + ) + super().__init__( + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=config.kv_lora_rank, + predicted_tokens_per_seq=predicted_tokens_per_seq, + max_position_embeddings=config.max_position_embeddings, + bias=False, + pos_embd_params=_deepseek_v4_pos_embd_params( + config, model_config, layer_idx, predicted_tokens_per_seq + ), + layer_idx=layer_idx, + dtype=config.torch_dtype, + config=model_config, + aux_stream=aux_stream, + num_groups=config.o_groups, + o_lora_rank=config.o_lora_rank, + mapping_with_cp=mapping_with_cp, + reduce_output=reduce_output, + ) + + self.indexer = getattr(self.mqa, "indexer", None) + self.compressor = getattr(self.mqa, "compressor", None) + + +class DeepseekV4Gate(nn.Module): + def __init__( + self, + hidden_size: int, + num_experts: int, + top_k: int, + n_group: int, + topk_group: int, + routed_scaling_factor: float, + is_hashed: bool, + vocab_size: Optional[int] = None, + dtype: Optional[torch.dtype] = None, + fuse_routing_kernel: bool = True, + apply_routing: bool = False, + moe_backend: str = "CUTLASS", + ): + super().__init__() + self.weight = nn.Parameter( + torch.empty((num_experts, hidden_size), dtype=dtype), requires_grad=False + ) + self.moe_backend = moe_backend + bias_dtype = torch.float32 + + self.is_hashed = is_hashed + self.top_k = top_k + self.n_group = n_group + self.topk_group = topk_group + self.routed_scaling_factor = routed_scaling_factor + + if self.is_hashed: + # self.tid2eid = nn.Parameter( + # torch.empty(vocab_size, top_k, dtype=torch.int32), requires_grad=False + # ) + # WAR to avoid illegal expert indexes in hashed gating + self.tid2eid = nn.Parameter( + torch.stack([torch.randperm(num_experts)[:top_k] for _ in range(vocab_size)]) + .to(torch.int32) + .contiguous(), + requires_grad=False, + ) + self.e_score_correction_bias = None + else: + self.e_score_correction_bias = nn.Parameter( + torch.empty(num_experts, dtype=bias_dtype), requires_grad=False + ) + + assert not apply_routing, "DeepseekV4Gate routing is called inside MoE" + + def fetch_e_score_correction_bias(): + if not self.is_hashed: + return self.e_score_correction_bias.to(torch.float32) + else: + return None + + self._routing_method = DeepSeekV4MoeRoutingMethod( + top_k=self.top_k, + n_group=self.n_group, + topk_group=self.topk_group, + routed_scaling_factor=self.routed_scaling_factor, + # Pass a callable to fetch the tensor from DeepseekV4Gate at runtime, ensuring it is on the correct device + callable_e_score_correction_bias=fetch_e_score_correction_bias, + callable_tid2eid=lambda: self.tid2eid, + is_hashed=self.is_hashed, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return torch.ops.trtllm.dsv3_router_gemm_op( + hidden_states, self.weight.t(), bias=None, out_dtype=torch.float32 + ) + + def load_weights(self, weights: List[Dict]): + assert len(weights) == 1 + + self.weight.copy_(weights[0]["weight"][:]) + + if self.is_hashed: + self.tid2eid.copy_(weights[0]["tid2eid"][:].to(self.tid2eid.dtype)) + else: + self.e_score_correction_bias.copy_( + weights[0]["e_score_correction_bias"][:].to(self.e_score_correction_bias.dtype) + ) + + @property + def routing_method(self) -> DeepSeekV4MoeRoutingMethod: + return self._routing_method + + def apply( + self, logits: torch.Tensor, input_ids: Optional[torch.Tensor] = None + ) -> tuple[torch.Tensor, torch.Tensor]: + # topk routing + return self.routing_method.apply(logits, input_ids) + + def get_experts_per_token(self): + return self.routing_method.top_k + + +class DeepseekV4MoE(nn.Module): + def __init__( + self, + *, + num_experts: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + shared_expert_intermediate_size: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + dtype: Optional[torch.dtype] = None, + model_config: ModelConfig = ModelConfig(), + override_quant_config: Optional[QuantConfig] = None, + layer_idx: Optional[int] = None, + ): + from ..distributed import AllReduce + + super().__init__() + config = model_config.pretrained_config + self.top_k = top_k + self.use_dp = model_config.mapping.enable_attention_dp + self.gate = DeepseekV4Gate( + hidden_size, + num_experts, + top_k=top_k, + n_group=config.n_group, + topk_group=config.topk_group, + routed_scaling_factor=config.routed_scaling_factor, + is_hashed=layer_idx < config.n_hash_layers, + vocab_size=config.vocab_size, + dtype=dtype, + fuse_routing_kernel=True, + apply_routing=False, + moe_backend=model_config.moe_backend, + ) + experts_quant_config = self._get_experts_quant_config(model_config, layer_idx) + if override_quant_config is not None and experts_quant_config is model_config.quant_config: + experts_quant_config = override_quant_config + + swiglu_limit = getattr(config, "swiglu_limit", None) + moe_swiglu_limit = None + if swiglu_limit is not None: + # `create_moe` only accepts swiglu_limit for these MoE classes; + # resolve via get_moe_cls so backend-string fallbacks (e.g. + # TRTLLM/CUTEDSL/DENSEGEMM dropping back to CutlassFusedMoE on + # unsupported quant) are handled correctly. + moe_cls = get_moe_cls(model_config, override_quant_config=experts_quant_config) + supports_swiglu_limit = moe_cls in ( + CutlassFusedMoE, + TritonFusedMoE, + TRTLLMGenFusedMoE, + WideEPMoE, + DeepGemmFusedMoE, + ) + # NVFP4 routed-expert path: the TRTLLM-Gen fp4-block-scale fused-MoE + # cubin produces near-zero accuracy without bias even when + # swiglu_limit is supplied; drop the limit there until the cubin + # gains a no-bias clamp variant. MXFP4 variants are unaffected. + kernel_requires_bias_for_swiglu_limit = ( + moe_cls in (TRTLLMGenFusedMoE, WideEPMoE) + and experts_quant_config.quant_mode.has_nvfp4() + ) + if supports_swiglu_limit and not kernel_requires_bias_for_swiglu_limit: + moe_load_balancer_config = getattr(model_config, "moe_load_balancer", None) + num_slots = ( + moe_load_balancer_config.num_slots + if moe_load_balancer_config and moe_load_balancer_config.num_slots + else num_experts + ) + local_num_slots = num_slots // model_config.mapping.moe_ep_size + device = "cuda" if torch.cuda.is_available() else "cpu" + moe_swiglu_limit = torch.full( + (local_num_slots,), float(swiglu_limit), dtype=torch.float32, device=device + ) + + self.experts = create_moe( + num_experts=num_experts, + routing_method=self.gate.routing_method, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + reduce_results=False, # In both low‑latency and attention‑DP modes, FusedMoE skips the in‑op all‑reduce. + model_config=model_config, + override_quant_config=experts_quant_config, + aux_stream_dict=aux_stream_dict, + layer_idx=layer_idx, + # DS-R1 W4A8 is only supported through custom quantization script from + # examples/quantization/quantize_mixed_precision_moe.py + weight_loading_mode=( + MoEWeightLoadingMode.W4A8_CUSTOM + if experts_quant_config.layer_quant_mode.is_int4_weight_only_per_group() + else MoEWeightLoadingMode.VANILLA + ), + swiglu_limit=moe_swiglu_limit, + swiglu_limit_scalar=(float(swiglu_limit) if swiglu_limit is not None else None), + ) + + self.mapping = model_config.mapping + + # FIXME: incompatible with mixed quantization mode (including excluding modules from quantization) + block_size = 1 + if model_config.quant_config and model_config.quant_config.group_size is not None: + block_size = model_config.quant_config.group_size + + shared_tp_size, self.shared_output_scale = self._compute_shared_expert_tp_size( + shared_expert_intermediate_size, block_size + ) + + self.shared_experts = GatedMLP( + hidden_size=hidden_size, + intermediate_size=shared_expert_intermediate_size, + bias=False, + dtype=dtype, + config=model_config, + overridden_tp_size=shared_tp_size, + reduce_output=False, + swiglu_limit=swiglu_limit, + ) + + self.allreduce = None + if not self.use_dp and self.mapping.tp_size > 1: + self.allreduce = AllReduce( + mapping=model_config.mapping, strategy=model_config.allreduce_strategy + ) + self.aux_stream = aux_stream_dict[AuxStreamType.MoeShared] + self.event_dict = {key: torch.cuda.Event() for key in [EventType.Main, EventType.MoeShared]} + + # Store config values for perfect routing. + self.model_config = model_config + self.dtype = dtype + + def _compute_shared_expert_tp_size( + self, intermediate_size: int, block_size: int + ) -> tuple[int, float | None]: + """ + In the case of Deepseek-R1, the TP size of MLP is capped by intermediate_size // block_size. + For example, when the intermediate_size is 2048 and block scaling size is 128, + TP sizes are limited to {1, 2, 4, 8, 16} because of 2048/128 = 16. + + Args: + intermediate_size (int): MLP intermediate size. + block_size (int): The quantization block scale size. In the case of Deepseek FP8 recipe, + it's 128. For NVFP4, it's 16. + + Returns: + tuple[int, float | None]: A tuple containing (shared_tp_size, shared_output_scale). + - shared_tp_size: The computed TP size. + - shared_output_scale: The output scale factor, or None if not needed. + """ + + assert intermediate_size % block_size == 0, ( + "intermediate_size must be divisible by block_size." + ) + + shared_output_scale = None + # The block scale size is 128, which requires shared_expert_intermediate_size to be divisible by 128. + if self.use_dp: + # If using attention DP, the shared experts also use DP instead of TP. + shared_tp_size = 1 + else: + # Due to the restriction of block scale size (i.e., 128), the supported + # TP sizes only include 1, 2, 4, 8, and 16. The math.gcd operation ensures that + # shared_tp_size falls in the supported TP sizes. + shared_tp_size = math.gcd( + intermediate_size // block_size, + self.mapping.tp_size, + ) + # If shared_tp_size has been overridden, the output of shared experts needs + # to be scaled down accordingly before all-reduce. + if shared_tp_size != self.mapping.tp_size: + shared_output_scale = shared_tp_size / self.mapping.tp_size + + return shared_tp_size, shared_output_scale + + @staticmethod + def _get_experts_quant_config(model_config, layer_idx: int) -> QuantConfig: + if getattr(model_config, "quant_config_dict", None) is None: + return model_config.quant_config + return model_config.quant_config_dict.get( + f"model.layers.{layer_idx}.mlp.experts", model_config.quant_config + ) + + def compute_routed_output( + self, hidden_states, hidden_states_fp4, input_ids, all_rank_num_tokens, do_finalize + ): + # max-throughput + use_dp_padding = False + # Add DP padding on SM120 for context comm performance + # TODO: Move this model-agonostic part to MoE + if self.use_dp and self.mapping.tp_size > 1 and get_sm_version() == 120: + use_dp_padding = True + hidden_states = torch.nn.functional.pad( + hidden_states, (0, 0, 0, max(all_rank_num_tokens) - hidden_states.shape[0]) + ) + input_ids = torch.nn.functional.pad( + input_ids, + (0, max(all_rank_num_tokens) - input_ids.shape[0]), + value=self.model_config.pretrained_config.pad_token_id, + ) + + router_logits = self.gate(hidden_states) + + # When ENABLE_PERFECT_ROUTER=1, the FusedMoE backend swaps router_logits + # for cached ideal logits inside its own forward (interface.py: + # MoE.forward -> _maybe_get_perfect_router_logits), so we don't need to + # do the substitution here. Replacing router_logits in this wrapper + # would duplicate the work and use the wrong ep_rank/routing_method. + + routed_output = self.experts( + hidden_states_fp4 if hidden_states_fp4 is not None else hidden_states, + router_logits, + input_ids=input_ids, + do_finalize=do_finalize, + output_dtype=hidden_states.dtype, + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=use_dp_padding, + **({"alltoall_result_do_sum": False} if isinstance(self.experts, WideEPMoE) else {}), + ) + + return routed_output + + def forward( + self, + hidden_states: torch.Tensor, + hidden_states_fp4: Optional[Fp4QuantizedTensor] = None, + input_ids: Optional[torch.Tensor] = None, + all_rank_num_tokens: Optional[list[int]] = None, + final_all_reduce_params: Optional[AllReduceParams] = None, + do_finalize: Optional[bool] = True, + ) -> torch.Tensor: + if not do_finalize: + assert not self.use_dp + + def _compute_shared_output(): + shared_output = self.shared_experts( + hidden_states_fp4 if hidden_states_fp4 is not None else hidden_states + ) + if self.shared_output_scale is not None: + shared_output *= self.shared_output_scale + return shared_output + + def _compute_routed_output(): + routed_output = self.compute_routed_output( + hidden_states, hidden_states_fp4, input_ids, all_rank_num_tokens, do_finalize + ) + return routed_output + + # NOTE: define compiled helpers at module scope to avoid defining decorators inside compiled frames + + routed_output, shared_output = maybe_execute_in_parallel( + _compute_routed_output, + _compute_shared_output, + self.event_dict[EventType.Main], + self.event_dict[EventType.MoeShared], + self.aux_stream, + ) + + if not do_finalize: + return [shared_output, *routed_output] + else: + if routed_output.dim() == 3: + assert shared_output.numel() * self.top_k == routed_output.numel(), ( + "unmatched tensor shape" + ) + final_hidden_states = moe_reduce_add_shared_output(routed_output, shared_output) + else: + assert shared_output.size() == routed_output.size(), "unmatched tensor shape" + final_hidden_states = shared_output + routed_output + + if not self.use_dp and self.mapping.tp_size > 1: + final_hidden_states = self.allreduce( + final_hidden_states, all_reduce_params=final_all_reduce_params + ) + + return final_hidden_states + + +class DeepseekV4DecoderLayer(DecoderLayer): + def __init__( + self, + model_config: ModelConfig[PretrainedConfig], + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + is_separate_draft_engine: bool = False, + mapping_with_cp: Optional[Mapping] = None, + disable_post_moe_fusion: bool = False, + ): + super().__init__() + self.model_config = model_config + self.config = model_config.pretrained_config + config = self.config + + self.hidden_size = config.hidden_size + self.moe_intermediate_size = config.moe_intermediate_size + self.num_experts = config.n_routed_experts + self.num_shared_experts = config.n_shared_experts + self.top_k = config.num_experts_per_tok + + self.mapping = model_config.mapping + mapping = self.mapping + self.enable_attention_dp = mapping.enable_attention_dp + self.mlp_tp_size = mapping.tp_size + self.is_p2p_supported = can_access_peer(mapping) + + self.hc_attn = mHC( + config.hc_mult, + config.hidden_size, + config.hc_sinkhorn_iters, + dtype=torch.float32, + post_mult_value=2.0, + ) + + layer_idx_for_attention = layer_idx + if is_separate_draft_engine: + # KVCacheManager only support 1 layer for separate draft engine + layer_idx_for_attention = layer_idx - model_config.pretrained_config.num_hidden_layers + + self.self_attn = DeepseekV4Attention( + model_config, + layer_idx=layer_idx_for_attention, + aux_stream=aux_stream_dict[AuxStreamType.Attention], + reduce_output=not self.enable_attention_dp and self.mapping.tp_size > 1, + ) + + self.fusion_config = EagerFusionConfig() + self.enable_fusion = os.environ.get("TRTLLM_DEEPSEEK_EAGER_FUSION_DISABLED", "0") == "0" + self.enable_fusion &= not self.enable_attention_dp + + self.hc_ffn = mHC( + config.hc_mult, + config.hidden_size, + config.hc_sinkhorn_iters, + dtype=torch.float32, + post_mult_value=2.0, + ) + + # FIXME: incompatible with mixed quantization mode + quant_config = self._get_decoder_layer_quant_config(model_config, layer_idx) + self.is_nvfp4 = quant_config.layer_quant_mode.has_nvfp4() + assert quant_config.quant_algo is not QuantAlgo.MIXED_PRECISION, ( + "MIXED_PRECISION is ambiguous" + ) + + self.allreduce = None + self.moe_allreduce = None + if not self.enable_attention_dp and self.mapping.tp_size > 1: + self.allreduce = AllReduce( + mapping=model_config.mapping, + strategy=model_config.allreduce_strategy, + dtype=config.torch_dtype, + ) + self.moe_allreduce = MoEAllReduce(self.mapping) + + has_tp = mapping.has_tp() + self.fusion_config.PRE_MOE_FUSION = self.enable_fusion and has_tp + # DeepSeek-V4 applies the next RMSNorm after mHC post_mapping and the + # next layer's mHC pre_mapping. Fusing the post-MoE all-reduce with the + # next RMSNorm would normalize the raw MoE output before mHC post_mapping, + # which is not equivalent. + self.fusion_config.POST_MOE_FUSION = False + + self.mlp = DeepseekV4MoE( + num_experts=self.num_experts, + top_k=self.top_k, + hidden_size=self.hidden_size, + intermediate_size=self.moe_intermediate_size, + shared_expert_intermediate_size=self.moe_intermediate_size * self.num_shared_experts, + dtype=config.torch_dtype, + model_config=model_config, + override_quant_config=quant_config, + aux_stream_dict=aux_stream_dict, + layer_idx=layer_idx, + ) + + self.input_layernorm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + + # When enable_attention_dp is True, we normally skip attention all-reduce since each + # DP rank works on different batch elements. However, with CP > 1, attention is split + # across CP ranks for the SAME batch element, so all-reduce is still needed. + has_cp = mapping_with_cp is not None and mapping_with_cp.cp_size > 1 + can_skip_for_attention_dp = self.enable_attention_dp and not has_cp + self.disable_attn_allreduce = ( + self.fusion_config.PRE_MOE_FUSION + or self.fusion_config.PRE_MLP_FUSION + or self.mapping.tp_size == 1 + or can_skip_for_attention_dp + ) + + self.post_attention_layernorm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + self.layer_idx = layer_idx + # is_first_layer is baked in at __init__ time so the Python-side branch in + # forward() resolves at CUDA-graph capture time. + self.is_first_layer = layer_idx == 0 + # fused_hc knob: pretrained-config attr `enable_fused_hc` controls whether + # the MHC boundary fusion (`mHC.fused_hc`) is used. When False, fall back + # to the unfused `post_mapping → pre_mapping` chain (same path engram + # layers already take). Env var TRTLLM_MHC_ENABLE_FUSED_HC overrides the + # config attr (set to "0" to force-disable for validation/rollback). + self.enable_fused_hc = _resolve_enable_fused_hc(config) + self.next_layer_layernorm: RMSNorm = None + # Finalized in DeepseekV4ForCausalLM.post_load_weights once the full layer + # list is visible: a layer may defer its hc_ffn.post_mapping only if + # the next layer is able to absorb it via fused_hc (i.e. the next + # layer has fused_hc enabled and no engram at its entry). Last layer + # never defers — hc_head consumes the residual directly. + self.defer_post_mapping: bool = False + + # Engram module (optional, for n-gram context augmentation) + self.engram: Optional[Engram] = None + _engram_config = getattr(config, "engram_config", None) + _engram_vocab_sizes_by_layer = getattr(config, "engram_vocab_sizes_by_layer", {}) + if _engram_config is not None and layer_idx in _engram_vocab_sizes_by_layer: + self.engram = Engram( + layer_id=layer_idx, + config=_engram_config, + vocab_sizes_flat=_engram_vocab_sizes_by_layer[layer_idx], + stream=aux_stream_dict[AuxStreamType.EngramPrecompute], + ) + + def _get_decoder_layer_quant_config( + self, model_config: ModelConfig[PretrainedConfig], layer_idx: int + ): + """ + The MTP layer in the nvfp4 checkpoint is unquantized. Because the TRTLLM + moe_backend only supports fp8/fp4 quantization, we need to override + the quant_config for the MTP layer. + """ + quant_config = model_config.quant_config + + layer_name = f"model.layers.{layer_idx}" + if quant_config.is_module_excluded_from_quantization(layer_name): + return QuantConfig( + quant_algo=None, + kv_cache_quant_algo=quant_config.kv_cache_quant_algo, + ) + else: + return model_config.quant_config + + def _compute_mlp_tp_size(self, intermediate_size: int, block_size: int) -> int: + """ + For DeepSeek‑R1, MLP TP size is limited by intermediate_size // block_size + and must also be multiples of gpus_per_node to avoid expensive inter‑node allreduce. + + Args: + intermediate_size (int): MLP intermediate size. + block_size (int): The quantization block scale size. In the case of Deepseek FP8 recipe, + it's 128. For NVFP4, it's 16. + + Returns: + int: The computed tp_size. + """ + + assert intermediate_size % block_size == 0, ( + "intermediate_size must be divisible by block_size." + ) + if self.enable_attention_dp: + # If using attention DP, the MLP also uses DP instead of TP. + mlp_tp_size = 1 + else: + # The two math.gcd operations ensure that mlp_tp_size falls in the candidate TP sizes. + tp = math.gcd( + intermediate_size // block_size, + self.mapping.tp_size, + ) + + if tp > self.mapping.gpus_per_node: + mlp_tp_size = math.gcd( + tp, + self.mapping.gpus_per_node, + ) # Avoid costly inter-node TP + else: + mlp_tp_size = tp + return mlp_tp_size + + def forward( + self, + position_ids: torch.IntTensor, + hc_state, + attn_metadata: DeepseekV4TrtllmAttentionMetadata, + spec_metadata: Optional[SpecMetadata] = None, + input_ids: Optional[torch.IntTensor] = None, + engram_embeddings=None, + **kwargs, + ): + """mHC-aware decoder layer with boundary fusion. + + ``hc_state`` carries the mHC pipeline state across layers: + + - ``is_first_layer=True``: ``hc_state`` is the initial residual tensor + ``[B, HC_MULT, hidden]`` (bf16). This layer bootstraps the stream + with ``hc_attn.pre_mapping`` (aka the standalone ``hc_pre``). + - ``is_first_layer=False``: ``hc_state`` is an ``HCState``. If + ``is_deferred`` (fused path), the 4 tensors feed the current + ``hc_attn.fused_hc``. Otherwise ``residual`` is already post-mapped + by the prior layer and this layer just runs ``pre_mapping``. + + Returns an ``HCState``. Fused mode returns a deferred state carrying + this layer's ``hc_ffn`` inputs, so the next layer absorbs + ``hc_ffn.post_mapping`` via its own ``fused_hc``. Engram / unfused + mode resolves the post_mapping in-layer and returns a resolved state. + + Engram (when enabled on this layer) injects a residual modification + between the previous block's post_mapping and this block's pre_mapping. + That mutation breaks the algebraic assumptions of ``fused_hc``, so + engram-enabled layers fall back to the unfused + ``post_mapping → +engram → pre_mapping`` chain at the entry boundary. + The mid-layer ``attn → MoE`` boundary is always safe to fuse. + """ + has_engram = self.engram is not None and engram_embeddings is not None + + # ------------------------------------------------------------------- + # Entry boundary: hc_pre (layer 0) or fused_hc / unfused chain. + # _entry_boundary returns layer_input already normalized by + # input_layernorm — fused via fused_hc on the deferred path, applied + # as a standalone RMSNorm on the layer-0 / engram / non-deferred path. + # ------------------------------------------------------------------- + residual, post_mix, comb_mix, layer_input = self._entry_boundary( + hc_state, engram_embeddings, has_engram + ) + + # ------------------------------------------------------------------- + # Attention block (layer_input is already input_layernorm-normed) + # ------------------------------------------------------------------- + x_attn = self.self_attn( + position_ids=position_ids, + hidden_states=layer_input, + attn_metadata=attn_metadata, + all_reduce_params=AllReduceParams(enable_allreduce=not (self.disable_attn_allreduce)), + **kwargs, + ) + + # ------------------------------------------------------------------- + # Mid-layer boundary: fuse hc_attn.post_mapping + hc_ffn.pre_mapping. + # No engram concern here because engram only fires at layer entry. + # When enable_fused_hc=False, fall back to the unfused chain. + # ------------------------------------------------------------------- + if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx): + self.fusion_config.POST_MOE_FUSION = False + if self.enable_fused_hc: + residual, post_mix, comb_mix, layer_input = self.hc_ffn.fused_hc( + x_prev=x_attn, + residual_prev=residual, + post_mix_prev=post_mix, + comb_mix_prev=comb_mix, + ) + else: + # Break fused_hc into post_mapping and pre_mapping as separate ops. + residual = self.hc_attn.post_mapping( + x=x_attn, + residual=residual, + post_layer_mix=post_mix, + comb_res_mix=comb_mix, + ) + post_mix, comb_mix, layer_input = self.hc_ffn.pre_mapping(residual) + + # ------------------------------------------------------------------- + # MoE block — returns x_ffn (post-MoE, already normed by + # next_layer_layernorm when POST_MOE_FUSION is on). + # ------------------------------------------------------------------- + x_ffn = self.forward_MoE( + hidden_states=layer_input, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + input_ids=input_ids, + ) + + # Defer this layer's hc_ffn.post_mapping only when the NEXT layer can + # absorb it via fused_hc (see post_load_weights). Otherwise resolve it + # here and hand the next layer a fully post-mapped residual. + if self.defer_post_mapping: + return HCState.deferred( + residual=residual, post_mix=post_mix, comb_mix=comb_mix, x_prev=x_ffn + ) + resolved_residual = self.hc_ffn.post_mapping( + x=x_ffn, + residual=residual, + post_layer_mix=post_mix, + comb_res_mix=comb_mix, + ) + return HCState.resolved(resolved_residual) + + def _entry_boundary(self, hc_state, engram_embeddings, has_engram): + """Resolve the per-layer entry into (residual, post_mix, comb_mix, layer_input). + + Two code paths: + 1. fused: previous layer deferred its post_mapping; fold it into this + layer's pre_mapping via hc_attn.fused_hc. The prev layer only + defers when post_load_weights has proved we can absorb it here + (no engram, fused_hc enabled), so a deferred state at entry is + guaranteed fusable. + 2. unfused: previous layer already resolved its post_mapping (or + this is layer 0); run pre_mapping on the residual directly. + """ + # Fused entry: prev layer deferred its post_mapping; fold it into + # hc_attn.fused_hc. By construction (post_load_weights), a deferred + # state at entry is guaranteed fusable. Layer 0 receives the raw + # residual tensor and has no HCState, so short-circuit on it. + # input_layernorm is folded into fused_hc's layer_input epilogue, so + # the returned layer_input is already RMSNorm-normalized. + if not self.is_first_layer and hc_state.is_deferred: + return self.hc_attn.fused_hc( + x_prev=hc_state.x_prev, + residual_prev=hc_state.residual, + post_mix_prev=hc_state.post_mix, + comb_mix_prev=hc_state.comb_mix, + norm_weight=self.input_layernorm.weight, + norm_eps=self.input_layernorm.variance_epsilon, + ) + + # Unfused entry: layer 0 hands us the initial residual tensor + # [B, HC_MULT, hidden] directly; post-layer-0 hands us a resolved + # HCState. Both collapse to "apply engram delta (if any) then run + # pre_mapping". Apply input_layernorm here so the caller always sees + # a normed layer_input regardless of which entry path ran. + residual = hc_state if self.is_first_layer else hc_state.residual + if has_engram: + residual = residual + self.engram(residual, engram_embeddings) + post_mix, comb_mix, layer_input = self.hc_attn.pre_mapping(residual) + layer_input = self.input_layernorm(layer_input) + return residual, post_mix, comb_mix, layer_input + + def forward_MoE( + self, + hidden_states: torch.Tensor, + attn_metadata: DeepseekV4TrtllmAttentionMetadata, + spec_metadata: Optional[SpecMetadata] = None, + input_ids: Optional[torch.IntTensor] = None, + ) -> torch.Tensor: + def _run_MoE(hidden_states, hidden_states_fp4, do_finalize, input_ids): + return self.mlp( + hidden_states, + hidden_states_fp4, + all_rank_num_tokens=attn_metadata.all_rank_num_tokens, + final_all_reduce_params=AllReduceParams( + enable_allreduce=not ( + self.fusion_config.POST_MOE_FUSION or self.mapping.tp_size == 1 + ) + ), + do_finalize=do_finalize, + input_ids=input_ids, + ) + + if self.fusion_config.PRE_MOE_FUSION: + # In DeepSeek-V4 the external residual connection is handled by mHC + # (hc_ffn.post_mapping), so there is no residual to add here. + # Use fused allreduce + RMSNorm (no residual addition). + hidden_states = self.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RMS_NORM, + norm_weight=self.post_attention_layernorm.weight, + eps=self.post_attention_layernorm.variance_epsilon, + trigger_completion_at_end=False, + ), + ) + else: + # No fusion: just normalize. + hidden_states = self.post_attention_layernorm(hidden_states) + + # Note: this fusion pattern is only supported for single-node TRTLLM-nvfp4 backend now + do_finalize = self.mapping.is_multi_node() or ( + not ( + self.fusion_config.POST_MOE_FUSION + and hidden_states.shape[0] <= self.moe_allreduce.max_token + and self.model_config.moe_backend == "TRTLLM" + and self.mlp.experts.has_nvfp4 + and self.is_p2p_supported + ) + ) + + hidden_states = _run_MoE( + hidden_states, hidden_states_fp4=None, input_ids=input_ids, do_finalize=do_finalize + ) + + if self.fusion_config.POST_MOE_FUSION: + if do_finalize: + # Fused allreduce + RMSNorm, no residual needed (mHC handles it). + hidden_states = self.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RMS_NORM, + norm_weight=self.next_layer_layernorm.weight, + eps=self.next_layer_layernorm.variance_epsilon, + trigger_completion_at_end=False, + ), + ) + else: + assert len(hidden_states) == 4, "hidden_states must have 4 elements" + + shared_output = hidden_states[0] + fc2_output = hidden_states[1] + expert_scale_factor = hidden_states[2] + expanded_idx_to_permuted_idx = hidden_states[3] + + moe_all_reduce_params = MoEAllReduceParams( + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + expert_scale_factor=expert_scale_factor, + shared_expert_output=shared_output, + residual=None, + norm_weight=self.next_layer_layernorm.weight, + eps=self.next_layer_layernorm.variance_epsilon, + is_cutlass_min_latency=False, + ) + (hidden_states,) = self.moe_allreduce( + fc2_output, all_reduce_params=moe_all_reduce_params + ) + else: + if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx): + spec_metadata.maybe_capture_hidden_states(self.layer_idx, hidden_states, None) + + return hidden_states + + +class DeepseekV4MTP(DeepseekV4DecoderLayer): + def __init__( + self, + model_config: ModelConfig[PretrainedConfig], + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + is_separate_draft_engine: bool = False, + ): + super().__init__( + model_config, + layer_idx, + aux_stream_dict, + is_separate_draft_engine, + disable_post_moe_fusion=True, + ) + config = model_config.pretrained_config + self.hidden_dim = config.hidden_size + self.moe_intermediate_size = config.moe_intermediate_size + self.num_experts = config.n_routed_experts + self.num_shared_experts = config.n_shared_experts + self.top_k = config.num_experts_per_tok + + self.aux_stream = aux_stream_dict[AuxStreamType.MoeShared] + self.event_dict = {key: torch.cuda.Event() for key in [EventType.Main, EventType.MoeShared]} + + self.enorm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + + self.hnorm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + self.hc_mult = config.hc_mult + if model_config.mapping.enable_attention_dp: + self.e_proj = Linear( + config.hidden_size, + config.hidden_size, + bias=False, + dtype=config.torch_dtype, + quant_config=model_config.get_quant_config(), + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + ) + self.h_proj = Linear( + config.hidden_size, + config.hidden_size, + bias=False, + dtype=config.torch_dtype, + quant_config=model_config.get_quant_config(), + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + ) + else: + self.e_proj = Linear( + config.hidden_size, + config.hidden_size, + bias=False, + dtype=config.torch_dtype, + tensor_parallel_mode=TensorParallelMode.ROW, + mapping=model_config.mapping, + reduce_output=True, + quant_config=model_config.get_quant_config(), + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + ) + self.h_proj = Linear( + config.hidden_size, + config.hidden_size, + bias=False, + dtype=config.torch_dtype, + tensor_parallel_mode=TensorParallelMode.ROW, + mapping=model_config.mapping, + reduce_output=True, + quant_config=model_config.get_quant_config(), + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + ) + + self.shared_head = DeepseekV4MTPHead(model_config) + + def forward( + self, + input_ids: torch.IntTensor, + position_ids: torch.IntTensor, + hidden_states: torch.Tensor, + embed_tokens: Embedding, + attn_metadata: DeepseekV4TrtllmAttentionMetadata, + all_rank_num_tokens: Optional[List[int]] = None, + spec_metadata: Optional[SpecMetadata] = None, + **kwargs, + ) -> torch.Tensor: + """Run an MTP layer. + + ``embed_tokens`` is injected by the one-model draft path and shared + with the target model. ``hidden_states`` is the flattened mHC residual + from the target or previous MTP layer: [num_tokens, hc_mult * hidden]. + """ + + def norm_embeds(): + return self.enorm(embed_tokens(input_ids)) # emdedding + + def norm_hidden(): + return self.hnorm(hidden_states.reshape(-1, self.hc_mult, self.hidden_dim)) + + inputs_embeds, hidden_states = maybe_execute_in_parallel( + norm_embeds, + norm_hidden, + self.event_dict[EventType.Main], + self.event_dict[EventType.MoeShared], + self.aux_stream, + ) + + # Split hidden_states columnwise based on TP + tp_size = self.model_config.mapping.tp_size + tp_rank = self.model_config.mapping.tp_rank + if tp_size > 1 and not (self.model_config.mapping.enable_attention_dp): + inputs_embeds = torch.chunk(inputs_embeds, tp_size, dim=-1)[tp_rank].contiguous() + hidden_states = torch.chunk(hidden_states, tp_size, dim=-1)[tp_rank].contiguous() + + inputs_embeds = self.e_proj(inputs_embeds).unsqueeze(1) + hidden_states = self.h_proj(hidden_states) + hidden_states = inputs_embeds + hidden_states + + original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = all_rank_num_tokens + try: + hc_state = super().forward( + position_ids=position_ids, + hc_state=HCState.resolved(hidden_states), + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + input_ids=input_ids, + **kwargs, + ) + finally: + attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens + hidden_states = hc_state.residual.flatten(1) + + return hidden_states + + +class DeepseekV4Model(DecoderModel): + def __init__( + self, model_config: ModelConfig[PretrainedConfig], mapping_with_cp: Optional[Mapping] = None + ): + super().__init__(model_config) + config = model_config.pretrained_config + self.vocab_size = config.vocab_size + self.num_hidden_layers = config.num_hidden_layers + self.hc_mult = config.hc_mult + aux_stream_list = [torch.cuda.Stream() for _ in range(5)] + self.aux_stream_dict = { + AuxStreamType.Attention: aux_stream_list[0], + AuxStreamType.MoeShared: aux_stream_list[0], + AuxStreamType.MoeChunkingOverlap: aux_stream_list[1], + AuxStreamType.MoeBalancer: aux_stream_list[2], + AuxStreamType.MoeOutputMemset: aux_stream_list[3], + AuxStreamType.EngramPrecompute: aux_stream_list[4], + } + + self.embed_tokens = Embedding( + config.vocab_size, + config.hidden_size, + dtype=config.torch_dtype, + ) + + self.hc_head = HCHead(config.hc_mult, config.hidden_size) + + # Engram hash provider (optional, for n-gram context augmentation) + # Must be created before layers so vocab sizes can be stored on config. + self.engram_hash_provider: Optional[EngramHashProvider] = None + self.use_engram = getattr(config, "has_engram", False) + if self.use_engram: + engram_config = EngramConfig( + tokenizer_name_or_path=getattr( + config, "tokenizer_name_or_path", "deepseek-ai/DeepSeek-V3" + ), + engram_vocab_size=config.engram_vocab_size, + max_ngram_size=config.engram_max_ngram_size, + n_embed_per_ngram=config.engram_n_embed_per_ngram, + n_head_per_ngram=config.engram_n_head_per_ngram, + layer_ids=config.engram_layer_ids, + pad_id=config.engram_pad_id, + seed=config.engram_seed, + kernel_size=config.engram_kernel_size, + hidden_size=config.hidden_size, + hc_mult=config.hc_mult, + norm_eps=config.rms_norm_eps, + dtype=config.torch_dtype, + ) + self.engram_hash_provider = EngramHashProvider(engram_config) + self.engram_layer_ids = engram_config.layer_ids + # Store engram config and per-layer vocab sizes on the pretrained config so + # DeepseekV4DecoderLayer can read them directly from model_config without extra params. + config.engram_config = engram_config + config.engram_vocab_sizes_by_layer = { + layer_id: [ + x + for y in self.engram_hash_provider.vocab_size_across_layers[layer_id] + for x in y + ] + for layer_id in engram_config.layer_ids + } + + self.layers = nn.ModuleList( + [ + DeepseekV4DecoderLayer( + model_config, + layer_idx, + self.aux_stream_dict, + mapping_with_cp=mapping_with_cp, + ) + for layer_idx in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + + def __pp_init__(self): + self.epilogue.append(self.hc_head) + super().__pp_init__() + + def forward( + self, + attn_metadata: DeepseekV4TrtllmAttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + **kwargs, + ) -> torch.Tensor: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" + ) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + # ----------------------------------------------------------------- + # Engram pre-computation (overlapped with main-stream layer forward) + # + # Hash provider internally applies CompressedTokenizer to input_ids, + # so no separate engram tokenizer input is needed. + # + # All precompute() calls are dispatched onto a dedicated engram CUDA + # stream so they run concurrently with the main stream processing the + # earlier transformer layers. A per-layer Event is recorded after + # each precompute; the main stream waits on the event just before it + # needs the result for that specific layer. + # ----------------------------------------------------------------- + engram_embeddings_cache: Optional[Dict] = None + engram_events: Dict[int, torch.cuda.Event] = {} + if self.use_engram and self.engram_hash_provider is not None and input_ids is not None: + hash_cache = self.engram_hash_provider.compute_hashes( + input_ids.view(-1), + seq_lens=attn_metadata.seq_lens_cuda, + ) + engram_embeddings_cache = {} + for layer_id in self.engram_layer_ids: + engram_mod = self.layers[layer_id].engram + if engram_mod is not None: + # precompute() dispatches onto the engram stream internally and + # records sync_event; the main stream will wait on it before use. + engram_embeddings_cache[layer_id] = engram_mod.precompute( + hash_cache[layer_id], dtype=inputs_embeds.dtype + ) + engram_events[layer_id] = engram_mod.sync_event + + mapping = self.model_config.mapping + if mapping.has_pp() and not mapping.is_first_pp_rank(): + hidden_states = torch.empty( + inputs_embeds.shape[0], + self.hc_mult, + inputs_embeds.shape[1], + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + else: + hidden_states = inputs_embeds + hidden_states = hidden_states.unsqueeze(1).repeat(1, self.hc_mult, 1) + + # ``hc_state`` carries the mHC pipeline state across layers. Layer 0 + # receives the initial residual tensor and bootstraps via + # hc_attn.pre_mapping ("hc_pre"). Every later layer receives an + # ``HCState``. In fused mode the state is "deferred" (prior layer's + # hc_ffn.post_mapping is folded into this layer's hc_attn.fused_hc); + # in unfused mode the state is "resolved" (residual already + # post-mapped). After the last layer, a deferred state is closed with + # a standalone hc_post; a resolved state feeds hc_head directly. + hc_state = hidden_states + + for idx, decoder_layer in enumerate(self.layers[: self.num_hidden_layers]): + engram_embeddings = None + if engram_embeddings_cache is not None and idx in engram_embeddings_cache: + # Sync: ensure the engram stream has finished precompute for this layer + # before the main stream reads the result. + engram_events[idx].wait(torch.cuda.current_stream()) + engram_embeddings = engram_embeddings_cache[idx] + + hc_state = decoder_layer( + position_ids=position_ids, + hc_state=hc_state, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + input_ids=input_ids, + engram_embeddings=engram_embeddings, + ) + + hidden_states = hc_state.residual.flatten(1) + + return hidden_states + + +@register_auto_model("DeepseekV4ForCausalLM") +class DeepseekV4ForCausalLM(SpecDecOneEngineForCausalLM[DeepseekV4Model, PretrainedConfig]): + @classmethod + def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: + return {"kv_cache_config": {"tokens_per_block": 128}} + + def __init__(self, model_config: ModelConfig[PretrainedConfig]): + self.mapping_with_cp = None + # Note: Currently the usage of mapping is all over the place making its usage brittle + # in this file. As a temporary WAR, we hold on to an original copy of mapping when CP + # is in action. This shall be passed on to attention which is the only layer that's + # affected by CP. For other layers, CP ranks are repurposed to TP. This shall be undone + # at the end of __init__. + if model_config.mapping.has_cp_helix(): + print( + "[DeepseekV4ForCausalLM::__init__] Repurposing KVP ranks to TP while keeping other details the same." + ) + self.mapping_with_cp = copy.deepcopy(model_config.mapping) + # Repurpose KVP ranks to TP while keeping other details the same. + model_config._frozen = False + model_config.mapping = model_config.mapping.repurpose_helix_cp_to_tp() + model_config._frozen = True + + # Rename some keys of quant_config_dict to support legacy checkpoints + if model_config.quant_config_dict is not None: + model_config = copy.deepcopy(model_config) + quant_config_dict = {} + for key, val in model_config.quant_config_dict.items(): + key_split = key.split(".") + if key_split[-1] == "fused_a": + key = ".".join(key_split[:-1] + ["kv_a_proj_with_mqa"]) + quant_config_dict[key] = val + model_config._frozen = False + model_config.quant_config_dict = quant_config_dict + model_config._frozen = True + + super().__init__( + model=DeepseekV4Model(model_config, mapping_with_cp=self.mapping_with_cp), + model_config=model_config, + ) + self.logits_processor = DeepseekV4LogitsProcessor( + model_config, self.model.hc_head, self.model.norm + ) + + # Exclude Engram weights from quantization. Engram embedding tables + # and small linear projections are not suited for NVFP4/FP8 quant. + if getattr(self.config, "has_engram", False): + if model_config.quant_config.exclude_modules is None: + model_config.quant_config.exclude_modules = [] + model_config.quant_config.exclude_modules.append("*engram*") + + self.model_nextn = 0 + if ( + model_config.spec_config is not None + and model_config.spec_config.spec_dec_mode.is_mtp_one_model() + ): + self.model_nextn = model_config.spec_config.num_nextn_predict_layers + ckpt_nextn = self.config.num_nextn_predict_layers + self.num_hidden_layers = self.config.num_hidden_layers + assert ckpt_nextn > 0, "There is not MTP modules in the checkpoint." + if ckpt_nextn == 1 and not model_config.spec_config.use_mtp_vanilla: + pass + else: + # modify the QuantConfig to support duplicated mtp layers + if model_config.quant_config.exclude_modules is not None: + extend_exclude_modules = [] + for model_mtp_idx in range( + self.num_hidden_layers, self.num_hidden_layers + self.model_nextn + ): + ckpt_mtp_idx = ( + model_mtp_idx - self.num_hidden_layers + ) % ckpt_nextn + self.num_hidden_layers + model_prefix = f"model.layers.{model_mtp_idx}" + ckpt_prefix = f"model.layers.{ckpt_mtp_idx}" + for exclude_module in model_config.quant_config.exclude_modules: + if ckpt_prefix in exclude_module and model_prefix not in exclude_module: + extend_exclude_modules.append( + exclude_module.replace(ckpt_prefix, model_prefix) + ) + self.model_config.quant_config.exclude_modules.extend(extend_exclude_modules) + self.model.layers.extend(self.draft_model.mtp_layers) + + # Undo any manipulations done to mapping. + if self.mapping_with_cp is not None: + print("[DeepseekV4ForCausalLM::__init__] Restoring original mapping.") + model_config._frozen = False + model_config.mapping = self.mapping_with_cp + model_config._frozen = True + + def forward( + self, + attn_metadata: DeepseekV4TrtllmAttentionMetadata, + input_ids: torch.IntTensor = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + return_context_logits: bool = False, + **kwargs, + ) -> torch.Tensor: + return super().forward( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + return_context_logits=return_context_logits, + **kwargs, + ) + + def load_weights(self, weights: Dict): + weight_loader = DeepseekV4WeightLoader(self) + weight_loader.load_weights(weights) + + def post_load_weights(self): + layers = self.model.layers[: self.config.num_hidden_layers] + last_idx = self.config.num_hidden_layers - 1 + for idx, layer in enumerate(layers): + if idx == last_idx: + # The V4 logits path is HCHead -> model.norm -> lm_head, so + # the final decoder layer must not fold model.norm into MoE. + layer.next_layer_layernorm = None + layer.fusion_config.POST_MOE_FUSION = False + layer.defer_post_mapping = False + else: + next_layer = layers[idx + 1] + layer.next_layer_layernorm = next_layer.input_layernorm + # Defer this layer's hc_ffn.post_mapping into the next layer's + # hc_attn.fused_hc only if that next layer can actually absorb + # it: fused_hc enabled on both sides, and no engram at the + # next layer's entry (engram needs the materialized residual + # before pre_mapping runs). + layer.defer_post_mapping = ( + layer.enable_fused_hc + and next_layer.enable_fused_hc + and next_layer.engram is None + ) diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 0bd2f8400157..60fba53aa376 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1457,6 +1457,9 @@ def __init__( case "step3p7" | "step3p5": from .modeling_step3p7 import Step3p7MTP mtp_layer = Step3p7MTP + case "deepseek_v4": + from .modeling_deepseekv4 import DeepseekV4MTP + mtp_layer = DeepseekV4MTP case _: raise ValueError( f"Model type {model_type} not supported for MTP") @@ -1516,6 +1519,12 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], elif model_type == "qwen3_next": from .modeling_qwen3_next import Qwen3NextMTP mtp_layer = Qwen3NextMTP(model_config, layer_idx, aux_stream_dict) + elif model_type == "deepseek_v4": + from .modeling_deepseekv4 import DeepseekV4MTP + mtp_layer = DeepseekV4MTP(model_config, + layer_idx, + aux_stream_dict, + is_separate_draft_engine=True) else: raise ValueError( f"MTPDraftModel does not support model_type: {model_type}") diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 03dbf36d0bc0..078c35397bdd 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -300,11 +300,25 @@ def __pp_init__(self): skip_forward(module) num_hidden_layers = self.model_config.pretrained_config.num_hidden_layers - assert num_hidden_layers >= mapping.pp_size, f"{num_hidden_layers} layers are not enough for PP{mapping.pp_size}" + assert num_hidden_layers >= mapping.pp_size, ( + f"{num_hidden_layers} layers are not enough for PP{mapping.pp_size}" + ) + pp_layer_list = mapping.pp_layers(num_hidden_layers) + total_num_layers = num_hidden_layers + spec_config = getattr(self.model_config, "spec_config", None) + if spec_config is not None: + from ..speculative.utils import get_num_spec_layers + + num_spec_layers = get_num_spec_layers(spec_config) or 0 + total_num_layers += num_spec_layers + if num_spec_layers > 0 and mapping.is_last_pp_rank(): + pp_layer_list.extend( + range(total_num_layers - num_spec_layers, total_num_layers)) + if len(pp_layer_list) == 0: + pp_layer_list.append(0) has_pp_layer = len(pp_layer_list) > 0 - for layer_idx in range(num_hidden_layers): - layer = self.layers[layer_idx] + for layer_idx, layer in enumerate(self.layers[:num_hidden_layers]): is_last_layer = (layer_idx == num_hidden_layers - 1) if layer_idx not in pp_layer_list: # keep next layer's input_layernorm's weights for fusion diff --git a/tensorrt_llm/_torch/modules/engram/__init__.py b/tensorrt_llm/_torch/modules/engram/__init__.py new file mode 100644 index 000000000000..ff1d218a8e4a --- /dev/null +++ b/tensorrt_llm/_torch/modules/engram/__init__.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from tensorrt_llm._torch.modules.engram.engram import Engram, EngramConfig, EngramHashProvider + +__all__ = ["Engram", "EngramConfig", "EngramHashProvider"] diff --git a/tensorrt_llm/_torch/modules/engram/engram.py b/tensorrt_llm/_torch/modules/engram/engram.py new file mode 100644 index 000000000000..41ddc5b4b6ab --- /dev/null +++ b/tensorrt_llm/_torch/modules/engram/engram.py @@ -0,0 +1,846 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Engram module implementation for TensorRT-LLM PyTorch backend. + +The Engram module provides n-gram based hash embeddings that augment +transformer hidden states with local context information. + +All operations run on device: + - Hash IDs (GPU) → embedding → flatten → linear projections (GEMMs) + → normed keys + projected value + - GPU main stream forward: SDP gating + short conv +""" + +import math +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np +import torch +from tokenizers import Regex, normalizers +from torch import nn +from transformers import AutoTokenizer + +# Default Engram embedding vocabulary size per n-gram level. +# Derived from the DeepSeek-V3 tokenizer vocab size (129280) scaled by 5. +_DEFAULT_ENGRAM_VOCAB_SIZE = 129280 * 5 + + +@dataclass +class EngramConfig: + """Configuration for the Engram module. + + Attributes: + tokenizer_name_or_path: Path or name of the HuggingFace tokenizer. + engram_vocab_size: List of vocabulary sizes for each n-gram level (2-gram, 3-gram, etc.). + max_ngram_size: Maximum n-gram size to compute hashes for. + n_embed_per_ngram: Embedding dimension for each n-gram. + n_head_per_ngram: Number of attention heads per n-gram. + layer_ids: List of layer indices where Engram modules are applied. + pad_id: Token ID used for padding. + seed: Random seed for hash multiplier generation. + kernel_size: Kernel size for the short convolution. + hidden_size: Hidden dimension of the backbone model. + hc_mult: Hyper-connection multiplier (number of residual streams). + norm_eps: Epsilon for RMSNorm. + dtype: Data type for model parameters (e.g., torch.float32, torch.bfloat16). + """ + + tokenizer_name_or_path: str = "deepseek-ai/DeepSeek-V3" + engram_vocab_size: List[int] = field( + default_factory=lambda: [_DEFAULT_ENGRAM_VOCAB_SIZE, _DEFAULT_ENGRAM_VOCAB_SIZE] + ) + max_ngram_size: int = 3 + n_embed_per_ngram: int = 512 + n_head_per_ngram: int = 8 + layer_ids: List[int] = field(default_factory=lambda: [1, 15]) + pad_id: int = 2 + seed: int = 0 + kernel_size: int = 4 + hidden_size: int = 1024 + hc_mult: int = 4 + norm_eps: float = 1e-5 + dtype: Optional[torch.dtype] = None + + +class CompressedTokenizer: + """Tokenizer wrapper that normalizes and compresses tokens. + + This class builds a lookup table mapping original token IDs to + normalized/compressed token IDs, reducing vocabulary size for + more efficient n-gram hashing. + + If a pre-built ``lookup_table`` (torch.Tensor of shape [vocab_size]) + and ``num_new_token`` count are provided, the tokenizer download is + skipped entirely — useful to avoid cold-start network fetches. + """ + + def __init__( + self, + tokenizer_name_or_path: str, + lookup_table: Optional[torch.Tensor] = None, + num_new_token: Optional[int] = None, + ): + if lookup_table is not None and num_new_token is not None: + self.lookup_table = lookup_table + self.num_new_token = num_new_token + return + + self.tokenizer = AutoTokenizer.from_pretrained( + tokenizer_name_or_path, trust_remote_code=True + ) + + SENTINEL = "\ue000" + self.normalizer = normalizers.Sequence( + [ + normalizers.NFKC(), + normalizers.NFD(), + normalizers.StripAccents(), + normalizers.Lowercase(), + normalizers.Replace(Regex(r"[ \t\r\n]+"), " "), + normalizers.Replace(Regex(r"^ $"), SENTINEL), + normalizers.Strip(), + normalizers.Replace(SENTINEL, " "), + ] + ) + + self.lookup_table, self.num_new_token = self._build_lookup_table() + + def __len__(self) -> int: + return self.num_new_token + + def _build_lookup_table(self): + old2new = {} + key2new = {} + new_tokens = [] + + vocab_size = len(self.tokenizer) + for tid in range(vocab_size): + text = self.tokenizer.decode([tid], skip_special_tokens=False) + + if "\ufffd" in text: + key = self.tokenizer.convert_ids_to_tokens(tid) + else: + norm = self.normalizer.normalize_str(text) + key = norm if norm else text + + nid = key2new.get(key) + if nid is None: + nid = len(new_tokens) + key2new[key] = nid + new_tokens.append(key) + old2new[tid] = nid + + lookup = torch.tensor([old2new[tid] for tid in range(vocab_size)], dtype=torch.long) + + return lookup, len(new_tokens) + + def _compress(self, input_ids: torch.Tensor) -> torch.Tensor: + if not isinstance(input_ids, torch.Tensor): + input_ids = torch.tensor(input_ids, dtype=torch.long) + ids = input_ids.long() + # Preserve negative padding IDs while compressing valid token IDs. + vocab_size = len(self.lookup_table) + compressed = self.lookup_table[ids.clamp(0, vocab_size - 1)] + return torch.where(ids < 0, ids, compressed) + + def __call__(self, input_ids): + return self._compress(input_ids) + + +def _is_prime(n: int) -> bool: + """Deterministic Miller-Rabin primality test for n < 3.3e24.""" + if n < 2: + return False + if n < 4: + return True + if n % 2 == 0 or n % 3 == 0: + return False + # Write n-1 as 2^r * d + d, r = n - 1, 0 + while d % 2 == 0: + d //= 2 + r += 1 + # Witnesses sufficient for n < 3.3e24 + for a in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37): + if a >= n: + continue + x = pow(a, d, n) + if x == 1 or x == n - 1: + continue + for _ in range(r - 1): + x = pow(x, 2, n) + if x == n - 1: + break + else: + return False + return True + + +def _find_next_prime(start: int, seen_primes: set) -> int: + """Find the next prime number greater than start that is not in seen_primes.""" + candidate = start + 1 + while True: + if _is_prime(candidate) and candidate not in seen_primes: + return candidate + candidate += 1 + + +class NgramHashMapping: + """Computes n-gram hash indices for embedding lookup. + + This class generates per-layer, per-head hash mappings for n-grams, + using prime moduli to reduce hash collisions across heads. + """ + + def __init__( + self, + engram_vocab_size: List[int], + max_ngram_size: int, + n_embed_per_ngram: int, + n_head_per_ngram: int, + layer_ids: List[int], + tokenizer_name_or_path: str, + pad_id: int, + seed: int, + ): + self.vocab_size_per_ngram = engram_vocab_size + self.max_ngram_size = max_ngram_size + self.n_embed_per_ngram = n_embed_per_ngram + self.n_head_per_ngram = n_head_per_ngram + self.pad_id = pad_id + self.layer_ids = layer_ids + + self.compressed_tokenizer = CompressedTokenizer( + tokenizer_name_or_path=tokenizer_name_or_path + ) + self.tokenizer_vocab_size = len(self.compressed_tokenizer) + if self.pad_id is not None: + self.pad_id = int(self.compressed_tokenizer.lookup_table[self.pad_id].item()) + + max_long = torch.iinfo(torch.long).max + M_max = int(max_long // self.tokenizer_vocab_size) + half_bound = max(1, M_max // 2) + PRIME_1 = 10007 + + self.layer_multipliers: Dict[int, torch.Tensor] = {} + + for layer_id in self.layer_ids: + base_seed = int(seed + PRIME_1 * int(layer_id)) + # Use numpy RNG for deterministic reproducibility with existing + # checkpoints — torch.Generator uses a different algorithm and + # would produce incompatible multipliers for the same seed. + g = np.random.default_rng(base_seed) + r = g.integers(low=0, high=half_bound, size=(self.max_ngram_size,), dtype=np.int64) + multipliers = torch.tensor(r * 2 + 1, dtype=torch.long) + self.layer_multipliers[layer_id] = multipliers + + self.vocab_size_across_layers = self._calculate_vocab_size_across_layers() + + def _calculate_vocab_size_across_layers(self) -> Dict[int, List[List[int]]]: + seen_primes = set() + vocab_size_across_layers = {} + + for layer_id in self.layer_ids: + all_ngram_vocab_sizes = [] + for ngram in range(2, self.max_ngram_size + 1): + current_ngram_heads_sizes = [] + + vocab_size = self.vocab_size_per_ngram[ngram - 2] + num_head = self.n_head_per_ngram + current_prime_search_start = vocab_size - 1 + + for _ in range(num_head): + found_prime = _find_next_prime(current_prime_search_start, seen_primes) + seen_primes.add(found_prime) + current_ngram_heads_sizes.append(found_prime) + current_prime_search_start = found_prime + + all_ngram_vocab_sizes.append(current_ngram_heads_sizes) + vocab_size_across_layers[layer_id] = all_ngram_vocab_sizes + + return vocab_size_across_layers + + def _get_ngram_hashes( + self, + input_ids: torch.Tensor, + layer_id: int, + ) -> torch.Tensor: + x = input_ids.long() + (T,) = x.shape + + multipliers = self.layer_multipliers[layer_id] + + def shift_k(k: int) -> torch.Tensor: + if k == 0: + return x + return torch.nn.functional.pad(x, (k, 0), value=self.pad_id)[:T] + + base_shifts = [shift_k(k) for k in range(self.max_ngram_size)] + + all_hashes: List[torch.Tensor] = [] + + for n in range(2, self.max_ngram_size + 1): + n_gram_index = n - 2 + tokens = base_shifts[:n] + mix = tokens[0] * multipliers[0] + for k in range(1, n): + mix = torch.bitwise_xor(mix, tokens[k] * multipliers[k]) + head_vocab_sizes = self.vocab_size_across_layers[layer_id][n_gram_index] + + for j in range(self.n_head_per_ngram): + mod = int(head_vocab_sizes[j]) + all_hashes.append(mix % mod) + + return torch.stack(all_hashes, dim=1) + + def hash(self, input_ids) -> Dict[int, torch.Tensor]: + """Compute hash indices for all configured layers. + + Args: + input_ids: Token IDs of shape ``[T]``. + + Returns: + Dictionary mapping layer_id to hash indices of shape ``[T, num_heads]``. + """ + input_ids = self.compressed_tokenizer(input_ids) + hash_ids_for_all_layers = {} + for layer_id in self.layer_ids: + hash_ids_for_all_layers[layer_id] = self._get_ngram_hashes(input_ids, layer_id=layer_id) + return hash_ids_for_all_layers + + def hash_single_layer(self, input_ids, layer_id: int) -> torch.Tensor: + """Compute hash indices for a single layer. + + Args: + input_ids: Token IDs of shape ``[T]`` (already compressed or raw). + layer_id: The layer to compute hashes for. + + Returns: + Hash indices of shape ``[T, num_heads]``. + """ + if not isinstance(input_ids, torch.Tensor): + input_ids = torch.tensor(input_ids, dtype=torch.long) + return self._get_ngram_hashes(input_ids, layer_id=layer_id) + + +class EngramHashProvider: + """Computes and caches n-gram hash indices for all Engram layers. + + All hash computation runs on GPU using PyTorch ops. CPU-side + ``NgramHashMapping`` is used only at init to derive constants + (lookup table, multipliers, moduli) and then discarded. + + Usage: + # At model initialization + hash_provider = EngramHashProvider(config) + + # At each forward pass + hash_cache = hash_provider.compute_hashes(input_ids) + + # In each Engram layer + precomputed = engram_layer.precompute(hash_cache[layer_id]) + output = engram_layer(hidden_states, precomputed=precomputed) + """ + + def __init__(self, config: EngramConfig): + self.config = config + + # Use NgramHashMapping only to derive constants, then discard it. + hash_mapping = NgramHashMapping( + engram_vocab_size=config.engram_vocab_size, + max_ngram_size=config.max_ngram_size, + n_embed_per_ngram=config.n_embed_per_ngram, + n_head_per_ngram=config.n_head_per_ngram, + layer_ids=config.layer_ids, + tokenizer_name_or_path=config.tokenizer_name_or_path, + pad_id=config.pad_id, + seed=config.seed, + ) + + # Store vocab sizes directly (needed by Engram layer init). + self._vocab_size_across_layers = hash_mapping.vocab_size_across_layers + + # GPU tensors for on-device hash computation. + # Created on CPU and lazily moved to GPU on first use. + self._lookup_table = hash_mapping.compressed_tokenizer.lookup_table.clone() + self._pad_id = hash_mapping.pad_id + + self._multipliers: Dict[int, torch.Tensor] = {} + self._modules: Dict[int, List[torch.Tensor]] = {} + for layer_id in config.layer_ids: + self._multipliers[layer_id] = hash_mapping.layer_multipliers[layer_id].clone() + self._modules[layer_id] = [ + torch.tensor(ngram_head_sizes, dtype=torch.long) + for ngram_head_sizes in hash_mapping.vocab_size_across_layers[layer_id] + ] + + self._device: Optional[torch.device] = None + + # Cache for CUDA graph capture - stores last computed hashes. + # _cached_hashes points to the most-recently used entry. + # _cached_hashes_store keeps ALL per-shape entries alive so that + # CUDA graph tensor addresses remain valid after subsequent compute calls + # change the active shape (e.g. warmup batch_size=4 → 3 → 2 → 1). + self._cached_hashes: Optional[Dict[int, torch.Tensor]] = None + self._cached_hashes_store: Dict[tuple, Dict[int, torch.Tensor]] = {} + + def _ensure_on_device(self, device: torch.device): + """Move hash tensors to the specified device (lazy, once).""" + if self._device == device: + return + self._lookup_table = self._lookup_table.to(device) + for layer_id in list(self._multipliers.keys()): + self._multipliers[layer_id] = self._multipliers[layer_id].to(device) + self._modules[layer_id] = [m.to(device) for m in self._modules[layer_id]] + self._device = device + + def compute_hashes( + self, + input_ids: torch.Tensor, + seq_lens: Optional[torch.Tensor] = None, + ) -> Dict[int, torch.Tensor]: + """Compute hash indices on GPU using PyTorch ops. + + Runs the compressed-tokenizer lookup, n-gram mixing, and modular + hashing entirely on the device where ``input_ids`` resides. + + During CUDA graph capture, returns cached hashes from the warmup pass. + + Args: + input_ids: Flattened token IDs of shape ``[T]`` on GPU. + When the tensor contains multiple packed sequences, pass + ``seq_lens`` so that n-gram hashing does not cross + sequence boundaries. + seq_lens: Optional 1-D tensor of per-sequence lengths. When + provided, shifted n-gram windows that would cross a + sequence boundary are replaced with ``pad_id``. + + Returns: + Dictionary mapping layer_id to hash indices as torch.Tensor + of shape ``[T, num_heads]`` on the same device as input_ids. + """ + if torch.cuda.is_current_stream_capturing(): + if self._cached_hashes is not None: + return self._cached_hashes + raise RuntimeError( + "EngramHashProvider.compute_hashes() called during CUDA " + "graph capture but no cached hashes available. " + "Ensure warmup runs before capture." + ) + + device = input_ids.device + self._ensure_on_device(device) + + # 1. Compress token ids via lookup table (1-D) + ids = input_ids.long().view(-1) + vocab_size = self._lookup_table.shape[0] + compressed = self._lookup_table[ids.clamp(0, vocab_size - 1)] + + (T,) = compressed.shape + + # 2. Pre-compute shifted versions (left-padded with pad_id) + shifts = [compressed] + for k in range(1, self.config.max_ngram_size): + padded = torch.nn.functional.pad(compressed, (k, 0), value=self._pad_id) + shifts.append(padded[:T]) + + # 2b. When input_ids is a packed/flattened tensor containing multiple + # sequences, mask out shifted positions that cross sequence boundaries + # so n-gram hashes stay within each sequence. + if seq_lens is not None and seq_lens.numel() > 1: + cum_lens = torch.cumsum(seq_lens, dim=0) + seq_starts = torch.cat( + [ + torch.zeros(1, device=device, dtype=cum_lens.dtype), + cum_lens[:-1], + ] + ) + positions = torch.arange(T, device=device) + # Map each position to its owning sequence + seq_idx = torch.searchsorted(cum_lens, positions, right=True) + seq_idx = seq_idx.clamp(max=seq_lens.numel() - 1) + pos_in_seq = positions - seq_starts[seq_idx] + + for k in range(1, self.config.max_ngram_size): + boundary_mask = pos_in_seq < k # [T] + shifts[k][boundary_mask] = self._pad_id + + # 3. Compute hashes per layer + result: Dict[int, torch.Tensor] = {} + for layer_id in self.config.layer_ids: + multipliers = self._multipliers[layer_id] + all_hashes: List[torch.Tensor] = [] + + for n in range(2, self.config.max_ngram_size + 1): + n_gram_index = n - 2 + mix = shifts[0] * multipliers[0] + for k in range(1, n): + mix = torch.bitwise_xor(mix, shifts[k] * multipliers[k]) + + moduli = self._modules[layer_id][n_gram_index] + for j in range(self.config.n_head_per_ngram): + head_hash = mix % int(moduli[j].item()) + all_hashes.append(head_hash) + + result[layer_id] = torch.stack(all_hashes, dim=1) + + # Cache for CUDA graph capture. + # Use a per-shape store so tensors captured by CUDA graphs are never + # garbage-collected when a subsequent call uses a different batch size. + # When the same shape is seen again, update the existing tensors + # in-place so CUDA graphs can read the freshly computed values from + # the same memory addresses used at capture time. + shape_key = tuple(tuple(v.shape) for v in result.values()) + if shape_key in self._cached_hashes_store: + cached = self._cached_hashes_store[shape_key] + for layer_id, hashes in result.items(): + cached[layer_id].copy_(hashes) + self._cached_hashes = cached + else: + self._cached_hashes_store[shape_key] = result + self._cached_hashes = result + return self._cached_hashes + + @property + def layer_ids(self) -> List[int]: + """Return the list of layer IDs that have Engram modules.""" + return self.config.layer_ids + + @property + def vocab_size_across_layers(self) -> Dict[int, List[List[int]]]: + """Return the vocabulary sizes for each layer and head.""" + return self._vocab_size_across_layers + + +class MultiHeadEmbedding(nn.Module): + """Multi-head embedding layer with per-head offset handling. + + Each head has its own embedding space, indexed by applying offsets + to the input indices before a single embedding lookup. + """ + + def __init__(self, list_of_N: List[int], D: int, dtype: Optional[torch.dtype] = None): + super().__init__() + self.num_heads = len(list_of_N) + self.embedding_dim = D + + offsets = [0] + for n in list_of_N[:-1]: + offsets.append(offsets[-1] + n) + + self.register_buffer("offsets", torch.tensor(offsets, dtype=torch.long)) + + total_N = sum(list_of_N) + self.embedding = nn.Embedding(num_embeddings=total_N, embedding_dim=D, dtype=dtype) + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + """Look up embeddings for multi-head input indices. + + Args: + input_ids: Indices of shape ``[T, num_heads]``. + + Returns: + Embeddings of shape ``[T, num_heads, D]``. + """ + shifted_input_ids = input_ids + self.offsets + # Clamp to valid embedding range to prevent CUDA kernel OOB assertion. + # head_hash values should be in [0, prime-1] by construction, but this + # acts as a defensive guard in case of unexpected inputs. + shifted_input_ids = shifted_input_ids.clamp(0, self.embedding.num_embeddings - 1) + output = self.embedding(shifted_input_ids) + return output + + +class ShortConv(nn.Module): + """Short depthwise convolution with RMSNorm and SiLU activation. + + Applies a causal depthwise convolution across the sequence dimension + with per-hyper-connection-stream normalization. + + Note: During token-by-token autoregressive generation (seq_len=1), + the Conv1d sees only the current token plus zero padding — no ring + buffer or conv state carries over from previous steps. This means + the short-range context capture is limited to prefill. A proper + conv state cache for generation is left as a future enhancement. + """ + + def __init__( + self, + hidden_size: int, + kernel_size: int = 4, + dilation: int = 1, + norm_eps: float = 1e-5, + hc_mult: int = 4, + activation: bool = True, + dtype: Optional[torch.dtype] = None, + ): + super().__init__() + self.hc_mult = hc_mult + self.activation = activation + + self.hidden_size = hidden_size + self.norm_eps = norm_eps + + total_channels = hidden_size * hc_mult + self.conv = nn.Conv1d( + in_channels=total_channels, + out_channels=total_channels, + kernel_size=kernel_size, + groups=total_channels, + bias=False, + padding=(kernel_size - 1) * dilation, + dilation=dilation, + dtype=dtype, + ) + + # Stacked RMSNorm weight [hc_mult, D] — single vectorised norm + # instead of hc_mult separate kernel launches. + self.norm_weight = nn.Parameter(torch.ones(hc_mult, hidden_size, dtype=dtype)) + + if self.activation: + self.act_fn = nn.SiLU() + + def load_weights(self, weights): + """Load weights, handling legacy ``norms.{i}.weight`` checkpoint keys. + + Legacy checkpoints store per-stream RMSNorm weights as + ``norms.0.weight``, ``norms.1.weight``, … This method stacks + them into the single ``norm_weight`` parameter. + """ + w = weights[0] if isinstance(weights, list) else weights + # Check for legacy per-stream norm keys + legacy_keys = [f"norms.{i}.weight" for i in range(self.hc_mult)] + if legacy_keys[0] in w: + stacked = torch.stack([w[k][:] for k in legacy_keys], dim=0) + self.norm_weight.data.copy_(stacked) + # Load remaining keys (e.g. conv.weight) via the generic path + for n, p in self.named_parameters(): + if n == "norm_weight": + continue + if n in w: + p.data.copy_(w[n][:]) + else: + for n, p in self.named_parameters(): + if n in w: + p.data.copy_(w[n][:]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply short convolution with normalization. + + Args: + x: Input tensor of shape ``[T, HC_MULT, D]``. + + Returns: + Output tensor of shape ``[T, HC_MULT, D]``. + """ + T, G, C = x.shape + + assert G == self.hc_mult, f"Input groups {G} != hc_mult {self.hc_mult}" + + rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.norm_eps).rsqrt() + x_normed = x * rms * self.norm_weight # [hc_mult, D] broadcasts over [T, G, D] + + x_norm = x_normed.reshape(T, G * C) + # Conv1d expects [N, C, L]; use N=1 so that L=T is the sequence dim. + x_nct = x_norm.unsqueeze(0).transpose(1, 2) # [1, G*C, T] + y_nct = self.conv(x_nct) + # Truncate to maintain causal masking + y_nct = y_nct[..., :T] + + if self.activation: + y_nct = self.act_fn(y_nct) + y = y_nct.transpose(1, 2).squeeze(0).view(T, G, C).contiguous() + + return y + + +class Engram(nn.Module): + """Engram module for n-gram based context augmentation. + + The Engram module computes n-gram hash embeddings from input tokens + and uses gated attention to augment the hidden states of the model. + + All operations run on device. ``precompute()`` is called on a separate + CUDA stream to overlap with other layers; the main stream performs + only SDP gating + short conv. + + Known limitations: + - **No TP sharding**: The embedding table is replicated on every rank. + For large-scale deployments with TP > 1 this means redundant memory. + - **No conv state for generation**: See ``ShortConv`` docstring. + """ + + def __init__( + self, + layer_id: int, + config: EngramConfig, + vocab_sizes_flat: Optional[List[int]] = None, + stream: Optional[torch.cuda.Stream] = None, + ): + super().__init__() + self.layer_id = layer_id + self.config = config + self.stream = stream + self.sync_event: Optional[torch.cuda.Event] = ( + torch.cuda.Event() if stream is not None else None + ) + + # If vocab_sizes not provided, compute them (for standalone usage) + if vocab_sizes_flat is None: + hash_mapping = NgramHashMapping( + engram_vocab_size=config.engram_vocab_size, + max_ngram_size=config.max_ngram_size, + n_embed_per_ngram=config.n_embed_per_ngram, + n_head_per_ngram=config.n_head_per_ngram, + layer_ids=config.layer_ids, + tokenizer_name_or_path=config.tokenizer_name_or_path, + pad_id=config.pad_id, + seed=config.seed, + ) + vocab_sizes_flat = [ + x for y in hash_mapping.vocab_size_across_layers[layer_id] for x in y + ] + self.hash_mapping = hash_mapping + else: + self.hash_mapping = None + + embed_dim_per_head = config.n_embed_per_ngram // config.n_head_per_ngram + + dtype = config.dtype + + self.multi_head_embedding = MultiHeadEmbedding( + list_of_N=vocab_sizes_flat, + D=embed_dim_per_head, + dtype=dtype, + ) + + self.short_conv = ShortConv( + hidden_size=config.hidden_size, + kernel_size=config.kernel_size, + dilation=config.max_ngram_size, + hc_mult=config.hc_mult, + norm_eps=config.norm_eps, + dtype=dtype, + ) + + engram_hidden_size = (config.max_ngram_size - 1) * config.n_embed_per_ngram + hc = config.hc_mult + D = config.hidden_size + + # Fused projection: one GEMM produces value (1 head) + all keys (hc_mult heads). + # Output layout: [value (D) | key_0 (D) | key_1 (D) | ... | key_{hc-1} (D)] + self.kv_proj = nn.Linear(engram_hidden_size, (1 + hc) * D, bias=False, dtype=dtype) + + # Per-HC-stream RMSNorm weights for keys and queries, stored as + # stacked parameters so we can apply a single vectorised norm + # instead of hc_mult separate kernel launches. + self.key_norm_weight = nn.Parameter(torch.ones(hc, D, dtype=dtype)) + self.query_norm_weight = nn.Parameter(torch.ones(hc, D, dtype=dtype)) + self.norm_eps = config.norm_eps + + # CUDA graph capture cache + self._cached_embeddings: Optional[torch.Tensor] = None + + def precompute( + self, + hash_indices: torch.Tensor, + dtype: Optional[torch.dtype] = None, + ) -> torch.Tensor: + """Pre-compute embeddings from hash indices. + + When a stream was supplied at construction time, the work is + dispatched onto that stream (overlapping with the main stream) and + ``sync_event`` is recorded so the caller can synchronize before + consuming the result. When no stream is set the computation runs + on the current stream synchronously and ``sync_event`` is None. + + Args: + hash_indices: Hash indices of shape ``[T, num_heads]`` on GPU. + dtype: Cast embeddings to this dtype after lookup. + If None, uses the embedding's native dtype. + + Returns: + Embedding tensor of shape ``[T, num_heads * embed_dim_per_head]``. + """ + if self.stream is not None: + # Fork: let the engram stream wait for any pending main-stream work + # (e.g. hash computation) before we start the embedding lookup. + self.stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self.stream): + embeddings = self.multi_head_embedding(hash_indices) + embeddings = embeddings.flatten(start_dim=-2) + if dtype is not None: + embeddings = embeddings.to(dtype) + self.sync_event.record() + else: + embeddings = self.multi_head_embedding(hash_indices) + embeddings = embeddings.flatten(start_dim=-2) + if dtype is not None: + embeddings = embeddings.to(dtype) + + return embeddings + + def forward( + self, + hidden_states: torch.Tensor, + embeddings: torch.Tensor, + conv_state: Optional[torch.Tensor] = None, + use_cache: bool = False, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """Forward pass of the Engram module. + + Args: + hidden_states: Hidden states of shape ``[T, HC_MULT, D]``. + embeddings: Pre-computed embeddings from ``precompute()``. + conv_state: Optional conv state from a previous decode step + (shape ``[1, C_total, conv_state_size]``). + use_cache: When ``True``, return ``(output, new_conv_state)``. + + Returns: + Output tensor of shape ``[T, HC_MULT, D]`` to be added as + residual, and optionally the updated conv state. + """ + # Fused key/value projection: single GEMM replaces hc_mult + 1 separate GEMMs. + # kv_proj output: [T, (1 + HC) * D] + D = self.config.hidden_size + HC = self.config.hc_mult + kv = self.kv_proj(embeddings) + value_raw, keys = kv.split([D, HC * D], dim=-1) + keys = keys.view(*keys.shape[:-1], HC, D) # [T, HC, D] + + # Vectorised RMSNorm for keys and queries (no per-HC kernel launches). + # rms_norm(x, w) = x / rms(x) * w where rms(x) = sqrt(mean(x^2) + eps) + key_rms = keys.pow(2).mean(dim=-1, keepdim=True).add(self.norm_eps).rsqrt() + normed_keys = keys * key_rms * self.key_norm_weight # [HC, D] broadcasts + + queries = hidden_states # [T, HC, D] + query_rms = queries.pow(2).mean(dim=-1, keepdim=True).add(self.norm_eps).rsqrt() + normed_queries = queries * query_rms * self.query_norm_weight + + # Gating: per-HC dot product between normed keys and queries. + gates = (normed_keys * normed_queries).sum(dim=-1) / math.sqrt(D) # [T, HC] + gates = gates.abs().clamp_min(1e-6).sqrt() * gates.sign() + gates = gates.sigmoid().unsqueeze(-1) # [T, HC, 1] + + value = gates * value_raw.unsqueeze(-2) # [T, HC, D] + + if use_cache: + conv_out, new_conv_state = self.short_conv(value, conv_state=conv_state, use_cache=True) + return value + conv_out, new_conv_state + + output = value + self.short_conv(value) + return output diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 5fdf56e3e121..c9cf887562c7 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -46,7 +46,7 @@ from ..distributed import Distributed from ..distributed.communicator import init_pp_comm from ..expert_statistic import ExpertStatistic -from ..memory_buffer_utils import with_shared_pool +from ..memory_buffer_utils import clear_memory_buffers, with_shared_pool from ..metadata import KVCacheParams from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader from ..models.modeling_multimodal_utils import filter_mm_token_from_input_ids @@ -704,6 +704,8 @@ def __init__( self.kv_cache_dtype_byte_size = self.get_kv_cache_dtype_byte_size() + self._prepare_inputs_event: Optional[torch.cuda.Event] = None + def register_forward_pass_callable(self, callable: Callable): self.forward_pass_callable = callable @@ -1123,8 +1125,36 @@ def _general_warmup_impl( logger.warning( f"OOM during general warmup with {num_tokens} tokens, " f"{num_gen_tokens} generation tokens. Skipping.") + # If the OOM aborted the forward between dispatch() and + # combine(), the MoE A2A state machines are stuck in + # ``dispatched`` and the next warmup will hit + # ``dispatch called twice``. Reset them before retrying a + # smaller shape. + self._reset_moe_alltoall_state() torch.cuda.empty_cache() + def _reset_moe_alltoall_state(self) -> None: + """Reset all MoE all-to-all state machines reachable from ``self.model``. + + Each MoE backend keeps a small dispatch/combine phase state per layer + (``MoeAlltoAll`` or ``NVLinkOneSided``). A forward that calls + ``dispatch`` but raises before reaching ``combine`` (e.g., a warmup + OOM mid-MoE) leaves that state in ``dispatched``, which fails the + invariant on the next ``dispatch`` call. This helper walks the model + and resets any A2A state found, so subsequent forwards start clean. + """ + for module in self.model.modules(): + for attr_name in ("moe_a2a", "comm"): + obj = getattr(module, attr_name, None) + reset = getattr(obj, "reset_state", None) + if callable(reset): + try: + reset() + except Exception as e: # noqa: BLE001 + logger.warning( + f"Failed to reset MoE A2A state on {type(module).__name__}.{attr_name}: {e}" + ) + def _run_attention_warmup(self, resource_manager: ResourceManager, can_run_general_warmup: bool = True) -> None: @@ -1227,6 +1257,14 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): ) AutoTuner.get().print_profiling_cache() + # Clear workspace buffers allocated during the autotuner forward pass. + # The autotuner runs a context-only forward with max_num_tokens, which + # causes the global Buffers pool to cache large MoE/GEMM workspaces. + # If not cleared, these inflate the memory baseline seen by the KV cache + # profiler, reducing memory available for activations during inference. + clear_memory_buffers() + torch.cuda.empty_cache() + def _compute_dynamic_draft_len_mapping(self) -> Optional[Dict[int, int]]: """Compute graph_bs → draft_len mapping for dynamic draft length feature. @@ -5038,6 +5076,8 @@ def forward(self, new_tensors_device, cache_indirection_buffer, num_accepted_tokens_device, req_id_to_old_request, resource_manager, can_run_graph) + self._prepare_inputs_event = torch.cuda.Event() + self._prepare_inputs_event.record() with with_shared_pool(self.cuda_graph_runner.get_graph_pool()): if not can_run_graph: @@ -5510,3 +5550,11 @@ def _execute_logit_post_processors(self, lp(request.py_request_id, logits_row, token_ids, None, None) logits_tensor[idx] = logits_row.view(-1) + + def wait_for_input_copy(self): + """ + Wait for input preparation and H2D copy of previous iteration before modifying host input, + otherwise the input of previous iteration will be overwritten. + """ + if self._prepare_inputs_event is not None: + self._prepare_inputs_event.synchronize() diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index e291fc905144..b72b94c94ed5 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -183,6 +183,16 @@ def get_spec_metadata(spec_config, return None +def get_mtp_hidden_size(model_config) -> int: + pretrained_config = getattr(model_config, "pretrained_config", model_config) + hidden_size = getattr(pretrained_config, "hidden_size", None) + if hidden_size is None: + hidden_size = getattr(model_config, "hidden_size") + if getattr(pretrained_config, "model_type", None) == "deepseek_v4": + return hidden_size * getattr(pretrained_config, "hc_mult", 1) + return hidden_size + + def get_spec_resource_manager(model_engine, draft_model_engine=None): spec_config = model_engine.spec_config if spec_config is None: @@ -205,7 +215,7 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): return Eagle3ResourceManager( spec_config, model_config.torch_dtype, - model_config.hidden_size, + get_mtp_hidden_size(model_config), max_num_requests, max_seq_len, max_num_tokens, @@ -222,7 +232,7 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): return MTPHiddenStatesManager( spec_config, model_config.torch_dtype, - model_config.hidden_size, + get_mtp_hidden_size(model_config), max_num_requests, sa_manager=sa_manager, ) diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index 3c5ebe993fdf..beb6c1c9f075 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -25,6 +25,7 @@ 'MoeBalancer', 'MoeOutputMemset', 'MoeFc2Alpha', + 'EngramPrecompute', ] AuxStreamType = Enum( 'AuxStreamType', diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index 0781efe2f1e2..9abe26037331 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -178,6 +178,7 @@ def torch_version(): int64=np.int64, int32=np.int32, int8=np.int8, + uint8=np.uint8, bool=np.bool_, bfloat16=np_bfloat16, fp8=np_float8, @@ -197,6 +198,7 @@ def str_dtype_to_np(dtype): int64=torch.int64, int32=torch.int32, int8=torch.int8, + uint8=torch.uint8, bool=torch.bool, fp8=torch.float8_e4m3fn, ) @@ -215,6 +217,7 @@ def str_dtype_to_torch(dtype): int64=DataType.INT64, int32=DataType.INT32, int8=DataType.INT8, + uint8=DataType.UINT8, bool=DataType.BOOL, fp8=DataType.FP8, ) diff --git a/tensorrt_llm/bench/build/dataclasses.py b/tensorrt_llm/bench/build/dataclasses.py index 0b27ca4cbf9b..37d63062f23a 100755 --- a/tensorrt_llm/bench/build/dataclasses.py +++ b/tensorrt_llm/bench/build/dataclasses.py @@ -207,6 +207,11 @@ def get_param_count_and_checkpoint_size(cls, model_hf_name, hf_model_path): if model_hf_name == "EleutherAI/gpt-j-6b": # GPT-J repo doesn't use safetensor format. param_count = 6053381344 checkpoint_size_in_gb = param_count * 2 / (1024**3) + # TODO: Remove this once HF supports DeepSeek-V4 + elif model_hf_name == "deepseek-ai/DeepSeek-V4": + param_count = 284347051735 + # DeepSeek-V4 ships as an FP8 checkpoint (~1 byte/param). + checkpoint_size_in_gb = param_count / (1024**3) else: model_name_or_path = hf_model_path or model_hf_name metadata = get_safetensors_metadata(model_name_or_path) diff --git a/tensorrt_llm/bench/dataclasses/configuration.py b/tensorrt_llm/bench/dataclasses/configuration.py index d88bd2b722fe..7025e372c986 100755 --- a/tensorrt_llm/bench/dataclasses/configuration.py +++ b/tensorrt_llm/bench/dataclasses/configuration.py @@ -82,7 +82,8 @@ def get_llm_args(self) -> Dict: if self.backend in backend_config_map: llm_args.update(backend_config_map[self.backend]()) - kv_cache_config = self.settings_config.get_kvcache_config().__dict__ + kv_cache_config = self.settings_config.get_kvcache_config().model_dump( + exclude_unset=True) backend_cache_config = llm_args.pop("kv_cache_config", {}) llm_args["kv_cache_config"] = backend_cache_config | kv_cache_config diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index ead3474a36f5..af5ee1759206 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -987,7 +987,7 @@ def serve( f"Cannot auto-detect reasoning parser for model '{model}'. " f"Supported model types for auto-detection: qwen3, qwen3_moe, " f"qwen3_5, qwen3_5_moe, qwen3_next, deepseek_v3 (R1 only), " - f"deepseek_v32 (R1 only), nemotron_h, gemma4, " + f"deepseek_v32 (R1 only), deepseek_v4, nemotron_h, gemma4, " f"kimi_k2, kimi_k25. " f"Please specify a parser explicitly: " f"{list(ReasoningParserFactory.keys())}", diff --git a/tensorrt_llm/inputs/utils.py b/tensorrt_llm/inputs/utils.py index eaa98cad44ef..fcc53dbc73f1 100644 --- a/tensorrt_llm/inputs/utils.py +++ b/tensorrt_llm/inputs/utils.py @@ -35,6 +35,7 @@ MultimodalPlaceholderPlacement) from tensorrt_llm.llmapi.llm_utils import ModelLoader from tensorrt_llm.tokenizer import TokenizerBase, TransformersTokenizer +from tensorrt_llm.tokenizer.deepseek_v4 import DeepseekV4Tokenizer from tensorrt_llm.tokenizer.deepseek_v32 import DeepseekV32Tokenizer logger = logging.get_logger(__name__) @@ -652,8 +653,8 @@ def apply_chat_template( - STRING: keeps flattened text with pre-inserted placeholders """ - # Handle DeepSeek V32 tokenizer with custom chat template - if isinstance(tokenizer, DeepseekV32Tokenizer): + # Handle DeepSeek tokenizers with custom chat templates. + if isinstance(tokenizer, (DeepseekV32Tokenizer, DeepseekV4Tokenizer)): prompt = tokenizer.apply_chat_template( messages=conversation, tools=tools, diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 78251d1ae978..96032c23c227 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -91,4 +91,5 @@ 'PrometheusMetricsConfig', 'ThinkingBudgetLogitsProcessor', 'add_thinking_budget_logits_processor', + 'DeepSeekV4SparseAttentionConfig', ] diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 12d749fcb39f..4290b82bfb76 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -764,15 +764,20 @@ class DeepSeekSparseAttentionConfig(SeqLenAwareSparseAttentionConfig): description= "Whether to enable Guess-Verify-Refine (GVR) Top-K for the DSA decode " "indexer. GVR reuses previous-step Top-K indices as hints to reduce " - "threshold search iterations. Currently supported for index_topk=2048 " - "on Blackwell (SM100+) and falls back to the production insertion/radix " - "Top-K path when prerequisites are not met.") + "threshold search iterations. Currently supported for index_topk ∈ " + "{512, 1024, 2048} on Blackwell (SM100+), with compress_ratio ∈ {1, 4} " + "(DSv3.2 + DSv4 indexers). Falls back to the production insertion/" + "radix Top-K path when prerequisites are not met.") indexer_k_dtype: Literal["fp8", "fp4"] = Field( default="fp8", description= - "Data type used for the indexer K cache. `fp4` requires Blackwell+ " - "(SM>=100) and index_head_dim=128, it can halve the indexer K cache " - "per-token footprint from 132 B to 68 B.") + "Data type used for the indexer K cache. `fp8` stores one FP8 E4M3 " + "byte per element with a per-128 float32 scale; `fp4` packs two FP4 " + "E2M1 codes per byte with a per-32 UE8M0 exponent, halving the " + "per-token indexer K footprint (132 B to 68 B at index_head_dim=128). " + "`fp4` requires Blackwell+ (SM>=100) at runtime and " + "index_head_dim=128.", + ) @model_validator(mode="after") def _validate_indexer_k_dtype(self): @@ -866,7 +871,6 @@ def _value(name: str, default=None): class DeepSeekV4SparseAttentionConfig(DeepSeekSparseAttentionConfig): """Configuration for DeepSeek-V4 Sparse Attention.""" - algorithm: Literal["deepseek_v4"] = "deepseek_v4" index_head_dim: Optional[int] = Field( default=128, @@ -875,6 +879,7 @@ class DeepSeekV4SparseAttentionConfig(DeepSeekSparseAttentionConfig): default=False, description= "Whether to skip the MQA and Top-K in the indexer for short sequences.") + compress_ratios: List[int] = Field( default_factory=lambda: [1, 1, 4, 128, 4, 128, 4], description="The compress ratios of each layer. DeepSeek-V4 uses 0 " @@ -910,6 +915,61 @@ def needs_separate_short_long_cuda_graphs(self) -> bool: # DeepSeek-V4 does not support short/long CUDA graph separation. return False + def to_sparse_params(self, **kwargs): + from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import \ + DeepSeekV4Params + + pretrained_config = kwargs.get("pretrained_config", None) + + def _value(name: str, default=None): + value = getattr(self, name) + if value is not None: + return value + if pretrained_config is not None: + return getattr(pretrained_config, name, default) + return default + + return DeepSeekV4Params( + index_n_heads=_value("index_n_heads"), + index_head_dim=_value("index_head_dim"), + index_topk=_value("index_topk"), + indexer_max_chunk_size=self.indexer_max_chunk_size, + skip_indexer_for_short_seqs=self.skip_indexer_for_short_seqs, + use_cute_dsl_topk=self.use_cute_dsl_topk, + use_cute_dsl_paged_mqa_logits=self.use_cute_dsl_paged_mqa_logits, + q_split_threshold=self.q_split_threshold, + indexer_rope_interleave=self.indexer_rope_interleave, + enable_heuristic_topk=self.enable_heuristic_topk, + indexer_k_dtype=self.indexer_k_dtype, + compress_ratios=self.compress_ratios, + window_size=self.window_size, + ) + + def to_sparse_metadata_params(self, **kwargs): + from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import \ + DeepSeekV4MetadataParams + + pretrained_config = kwargs.get("pretrained_config", None) + + def _value(name: str, default=None): + value = getattr(self, name) + if value is not None: + return value + if pretrained_config is not None: + return getattr(pretrained_config, name, default) + return default + + return DeepSeekV4MetadataParams( + indexer_max_chunk_size=self.indexer_max_chunk_size or 32768, + max_sparse_topk=_value("index_topk"), + enable_indexer_skip=self.skip_indexer_for_short_seqs, + enable_heuristic_topk=self.enable_heuristic_topk, + use_cute_dsl_paged_mqa_logits=(self.use_cute_dsl_paged_mqa_logits), + q_split_threshold=self.q_split_threshold, + compress_ratios=self.compress_ratios, + window_size=self.window_size, + ) + class SkipSoftmaxAttentionConfig(BaseSparseAttentionConfig): """Configuration for skip softmax attention.""" @@ -1108,7 +1168,7 @@ class MoeConfig(StrictBaseModel): """Configuration for MoE.""" backend: Literal[ "AUTO", "CUTLASS", "CUTEDSL", "WIDEEP", "TRTLLM", "DEEPGEMM", - "DENSEGEMM", "VANILLA", "TRITON", "MARLIN"] = Field( + "DENSEGEMM", "VANILLA", "TRITON", "MARLIN", "MEGAMOE_DEEPGEMM"] = Field( default='AUTO', description="MoE backend to use. " "AUTO selects default backend based on model. It currently doesn\'t always give the best choice for all scenarios. The capabilities of auto selection will be improved in future releases." @@ -1144,6 +1204,7 @@ class MoeConfig(StrictBaseModel): # Maps alias → full import path (module.ClassName). TOKENIZER_ALIASES = { 'deepseek_v32': 'tensorrt_llm.tokenizer.deepseek_v32.DeepseekV32Tokenizer', + 'deepseek_v4': 'tensorrt_llm.tokenizer.deepseek_v4.DeepseekV4Tokenizer', } @@ -3911,23 +3972,39 @@ def validate_and_init_tokenizer(self): if self.skip_tokenizer_init: self.tokenizer = None elif self.custom_tokenizer: - # If tokenizer is already a tokenizer object, custom_tokenizer is not compatible - if isinstance(self.tokenizer, - (TokenizerBase, PreTrainedTokenizerBase)): + # IPC workers receive the tokenizer object that was already loaded + # in the parent LLM process. Reuse TRT-LLM tokenizer wrappers as-is. + if isinstance(self.tokenizer, TokenizerBase): + return self + # A raw HF tokenizer object would bypass the requested custom + # wrapper, so keep rejecting that combination. + if isinstance(self.tokenizer, PreTrainedTokenizerBase): raise ValueError( "Cannot use custom_tokenizer when tokenizer is already a tokenizer object. " "Please specify a tokenizer path or leave it as None to load from model path." ) - from tensorrt_llm.tokenizer import load_custom_tokenizer + # Resolve short aliases via the module-level TOKENIZER_ALIASES. + tokenizer_path = TOKENIZER_ALIASES.get(self.custom_tokenizer, + self.custom_tokenizer) - # Use tokenizer path if specified, otherwise use model path - load_path = self.tokenizer if self.tokenizer else self.model - self.tokenizer = load_custom_tokenizer( - self.custom_tokenizer, - load_path, - trust_remote_code=self.trust_remote_code, - use_fast=self.tokenizer_mode != 'slow') + # Dynamically import and use custom tokenizer + from importlib import import_module + try: + module_path, class_name = tokenizer_path.rsplit('.', 1) + module = import_module(module_path) + tokenizer_class = getattr(module, class_name) + # Use tokenizer path if specified, otherwise use model path + load_path = self.tokenizer if self.tokenizer else self.model + self.tokenizer = tokenizer_class.from_pretrained( + load_path, + trust_remote_code=self.trust_remote_code, + use_fast=self.tokenizer_mode != 'slow') + except (ValueError, ImportError, AttributeError) as e: + raise ValueError( + f"Failed to load custom tokenizer '{self.custom_tokenizer}': {e}. " + "Expected format: 'module.path.ClassName' or a recognized alias." + ) from e else: self.tokenizer = tokenizer_factory( self.tokenizer, diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index addcd5674ae4..e01043406944 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -403,6 +403,7 @@ def _update_from_hf_quant_config(self) -> bool: "fp8": QuantAlgo.FP8, "nvfp4": QuantAlgo.NVFP4, }.get(kv_cache_dtype) + requires_global_quant_config_fallback = False hf_quant_config_path = f"{self._model_dir}/hf_quant_config.json" if os.path.exists(hf_quant_config_path): @@ -423,9 +424,20 @@ def _update_from_hf_quant_config(self) -> bool: ) except FileNotFoundError: pass - self._apply_modelopt_quant_config(normalized, - explicit_kv_cache_quant_algo) - return True + if normalized.get("quant_algo") is None: + if normalized.get("quantized_layers") is not None: + requires_global_quant_config_fallback = True + logger.info( + "hf_quant_config.json does not set a global quant_algo; " + "falling back to config.json or model_kwargs for global " + "quantization.") + else: + raise ValueError( + "Pre-quantized checkpoint must have quant_algo.") + else: + self._apply_modelopt_quant_config(normalized, + explicit_kv_cache_quant_algo) + return True hf_config_path = f"{self._model_dir}/config.json" hf_quant_config = None @@ -447,6 +459,11 @@ def _update_from_hf_quant_config(self) -> bool: f"Use quantization_config from {hf_config_path}: quantization_config={hf_quant_config}" ) + if requires_global_quant_config_fallback and hf_quant_config is None: + raise ValueError( + "hf_quant_config.json does not set a global quant_algo and no " + "quantization_config fallback was found.") + if hf_quant_config is not None: if is_modelopt_quant_config(hf_quant_config): self._apply_modelopt_quant_config( @@ -458,6 +475,8 @@ def _update_from_hf_quant_config(self) -> bool: if hf_quant_config.get("quant_method") == "fp8": if hf_quant_config.get("weight_block_size") is not None: quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES + quant_config.group_size = hf_quant_config[ + "weight_block_size"][0] quant_config.exclude_modules = ["*eh_proj"] else: # Ministral 3 static quant diff --git a/tensorrt_llm/llmapi/reasoning_parser.py b/tensorrt_llm/llmapi/reasoning_parser.py index 0b2af87978b4..ccc81e2e05e2 100644 --- a/tensorrt_llm/llmapi/reasoning_parser.py +++ b/tensorrt_llm/llmapi/reasoning_parser.py @@ -95,6 +95,19 @@ def finish(self) -> ReasoningParserResult: return ReasoningParserResult() +class IdentityReasoningParser(BaseReasoningParser): + """Reasoning parser that treats all model output as visible content.""" + + reasoning_start = "" + reasoning_end = "" + + def parse(self, text: str) -> ReasoningParserResult: + return ReasoningParserResult(content=text) + + def parse_delta(self, delta_text: str) -> ReasoningParserResult: + return ReasoningParserResult(content=delta_text) + + @register_reasoning_parser("deepseek-r1", reasoning_at_start=True) @register_reasoning_parser("laguna") @register_reasoning_parser("qwen3") @@ -198,6 +211,42 @@ def parse_delta(self, delta_text: str) -> ReasoningParserResult: "Unreachable code reached in `DeepSeekR1Parser.parse_delta`") +@register_reasoning_parser("deepseek_v4") +class DeepSeekV4ReasoningParser(BaseReasoningParser): + """DeepSeek-V4 parser selected by thinking-mode chat template kwargs.""" + + reasoning_start = "" + reasoning_end = "" + + def __init__( + self, + *, + chat_template_kwargs: Optional[dict[str, Any]] = None, + ) -> None: + super().__init__(chat_template_kwargs=chat_template_kwargs) + chat_template_kwargs = chat_template_kwargs or {} + thinking = bool( + chat_template_kwargs.get("thinking", False) + or chat_template_kwargs.get("enable_thinking", False)) + if thinking: + self._parser = DeepSeekR1Parser( + reasoning_at_start=True, + chat_template_kwargs=chat_template_kwargs, + ) + else: + self._parser = IdentityReasoningParser( + chat_template_kwargs=chat_template_kwargs) + + def parse(self, text: str) -> ReasoningParserResult: + return self._parser.parse(text) + + def parse_delta(self, delta_text: str) -> ReasoningParserResult: + return self._parser.parse_delta(delta_text) + + def finish(self) -> ReasoningParserResult: + return self._parser.finish() + + @register_reasoning_parser("minimax_m3") class MiniMaxM3ReasoningParser(DeepSeekR1Parser): """Reasoning parser for MiniMax-M3. @@ -254,6 +303,7 @@ def parse(self, text: str) -> ReasoningParserResult: "deepseek_v3": "deepseek-r1", "deepseek_v32": "deepseek-r1", "laguna": "laguna", + "deepseek_v4": "deepseek_v4", "nemotron_h": "nemotron-v3", "nemotron_h_puzzle": "nemotron-v3", "gemma4": "gemma4", diff --git a/tensorrt_llm/llmapi/tokenizer.py b/tensorrt_llm/llmapi/tokenizer.py index 214ca86f9780..67a7bf5c227f 100644 --- a/tensorrt_llm/llmapi/tokenizer.py +++ b/tensorrt_llm/llmapi/tokenizer.py @@ -6,6 +6,7 @@ _llguidance_tokenizer_info, _xgrammar_tokenizer_info, load_hf_tokenizer, tokenizer_factory) +from tensorrt_llm.tokenizer.deepseek_v4 import DeepseekV4Tokenizer from tensorrt_llm.tokenizer.deepseek_v32 import DeepseekV32Tokenizer __all__ = [ @@ -14,6 +15,7 @@ "TokenizerBase", "TransformersTokenizer", "DeepseekV32Tokenizer", + "DeepseekV4Tokenizer", "tokenizer_factory", "_xgrammar_tokenizer_info", "_llguidance_tokenizer_info", diff --git a/tensorrt_llm/serve/tool_parser/deepseekv32_parser.py b/tensorrt_llm/serve/tool_parser/deepseekv32_parser.py index 25c49ae2a2cb..3362ab7cb68a 100644 --- a/tensorrt_llm/serve/tool_parser/deepseekv32_parser.py +++ b/tensorrt_llm/serve/tool_parser/deepseekv32_parser.py @@ -133,7 +133,7 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult try: # Extract content between function_calls tags function_calls_match = re.search( - r"<｜DSML｜function_calls>(.*?)", + re.escape(self.bot_token) + r"(.*?)" + re.escape(self.eot_token), text, re.DOTALL, ) diff --git a/tensorrt_llm/serve/tool_parser/deepseekv4_parser.py b/tensorrt_llm/serve/tool_parser/deepseekv4_parser.py new file mode 100644 index 000000000000..05896561f839 --- /dev/null +++ b/tensorrt_llm/serve/tool_parser/deepseekv4_parser.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .deepseekv32_parser import DeepSeekV32Parser + + +class DeepSeekV4Parser(DeepSeekV32Parser): + """Tool parser for the DeepSeek V4 DSML tool call format.""" + + def __init__(self) -> None: + super().__init__() + self.bot_token = "<｜DSML｜tool_calls>" # nosec B105 + self.eot_token = "" # nosec B105 diff --git a/tensorrt_llm/serve/tool_parser/tool_parser_factory.py b/tensorrt_llm/serve/tool_parser/tool_parser_factory.py index 0258556a986b..87bc6d090b48 100644 --- a/tensorrt_llm/serve/tool_parser/tool_parser_factory.py +++ b/tensorrt_llm/serve/tool_parser/tool_parser_factory.py @@ -4,6 +4,7 @@ from .base_tool_parser import BaseToolParser from .deepseekv3_parser import DeepSeekV3Parser +from .deepseekv4_parser import DeepSeekV4Parser from .deepseekv31_parser import DeepSeekV31Parser from .deepseekv32_parser import DeepSeekV32Parser from .gemma4_parser import Gemma4ToolParser @@ -25,6 +26,7 @@ "qwen3_next": "qwen3", "deepseek_v3": "deepseek_v3", "deepseek_v32": "deepseek_v32", + "deepseek_v4": "deepseek_v4", "kimi_k2": "kimi_k2", "kimi_k25": "kimi_k2", "glm4": "glm4", @@ -59,6 +61,7 @@ class ToolParserFactory: "deepseek_v3": DeepSeekV3Parser, "deepseek_v31": DeepSeekV31Parser, "deepseek_v32": DeepSeekV32Parser, + "deepseek_v4": DeepSeekV4Parser, "gemma4": Gemma4ToolParser, "glm4": Glm4ToolParser, "glm47": Glm47ToolParser, diff --git a/tensorrt_llm/tokenizer/deepseek_v4/__init__.py b/tensorrt_llm/tokenizer/deepseek_v4/__init__.py new file mode 100644 index 000000000000..30c51f77d64d --- /dev/null +++ b/tensorrt_llm/tokenizer/deepseek_v4/__init__.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .tokenizer import DeepseekV4Tokenizer + +__all__ = ["DeepseekV4Tokenizer"] diff --git a/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py b/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py new file mode 100644 index 000000000000..a5600816b072 --- /dev/null +++ b/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py @@ -0,0 +1,457 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ruff: noqa: E501 + +import copy +import json +from pathlib import Path +from typing import Any + +from transformers import AutoTokenizer + +from ..tokenizer import TransformersTokenizer + +BOS_TOKEN = "<｜begin▁of▁sentence｜>" # nosec B105 +EOS_TOKEN = "<｜end▁of▁sentence｜>" # nosec B105 +USER_TOKEN = "<｜User｜>" # nosec B105 +ASSISTANT_TOKEN = "<｜Assistant｜>" # nosec B105 +LATEST_REMINDER_TOKEN = "<｜latest_reminder｜>" # nosec B105 +THINKING_START_TOKEN = "" # nosec B105 +THINKING_END_TOKEN = "" # nosec B105 +DSML_TOKEN = "｜DSML｜" # nosec B105 + +TOOL_CALLS_BLOCK_NAME = "tool_calls" +VALID_TASKS = { + "action": "<｜action｜>", + "query": "<｜query｜>", + "authority": "<｜authority｜>", + "domain": "<｜domain｜>", + "title": "<｜title｜>", + "read_url": "<｜read_url｜>", +} + +REASONING_EFFORT_MAX = ( + "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n" + "You MUST be very thorough in your thinking and comprehensively decompose " + "the problem to resolve the root cause, rigorously stress-testing your " + "logic against all potential paths, edge cases, and adversarial scenarios.\n" + "Explicitly write out your entire deliberation process, documenting every " + "intermediate step, considered alternative, and rejected hypothesis to " + "ensure absolutely no assumption is left unchecked.\n\n" +) + +TOOLS_TEMPLATE = """## Tools + +You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following: + +<{dsml_token}tool_calls> +<{dsml_token}invoke name="$TOOL_NAME"> +<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE +... + +<{dsml_token}invoke name="$TOOL_NAME2"> +... + + + +String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`. + +If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response. + +Otherwise, output directly after {thinking_end_token} with tool calls or final response. + +### Available Tool Schemas + +{tool_schemas} + +You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls. +""" + +RESPONSE_FORMAT_TEMPLATE = ( + "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}" +) +TOOL_CALL_TEMPLATE = '<{dsml_token}invoke name="{name}">\n{arguments}\n' +TOOL_CALLS_TEMPLATE = "<{dsml_token}{block_name}>\n{tool_calls}\n" +TOOL_OUTPUT_TEMPLATE = "{content}" + + +def _message_content_to_text(content: Any) -> str: + if content is None: + return "" + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for block in content: + if isinstance(block, dict) and block.get("type") == "text": + parts.append(str(block.get("text", ""))) + else: + parts.append(str(block)) + return "\n\n".join(parts) + return str(content) + + +def _to_json(value: Any) -> str: + try: + return json.dumps(value, ensure_ascii=False) + except TypeError: + return json.dumps(value, ensure_ascii=True) + + +def _tools_from_openai_format(tools: list[dict[str, Any]]) -> list[dict[str, Any]]: + return [tool["function"] for tool in tools] + + +def _tool_calls_from_openai_format(tool_calls: list[dict[str, Any]]) -> list[dict[str, Any]]: + return [ + { + "name": tool_call["function"]["name"], + "arguments": tool_call["function"]["arguments"], + } + for tool_call in tool_calls + ] + + +def _encode_arguments_to_dsml(tool_call: dict[str, Any]) -> str: + raw_args = tool_call["arguments"] + arguments = json.loads(raw_args) if isinstance(raw_args, str) else raw_args + if not isinstance(arguments, dict): + raise ValueError("DeepSeek-V4 tool call arguments must be a JSON object.") + + parameters = [] + for key, value in arguments.items(): + parameters.append( + f'<{DSML_TOKEN}parameter name="{key}" ' + f'string="{"true" if isinstance(value, str) else "false"}">' + f"{value if isinstance(value, str) else _to_json(value)}" + ) + return "\n".join(parameters) + + +def _render_tools(tools: list[dict[str, Any]]) -> str: + return TOOLS_TEMPLATE.format( + tool_schemas="\n".join(_to_json(tool) for tool in tools), + dsml_token=DSML_TOKEN, + thinking_start_token=THINKING_START_TOKEN, + thinking_end_token=THINKING_END_TOKEN, + ) + + +def _find_last_user_index(messages: list[dict[str, Any]]) -> int: + for index in range(len(messages) - 1, -1, -1): + if messages[index].get("role") in ("user", "developer"): + return index + return -1 + + +def _merge_tool_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + merged: list[dict[str, Any]] = [] + for message in messages: + message = copy.deepcopy(message) + role = message.get("role") + + if role == "tool": + tool_block = { + "type": "tool_result", + "tool_use_id": message.get("tool_call_id", ""), + "content": message.get("content", ""), + } + if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1]: + merged[-1]["content_blocks"].append(tool_block) + else: + merged.append({"role": "user", "content_blocks": [tool_block]}) + elif role == "user": + text_block = {"type": "text", "text": _message_content_to_text(message.get("content"))} + if ( + merged + and merged[-1].get("role") == "user" + and "content_blocks" in merged[-1] + and merged[-1].get("task") is None + ): + merged[-1]["content_blocks"].append(text_block) + else: + new_message = { + "role": "user", + "content": message.get("content", ""), + "content_blocks": [text_block], + } + for key in ("task", "wo_eos", "mask"): + if key in message: + new_message[key] = message[key] + merged.append(new_message) + else: + merged.append(message) + + return merged + + +def _sort_tool_results_by_call_order(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + last_tool_call_order: dict[str, int] = {} + + for message in messages: + if message.get("role") == "assistant" and message.get("tool_calls"): + last_tool_call_order = {} + for index, tool_call in enumerate(message["tool_calls"]): + tool_call_id = tool_call.get("id") or tool_call.get("function", {}).get("id", "") + if tool_call_id: + last_tool_call_order[tool_call_id] = index + elif message.get("role") == "user" and message.get("content_blocks"): + tool_blocks = [ + block for block in message["content_blocks"] if block.get("type") == "tool_result" + ] + if len(tool_blocks) > 1 and last_tool_call_order: + sorted_blocks = sorted( + tool_blocks, + key=lambda block: last_tool_call_order.get(block.get("tool_use_id", ""), 0), + ) + sorted_index = 0 + new_blocks = [] + for block in message["content_blocks"]: + if block.get("type") == "tool_result": + new_blocks.append(sorted_blocks[sorted_index]) + sorted_index += 1 + else: + new_blocks.append(block) + message["content_blocks"] = new_blocks + + return messages + + +def _drop_thinking_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + last_user_index = _find_last_user_index(messages) + result = [] + keep_roles = {"user", "system", "tool", "latest_reminder", "direct_search_results"} + + for index, message in enumerate(messages): + role = message.get("role") + if role in keep_roles or index >= last_user_index: + result.append(message) + elif role == "assistant": + message_without_reasoning = copy.copy(message) + message_without_reasoning.pop("reasoning", None) + message_without_reasoning.pop("reasoning_content", None) + result.append(message_without_reasoning) + + return result + + +def _render_user_content(message: dict[str, Any]) -> str: + content_blocks = message.get("content_blocks") + if not content_blocks: + return _message_content_to_text(message.get("content")) + + parts = [] + for block in content_blocks: + block_type = block.get("type") + if block_type == "text": + parts.append(str(block.get("text", ""))) + elif block_type == "tool_result": + tool_content = block.get("content", "") + if isinstance(tool_content, list): + text_parts = [] + for item in tool_content: + if isinstance(item, dict) and item.get("type") == "text": + text_parts.append(str(item.get("text", ""))) + elif isinstance(item, dict): + text_parts.append(f"[Unsupported {item.get('type')}]") + else: + text_parts.append(str(item)) + tool_content = "\n\n".join(text_parts) + parts.append(TOOL_OUTPUT_TEMPLATE.format(content=tool_content)) + else: + parts.append(f"[Unsupported {block_type}]") + return "\n\n".join(parts) + + +def _render_message( + index: int, + messages: list[dict[str, Any]], + thinking_mode: str, + drop_thinking: bool, + add_generation_prompt: bool, + reasoning_effort: str | None, +) -> str: + if thinking_mode not in ("chat", "thinking"): + raise ValueError(f"Invalid thinking_mode: {thinking_mode}") + + message = messages[index] + last_user_index = _find_last_user_index(messages) + role = message.get("role") + content = _message_content_to_text(message.get("content")) + tools = message.get("tools") + response_format = message.get("response_format") + tool_calls = message.get("tool_calls") + reasoning = message.get("reasoning") or message.get("reasoning_content") or "" + prompt = "" + + if tools: + tools = _tools_from_openai_format(tools) + if tool_calls: + tool_calls = _tool_calls_from_openai_format(tool_calls) + + if index == 0 and thinking_mode == "thinking" and reasoning_effort == "max": + prompt += REASONING_EFFORT_MAX + + if role == "system": + prompt += content + if tools: + prompt += "\n\n" + _render_tools(tools) + if response_format: + prompt += "\n\n" + RESPONSE_FORMAT_TEMPLATE.format(schema=_to_json(response_format)) + elif role == "developer": + prompt += USER_TOKEN + content + if tools: + prompt += "\n\n" + _render_tools(tools) + if response_format: + prompt += "\n\n" + RESPONSE_FORMAT_TEMPLATE.format(schema=_to_json(response_format)) + elif role == "user": + prompt += USER_TOKEN + _render_user_content(message) + elif role == "latest_reminder": + prompt += LATEST_REMINDER_TOKEN + content + elif role == "tool": + raise NotImplementedError( + "DeepSeek-V4 merges tool messages into user messages; " + "preprocess with _merge_tool_messages()." + ) + elif role == "assistant": + tool_calls_content = "" + if tool_calls: + rendered_tool_calls = [ + TOOL_CALL_TEMPLATE.format( + dsml_token=DSML_TOKEN, + name=tool_call.get("name"), + arguments=_encode_arguments_to_dsml(tool_call), + ) + for tool_call in tool_calls + ] + tool_calls_content += "\n\n" + TOOL_CALLS_TEMPLATE.format( + dsml_token=DSML_TOKEN, + block_name=TOOL_CALLS_BLOCK_NAME, + tool_calls="\n".join(rendered_tool_calls), + ) + + thinking_part = "" + prev_has_task = index - 1 >= 0 and messages[index - 1].get("task") is not None + if thinking_mode == "thinking" and not prev_has_task: + if not drop_thinking or index > last_user_index: + thinking_part = reasoning + THINKING_END_TOKEN + + if message.get("wo_eos", False): + prompt += thinking_part + content + tool_calls_content + else: + prompt += thinking_part + content + tool_calls_content + EOS_TOKEN + else: + raise NotImplementedError(f"Unsupported DeepSeek-V4 message role: {role}") + + next_role = messages[index + 1].get("role") if index + 1 < len(messages) else None + if next_role is not None and next_role not in ("assistant", "latest_reminder"): + return prompt + + task = message.get("task") + if task is not None: + if task not in VALID_TASKS: + raise ValueError(f"Invalid DeepSeek-V4 task: {task}") + if task == "action": + prompt += ASSISTANT_TOKEN + prompt += THINKING_START_TOKEN if thinking_mode == "thinking" else THINKING_END_TOKEN + prompt += VALID_TASKS[task] + elif role in ("user", "developer") and (next_role == "assistant" or add_generation_prompt): + prompt += ASSISTANT_TOKEN + if thinking_mode == "thinking" and (not drop_thinking or index >= last_user_index): + prompt += THINKING_START_TOKEN + else: + prompt += THINKING_END_TOKEN + + return prompt + + +def _encode_messages( + messages: list[dict[str, Any]], + thinking_mode: str, + drop_thinking: bool, + add_generation_prompt: bool, + reasoning_effort: str | None, +) -> str: + messages = _merge_tool_messages(messages) + messages = _sort_tool_results_by_call_order(messages) + + effective_drop_thinking = drop_thinking + if any(message.get("tools") for message in messages): + effective_drop_thinking = False + + if thinking_mode == "thinking" and effective_drop_thinking: + messages = _drop_thinking_messages(messages) + + prompt = BOS_TOKEN + for index in range(len(messages)): + prompt += _render_message( + index, + messages, + thinking_mode=thinking_mode, + drop_thinking=effective_drop_thinking, + add_generation_prompt=add_generation_prompt, + reasoning_effort=reasoning_effort, + ) + return prompt + + +class DeepseekV4Tokenizer(TransformersTokenizer): + """DeepSeek-V4 tokenizer with the checkpoint reference chat format.""" + + @classmethod + def from_pretrained( + cls, + path_or_repo_id: str | Path, + *args, + trust_remote_code: bool = False, + revision: str | None = None, + **kwargs, + ) -> "DeepseekV4Tokenizer": + tokenizer = AutoTokenizer.from_pretrained( + path_or_repo_id, + *args, + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + return cls(tokenizer) + + def apply_chat_template(self, messages, tools=None, **kwargs): + tokenize = kwargs.get("tokenize", False) + thinking = kwargs.get("thinking", False) or kwargs.get("enable_thinking", False) + thinking_mode = "thinking" if thinking else "chat" + reasoning_effort = kwargs.get("reasoning_effort") + if reasoning_effort not in ("max", "high"): + reasoning_effort = None + + conversation = kwargs.get("conversation", messages) + messages = list(conversation) + if tools: + messages.insert(0, {"role": "system", "tools": tools}) + + rendered = _encode_messages( + messages=messages, + thinking_mode=thinking_mode, + drop_thinking=kwargs.get("drop_thinking", True), + add_generation_prompt=True, + reasoning_effort=reasoning_effort, + ) + + if tokenize: + tokenizer_kwargs = { + key: kwargs[key] for key in ("truncation", "max_length") if key in kwargs + } + return self.encode(rendered, add_special_tokens=False, **tokenizer_kwargs) + return rendered diff --git a/tensorrt_llm/tokenizer/tokenizer.py b/tensorrt_llm/tokenizer/tokenizer.py index 3a44d77543e8..ad09cc316a2b 100644 --- a/tensorrt_llm/tokenizer/tokenizer.py +++ b/tensorrt_llm/tokenizer/tokenizer.py @@ -26,6 +26,7 @@ # Aliases for built-in custom tokenizers. TOKENIZER_ALIASES = { "deepseek_v32": "tensorrt_llm.tokenizer.deepseek_v32.DeepseekV32Tokenizer", + "deepseek_v4": "tensorrt_llm.tokenizer.deepseek_v4.DeepseekV4Tokenizer", } TLLM_INCREMENTAL_DETOKENIZATION_BACKEND = os.environ.get( diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index 1983cf140e4a..f5d47675b8b8 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -130,6 +130,16 @@ deepseek-ai/DeepSeek-V3.2-Exp: kv_cache_quant_algo: FP8 spec_dec_algo: MTP accuracy: 95.6 +deepseek-ai/DeepSeek-V4-Flash: + # GSM8K measurements: + # * 95.11 on 8x B200 178GB at TP=8 with tightened config (fraction=0.15, + # cuda_graph=None, max_num_tokens=2048) — original measurement. + # * 95.38 on 4x B300 275GB at TP=4 with default config (fraction=0.5, + # CudaGraphConfig()) — current CI path (test_nvfp4_4gpus_static_eplb + # in l0_dgx_b300_ds.yml). Drift +0.27 from 95.11 is within sigma; the + # 95.11 reference still holds for the hypothesis test. + - quant_algo: FP8_BLOCK_SCALES + accuracy: 95.11 Qwen3/Qwen3-4B: - spec_dec_algo: Eagle accuracy: 85.823 diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 32df5bb9f15b..159eb9e13fa5 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -1416,6 +1416,25 @@ def test_bfloat16_python_scheduler(self, mtp_nextn, attention_dp, task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + @pytest.mark.skip_less_device_memory(60000) + def test_bfloat16_2_model_mtp(self): + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.3) + pytorch_config = dict( + disable_overlap_scheduler=True, + cuda_graph_config=CudaGraphConfig(), + ) + mtp_config = MTPDecodingConfig(max_draft_len=3, + mtp_eagle_one_model=False, + speculative_model=self.MODEL_PATH) + with LLM(self.MODEL_PATH, + kv_cache_config=kv_cache_config, + enable_chunked_prefill=False, + max_num_tokens=8192, + **pytorch_config, + speculative_config=mtp_config) as llm: + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + @pytest.mark.skip_less_device_memory(60000) def test_bfloat16_mtp_sa(self): """Accuracy test for MTP + Suffix Automaton (MTP+SA) speculative decoding.""" @@ -3405,6 +3424,201 @@ def test_nvfp4_multi_gpus_chunked_prefill(self, tp_size, pp_size, ep_size, task.evaluate(llm) +def _make_deepseekv4_eplb_config(model_path, layer_updates_per_iter, ep_size=8): + """Build a MoeLoadBalancerConfig for DeepSeek V4 from the HF config. + + All V4 layers run MoE (no first_k_dense_replace prefix), so static + assignments cover every layer in 0..num_hidden_layers-1. Extra slots + per rank match the TestNemotronV3Super pattern (16 redundant per rank). + """ + with open(f"{model_path}/config.json") as f: + cfg = json.load(f) + num_experts = cfg["n_routed_experts"] + num_slots = num_experts + 16 * ep_size + if layer_updates_per_iter > 0: + return MoeLoadBalancerConfig( + num_slots=num_slots, layer_updates_per_iter=layer_updates_per_iter) + num_hidden_layers = cfg["num_hidden_layers"] + initial_global_assignments = { + i: [(i + j) % num_experts for j in range(num_slots)] + for i in range(num_hidden_layers) + } + return MoeLoadBalancerConfig( + num_slots=num_slots, + initial_global_assignments=initial_global_assignments, + layer_updates_per_iter=0) + + +def _run_deepseekv4_eplb(model_name, + model_path, + moe_backend, + eplb_config, + mtp_nextn=0, + tensor_parallel_size=4): + # Default config targets 4x B300 (~288 GB/GPU): plenty of headroom at + # TP=4 for V4-Flash NVFP4 (~36 GB/rank weights) or V4-Flash-Base FP8 + # (~71 GB/rank weights), even with EPLB redundancy + DeepGemm MoE + # workspace + cuda graph capture. + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.5) + pytorch_config = dict(cuda_graph_config=CudaGraphConfig(), + moe_config=MoeConfig(backend=moe_backend, + load_balancer=eplb_config)) + mtp_config = (MTPDecodingConfig( + num_nextn_predict_layers=mtp_nextn) if mtp_nextn > 0 else None) + with LLM(model_path, + tensor_parallel_size=tensor_parallel_size, + moe_expert_parallel_size=tensor_parallel_size, + kv_cache_config=kv_cache_config, + enable_attention_dp=True, + max_seq_len=4096, + **pytorch_config, + speculative_config=mtp_config) as llm: + task = GSM8K(model_name) + task.evaluate(llm) + + +@pytest.mark.timeout(14400) +@pytest.mark.skip_less_device_memory(140000) +@skip_pre_blackwell +class TestDeepSeekV4Flash(LlmapiAccuracyTestHarness): + MODEL_NAME = "deepseek-ai/DeepSeek-V4-Flash" + MODEL_PATH = f"{llm_models_root()}/DeepSeek-V4-Flash" + + @pytest.mark.skip_less_mpi_world_size(4) + def test_auto_dtype(self): + # Aggregate (non-disagg, non-EPLB) smoke test. NVFP4 weights are ~71 + # GB/rank at TP=2, ~36 GB/rank at TP=4 — TP=4 fits comfortably on + # 4x B200 178GB. TRTLLM backend required because V4-Flash MXFP4 + # routed experts are unsupported by WIDEEP (raises "Unsupported + # quantization mode: [65536]"). is_integration_test=True keeps this + # to a 1-sample smoke. + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.5) + with LLM(self.MODEL_PATH, + tensor_parallel_size=4, + moe_expert_parallel_size=4, + moe_config=MoeConfig(backend="TRTLLM"), + enable_attention_dp=True, + max_batch_size=4, + max_seq_len=4096, + max_num_tokens=4096, + kv_cache_config=kv_cache_config) as llm: + task = MMLU(self.MODEL_NAME) + task.evaluate(llm, is_integration_test=True) + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm, is_integration_test=True) + + @pytest.mark.skip_less_mpi_world_size(4) + @parametrize_with_ids("moe_backend", [ + pytest.param( + "WIDEEP", + marks=pytest.mark.skip( + reason= + "V4-Flash MXFP4 routed experts: WIDEEP _get_quant_method has " + "no MXFP4 branch (raises 'Unsupported quantization mode: " + "[65536]'). Re-enable once fused_moe_wide_ep.py supports MXFP4." + )), + "TRTLLM", + "MEGAMOE_DEEPGEMM", + ]) + def test_nvfp4_4gpus_static_eplb(self, moe_backend): + eplb_config = _make_deepseekv4_eplb_config(self.MODEL_PATH, + layer_updates_per_iter=0, + ep_size=4) + _run_deepseekv4_eplb(self.MODEL_NAME, self.MODEL_PATH, moe_backend, + eplb_config) + + @pytest.mark.skip_less_mpi_world_size(4) + @parametrize_with_ids("moe_backend", [ + pytest.param( + "WIDEEP", + marks=pytest.mark.skip( + reason= + "V4-Flash MXFP4 routed experts: WIDEEP _get_quant_method has " + "no MXFP4 branch (raises 'Unsupported quantization mode: " + "[65536]'). Re-enable once fused_moe_wide_ep.py supports MXFP4." + )), + "TRTLLM", + ]) + @parametrize_with_ids("mtp_nextn", [0, 1]) + def test_nvfp4_4gpus_online_eplb(self, moe_backend, mtp_nextn): + eplb_config = _make_deepseekv4_eplb_config(self.MODEL_PATH, + layer_updates_per_iter=2, + ep_size=4) + _run_deepseekv4_eplb(self.MODEL_NAME, + self.MODEL_PATH, + moe_backend, + eplb_config, + mtp_nextn=mtp_nextn) + + +@pytest.mark.timeout(14400) +@pytest.mark.skip_less_device_memory(140000) +@skip_pre_blackwell +class TestDeepSeekV4FlashBase(LlmapiAccuracyTestHarness): + MODEL_NAME = "deepseek-ai/DeepSeek-V4-Flash-Base" + MODEL_PATH = f"{llm_models_root()}/DeepSeek-V4-Flash-Base" + + @pytest.mark.skip_less_mpi_world_size(4) + @parametrize_with_ids("moe_backend", ["WIDEEP", "TRTLLM"]) + def test_auto_dtype(self, moe_backend): + # Aggregate (non-disagg, non-EPLB) smoke test. FP8 weights ~71 GB/rank + # at TP=4 — fits on 4x B300 (~288 GB/GPU). 1-sample smoke. CUTLASS is + # Hopper-only on Blackwell and skipped here. + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.5) + with LLM(self.MODEL_PATH, + tensor_parallel_size=4, + moe_expert_parallel_size=4, + moe_config=MoeConfig(backend=moe_backend), + enable_attention_dp=True, + max_seq_len=4096, + kv_cache_config=kv_cache_config) as llm: + task = MMLU(self.MODEL_NAME) + task.evaluate(llm, is_integration_test=True) + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm, is_integration_test=True) + + @pytest.mark.skip_less_mpi_world_size(4) + def test_fp8_chunked_prefill(self): + kv_cache_config = KvCacheConfig(dtype="fp8", + free_gpu_memory_fraction=0.5, + enable_block_reuse=False) + with LLM(self.MODEL_PATH, + tensor_parallel_size=4, + moe_expert_parallel_size=4, + moe_config=MoeConfig(backend="WIDEEP"), + cuda_graph_config=CudaGraphConfig(max_batch_size=16, + enable_padding=True), + enable_attention_dp=True, + enable_chunked_prefill=True, + max_batch_size=16, + max_num_tokens=128, + max_seq_len=4096, + kv_cache_config=kv_cache_config) as llm: + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm, is_integration_test=True) + + # CUTLASS is omitted: V4 Flash-Base FP8 block-scale weights take a + # Hopper-only kernel path (CutlassFp8BlockScaleGemmRunner::moeGemm) that + # fails on Blackwell. WIDEEP avoids that path and works on B200/B300. + @pytest.mark.skip_less_mpi_world_size(4) + @parametrize_with_ids("moe_backend", ["WIDEEP"]) + def test_fp8_4gpus_static_eplb(self, moe_backend): + eplb_config = _make_deepseekv4_eplb_config(self.MODEL_PATH, + layer_updates_per_iter=0, + ep_size=4) + _run_deepseekv4_eplb(self.MODEL_NAME, self.MODEL_PATH, moe_backend, + eplb_config) + + @pytest.mark.skip_less_mpi_world_size(4) + @parametrize_with_ids("moe_backend", ["WIDEEP"]) + def test_fp8_4gpus_online_eplb(self, moe_backend): + eplb_config = _make_deepseekv4_eplb_config(self.MODEL_PATH, + layer_updates_per_iter=2, + ep_size=4) + _run_deepseekv4_eplb(self.MODEL_NAME, self.MODEL_PATH, moe_backend, + eplb_config) + + @pytest.mark.timeout(10800) @pytest.mark.skip_less_device_memory(100000) class TestKimiK2(LlmapiAccuracyTestHarness): diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py index 59598f11c9de..14d0e4c97839 100644 --- a/tests/integration/defs/conftest.py +++ b/tests/integration/defs/conftest.py @@ -1776,7 +1776,12 @@ def skip_by_device_memory(request): def get_sm_version(): "get compute capability" - prop = torch.cuda.get_device_properties(0) + if not torch.cuda.is_available(): + return 0 + try: + prop = torch.cuda.get_device_properties(0) + except RuntimeError: + return 0 return prop.major * 10 + prop.minor @@ -1802,6 +1807,15 @@ def check_device_contain(keyword_list): return any(keyword in device for keyword in keyword_list) +def is_ipc_nvls_supported(): + if not torch.cuda.is_available(): + return False + try: + return ipc_nvls_supported() + except RuntimeError: + return False + + skip_pre_ada = pytest.mark.skipif( get_sm_version() < 89, reason="This test is not supported in pre-Ada architecture") @@ -1836,7 +1850,7 @@ def check_device_contain(keyword_list): reason="This test is not supported on GB200 or GB100", ) -skip_no_nvls = pytest.mark.skipif(not ipc_nvls_supported(), +skip_no_nvls = pytest.mark.skipif(not is_ipc_nvls_supported(), reason="NVLS is not supported") skip_no_hopper = pytest.mark.skipif( get_sm_version() != 90, diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 8193517eb8a2..eea20ae0beb5 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -103,6 +103,20 @@ l0_b200: - unittest/_torch/attention - unittest/_torch/compilation - unittest/_torch/debugger + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_2_model_mtp + - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py TIMEOUT (60) + - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py TIMEOUT (60) + - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py TIMEOUT (60) + - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py TIMEOUT (30) + - unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py TIMEOUT (90) + - unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py TIMEOUT (90) + - unittest/_torch/attention/sparse/deepseek_v4/test_compressor_tf32.py TIMEOUT (15) + - unittest/_torch/attention/sparse/test_sparse_mla_forward.py TIMEOUT (60) + - unittest/_torch/modeling/test_modeling_deepseekv4.py + - unittest/llmapi/test_deepseek_v4_tokenizer.py + - unittest/_torch/modules/test_mhc.py + - unittest/_torch/modules/test_engram.py + - unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py TIMEOUT (15) # ------------- modules (non-MoE) --------------- - unittest/_torch/modules/test_mla_helix.py - unittest/_torch/modules/test_fused_add_rms_norm_quant.py diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index a877c051a7fe..150bbd614b15 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -50,8 +50,18 @@ l0_dgx_b200: - disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b] - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[latency_adp_lmtp_tp4] - accuracy/test_llm_api_pytorch.py::TestMiniMaxM2::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4] TIMEOUT (60) + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] TIMEOUT (60) + - unittest/_torch/modeling/test_modeling_deepseekv4.py + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_auto_dtype TIMEOUT (60) # ------------- NVBug 6025177: trtllm-serve cross-request KV contamination (OpenAI) --------------- - test_e2e.py::test_openai_kv_cache_contamination TIMEOUT (120) + # ------------- DSA FP4 indexer (Blackwell-only) --------------- + - unittest/_torch/attention/sparse/test_cpp_custom_ops.py::test_fused_cat_fp4_matches_deepgemm + - unittest/_torch/attention/sparse/test_cpp_custom_ops.py::test_fused_cat_fp4_shape_dispatch + - unittest/_torch/attention/sparse/test_cpp_custom_ops.py::test_fused_cat_fp4_noncontiguous_split + - unittest/_torch/attention/sparse/test_cpp_custom_ops.py::test_fused_cat_fp4_dsv32_prefill_shape + - unittest/_torch/attention/sparse/test_cpp_custom_ops.py::test_indexer_k_cache_gather_contiguous_fp4 + - unittest/_torch/attention/sparse/dsa/test_dsa_fp4_indexer.py - condition: ranges: system_gpu_count: @@ -166,6 +176,10 @@ l0_dgx_b200: - accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=0-block_reuse=False-use_py_transceiver=True] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=3-block_reuse=True-use_py_transceiver=False] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=0-block_reuse=False-use_py_transceiver=False] TIMEOUT (60) + # DeepSeek-V4 EPLB pre-merge sanity (uncomment once DeepSeek-V4-Flash/Flash-Base + # checkpoints are staged under llm_models_root()). + # - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_nvfp4_8gpus_static_eplb[moe_backend=WIDEEP] TIMEOUT (120) + # - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4FlashBase::test_fp8_8gpus_static_eplb[moe_backend=WIDEEP] TIMEOUT (120) - accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_ctx_dp2_gen_tp4 TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestQwen3NextInstruct::test_auto_dtype[use_py_transceiver=True] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestQwen3NextInstruct::test_auto_dtype[use_py_transceiver=False] TIMEOUT (60) diff --git a/tests/integration/test_lists/test-db/l0_dgx_b300.yml b/tests/integration/test_lists/test-db/l0_dgx_b300.yml index b608a2e2221f..ab0ff8adc2ae 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b300.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b300.yml @@ -122,3 +122,6 @@ l0_dgx_b300: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=TRTLLM-mtp_nextn=0-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-tp2pp2-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=TRTLLM-mtp_nextn=0-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=True-torch_compile=False] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_auto_dtype TIMEOUT (60) + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_nvfp4_4gpus_static_eplb[moe_backend=TRTLLM] TIMEOUT (120) + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_nvfp4_4gpus_static_eplb[moe_backend=MEGAMOE_DEEPGEMM] TIMEOUT (120) diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml index 5c4d59381418..e90d96ddef20 100644 --- a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml @@ -15,6 +15,8 @@ l0_gb200_multi_gpus: backend: pytorch tests: - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=TRTLLM-torch_compile=True] + - unittest/_torch/modules/test_engram.py + - unittest/_torch/modules/test_mhc.py - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=FLASHINFER-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index d18338e1e340..d626c49ad94b 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -716,6 +716,57 @@ def validate_topk_indices(topk_indices_0, topk_indices_1, total_tokens): return num_exact_matches, total_similarity, min_similarity +@pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") +@skip_pre_hopper +def test_recompute_slot_mappings_matches_prepare_with_cached_tokens(): + """Recompute slot mappings without re-running full Indexer.prepare().""" + head_dim = 128 + block_size = 64 + request_ids = [0, 1] + cached_tokens = [64, 128] + seq_lens = torch.tensor([32, 4], dtype=torch.int32) + kv_lens = torch.tensor([96, 132], dtype=torch.int32) + num_ctx_tokens = seq_lens[0].item() + num_tokens = seq_lens.sum().item() + + cache_manager, _ = create_dsa_cache_manager( + batch_size=len(request_ids), + head_dim=head_dim, + tokens_per_block=block_size, + max_seq_len=256, + num_layers=1, + ) + cache_manager.add_dummy_requests( + request_ids, kv_lens.tolist(), is_gen=False, prepare_resource=True + ) + metadata = _create_mock_metadata( + request_ids, + len(request_ids), + num_contexts=1, + num_generations=1, + seq_lens=seq_lens, + kv_lens=kv_lens, + num_cached_tokens=cached_tokens, + cache_manager=cache_manager, + num_ctx_tokens=num_ctx_tokens, + num_tokens=num_tokens, + indexer_head_dim=head_dim, + ) + + Indexer.prepare(metadata) + expected_fp8 = metadata.slot_mapping_fp8[:num_tokens].clone() + expected_scale = metadata.slot_mapping_scale[:num_tokens].clone() + + metadata.host_slot_mapping_fp8.zero_() + metadata.host_slot_mapping_scale.zero_() + metadata.slot_mapping_fp8.zero_() + metadata.slot_mapping_scale.zero_() + Indexer.recompute_slot_mappings(metadata) + + torch.testing.assert_close(metadata.slot_mapping_fp8[:num_tokens], expected_fp8) + torch.testing.assert_close(metadata.slot_mapping_scale[:num_tokens], expected_scale) + + @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @skip_pre_hopper def test_indexer_k_cache_scatter_custom_op(): diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py new file mode 100644 index 000000000000..a21a035f0202 --- /dev/null +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -0,0 +1,869 @@ +import ast +import inspect +import json +import struct +import textwrap +import weakref +from copy import deepcopy + +import pytest +import torch +from transformers import PretrainedConfig + +# from utils.util import default_dtype +import tensorrt_llm +from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.cache_manager import ( + DeepseekV4CacheManager, +) +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.compressor import Compressor +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import ( + DeepseekV4Indexer, + DeepseekV4TrtllmAttention, + DeepseekV4TrtllmAttentionMetadata, +) +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention +from tensorrt_llm._torch.configs.deepseekv4 import DeepseekV4Config +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_deepseekv4 import ( + DeepseekV4DecoderLayer, + DeepseekV4ForCausalLM, + DeepseekV4Gate, + DeepseekV4MTP, + _copy_deepseek_v4_fused_a_weight_scale, + _deepseek_v4_pos_embd_params, + _remap_deepseek_v4_checkpoint_keys, + _resolve_enable_fused_hc, +) +from tensorrt_llm._torch.modules.linear import TensorParallelMode +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.utils import AuxStreamType, model_extra_attrs +from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType +from tensorrt_llm.llmapi.llm_args import ( + DeepSeekV4SparseAttentionConfig, + KvCacheConfig, + MTPDecodingConfig, +) +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +DEEPSEEK_V4_TINY_CONFIG = { + "architectures": ["DeepseekV4ForCausalLM"], + "model_type": "deepseek_v4", + "hidden_size": 4096, + "num_attention_heads": 64, + "num_key_value_heads": 1, + "qk_nope_head_dim": 448, + "qk_rope_head_dim": 64, + "v_head_dim": 512, + "q_lora_rank": 1024, + "kv_lora_rank": 448, + "o_groups": 8, + "o_lora_rank": 1024, + "max_position_embeddings": 65536, + "rms_norm_eps": 1e-6, + "dtype": "bfloat16", + "vocab_size": 129280, + "num_hidden_layers": 7, + "n_hash_layers": 3, + "moe_intermediate_size": 2048, + "n_routed_experts": 256, + "n_shared_experts": 1, + "num_experts_per_tok": 6, + "n_group": 1, + "topk_group": 1, + "routed_scaling_factor": 1.5, + "score_func": "sqrtsoftplus", + "hc_mult": 4, + "hc_sinkhorn_iters": 20, + "hc_eps": 1e-6, + "compress_rope_theta": 40000.0, + "rope_theta": 10000.0, + "rope_scaling": { + "type": "yarn", + "factor": 4.0, + "original_max_position_embeddings": 65536, + "beta_fast": 32, + "beta_slow": 1, + }, + "quantization_config": { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + "scale_fmt": "ue8m0", + "weight_block_size": [128, 128], + }, +} + + +def _source_calls(source): + return { + ast.unparse(node) + for node in ast.walk(ast.parse(textwrap.dedent(source))) + if isinstance(node, ast.Call) + } + + +def _write_safetensors_header(path, tensor_name, dtype, shape): + header = { + tensor_name: { + "dtype": dtype, + "shape": shape, + "data_offsets": [0, 0], + } + } + payload = json.dumps(header).encode("utf-8") + path.write_bytes(struct.pack("ls<" in result +# ============================================================================ +# DeepSeekV4Parser Tests +# ============================================================================ + + +class TestDeepSeekV4Parser(BaseToolParserTestClass): + """Test suite for DeepSeekV4Parser class.""" + + def make_parser(self): + return DeepSeekV4Parser() + + def make_tool_parser_test_cases(self): + return ToolParserTestCases( + has_tool_call_true= + ('Some text <｜DSML｜tool_calls> <｜DSML｜invoke name="get_weather"> ' + '<｜DSML｜parameter name="location" string="true">NYC ' + " "), + detect_and_parse_single_tool=( + ('Normal text<｜DSML｜tool_calls> <｜DSML｜invoke name="get_weather"> ' + '<｜DSML｜parameter name="location" string="true">NYC ' + " "), + "Normal text", + "get_weather", + { + "location": "NYC" + }, + ), + detect_and_parse_multiple_tools=( + ('<｜DSML｜tool_calls> <｜DSML｜invoke name="get_weather"> ' + '<｜DSML｜parameter name="location" string="true">NYC ' + ' <｜DSML｜invoke name="search_web"> ' + '{ "query": "AI" } '), + ("get_weather", "search_web"), + ), + detect_and_parse_malformed_tool= + ('<|DSML|tool_calls> <|DSML|invoke name="get_weather"> ' + '<|DSML|parameter name="location" string="true">NYC ' + " "), + detect_and_parse_with_parameters_key=( + ('<｜DSML｜tool_calls> <｜DSML｜invoke name="search_web"> ' + '{ "query": "test" } '), + "search_web", + { + "query": "test" + }, + ), + parse_streaming_increment_partial_bot_token="<｜DSML｜tool", + undefined_tool= + ('<｜DSML｜tool_calls> <｜DSML｜invoke name="undefined_func"> ' + '<｜DSML｜parameter name="arg" string="true">value ' + " "), + ) + + # ============================================================================ # Glm4ToolParser Tests # ============================================================================ diff --git a/tests/unittest/llmapi/test_deepseek_v4_tokenizer.py b/tests/unittest/llmapi/test_deepseek_v4_tokenizer.py new file mode 100644 index 000000000000..dec508d49cbc --- /dev/null +++ b/tests/unittest/llmapi/test_deepseek_v4_tokenizer.py @@ -0,0 +1,449 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from tensorrt_llm.inputs.utils import apply_chat_template +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs +from tensorrt_llm.tokenizer.deepseek_v4 import DeepseekV4Tokenizer + + +class _DummyTokenizer: + all_special_tokens = [] + eos_token_id = 1 + pad_token_id = 0 + name_or_path = "dummy" + + def encode(self, text, *args, **kwargs): + self.last_encoded_text = text + return [1, 2, 3] + + +def test_deepseek_v4_chat_template_matches_reference_single_user_prompt(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "Question: 1+1?\nAnswer:", + } + ], + tokenize=False, + add_generation_prompt=True, + ) + + assert prompt == ( + "<｜begin▁of▁sentence｜><｜User｜>Question: 1+1?\nAnswer:<｜Assistant｜>" + ) + + +def test_deepseek_v4_chat_template_tokenize_uses_rendered_prompt(): + dummy = _DummyTokenizer() + tokenizer = DeepseekV4Tokenizer(dummy) + + token_ids = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "hello", + } + ], + tokenize=True, + add_generation_prompt=True, + ) + + assert token_ids == [1, 2, 3] + assert dummy.last_encoded_text == ( + "<｜begin▁of▁sentence｜><｜User｜>hello<｜Assistant｜>" + ) + + +def test_deepseek_v4_chat_template_supports_thinking_mode(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "hello", + } + ], + tokenize=False, + enable_thinking=True, + ) + + assert prompt == ("<｜begin▁of▁sentence｜><｜User｜>hello<｜Assistant｜>") + + +def test_deepseek_v4_chat_template_supports_thinking_alias(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "hello", + } + ], + tokenize=False, + thinking=True, + ) + + assert prompt == ("<｜begin▁of▁sentence｜><｜User｜>hello<｜Assistant｜>") + + +def test_deepseek_v4_chat_template_matches_vllm_add_generation_prompt_behavior(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "hello", + } + ], + tokenize=False, + add_generation_prompt=False, + ) + + assert prompt == ("<｜begin▁of▁sentence｜><｜User｜>hello<｜Assistant｜>") + + +def test_deepseek_v4_chat_template_accepts_openai_reasoning_effort_values(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + for reasoning_effort in ("none", "low", "medium", "high"): + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "hello", + } + ], + tokenize=False, + enable_thinking=True, + reasoning_effort=reasoning_effort, + ) + + assert prompt.endswith("<｜Assistant｜>") + assert "Reasoning Effort: Absolute maximum" not in prompt + + +def test_deepseek_v4_chat_template_preserves_reference_max_reasoning_effort(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "hello", + } + ], + tokenize=False, + enable_thinking=True, + reasoning_effort="max", + ) + + assert prompt.startswith("<｜begin▁of▁sentence｜>Reasoning Effort: Absolute maximum") + assert prompt.endswith("<｜User｜>hello<｜Assistant｜>") + + +def test_deepseek_v4_chat_template_drops_historical_thinking_without_tools(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "first", + }, + { + "role": "assistant", + "reasoning": "hidden chain", + "content": "answer", + }, + { + "role": "user", + "content": "second", + }, + ], + tokenize=False, + enable_thinking=True, + ) + + assert "hidden chain" not in prompt + assert "answer<｜end▁of▁sentence｜>" in prompt + assert prompt.endswith("<｜User｜>second<｜Assistant｜>") + + +def test_deepseek_v4_chat_template_keeps_historical_thinking_with_tools(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + tools = [ + { + "type": "function", + "function": { + "name": "search", + "description": "Search", + "parameters": {"type": "object"}, + }, + } + ] + + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "first", + }, + { + "role": "assistant", + "reasoning": "kept chain", + "content": "answer", + }, + { + "role": "user", + "content": "second", + }, + ], + tools=tools, + tokenize=False, + enable_thinking=True, + ) + + assert "kept chainanswer<｜end▁of▁sentence｜>" in prompt + assert prompt.endswith("<｜User｜>second<｜Assistant｜>") + + +def test_deepseek_v4_chat_template_renders_developer_tools_and_latest_reminder(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + tools = [ + { + "type": "function", + "function": { + "name": "search", + "description": "Search", + "parameters": {"type": "object"}, + }, + } + ] + messages = [ + { + "role": "system", + "content": "sys", + }, + { + "role": "latest_reminder", + "content": "today", + }, + { + "role": "developer", + "content": "dev", + "tools": tools, + }, + { + "role": "assistant", + "reasoning": "need search", + "content": "", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "search", + "arguments": '{"query": "x"}', + }, + } + ], + }, + { + "role": "tool", + "content": "[0]", + }, + ] + + prompt = tokenizer.apply_chat_template( + messages, + tokenize=False, + enable_thinking=True, + ) + + assert prompt.startswith("<｜begin▁of▁sentence｜>sys<｜latest_reminder｜>today<｜User｜>dev") + assert "## Tools" in prompt + assert '<｜DSML｜invoke name="search">' in prompt + assert "need search" in prompt + assert "<｜User｜>[0]<｜Assistant｜>" in prompt + + +def test_deepseek_v4_chat_template_renders_action_task_token(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + prompt = tokenizer.apply_chat_template( + [ + { + "role": "system", + "content": "sys", + }, + { + "role": "latest_reminder", + "content": "today", + }, + { + "role": "user", + "content": "search this", + "task": "action", + }, + { + "role": "assistant", + "content": "Search", + }, + ], + tokenize=False, + ) + + assert prompt == ( + "<｜begin▁of▁sentence｜>sys<｜latest_reminder｜>today" + "<｜User｜>search this<｜Assistant｜><｜action｜>" + "Search<｜end▁of▁sentence｜>" + ) + + +def test_deepseek_v4_chat_template_uses_v4_tool_prompt_from_request_tools(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ] + + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "Weather?", + } + ], + tools=tools, + tokenize=False, + ) + + assert "## Tools" in prompt + assert "<｜DSML｜tool_calls>" in prompt + assert "" in prompt + assert "function_calls" not in prompt + assert '"name": "get_weather"' in prompt + assert prompt.endswith("<｜User｜>Weather?<｜Assistant｜>") + + +def test_deepseek_v4_chat_template_renders_tool_call_history(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + messages = [ + { + "role": "user", + "content": "List the repo", + }, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "str_replace_editor", + "arguments": '{"command": "view", "path": "/testbed"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "file list", + }, + ] + + prompt = tokenizer.apply_chat_template(messages, tokenize=False) + + assert '<｜DSML｜invoke name="str_replace_editor">' in prompt + assert '<｜DSML｜parameter name="command" string="true">view' in prompt + assert '<｜DSML｜parameter name="path" string="true">/testbed' in prompt + assert "<｜User｜>file list<｜Assistant｜>" in prompt + assert 'parameter name="arguments"' not in prompt + + +def test_deepseek_v4_custom_tokenizer_reuses_loaded_wrapper(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + args = TorchLlmArgs(model="dummy", tokenizer=tokenizer, custom_tokenizer="deepseek_v4") + + assert args.tokenizer is tokenizer + + +def test_deepseek_v4_server_chat_template_path_uses_custom_tokenizer(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + prompt = apply_chat_template( + model_type="deepseek_v4", + tokenizer=tokenizer, + processor=None, + conversation=[ + { + "role": "user", + "content": "hello", + } + ], + add_generation_prompt=True, + mm_placeholder_counts=[{}], + ) + + assert prompt == ("<｜begin▁of▁sentence｜><｜User｜>hello<｜Assistant｜>") + + +def test_deepseek_v4_server_chat_template_path_forwards_tools(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + tools = [ + { + "type": "function", + "function": { + "name": "search", + "description": "Search", + "parameters": {"type": "object"}, + }, + } + ] + + prompt = apply_chat_template( + model_type="deepseek_v4", + tokenizer=tokenizer, + processor=None, + conversation=[ + { + "role": "user", + "content": "hello", + } + ], + add_generation_prompt=True, + mm_placeholder_counts=[{}], + tools=tools, + ) + + assert "<｜DSML｜tool_calls>" in prompt + assert '"name": "search"' in prompt diff --git a/tests/unittest/llmapi/test_reasoning_parser.py b/tests/unittest/llmapi/test_reasoning_parser.py index 6806428ec5ea..68543c247a5f 100644 --- a/tests/unittest/llmapi/test_reasoning_parser.py +++ b/tests/unittest/llmapi/test_reasoning_parser.py @@ -60,6 +60,34 @@ def test_deepseek_r1_reasoning_parser_stream(delta_texts: list, content: list, assert result.reasoning_content == reasoning_context[i] +@pytest.mark.parametrize("chat_template_kwargs", [{ + "thinking": True +}, { + "enable_thinking": True +}]) +def test_deepseek_v4_reasoning_parser_extracts_when_thinking( + chat_template_kwargs: dict): + reasoning_parser = ReasoningParserFactory.create_reasoning_parser( + "deepseek_v4", chat_template_kwargs) + + result = reasoning_parser.parse(f"hidden{R1_END}visible") + + assert result.content == "visible" + assert result.reasoning_content == "hidden" + + +def test_deepseek_v4_reasoning_parser_streams_when_thinking(): + reasoning_parser = ReasoningParserFactory.create_reasoning_parser( + "deepseek_v4", {"enable_thinking": True}) + + deltas = ["hid", f"den{R1_END}visible", " tail"] + results = [reasoning_parser.parse_delta(delta) for delta in deltas] + + assert [result.content for result in results] == ["", "visible", " tail"] + assert [result.reasoning_content + for result in results] == ["hid", "den", ""] + + TOOL_START = "<|tool_calls_section_begin|>"