[VALLIS-NERIA]

Description

Extend NVIDIA/TensorRT-LLM#19991 with isolated NVFP4 Kimi K3 targets, fused FP8 decode attention, and independent KV/SSM storage choices.

• Keep both original MXFP4/SM100 target trees byte-for-byte identical to the parent PR (7b203f8c80), and leave the original shared KDA verify kernel unchanged.
• Share the separate NVFP4 modeling, weight-loading, communication and attention implementation between thin SM103 and SM107 targets.
• Retain FP8 projection weights and activations. KDA fuses the f_b projection, convolution, FP32 recurrence and gated output normalization. MLA fuses query normalization/Q-B/absorption and V-B/output gating; Q-B follows the selected GEMM backend's FP8/UE8M0 recipe.
• Support BF16 KV (64-token pages), FP8 KV, and NVFP4 KV (128-token pages), independently of FP32 or BF16 KDA recurrent-state storage. Compose the existing native FP4 MLA layout policy with the hybrid Mamba V2 lifecycle so the complete model can allocate FP4 KV, HP tails and recurrent states together.
• Keep FP32 recurrence and FP32 per-token draft records with either SSM storage dtype. Use an NVFP4-owned verify kernel with explicit FP32 loads and typed state stores.

NVFP4 KV uses the TRTLLM attention backend and V2 cache manager, with chunked prefill disabled. Existing native restrictions on cached context, context parallelism and disaggregated FP4 MLA remain in place.

The head retains the updated parent-PR stack. Supersedes NVIDIA/TensorRT-LLM#20532.

Validation

Validated head: fbde4c91b39c13c54a8736e400751428bb790508.

• Clean SM103 Release C++ SDK and wheel build passed; final Python/Triton sources were tested against that matching SDK (PyTorch 2.14.0a0, CUDA 13.4, Triton 3.8). The public build image needed an official CMake family-architecture validator backport for NVSHMEM's 100f-real setting.
• 397 tests passed with no skips: 126 CPU construction/routing/claims/config checks, 57 GPU precision/kernel cases (including real hybrid FP4 MLA allocation, CUDA Graph replay and BF16 verify/replay), and all 214 existing hybrid Mamba cache-manager cases.
• Full 93-layer modeling_v2 NVFP4 functionality passed on 16 B300 GPUs (SM103, TP16, MoE TP4/EP4, attention DP off): 11 deterministic generation requests submitted in groups of 1, 2 and 4, CUDA Graph on (max batch 4), native NVFP4 KV with 128-token pages and BF16 SSM storage. All 16 ranks verified the target class, actual cache/state dtypes, and eager/captured fused KDA/MLA calls. The functional run used enable_autotuner=False to avoid compiling the complete performance-tuning search space; it retained the 4096-token prefill warmup. This is a functionality check, not a full-dataset accuracy or performance result.
• Applicable pre-commit checks passed, including formatting, lint and test-list validation. All commits carry DCO sign-offs.

Earlier full-model accuracy baseline

The earlier stock-attention code at 5c870f584c completed a clean SM103 Release build, 312 CPU tests, 117 B200 component tests, and full 93-layer inference on 16 B300 GPUs (TP16, MoE TP4/EP4, attention DP off, CUDA Graph on).

Max batch	GSM8K questions	Strict match	Flexible extract
1	1,319	96.3609%	96.3609%
2	1,319	97.0432%	97.0432%
4	1,319	96.3609%	96.4367%

These scores apply to that earlier code, not the fused FP8 attention/FP4 KV/BF16 SSM implementation. Original wrapper jobs exited 1 during summary postprocessing after per-question outputs were saved; an independent audit recovered all 1,319 unique questions per filter and reproduced the scores. The external harness issues were duplicate per-filter counting and NumPy-boolean JSON serialization.

Draft limitations

• Full GSM8K accuracy and performance have not been rerun for the new implementation; the new full-model e2e is a functionality gate.
• SM107 hardware has not been tested.
• The earlier full-model batch=8 CUDA Graph warmup failure has not been confirmed resolved; batch=1/2/4 validation does not establish batch=8 correctness.