[None][fix] Qwen3.8-Flash-Next: mask FP8 PLE rows by bytes under attention DP

tensorrt_llm/_torch/modules/qwen4_exp/ple.py

@@ -930,7 +930,13 @@ def embed(
             torch.zeros_like(lookup_ids),
         )
         partial = F.embedding(local_ids, self.ngram_embedding.weight)
-        partial.masked_fill_(~owned.unsqueeze(-1), 0)
+        # FP8 rows stay quantized until _finish_embedding_lookup, and FP8 has no
+        # masked_fill_ kernel. An all-zero byte is +0.0 in both FP8 formats, so
+        # clearing the raw bytes matches masking after dequantization.
+        if partial.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+            partial.view(torch.uint8).masked_fill_(~owned.unsqueeze(-1), 0)
+        else:
+            partial.masked_fill_(~owned.unsqueeze(-1), 0)
         return self._finish_embedding_lookup(
             partial,
             semantic_tokens,

tests/integration/defs/.test_durations

@@ -505,7 +505,7 @@
    "accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_fp8_adp4_mtp3_prefix_cache": 904.0571428571428,
    "accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_fp8_adp4_mtp3_trtllm_ple_offload": 690.0651202531645,
    "accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_megamoe_cutedsl_ple_offload": 427.0257848101266,
-   "accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_mtp3_trtllm_ple_offload": 438.4012215189873,
+   "accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_mtp3_trtllm": 438.4012215189873,
    "accuracy/test_llm_api_pytorch.py::TestSeedOss_36B::test_auto_dtype": 552.3427058823529,
    "accuracy/test_llm_api_pytorch.py::TestStep3p7SpeculativeDecoding::test_step3p7_mtp": 325.6583,
    "accuracy/test_llm_api_pytorch_encode.py::TestDecoderEncode::test_decoder_encode_cuda_graph_matches_eager_logits[tinyllama-1.1b]": 15.013406576980568,

tests/integration/defs/accuracy/references/acceptance_length.yaml

@@ -68,7 +68,7 @@ TestQwen3_5_35B_A3B::test_bf16_mtp_flashinfer_ring_replay:
 TestQwen3_8_Flash_Next::test_fp8_adp4_mtp3_trtllm_ple_offload:
   ref_al: 1.890148
   min_al: 1.7956406
-TestQwen3_8_Flash_Next::test_nvfp4_adp4_mtp3_trtllm_ple_offload:
+TestQwen3_8_Flash_Next::test_nvfp4_adp4_mtp3_trtllm:
   ref_al: 3.297978021978022
   min_al: 3.133079120879121
 # The usual 95% tolerance for these low values is below the no-acceptance

tests/integration/defs/accuracy/test_llm_api_pytorch.py

@@ -6085,11 +6085,13 @@ def _run_evals(
             enable_block_reuse: bool = False,
             sparse_attention_config: Optional[QSASparseAttentionConfig] = None,
             acceptance_length_key: Optional[str] = None,
-            kv_cache_dtype: str = "auto") -> None:
+            kv_cache_dtype: str = "auto",
+            ple_host_offload: bool = True) -> None:
         if not os.path.exists(model_path):
             pytest.skip(f"Model directory {model_path} does not exist")
 
-        monkeypatch.setenv("TRTLLM_QWEN4_EXP_PLE_HOST_OFFLOAD", "1")
+        monkeypatch.setenv("TRTLLM_QWEN4_EXP_PLE_HOST_OFFLOAD",
+                           "1" if ple_host_offload else "0")
 
         # A case opts into the acceptance-length guard by naming its own entry
         # in references/acceptance_length.yaml.
