[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))
 
@@ -2073,16 +2076,22 @@ def _fused_ob_gemm(self, o_lora_fp8: torch.Tensor, sf_out: torch.Tensor,
         SK = _select_dsv4_ob_split_k(M, getattr(self, "ob_split_k", None))
         packed_k = K_ob // _DSV4_OB_PACKED_SCALE_K
         # CuTe specializes the DSV4-Pro O_b shape.
-        can_try_cute = (IS_CUTLASS_DSL_AVAILABLE and M > 0
-                        and (M >= _DSV4_OB_CUTE_SK1_MIN_M or SK > 1)
-                        and N == _DSV4_PRO_OB_N and K_ob == _DSV4_PRO_OB_K
-                        and K_ob % (_DSV4_OB_PACKED_SCALE_K * SK) == 0
-                        and sf_out.dtype == torch.int32 and sf_out.dim() == 2
-                        and sf_out.shape == (M, packed_k)
-                        and sf_out.stride(0) == 1 and sf_out.stride(1) >= M
-                        and sf_out.stride(1) % _UE8M0_SCALES_PER_WORD == 0)
+        use_cute_tactic = (IS_CUTLASS_DSL_AVAILABLE and M > 0
+                           and (M >= _DSV4_OB_CUTE_SK1_MIN_M or SK > 1)
+                           and N == _DSV4_PRO_OB_N and K_ob == _DSV4_PRO_OB_K
+                           and K_ob % (_DSV4_OB_PACKED_SCALE_K * SK) == 0)
+        activation_layout_ok = (sf_out.dtype == torch.int32
+                                and sf_out.dim() == 2
+                                and sf_out.shape == (M, packed_k)
+                                and sf_out.stride(0) == 1
+                                and sf_out.stride(1) >= M and
+                                sf_out.stride(1) % _UE8M0_SCALES_PER_WORD == 0)
+        if use_cute_tactic and not activation_layout_ok:
+            raise RuntimeError(
+                "DSV4-Pro O_b CuTe tactic requires packed activation scales.")
+
         cute_wsf = wsf
-        if can_try_cute and cute_wsf.dtype != torch.int32:
+        if use_cute_tactic and cute_wsf.dtype != torch.int32:
             from tensorrt_llm.quantization.utils.fp8_utils import \
                 transform_sf_into_required_layout
             cute_wsf = getattr(self.o_b_proj, "_ob_wsf_int", None)
@@ -2100,12 +2109,17 @@ def _fused_ob_gemm(self, o_lora_fp8: torch.Tensor, sf_out: torch.Tensor,
                     is_sfa=False)
                 self.o_b_proj._ob_wsf_int = cute_wsf
 
-        can_use_cute = (can_try_cute and cute_wsf.dtype == torch.int32
-                        and cute_wsf.dim() == 2
-                        and cute_wsf.shape == (N, packed_k)
-                        and cute_wsf.stride(0) == 1 and cute_wsf.stride(1) == N)
-        if not can_use_cute:
-            # Keep DeepGEMM for unsupported shapes and layouts.
+        weight_layout_ok = (cute_wsf.dtype == torch.int32
+                            and cute_wsf.dim() == 2
+                            and cute_wsf.shape == (N, packed_k)
+                            and cute_wsf.stride(0) == 1
+                            and cute_wsf.stride(1) == N)
+        if use_cute_tactic and not weight_layout_ok:
+            raise RuntimeError(
+                "DSV4-Pro O_b CuTe tactic requires packed weight scales.")
+
+        if not use_cute_tactic:
+            # Keep DeepGEMM for shapes outside the tuned CuTe buckets.
             hidden = torch.empty([M, N],
                                  device=o_lora_fp8.device,
                                  dtype=self.dtype)

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

@@ -368,7 +368,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 +622,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 +643,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)