[TRTLLM-16621][feat] Support DeepSeek-V4.1-Flash

[ziyixiong-nv]

Description

Add DeepSeek-V4.1 inference support to the PyTorch backend for:

• DeepSeek-V4.1-Flash
• DeepSeek-V4.1-Flash-NVFP4

The implementation includes CSA2 sparse attention, Engram, mHC, model-specific checkpoint loading, tokenizer/tool parsing, and registration of text-decoder and conditional-generation architectures.

Key Features

• Encoder and decoder bounded replay: optional encoder-state recovery retains reusable GLOBAL prefixes; bounded decoder replay preserves complete GLOBAL/index KV while restricting decoder prefill to a prompt suffix. The decoder path is approximate, enabled by default, and can be disabled for full decoder prefill.
• DSpark speculative decoding: supports the embedded draft model, captured target states, and acceptance-length tracking.
• Conditional disaggregation: eligible requests with a short uncached suffix can bypass context and run on the selected generation worker, using that worker's prefix-cache match. Experimental generation-side bounded replay is separately opt-in, subject to the documented configuration constraints.
• DeepSeek-V4 FP4 indexer scale addressing: indexer slot maps use the index cache's 128-channel quantization blocks, independently of the main FP4 cache's 32-channel blocks. This corrects the indexer scale stride from 16 to 4 bytes per row; V3.2 and V4 FP8 addressing are unchanged. Existing deepseek_v4/test_indexer_slot_mapping.py coverage includes V4 FP4 context, eager decode and fused decode against the physical cache payload and scale bytes.

Validation

GSM8K validation uses the NVFP4 checkpoint on 4 NVIDIA GB300 GPUs with TP4/EP4, attention DP, CUTEDSL MoE, and DSpark (max_draft_len=5, block_size=5).

GSM8K: 96.36% (1,271/1,319), with DSpark acceptance length 4.0000. Full evaluation uses batch size 32, 32-shot non-thinking prompts, greedy decoding, and a 1,024-token output limit.

PR Checklist

• Reviewed the PR checklist items as applicable to this PR. Required CODEOWNER approval remains pending.

Dev Engineer Review

The available diff changes only import formatting in tests/unittest/_torch/sampler/test_beam_search.py. It does not implement the DeepSeek-V4.1 features described in the PR objectives. No source or configuration behavior changes are present in the inspected diff.

QA Engineer Review

One test file changed: tests/unittest/_torch/sampler/test_beam_search.py. The changes reformat imports and local import statements; no test functions, selectors, or waivers changed. No integration tests changed, so no test-list registration applies. Coverage verdict: sufficient for the formatting-only changes. Test execution results are unavailable.

Per-File QA Perspective

• tests/unittest/_torch/sampler/test_beam_search.py — Existing beam-search coverage is unchanged; this diff only reformats imports. The file is not an integration test and does not require listing in the CI or manual-QA integration test lists.

[ziyixiong-nv]

/bot run

[tensorrt-cicd]

PR_Github #75832 [ run ] triggered by Bot. Commit: 329ce16 Link to invocation

[tensorrt-cicd]

PR_Github #75832 [ run ] completed with state SUCCESS. Commit: 329ce16
/LLM/main/L0_MergeRequest_PR pipeline #62523 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[tensorrt-cicd]

PR_Github #77137 [ run ] completed with state SUCCESS. Commit: e8a9ebe
/LLM/main/L0_MergeRequest_PR pipeline #63593 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

[ziyixiong-nv]

/bot reuse-pipeline

[tensorrt-cicd]

PR_Github #77144 [ reuse-pipeline ] triggered by Bot. Commit: 497e281 Link to invocation

[tensorrt-cicd]

PR_Github #77144 [ reuse-pipeline ] completed with state SUCCESS. Commit: 497e281
Reusing PR_Github #77137 for commit 497e281

Link to invocation