@@ -6185,29 +6187,32 @@ def test_fp8_adp4_mtp3_prefix_cache(self, monkeypatch: pytest.MonkeyPatch,
     @skip_pre_blackwell
     @pytest.mark.skip_less_device(4)
     @pytest.mark.skip_less_device_memory(70000)
-    @pytest.mark.skip_less_host_memory(131072)
-    def test_nvfp4_adp4_mtp3_trtllm_ple_offload(self,
-                                                monkeypatch: pytest.MonkeyPatch,
-                                                mocker) -> None:
-        """NVFP4 on four GPUs with attention DP, MTP3 and PLE offload.
+    def test_nvfp4_adp4_mtp3_trtllm(self, monkeypatch: pytest.MonkeyPatch,
+                                    mocker) -> None:
+        """NVFP4 on four GPUs with attention DP, MTP3 and the PLE table on GPU.
 
         Blackwell-only, so the GVR Top-K prerequisites hold. This is the case
         that keeps both indexer options under accuracy guard: the heuristic
         Top-K, and the draft loop reusing one selection.
+
+        The checkpoint stores the PLE n-gram table as scaled FP8. Under
+        attention DP each rank owns a row slice of it and masks the rows it
+        does not own before the reduce-scatter, while the rows are still FP8.
         """
-        self._run_evals(
-            f"{llm_models_root()}/Qwen3.8-Flash-Next-NVFP4",
-            tensor_parallel_size=4,
-            moe_backend="TRTLLM",
-            max_draft_len=3,
-            expected_quant_algo=QuantAlgo.MIXED_PRECISION,
-            monkeypatch=monkeypatch,
-            mocker=mocker,
-            moe_expert_parallel_size=4,
-            enable_attention_dp=True,
-            sparse_attention_config=QSASparseAttentionConfig(
-                enable_heuristic_topk=True, index_share_for_mtp_iteration=True),
-            acceptance_length_key="test_nvfp4_adp4_mtp3_trtllm_ple_offload")
+        self._run_evals(f"{llm_models_root()}/Qwen3.8-Flash-Next-NVFP4",
+                        tensor_parallel_size=4,
+                        moe_backend="TRTLLM",
+                        max_draft_len=3,
+                        expected_quant_algo=QuantAlgo.MIXED_PRECISION,
+                        monkeypatch=monkeypatch,
+                        mocker=mocker,
+                        moe_expert_parallel_size=4,
+                        enable_attention_dp=True,
+                        sparse_attention_config=QSASparseAttentionConfig(
+                            enable_heuristic_topk=True,
+                            index_share_for_mtp_iteration=True),
+                        acceptance_length_key="test_nvfp4_adp4_mtp3_trtllm",
+                        ple_host_offload=False)
 
     @skip_pre_blackwell
     @pytest.mark.skip_less_device_memory(100000)

tests/integration/test_lists/qa/llm_function_core.txt

@@ -621,7 +621,7 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_bf16_tep4_cutlass
 accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_fp8_adp4_mtp3_trtllm_ple_offload
 accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_fp8_adp4_mtp3_prefix_cache
 accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_1gpu_cutedsl_ple_offload
-accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_mtp3_trtllm_ple_offload
+accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_mtp3_trtllm
 accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_megamoe_cutedsl_ple_offload
 accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_nvfp4kv_trtllm_ple_offload
 accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_tep4_fp8kv_trtllm_ple_offload

tests/integration/test_lists/test-db/l0_gb300.yml

@@ -63,7 +63,7 @@ l0_gb300:
   - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v2_kv_cache-True-True-cutlass-auto]
   - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_fp8_block_scales[latency] # Cover nvbugs 5461712 and 5505402
   - accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_fp8_adp4_mtp3_trtllm_ple_offload
-  - accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_mtp3_trtllm_ple_offload
+  - accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_mtp3_trtllm
   - accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_megamoe_cutedsl_ple_offload
   - accuracy/test_disaggregated_serving.py::TestQwen3_8_Flash_Next::test_fp8_nixl_python[prefix_cache]
   - unittest/_torch/moe/kernels/parallel

tests/unittest/_torch/modules/test_qwen4_exp_ple.py

@@ -765,3 +765,85 @@ def fake_reducescatter(partial, actual_mapping, dim, sizes):
     )
 
     torch.testing.assert_close(output, full_weight[local_ids])
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="FP8 masking is checked on CUDA")
+def test_ple_attention_dp_row_shard_masks_fp8_table(monkeypatch: pytest.MonkeyPatch) -> None:
+    from tensorrt_llm._torch.modules.qwen4_exp import ple as qwen4_exp_ple
+
+    # Host offload would build a BF16 pinned table instead of the resident one under test.
+    monkeypatch.setenv("TRTLLM_QWEN4_EXP_PLE_HOST_OFFLOAD", "0")
+    config = _make_config()
+    mapping = SimpleNamespace(tp_size=2, tp_rank=0, cp_size=1, enable_attention_dp=True)
+    module = Qwen4ExpNGramEmbedding(
+        config,
+        embedding_dim=32,
+        dtype=torch.bfloat16,
+        mapping=mapping,
+    )
+    module.configure_fp8_weight_storage(torch.tensor(0.5), torch.float8_e4m3fn)
+    module = module.cuda()
+    # Small integers are exact in e4m3, so the dequantized rows are exact too.
+    full_weight = (
+        torch.arange(
+            module.padded_vocab_size * module.head_dim_per_ngram,
+            dtype=torch.float32,
+            device="cuda",
+        ).remainder(15)
+        - 7
+    ).reshape(module.padded_vocab_size, module.head_dim_per_ngram)
+    with torch.no_grad():
+        module.ngram_embedding.weight.copy_(
+            full_weight[module.vocab_start_index : module.vocab_end_index].to(torch.float8_e4m3fn)
+        )
+
+    local_ids = torch.tensor(
+        [
+            [1, module.vocab_end_index + 1] * (NGRAM_HEADS // 2),
+            [module.vocab_end_index - 1, module.padded_vocab_size - 1] * (NGRAM_HEADS // 2),
+        ],
+        dtype=torch.long,
+        device="cuda",
+    )
+    remote_ids = torch.tensor(
+        [[module.vocab_end_index + 3, 7] * (NGRAM_HEADS // 2)],
+        dtype=torch.long,
+        device="cuda",
+    )
+    gathered_ids = torch.cat((local_ids, torch.zeros_like(local_ids[:1]), remote_ids))
+    owned = gathered_ids < module.vocab_end_index
+    expected_partial = torch.zeros(
+        (*gathered_ids.shape, module.head_dim_per_ngram),
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    expected_partial[owned] = (full_weight[gathered_ids[owned]] * 0.5).to(torch.bfloat16)
+
+    def fake_allgather(
+        input_ids: torch.Tensor,
+        actual_mapping: object,
+        dim: int,
+        sizes: list[int] | None,
+    ) -> torch.Tensor:
+        return gathered_ids
+
+    def fake_reducescatter(
+        partial: torch.Tensor,
+        actual_mapping: object,
+        dim: int,
+        sizes: list[int] | None,
+    ) -> torch.Tensor:
+        assert partial.dtype == torch.bfloat16
+        torch.testing.assert_close(partial, expected_partial, rtol=0, atol=0)
+        return partial[:3]
+
+    monkeypatch.setattr(qwen4_exp_ple, "allgather", fake_allgather)
+    monkeypatch.setattr(qwen4_exp_ple, "reducescatter", fake_reducescatter)
+
+    output = module.embed(
+        local_ids,
+        physical_tokens=3,
+        all_rank_num_tokens=[3, 1],
+    )
+
+    torch.testing.assert_close(output, expected_partial[:2], rtol=0, atol=0)