[https://nvbugs/6796296][fix] Pad DSv4 NVFP4 compressed FP8 scratch by one FMHA KV tile

tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/backend.py

@@ -43,6 +43,31 @@
 if TYPE_CHECKING:
     from tensorrt_llm.llmapi.llm_args import SparseAttentionConfig
 
+# The trtllm-gen DSv4 dynamic sparse MLA FMHA can read its aux KV pool up to one KV tile
+# (tileSizeKv, 128 rows) past the rows named in the sparse indices. With FP8 KV the aux pool
+# is the whole compressed KV pool, so those reads stay inside it; with NVFP4 KV it is the
+# per-call FP8 scratch below, which is sized to the batch.
+_SPARSE_MLA_FMHA_KV_TILE_SIZE = 128
+
+
+def _alloc_nvfp4_compress_fp8_scratch(
+    num_rows: int, head_dim: int, device: torch.device
+) -> torch.Tensor:
+    """Allocate the FP8 scratch the sparse MLA FMHA reads as its aux KV pool.
+
+    Rows [0, num_rows) receive the dequantized compressed KV. A zeroed tail of one FMHA KV tile
+    keeps the kernel's reads past the last indexed row inside the allocation; without it the
+    FMHA faults whenever the caching allocator places the scratch at the end of a mapped
+    segment (https://nvbugs/6796296).
+    """
+    scratch = torch.empty(
+        (num_rows + _SPARSE_MLA_FMHA_KV_TILE_SIZE, head_dim),
+        dtype=torch.float8_e4m3fn,
+        device=device,
+    )
+    scratch[num_rows:].zero_()
+    return scratch
+
 
 class DeepseekV4TrtllmAttention(TrtllmAttention):
     Metadata = DeepseekV4TrtllmAttentionMetadata
@@ -399,16 +424,15 @@ def sparse_attn_predict(
             # the caching allocator can reuse the storage in stream order.
             metadata.nvfp4_compress_fp8_scratch = None
             if attention_input_type == AttentionInputType.generation_only:
-                scratch = torch.empty(
-                    (*compressed_indices.shape, self.head_dim),
-                    dtype=torch.float8_e4m3fn,
-                    device=global_indices.device,
+                num_scratch_rows = compressed_indices.numel()
+                scratch = _alloc_nvfp4_compress_fp8_scratch(
+                    num_scratch_rows, self.head_dim, global_indices.device
                 )
                 torch.ops.trtllm.nvfp4_mla_kv_cache_gather_direct(
                     data_pool,
                     scale_pool,
                     compressed_indices,
-                    scratch,
+                    scratch[:num_scratch_rows].view(*compressed_indices.shape, self.head_dim),
                     self._nvfp4_compress_scale_quant_orig,
                     kv_cache_manager.nvfp4_residual_dim,
                     num_pool_tokens,
@@ -437,10 +461,8 @@ def sparse_attn_predict(
                     compressed_indices.shape[1],
                     min(max_compressed_kv_tokens, compressed_indices.numel()),
                 )
-                scratch = torch.empty(
-                    (scratch_capacity, 1, self.head_dim),
-                    dtype=torch.float8_e4m3fn,
-                    device=global_indices.device,
+                scratch = _alloc_nvfp4_compress_fp8_scratch(
+                    scratch_capacity, self.head_dim, global_indices.device
                 )
                 if max_compressed_kv_tokens > 0:
                     torch.ops.trtllm.nvfp4_mla_context_kv_cache_gather_direct(
@@ -450,7 +472,7 @@ def sparse_attn_predict(
                         req_id.contiguous(),
                         compressed_kv_lengths,
                         compressed_indices,
-                        scratch,
+                        scratch[:scratch_capacity].unsqueeze(1),
                         self._nvfp4_compress_scale_quant_orig,
                         kv_cache_manager.nvfp4_residual_dim,
                         max_compressed_kv_tokens,

tests/integration/test_lists/waives.txt

@@ -6,7 +6,6 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp2pp2
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6198774)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-enable_chunked_prefill=True] SKIP (https://nvbugs/6388139)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-enable_chunked_prefill=True] SKIP (https://nvbugs/6507095)
-accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_nvfp4_kv SKIP (https://nvbugs/6796296)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Pro::test_gsm8k_full_accuracy SKIP (https://nvbugs/6939213)
 accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[cutlass-one_model-overlap_scheduler] SKIP (https://nvbugs/6762288)
 accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[triton-one_model-overlap_scheduler] SKIP (https://nvbugs/6762288)