[None][fix] Enable DSV4 fused O-projection by default

Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>

Author: mingyangHao

tensorrt_llm/_torch/modules/attention.py

@@ -2033,9 +2033,12 @@ def _validate_dsv4_epilogue_buffers(
         return fp8_o, output_sf
 
     def _should_use_fused_oproj(self) -> bool:
-        return (os.environ.get("DSV4_FUSE_OPROJ", "0") == "1"
+        # Eligible Blackwell paths are on by default; keep an emergency fallback.
+        return (not _is_env_truthy("TRTLLM_DSV4_DISABLE_FUSED_OPROJ")
+                and is_sm_100f() and IS_CUTLASS_DSL_AVAILABLE
                 and self.n_local_groups == self.num_groups
                 and getattr(self.o_b_proj, "tp_size", 1) == 1
+                and self.o_a_proj.dtype == torch.float8_e4m3fn
                 and self.o_b_proj.has_fp8_block_scales and not getattr(
                     self.o_b_proj, "use_cute_dsl_blockscaling_mm", False))
 

tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py

@@ -44,6 +44,27 @@ def test_select_dsv4_ob_split_k_overrides(monkeypatch):
         _select_dsv4_ob_split_k(32, configured_split=3)
 
 
+def test_dsv4_fused_oproj_is_enabled_by_default(monkeypatch):
+    module = SimpleNamespace(
+        n_local_groups=16,
+        num_groups=16,
+        o_a_proj=torch.empty((), device="meta", dtype=torch.float8_e4m3fn),
+        o_b_proj=SimpleNamespace(
+            tp_size=1,
+            has_fp8_block_scales=True,
+            use_cute_dsl_blockscaling_mm=False,
+        ),
+    )
+    monkeypatch.setattr(attention_module, "IS_CUTLASS_DSL_AVAILABLE", True)
+    monkeypatch.setattr(attention_module, "is_sm_100f", lambda: True)
+    monkeypatch.delenv("TRTLLM_DSV4_DISABLE_FUSED_OPROJ", raising=False)
+
+    assert MLA._should_use_fused_oproj(module)
+
+    monkeypatch.setenv("TRTLLM_DSV4_DISABLE_FUSED_OPROJ", "1")
+    assert not MLA._should_use_fused_oproj(module)
+
+
 def test_dsv4_fmha_epilogue_output_uses_fused_oproj():
     attn_fp8 = torch.empty((16, 4, 4096), device="meta", dtype=torch.float8_e4m3fn)
     attn_scale = torch.empty((16, 32, 4), device="meta")
@@ -368,7 +389,7 @@ def calculate_reference_deepseek_v4_o_proj(
 
 def _build_dsv4_o_proj_case(num_tokens: int, dtype_str: str, device: torch.device):
     """Build an MLA module, inputs, and reference-path tensors for the DeepSeek-V4
-    o_proj tests. Shared by the correctness test and the DSV4_FUSE_OPROJ fused
+    o_proj tests. Shared by the correctness test and the default fused
     fp8-equivalence test so both exercise an identical setup.
 
     Returns:
@@ -622,18 +643,16 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str):
 @pytest.mark.skip_less_device_memory(80000)
 @pytest.mark.parametrize("num_tokens", [1, 16, 32, 128, 256])
 def test_deepseek_v4_o_proj_fused_fp8_equivalence(num_tokens: int, monkeypatch):
-    """The opt-in DSV4_FUSE_OPROJ fused fp8 epilogue must be numerically
-    equivalent to the default unfused path.
+    """The default fused FP8 epilogue must match the disabled baseline.
 
-    Fused (DSV4_FUSE_OPROJ=1): o_a's CuTe-DSL GEMM emits o_lora directly as fp8
+    Fused (default): o_a's CuTe-DSL GEMM emits o_lora directly as fp8
     e4m3 + packed-UE8M0 1x128 scale factors, fed straight to o_b.
-    Unfused (default): o_a emits bf16 o_lora, then ``o_b_proj`` runs the separate
-    1x128 quant + DeepGEMM. Since the fusion only folds the *same* quant into o_a's
-    epilogue, the two production paths must match far tighter than the fp8-vs-bf16
-    reference bar.
+    Disabled: o_a emits bf16 o_lora, then ``o_b_proj`` runs the separate 1x128
+    quant + DeepGEMM. Since the fusion only folds the *same* quant into o_a's
+    epilogue, the two paths must match far tighter than the fp8-vs-bf16 reference.
     """
     if get_sm_version() < 100:
-        pytest.skip("DSV4_FUSE_OPROJ fp8 fusion requires Blackwell (SM100+)")
+        pytest.skip("fused DeepSeek-V4 FP8 O-projection requires Blackwell (SM100+)")
 
     device = torch.device("cuda")
     mla, attn_out_latent, position_ids, refs = _build_dsv4_o_proj_case(num_tokens, "fp8", device)
@@ -645,12 +664,12 @@ def test_deepseek_v4_o_proj_fused_fp8_equivalence(num_tokens: int, monkeypatch):
     assert mla.o_b_proj.has_fp8_block_scales
     assert not getattr(mla.o_b_proj, "use_cute_dsl_blockscaling_mm", False)
 
-    # Unfused (default): DSV4_FUSE_OPROJ unset.
-    monkeypatch.delenv("DSV4_FUSE_OPROJ", raising=False)
+    # Explicit kill switch retains the unfused fallback.
+    monkeypatch.setenv("TRTLLM_DSV4_DISABLE_FUSED_OPROJ", "1")
     out_unfused = mla._deepseek_v4_o_proj(attn_out_latent.clone(), position_ids)
 
-    # Fused: opt in.
-    monkeypatch.setenv("DSV4_FUSE_OPROJ", "1")
+    # Unset is the production default and must take the fused path.
+    monkeypatch.delenv("TRTLLM_DSV4_DISABLE_FUSED_OPROJ", raising=False)
     mla.ob_split_k = 1
     out_fused = mla._deepseek_v4_o_proj(attn_out_latent.clone(), position_ids)