FROST SDPA fp8: block-scaled O epilogues (NVFP4 / MXFP8) with an sf_o output

benchmark/attention_training/README.md

@@ -209,6 +209,17 @@ python benchmark_single_sdpa.py \
     --sdpa_backend cudnn --data_type fp8 \
     --attn_mask top_left --fwd_bwd
 
+# cuDNN OSS FROST engines, FP8 inputs with a block-scaled O epilogue
+# (forward-only inference): `fp8_nvfp4` writes O as FP4 (E2M1, two per byte)
+# with one E4M3 scale per 16 head-dim elements, `fp8_mxfp8` writes E4M3 O with
+# one UE8M0 scale per 32, both into the extra `sf_o` output of `sdpa_fp8` in
+# the F8_128x4 atom order a block-scaled GEMM consumes directly. d = 128 only.
+python benchmark_single_sdpa.py \
+    --batch_size 1 --q_seqlen 8192 --kv_seqlen 8192 \
+    --num_q_heads 64 --num_kv_heads 8 --head_dim 128 \
+    --sdpa_backend cudnn_oss --data_type fp8_nvfp4 \
+    --attn_mask no_mask --profile_pass fwd
+
 # FlashAttention 4
 python benchmark_single_sdpa.py \
     --batch_size 1 --q_seqlen 8192 --kv_seqlen 8192 \

benchmark/attention_training/benchmark_single_sdpa.py

@@ -101,8 +101,10 @@ class UnsupportedConfigError(RuntimeError):
     # (No fp8/mxfp8 entries: Ampere has neither datapath.)
     8: {"bfloat16": 2048, "float16": 2048},
     9: {"bfloat16": 4096, "float16": 4096, "fp8": 8192},
-    10: {"bfloat16": 8192, "float16": 8192, "fp8": 16384, "mxfp8": 16384},
-    12: {"bfloat16": 1024, "float16": 1024, "fp8": 2048, "mxfp8": 2048},
+    # fp8_nvfp4 / fp8_mxfp8: fp8 inputs with a block-scaled O epilogue -- the
+    # MMAs are the fp8 ones, so the peak is the fp8 peak.
+    10: {"bfloat16": 8192, "float16": 8192, "fp8": 16384, "mxfp8": 16384, "fp8_nvfp4": 16384, "fp8_mxfp8": 16384},
+    12: {"bfloat16": 1024, "float16": 1024, "fp8": 2048, "mxfp8": 2048, "fp8_nvfp4": 2048, "fp8_mxfp8": 2048},
 }
 
 
