[PyTorch] Compile-friendly attention API: pointer-free qkv layout detection

transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py

@@ -1009,10 +1009,11 @@ def get_quantizer_roles(
     def forward(
         self,
         query_layer: torch.Tensor,
-        key_layer: torch.Tensor,
-        value_layer: torch.Tensor,
+        key_layer: Optional[torch.Tensor] = None,
+        value_layer: Optional[torch.Tensor] = None,
         attention_mask: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] = None,
         qkv_format: str = None,
+        qkv_layout: Optional[str] = None,
         cu_seqlens_q: torch.Tensor = None,
         cu_seqlens_kv: torch.Tensor = None,
         cu_seqlens_q_padded: torch.Tensor = None,
@@ -1253,6 +1254,38 @@ def forward(
             allow_non_contiguous=True,
             allow_different_data_and_param_types=self.softmax_type != "vanilla",
         ) as query_layer:
+            # ------------------------------------------------------------------
+            # JAX-style, pointer-free qkv layout signalling.
+            #
+            # Packing is derived from *how many tensors are passed* (argument
+            # presence) plus `qkv_format`, never from storage data pointers.
+            #   - key_layer is None and value_layer is None:
+            #         query_layer is a single qkv-packed tensor [.., 3, h, d]
+            #   - value_layer is None (key_layer given):
+            #         key_layer is a kv-packed tensor [.., 2, h, d], q separate
+            #   - all three given: three separate tensors
+            # The packed tensor(s) are unbound into q, k, v *views* (traceable,
+            # no data_ptr), so the rest of forward is unchanged. `qkv_layout`
+            # (the layout string) may also be passed explicitly to skip
+            # detection entirely.
+            if qkv_format is None:
+                qkv_format = self.qkv_format
+            qkv_packing = None  # one of {None (separate), "qkv", "kv"}
+            if key_layer is None and value_layer is None:
+                qkv_packing = "qkv"
+                query_layer, key_layer, value_layer = query_layer.unbind(dim=-3)
+            elif value_layer is None:
+                qkv_packing = "kv"
+                key_layer, value_layer = key_layer.unbind(dim=-3)
+            # Use the pointer-free layout path when packing is signalled via
+            # argument presence, an explicit layout string is given, or we are
+            # tracing under torch.compile (where data_ptr graph-breaks).
+            pointer_free_layout = (
+                qkv_packing is not None
+                or qkv_layout is not None
+                or torch.compiler.is_compiling()
+            )
+
             # checks for RNG
             if self.rng_states_tracker is not None and is_graph_capturing():
                 assert isinstance(
@@ -1432,8 +1465,55 @@ def forward(
                 cu_seqlens_kv_padded = None
 
             # get qkv's memory layout
-            if all(
-                isinstance(x, Float8TensorStorage) for x in [query_layer, key_layer, value_layer]
+            if pointer_free_layout:
+                # Pointer-free (JAX-style) layout: derive the layout string from
+                # `qkv_format` + packing (or use the explicit `qkv_layout`), with
+                # no data_ptr / storage_offset / stride inspection. For three
+                # separate tensors under torch.compile this yields the separate
+                # layout for the given format, which is functionally correct for
+                # flash/fused attention (it forgoes only the packed-storage
+                # micro-optimization that pointer inspection would detect).
+                is_fp8 = all(
+                    isinstance(x, Float8TensorStorage)
+                    for x in [query_layer, key_layer, value_layer]
+                )
+                if is_fp8:
+                    (
+                        qkv_layout,
+                        query_layer._data,
+                        key_layer._data,
+                        value_layer._data,
+                        q_format,
+                        kv_format,
+                    ) = dpa_utils.get_qkv_layout_pointer_free(
+                        query_layer._data,
+                        key_layer._data,
+                        value_layer._data,
+                        qkv_format=qkv_format,
+                        qkv_packing=qkv_packing,
+                        qkv_layout=qkv_layout,
+                        inference_params=inference_params,
+                    )
+                else:
+                    (
+                        qkv_layout,
+                        query_layer,
+                        key_layer,
+                        value_layer,
+                        q_format,
+                        kv_format,
+                    ) = dpa_utils.get_qkv_layout_pointer_free(
+                        query_layer,
+                        key_layer,
+                        value_layer,
+                        qkv_format=qkv_format,
+                        qkv_packing=qkv_packing,
+                        qkv_layout=qkv_layout,
+                        inference_params=inference_params,
+                    )
+            elif all(
+                isinstance(x, Float8TensorStorage)
+                for x in [query_layer, key_layer, value_layer]
             ):
                 (
                     qkv_layout,

transformer_engine/pytorch/attention/dot_product_attention/utils.py

@@ -2308,6 +2308,18 @@ def get_qkv_layout(
         Format of the key and value tensors, {`bshd`, `sbhd`, `thd`}.
     """
 
+    # Under torch.compile, pointer inspection (untyped_storage().data_ptr(),
+    # storage_offset()) graph-breaks. The gate in DotProductAttention.forward
+    # normally routes compiled calls to the pointer-free path, but if an
+    # earlier graph break makes Dynamo skip the forward frame, this function
+    # is compiled as its own frame and the forward-level gate never fired.
+    # Checking here folds to a constant during tracing, so a traced call
+    # always takes the pointer-free (separate-layout) path.
+    if torch.compiler.is_compiling():
+        return get_qkv_layout_pointer_free(
+            q, k, v, qkv_format=qkv_format, inference_params=inference_params
+        )
+
     check_last_dim_contiguous = all(x.stride(-1) == 1 for x in [q, k, v])
     assert check_last_dim_contiguous, "q, k and v must have stride 1 in their last dimension!"
     if "_2" in qkv_format:
@@ -2434,6 +2446,91 @@ def run_iteratively(q, k, v):
     return qkv_layout, q, k, v, q_format, kv_format
 
 
+def get_qkv_layout_pointer_free(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    qkv_format: str = "sbhd",
+    qkv_packing: Optional[str] = None,
+    qkv_layout: Optional[str] = None,
+    inference_params: InferenceParams = None,
+) -> str:
+    """Get qkv layout WITHOUT inspecting storage (JAX-style, torch.compile-safe).
+
+    Unlike :func:`get_qkv_layout`, this never calls ``data_ptr``,
+    ``storage_offset`` or ``stride`` (all of which graph-break under
+    ``torch.compile``). Instead the layout is derived from ``qkv_format`` plus
+    how the tensors were packed (``qkv_packing``), or taken directly from an
+    explicit ``qkv_layout`` string.
+
+    Parameters
+    ----------
+    q, k, v : torch.Tensor
+        Query/key/value tensors. When ``qkv_packing`` is ``"qkv"`` or ``"kv"``,
+        these are already the unbound *views* of the packed tensor(s).
+    qkv_format : str, default = "sbhd"
+        Dimension format, one of {``sbhd``, ``bshd``, ``thd``} (or a
+        ``q_2kv``-style format for cross attention with KV caching).
+    qkv_packing : Optional[str], default = None
+        How the inputs were packed, derived from argument presence in
+        ``DotProductAttention.forward``:
+          ``None``  -> q, k, v are three separate tensors.
+          ``"qkv"`` -> a single qkv-packed tensor ``[.., 3, h, d]``.
+          ``"kv"``  -> q separate, k/v from a kv-packed tensor ``[.., 2, h, d]``.
+        The ``3``/``2`` are assumed to sit at dim ``-3`` (i.e. the ``3hd``/``2hd``
+        interleaving, matching the JAX convention). The ``h3d``/``h2d`` variants
+        cannot be signalled this way and require an explicit ``qkv_layout``.
+    qkv_layout : Optional[str], default = None
+        Explicit layout string. When provided it is used verbatim and packing
+        detection is skipped.
+    inference_params : InferenceParams, default = None
+        InferenceParams related to KV caching (used only to add the
+        ``paged_kv_`` prefix).
+
+    Returns
+    ----------
+    qkv_layout : str
+        Memory layout string (same vocabulary as :func:`get_qkv_layout`).
+    q, k, v : torch.Tensor
+        Returned unchanged (kept for a matching return contract).
+    q_format, kv_format : str
+        Formats of q and of k/v, in {``bshd``, ``sbhd``, ``thd``}.
+    """
+    check_last_dim_contiguous = all(x.stride(-1) == 1 for x in [q, k, v])
+    assert (
+        check_last_dim_contiguous
+    ), "q, k and v must have stride 1 in their last dimension!"
+
+    if "_2" in qkv_format:
+        q_format, kv_format = qkv_format.split("_2")
+        is_same_q_kv_format = False
+    else:
+        q_format = qkv_format
+        kv_format = qkv_format
+        is_same_q_kv_format = True
+
+    if qkv_layout is None:
+        if qkv_packing == "qkv":
+            # sb3hd, bs3hd, t3hd -- one chunk, q/k/v interleaved at dim=-3.
+            assert is_same_q_kv_format, "qkv-packing requires q_format == kv_format!"
+            qkv_layout = q_format[:-2] + "3" + q_format[-2:]
+        elif qkv_packing == "kv":
+            # sbhd_sb2hd, bshd_bs2hd, thd_t2hd -- q separate, k/v interleaved at dim=-3.
+            qkv_layout = q_format + "_" + kv_format[:-2] + "2" + kv_format[-2:]
+        else:
+            # sbhd_sbhd_sbhd, bshd_bshd_bshd, thd_thd_thd (or the mixed q_2kv forms).
+            if is_same_q_kv_format:
+                qkv_layout = "_".join([qkv_format] * 3)
+            else:
+                qkv_layout = q_format + "_" + kv_format + "_" + kv_format
+
+    if inference_params is not None and inference_params.is_paged:
+        if not qkv_layout.startswith("paged_kv_"):
+            qkv_layout = "paged_kv_" + qkv_layout
+
+    return qkv_layout, q, k, v, q_format, kv_format
+
+
 def check_set_window_size(
     attn_mask_type: str,
     window_size: Tuple[int, int] = None,