@@ -271,8 +273,12 @@ def parse_args():
     parser.add_argument(
         "--data_type",
         default="bfloat16",
-        choices=["bfloat16", "float16", "float", "fp8", "mxfp8"],
-        help="Data type. Can be bfloat16, float16, fp8, or mxfp8",
+        choices=["bfloat16", "float16", "float", "fp8", "mxfp8", "fp8_nvfp4", "fp8_mxfp8"],
+        help=(
+            "Data type: bfloat16, float16, float, fp8, mxfp8, or the fp8-input block-scaled-output modes "
+            "fp8_nvfp4 (FP4 O + E4M3 scale per 16 d) and fp8_mxfp8 (E4M3 O + UE8M0 scale per 32 d); "
+            "those two are forward-only and served by the cudnn_oss backend"
+        ),
     )
     parser.add_argument(
         "--num_iterations",
@@ -582,6 +588,15 @@ def exit_unsupported(reason):
         print(f"UNSUPPORTED_CONFIG: {reason}")
         sys.exit(UNSUPPORTED_CONFIG_RETURN_CODE)
 
+    # Block-scaled O modes ride the fp8 graph (fp8 Q/K/V, sdpa_fp8) with an
+    # extra sf_o output: 16 = FP4 O + E4M3 scale per 16 d, 32 = E4M3 O + UE8M0
+    # scale per 32 d. Forward-only, cudnn_oss only.
+    o_block_scale = {"fp8_nvfp4": 16, "fp8_mxfp8": 32}.get(args.data_type, 0)
+    if o_block_scale:
+        if args.sdpa_backend != "cudnn_oss":
+            exit_unsupported(f"{args.data_type} (block-scaled O) is served by the cudnn_oss backend only")
+        args.data_type = "fp8"
+
     if args.data_type == "bfloat16":
         target_dtype = torch.bfloat16
     elif args.data_type == "float16":
@@ -626,6 +641,10 @@ def exit_unsupported(reason):
     else:
         run_fwd = True
         run_bwd = False
+    if o_block_scale and run_bwd:
+        exit_unsupported("block-scaled O (fp8_nvfp4 / fp8_mxfp8) is a forward-only inference epilogue")
+    if o_block_scale and (head_dim_qk != 128 or head_dim_vo != 128):
+        exit_unsupported("block-scaled O is served by the d128 fp8 kernel flavor only")
     enable_gqa = num_q_heads != num_kv_heads
     if args.data_type == "mxfp8":
         if not is_cudnn_fe:
@@ -684,8 +703,21 @@ def convert_to_cudnn_type(torch_type):
         value = torch.randn(batch_size, kv_seqlen, num_kv_heads, head_dim_vo, dtype=randn_dtype, device=device).to(target_dtype).transpose(1, 2)
         if args.data_type == "mxfp8":
             output = torch.empty(batch_size, q_seqlen, num_q_heads, head_dim_vo, dtype=output_dtype, device=device).transpose(1, 2)
+        elif o_block_scale == 16:
+            # FP4 O: the packed byte container (two E2M1 per byte).
+            output = (
+                torch.empty(batch_size, q_seqlen, num_q_heads, head_dim_vo // 2, dtype=torch.uint8, device=device).view(torch.float4_e2m1fn_x2).transpose(1, 2)
+            )
         else:
             output = torch.empty(batch_size, q_seqlen, num_q_heads, head_dim_vo, dtype=target_dtype, device=device).transpose(1, 2)
+        sf_o_gpu = None
+        if o_block_scale:
+            # Per-(b,h) planes of block scales in the F8_128x4 atom order:
+            # [B, H, S rounded to 128, d/block rounded to 4]; zero-initialized
+            # (the epilogue owns the rows it writes, the caller the rest).
+            sf_rows = -(-q_seqlen // 128) * 128
+            sf_cols = max(4, -(-(head_dim_vo // o_block_scale) // 4) * 4)
+            sf_o_gpu = torch.zeros(batch_size, num_q_heads, sf_rows, sf_cols, dtype=torch.uint8, device=device)
 
         # FP8-specific descale/scale/amax tensors
         if args.data_type == "fp8":
@@ -786,6 +818,9 @@ def convert_to_cudnn_type(torch_type):
             scale_s_fwd = graph_fwd.tensor_like(scale_s_gpu)
             scale_o_fwd = graph_fwd.tensor_like(scale_o_gpu)
 
+            sf_o_fwd = None
+            if o_block_scale:
+                sf_o_fwd = graph_fwd.tensor_like(sf_o_gpu).set_data_type(cudnn.data_type.FP8_E4M3 if o_block_scale == 16 else cudnn.data_type.FP8_E8M0)
             # Amax_S is returned unconditionally but never requested: it only becomes a
             # real graph output if set_output(True) is called on it, and nothing here
             # consumes it. Leaving it virtual also keeps the graph servable by engines
@@ -807,6 +842,7 @@ def convert_to_cudnn_type(torch_type):
                 left_bound=left_bound,
                 right_bound=right_bound,
                 # dropout=dropout_tuple if is_dropout else None,
+                **({"sf_o": sf_o_fwd} if o_block_scale else {}),
             )
         elif args.data_type == "mxfp8":
             q_fwd = graph_fwd.tensor_like(query)
@@ -879,7 +915,16 @@ def convert_to_cudnn_type(torch_type):
             # unmarked tensor, which the FROST engines' variant-pack binding
             # does not do). The stats buffer is already in every variant pack.
             if args.data_type == "fp8":
-                o_fwd.set_output(True).set_dim(output.size()).set_stride(output.stride()).set_data_type(cudnn.data_type.FP8_E4M3)
+                if o_block_scale == 16:
+                    # FP4 O: logical (B, H, S, d) dims over the packed byte container.
+                    o_dims = list(output.size())
+                    o_dims[-1] *= 2
+                    o_strides = [s * 2 if s != 1 else 1 for s in output.stride()]
+                    o_fwd.set_output(True).set_dim(o_dims).set_stride(o_strides).set_data_type(cudnn.data_type.FP4_E2M1)
+                else:
+                    o_fwd.set_output(True).set_dim(output.size()).set_stride(output.stride()).set_data_type(cudnn.data_type.FP8_E4M3)
+                if o_block_scale:
+                    sf_o_fwd.set_output(True)
                 (stats_fwd.set_output(True).set_dim(stats.size()).set_stride(stats.stride()).set_data_type(cudnn.data_type.FLOAT) if not is_infer else None)
             elif args.data_type == "mxfp8":
                 o_fwd.set_output(True).set_dim(output.size()).set_stride(output.stride()).set_data_type(convert_to_cudnn_type(output_dtype))
@@ -1329,6 +1374,8 @@ def pin_frost_engine_or_skip(graph, pass_name, required):
                     scale_o_fwd: scale_o_gpu,
                     amax_o_fwd: amax_o_gpu,
                 }
+                if o_block_scale:
+                    variant_pack_fwd[sf_o_fwd] = sf_o_gpu
                 workspace = torch.empty(graph_fwd.get_workspace_size(), device="cuda", dtype=torch.uint8)
             elif args.data_type == "mxfp8":
                 variant_pack_fwd = {
@@ -1764,6 +1811,8 @@ def tflops_per_sec(
                         scale_o_fwd: scale_o_gpu,
                         amax_o_fwd: amax_o_gpu,
                     }
+                    if o_block_scale:
+                        variant_pack_fwd[sf_o_fwd] = sf_o_gpu
                 elif args.data_type == "mxfp8":
                     variant_pack_fwd = {
                         q_fwd: query,

benchmark/attention_training/charts.py

@@ -46,6 +46,9 @@
     "float16": 0,  # same bucket as bf16
     "mxfp8": 1,
     "fp8": 2,
+    # fp8 inputs with a block-scaled O epilogue (FROST OSS, forward-only)
+    "fp8_mxfp8": 3,
+    "fp8_nvfp4": 4,
 }
 
 
@@ -96,6 +99,10 @@ def get_backend_display_name(backend: str, data_type: str, cudnn_backend_version
         return f"{base_name} (FP8)"
     elif data_type == "mxfp8":
         return f"{base_name} (MXFP8)"
+    elif data_type == "fp8_nvfp4":
+        return f"{base_name} (FP8 in, NVFP4 out)"
+    elif data_type == "fp8_mxfp8":
+        return f"{base_name} (FP8 in, MXFP8 out)"
     elif data_type == "float16":
         return f"{base_name} (FP16)"
     elif data_type == "bfloat16":
@@ -119,6 +126,9 @@ def get_backend_color(backend: str, data_type: str) -> str:
         return config["color_fp8"]
     if data_type == "mxfp8" and "color_mxfp8" in config:
         return config["color_mxfp8"]
+    if data_type in ("fp8_nvfp4", "fp8_mxfp8") and "color_fp8" in config:
+        # Same fp8 datapath as "fp8"; the hatch (set by the caller) tells the O modes apart.
+        return config["color_fp8"]
     return config.get("color", "gray")
 
 
@@ -129,6 +139,9 @@ def get_backend_color(backend: str, data_type: str) -> str:
     "float16": {"color": "#333333", "label": "FP16 peak (dense MMA)"},
     "fp8": {"color": "#8b0000", "label": "FP8/MXFP8 peak (dense MMA)"},
     "mxfp8": {"color": "#8b0000", "label": "FP8/MXFP8 peak (dense MMA)"},
+    # block-scaled-O modes run the fp8 MMAs: same peak line
+    "fp8_nvfp4": {"color": "#8b0000", "label": "FP8/MXFP8 peak (dense MMA)"},
+    "fp8_mxfp8": {"color": "#8b0000", "label": "FP8/MXFP8 peak (dense MMA)"},
 }
 
 

benchmark/attention_training/config_types.py

@@ -103,6 +103,10 @@ class BenchmarkConfig:
     models: List[ModelPreset]
     seqlens: List[Tuple[int, int]]
     backends: List[str] = field(default_factory=lambda: ["cudnn"])
+    # "bfloat16" / "float16" / "fp8" / "mxfp8", plus the fp8-input block-scaled-
+    # output modes "fp8_nvfp4" (FP4 O + E4M3 scale per 16 d) and "fp8_mxfp8"
+    # (E4M3 O + UE8M0 scale per 32 d) -- forward-only, cudnn_oss backend only
+    # (the runner drops other backend / pass combinations for them).
     data_types: List[str] = field(default_factory=lambda: ["bfloat16"])
     attn_masks: List[str] = field(default_factory=lambda: ["top_left"])
     profile_pass: str = "fwd"

benchmark/attention_training/configs/wan22.py

@@ -45,14 +45,19 @@
     name="wan22",
     models=[WAN22_A14B],
     seqlens=[
+        (2100, 2100),  # 720p, 81 frames, sequence-parallel shard (inference serving)
         (7800, 7800),  # 480p, 17 frames
         (17160, 17160),  # 480p, 41 frames
         (32760, 32760),  # 480p, 81 frames (default)
         (48360, 48360),  # 480p, 121 frames
         (75600, 75600),  # 720p, 81 frames
     ],
     backends=["cudnn", "cudnn_oss", "flash_attention_4"] + fa2_on_ampere(),
-    data_types=["bfloat16"],  # Wan 2.2 is released/trained in bf16; no official FP8 checkpoint
+    # bf16 = the released/trained precision. The fp8 rows model quantized
+    # inference: per-tensor fp8 in/out, and the fp8-input block-scaled-output
+    # modes (FP4 O + E4M3/16 scales, E4M3 O + UE8M0/32 scales) that feed a
+    # block-scaled projection GEMM directly (forward-only, cudnn_oss).
+    data_types=["bfloat16", "fp8", "fp8_nvfp4", "fp8_mxfp8"],
     attn_masks=["no_mask"],  # bidirectional diffusion DiT, no causal mask
     profile_pass="both",  # forward + backward for training
     deterministic_bwd=[False, True],

benchmark/attention_training/runner.py

@@ -151,6 +151,11 @@ def expand_config(self, config: BenchmarkConfig) -> Iterator[Dict[str, Any]]:
                 if data_type in ("fp8", "mxfp8"):
                     if backend not in ("cudnn", "cudnn_oss"):
                         continue
+                # Block-scaled O modes (fp8 inputs, FP4 / MXFP8 O + sf_o): a
+                # forward-only inference epilogue of the FROST OSS fp8 kernels.
+                if data_type in ("fp8_nvfp4", "fp8_mxfp8"):
+                    if backend != "cudnn_oss" or profile_pass != "fwd":
+                        continue
 
                 yield {
                     "config_name": config.name,

python/cudnn/_pygraph.py

@@ -1138,6 +1138,8 @@ def _unlowerable_node(self) -> Optional[Node]:
                 return node  # declared python-only: lowering raises by design
             if any(node.params.get(attr) is not None for attr in spec_entry[1].get("python_only_attrs", ())):
                 return node  # an op attribute the backend has no field for is SET: python engines only
+            if any(node.outputs.get(port) is not None for port in spec_entry[1].get("python_only_out_kwargs", ())):
+                return node  # an output the backend cannot produce (sf_o) is requested: python engines only
         return None
 
     def _backend_lowerable(self) -> bool:
@@ -2467,8 +2469,9 @@ def push_output_attrs(out_t: Tensor, cpp_t: Any) -> None:
                 for port, t in node.inputs.items():
                     if not port.startswith("dropout_"):
                         kw[port] = tensor_map[t.uid]
+                python_only_outs = spec.get("python_only_out_kwargs", ())
                 for port in spec.get("out_kwargs", ()):
-                    if port in node.outputs:  # classic passes these descriptors as args
+                    if port in node.outputs and port not in python_only_outs:  # classic passes these descriptors as args
                         kw[port] = lower_tensor(node.outputs[port])
                 n_drop = node.params.get("_dropout_n")
                 if n_drop:
@@ -3515,7 +3518,13 @@ def _sdpa_stats_dims(node):  # Stats: q dims with last dim 1
         node_type=NodeType.SDPA_FP8,
         pos=("q", "k", "v", "descale_q", "descale_k", "descale_v", "descale_s", "scale_s", "scale_o"),
         outputs=("O", "Stats", "Amax_S", "Amax_O"),
-        out_kwargs=("rng_dump", "score_max", "score_sum_exp"),
+        # ``sf_o``: block-scaled O scale factors (O declared FP4_E2M1 -> one
+        # E4M3 scale per 16 d elements; O FP8_E4M3 + sf_o -> one UE8M0 scale
+        # per 32). The caller passes its descriptor like rng_dump; the cuDNN
+        # backend has no field for it, so a graph that sets it is served by
+        # python engines only (see python_only_out_kwargs / _unlowerable_node).
+        out_kwargs=("rng_dump", "score_max", "score_sum_exp", "sf_o"),
+        python_only_out_kwargs=("sf_o",),
         maybe={"Stats": _stats_expected},
         infer={"O": _sdpa_o_dims, "Stats": _sdpa_stats_dims, "Amax_S": _AMAX, "Amax_O": _AMAX},
         python_only_attrs=("softmax_precision",),  # see "sdpa"

python/cudnn/_sdpa_validate.py

@@ -340,6 +340,24 @@ def _validate_forward(node) -> None:
             raise _not_supported("SDPA FP8 does not support bias")
         if (d_qk % 16 != 0) or (d_v % 16 != 0):
             raise _not_supported("hidden_dim should be multiple of 16")
+        # Block-scaled O: FP4_E2M1 O carries one E4M3 scale per 16 d elements in
+        # sf_o; an FP8_E4M3 O with sf_o carries one UE8M0 scale per 32 (MXFP8 output).
+        sf_o = node.outputs.get("sf_o")
+        o_dtype = _dtype_name(o)
+        if o_dtype == "FP4_E2M1":
+            if sf_o is None:
+                raise _not_supported("An FP4_E2M1 O requires the sf_o output (E4M3 scale factors, one per 16 d elements).")
+            if _dtype_name(sf_o) not in (None, "FP8_E4M3"):
+                raise _not_supported("sf_o must be FP8_E4M3 for an FP4_E2M1 O.")
+        elif sf_o is not None:
+            if o_dtype not in (None, "FP8_E4M3"):
+                raise _not_supported("sf_o with a non-FP4 O requires an FP8_E4M3 O (MXFP8 output).")
+            if _dtype_name(sf_o) not in (None, "FP8_E8M0"):
+                raise _not_supported("sf_o must be FP8_E8M0 for an FP8_E4M3 O (MXFP8 output).")
+            if d_v % 32 != 0:
+                raise _not_supported("MXFP8 output needs d_v to be a multiple of 32.")
+        if sf_o is not None:
+            _check_dim_stride(node, "sf_o", sf_o)
 
     if node.node_type == NodeType.SDPA_MXFP8:
         _validate_mxfp8_descales(node, q, k, v, s_kv if s_kv is not None else k.get_dim()[2])

python/cudnn/frost/tile_dsl/constants.py

@@ -16,3 +16,10 @@
 DTYPE_E5M2 = 1
 DTYPE_BF16 = 2
 DTYPE_FP16 = 3
+# Output-only block-scaled codes (per-tensor FP8 fprop epilogue): the data
+# dtype plus a scale-factor tensor written alongside O.
+DTYPE_O_NVFP4 = 4  # E2M1 data (2 per byte) + E4M3 scale per 16 d-elements
+DTYPE_O_MXFP8 = 5  # E4M3 data + UE8M0 scale per 32 d-elements
+
+# Scale-factor block along d for each output code; 0 = no block scaling.
+O_BLOCK_SCALE_BY_DTYPE = {DTYPE_E4M3: 0, DTYPE_E5M2: 0, DTYPE_BF16: 0, DTYPE_FP16: 0, DTYPE_O_NVFP4: 16, DTYPE_O_MXFP8: 32}