diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index f2ab4a8495..7bcdd10daf 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -132,6 +132,7 @@ def test_dot_product_attention( qkv_layout, swa, pad_between_seqs, + declarative_packed=False, ): """Test DotProductAttention module""" @@ -149,6 +150,8 @@ def test_dot_product_attention( qkv_layout = "bshd_bs2hd" if not is_mla and not is_mqa_gqa else "bshd_bshd_bshd" if "3" in qkv_layout and config.attn_type == "cross": pytest.skip("No need to test this layout for cross attention") + if declarative_packed and not any(c.isdigit() for c in qkv_layout): + pytest.skip("Declarative packed inputs only apply to packed qkv layouts.") if config.window_size == (-1, -1) and swa: config.window_size = [2, 2] @@ -221,6 +224,7 @@ def test_dot_product_attention( qkv_layout, pad_between_seqs, is_training, + declarative_packed=declarative_packed, ) # FlashAttention backend @@ -233,6 +237,7 @@ def test_dot_product_attention( qkv_layout, pad_between_seqs, is_training, + declarative_packed=declarative_packed, ) # Compare results @@ -924,6 +929,26 @@ def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout): test_dot_product_attention(dtype, model_configs, model, False, qkv_layout, False, False) +qkv_layouts_packed = [l for l in qkv_layouts if any(c.isdigit() for c in l)] + + +@pytest.mark.skipif(get_cudnn_version() < (8, 9, 5), reason="cuDNN 8.9.5+ is required.") +@pytest.mark.parametrize("dtype", param_types_lean) +@pytest.mark.parametrize("model_configs", [model_configs_layout]) +@pytest.mark.parametrize("model", ["layout_1_1", "layout_1_2"]) +@pytest.mark.parametrize("qkv_layout", qkv_layouts_packed) +def test_dpa_qkv_layout_declarative(dtype, model_configs, model, qkv_layout): + """Declarative packed inputs: the packed buffer is passed to + DotProductAttention via qkv_layer/kv_layer (declared layout, gradients read + off the packed buffer) instead of q/k/v views + pointer-based detection. + Layout coverage is complete; the model-config dimension is trimmed to one + self-attention and one cross-attention config, since past the input + handling the backend code is identical to test_dpa_qkv_layout.""" + test_dot_product_attention( + dtype, model_configs, model, False, qkv_layout, False, False, declarative_packed=True + ) + + qkv_layouts_thd = ["t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"] model_configs_layout_thd = { # test: ModelConfig(b, sq, hq, dqk) @@ -988,7 +1013,7 @@ def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout): @pytest.mark.parametrize("model_configs", [model_configs_layout_thd]) @pytest.mark.parametrize("model", model_configs_layout_thd.keys()) @pytest.mark.parametrize("qkv_layout", qkv_layouts_thd) -def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout): +def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, declarative_packed=False): """Test DotProductAttention module with different QKV layouts""" config = model_configs[model] if config.num_heads != config.num_gqa_groups and "3" in qkv_layout: @@ -996,17 +1021,47 @@ def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout): logging.info("[test_dpa_qkv_layout_thd]: pad_between_seqs = True") pad_between_seqs = True test_dot_product_attention( - dtype, model_configs, model, False, qkv_layout, False, pad_between_seqs + dtype, + model_configs, + model, + False, + qkv_layout, + False, + pad_between_seqs, + declarative_packed=declarative_packed, ) if get_cudnn_version() >= (9, 3, 0): logging.info("[test_dpa_qkv_layout_thd]: pad_between_seqs = False") # cuDNN 9.3.0+ is required to run pad_between_seqs = False/True in the same run pad_between_seqs = False test_dot_product_attention( - dtype, model_configs, model, False, qkv_layout, False, pad_between_seqs + dtype, + model_configs, + model, + False, + qkv_layout, + False, + pad_between_seqs, + declarative_packed=declarative_packed, ) +qkv_layouts_thd_packed = [l for l in qkv_layouts_thd if any(c.isdigit() for c in l)] + + +@pytest.mark.skipif(get_cudnn_version() < (9, 0, 0), reason="cuDNN 9.0.0+ is required.") +@pytest.mark.skipif( + get_device_compute_capability() < (9, 0), reason="THD is only supported on Hopper+." +) +@pytest.mark.parametrize("dtype", param_types_lean) +@pytest.mark.parametrize("model_configs", [model_configs_layout_thd]) +@pytest.mark.parametrize("model", ["layout_0_0"]) +@pytest.mark.parametrize("qkv_layout", qkv_layouts_thd_packed) +def test_dpa_qkv_layout_thd_declarative(dtype, model_configs, model, qkv_layout): + """Declarative packed thd inputs, see test_dpa_qkv_layout_declarative.""" + test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, declarative_packed=True) + + def _run_dot_product_attention( dtype: torch.dtype, config: ModelConfig, @@ -1015,8 +1070,14 @@ def _run_dot_product_attention( qkv_layout: str, pad_between_seqs: bool, is_training: bool, + declarative_packed: bool = False, ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: - """Run DotProductAttention module with one forward pass and one backward pass""" + """Run DotProductAttention module with one forward pass and one backward pass. + + With declarative_packed=True (packed qkv_layout only), the packed buffer is + passed to DotProductAttention directly via qkv_layer/kv_layer instead of + slicing it into q/k/v views, and input gradients are read off the packed + buffer itself.""" # Set RNG and environment varables reset_rng_states() os.environ["NVTE_FLASH_ATTN"] = "0" @@ -1213,6 +1274,12 @@ def _run_dot_product_attention( tensor_count = int(l) split_dim = dim break + if declarative_packed and split_dim != 0: + # The packed buffer is the autograd leaf; q/k/v below are non-leaf + # views of it, and DPA receives the buffer via qkv_layer/kv_layer. + tensor.requires_grad_() + packed_tensor = tensor + packed_interleave_dim = split_dim - tensor.dim() tensors = torch.split(tensor, 1, dim=split_dim) if split_dim != 0 else [tensor] tensors_orig = ( torch.split(tensor_orig, 1, dim=split_dim) if split_dim != 0 else [tensor_orig] @@ -1225,8 +1292,10 @@ def _run_dot_product_attention( inp.append(tensors[j]) inp_orig.append(tensors_orig[j]) for i in range(3): - inp[i].requires_grad = True - inp_orig[i].requires_grad = True + if inp[i].is_leaf: + inp[i].requires_grad = True + if inp_orig[i].is_leaf: + inp_orig[i].requires_grad = True # Create output gradient qkv_format_kv = "_".join(qkv_format) @@ -1308,10 +1377,21 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: k = inp[1] v = inp[2] d_out = out_grad + packed_kwargs = {} + if declarative_packed: + assert backend in ["FusedAttention", "FlashAttention"] + packed_kwargs["qkv_interleave_dim"] = packed_interleave_dim + if len(qkv_layout.split("_")) == 1: + packed_kwargs["qkv_layer"] = packed_tensor + q, k, v = None, None, None + else: + packed_kwargs["kv_layer"] = packed_tensor + k, v = None, None out = block( q, k, v, + **packed_kwargs, window_size=config.window_size, attention_mask=attention_mask, qkv_format=qkv_format, @@ -1341,15 +1421,33 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: if is_training: out.backward(d_out) + q_grad, k_grad, v_grad = None, None, None + if is_training: + if declarative_packed: + # Input gradients live on the packed buffer; slice them back out so + # the cross-backend comparisons below stay uniform with the + # separate-q/k/v path. + assert ( + packed_tensor.grad is not None and packed_tensor.grad.shape == packed_tensor.shape + ) + packed_grads = [ + packed_tensor.grad.select(packed_interleave_dim, j) + for j in range(packed_tensor.shape[packed_interleave_dim]) + ] + if len(qkv_layout.split("_")) == 1: + q_grad, k_grad, v_grad = packed_grads + else: + q_grad = q.grad + k_grad, v_grad = packed_grads + else: + q_grad, k_grad, v_grad = q.grad, k.grad, v.grad + d_softmax_offset = None if is_training and config.softmax_type != "vanilla": d_softmax_offset = block.softmax_offset.grad if backend in ["UnfusedDotProductAttention"]: - if is_training: - return out, max_logit, (q.grad, k.grad, v.grad, d_softmax_offset) - else: - return out, max_logit, (None, None, None, d_softmax_offset) + return out, max_logit, (q_grad, k_grad, v_grad, d_softmax_offset) if backend in ["FusedAttention", "FlashAttention"]: if qkv_format == "thd" and pad_between_seqs: out_orig = torch.Tensor([]).to(device="cuda", dtype=dtype) @@ -1369,13 +1467,13 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: out_orig = torch.cat([out_orig, out[valid_range_q[0] : valid_range_q[1]]], dim=0) if is_training: q_grad_orig = torch.cat( - [q_grad_orig, q.grad[valid_range_q[0] : valid_range_q[1]]], dim=0 + [q_grad_orig, q_grad[valid_range_q[0] : valid_range_q[1]]], dim=0 ) k_grad_orig = torch.cat( - [k_grad_orig, k.grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 + [k_grad_orig, k_grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 ) v_grad_orig = torch.cat( - [v_grad_orig, v.grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 + [v_grad_orig, v_grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 ) if is_training: return ( @@ -1386,10 +1484,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: else: return out_orig, max_logit, (None, None, None, d_softmax_offset) else: - if is_training: - return out, max_logit, (q.grad, k.grad, v.grad, d_softmax_offset) - else: - return out, max_logit, (None, None, None, d_softmax_offset) + return out, max_logit, (q_grad, k_grad, v_grad, d_softmax_offset) model_configs_te_layer = { diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 785e438cda..d6acc1aaf7 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -1328,6 +1328,8 @@ def forward( fp8_output, layer_number, return_max_logit, + packed_qkv=None, + packed_kv=None, ): # pylint: disable=missing-function-docstring @@ -1390,7 +1392,14 @@ def forward( q_fp8, k_fp8, v_fp8 = q, k, v else: q_fp8, k_fp8, v_fp8, qkv_layout, qkv_scale_inv_format = combine_and_quantize( - qkv_layout, q, k, v, QKV_quantizer, used_in_backward=is_training + qkv_layout, + q, + k, + v, + QKV_quantizer, + used_in_backward=is_training, + combined_qkv=packed_qkv, + combined_kv=packed_kv, ) # print quantizers @@ -1920,6 +1929,8 @@ def backward(ctx, d_out, *_args): None, None, None, + None, # packed_qkv + None, # packed_kv ) @@ -2014,6 +2025,8 @@ def forward( score_mod_bprop: Optional[Callable] = None, score_mod_tensors: Optional[Dict[str, torch.Tensor]] = None, score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]] = None, + packed_qkv: Optional[torch.Tensor] = None, + packed_kv: Optional[torch.Tensor] = None, ) -> torch.Tensor: """fused attention fprop""" assert ( @@ -2233,6 +2246,8 @@ def forward( fp8_output, self.layer_number, self.return_max_logit, + packed_qkv, + packed_kv, ) if self.return_max_logit: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index dd812f15e9..82a23ff021 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -196,6 +196,111 @@ def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim return attn_out[..., :orig_head_dim_v].reshape(*out_shape, -1) +def _unpack_packed_qkv( + qkv_layer: Optional[torch.Tensor], + kv_layer: Optional[torch.Tensor], + query_layer: Optional[torch.Tensor], + key_layer: Optional[torch.Tensor], + value_layer: Optional[torch.Tensor], + qkv_format: str, + qkv_interleave_dim: int, + inference_params: Optional[InferenceParams], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[str]]: + """Resolve declarative packed inputs into q/k/v. + + Derives q/k/v as zero-copy views of the packed buffer (``qkv_layer`` or + ``kv_layer``) and constructs the exact layout string from the declaration. + The layout enum is truthful by construction, so no pointer-based detection + is needed downstream (this also covers thd and FP8 DPA). + + Returns ``(query_layer, key_layer, value_layer, declared_qkv_layout)``; + the layout is ``None`` when no packed input is given. + """ + if qkv_layer is None and kv_layer is None: + if query_layer is None or key_layer is None or value_layer is None: + raise ValueError( + "query_layer, key_layer and value_layer are required unless packed" + " inputs (qkv_layer or query_layer + kv_layer) are provided." + ) + return query_layer, key_layer, value_layer, None + + if qkv_layer is not None and kv_layer is not None: + raise ValueError("qkv_layer and kv_layer are mutually exclusive.") + if inference_params is not None: + raise ValueError( + "Packed inputs (qkv_layer/kv_layer) are not supported with KV caching" + " (inference_params); pass separate query/key/value tensors instead." + ) + if qkv_interleave_dim not in (-3, -2): + raise ValueError( + "qkv_interleave_dim must be -3 (e.g. bs3hd) or -2 (e.g. bsh3d), got" + f" {qkv_interleave_dim}." + ) + packed = qkv_layer if qkv_layer is not None else kv_layer + # The declared layout describes the packed buffer's memory, so it must have + # stride 1 in its last dimension (the check get_qkv_layout would otherwise + # perform on the derived q/k/v views). + if packed.stride(-1) != 1: + raise ValueError( + "The packed tensor (qkv_layer/kv_layer) must have stride 1 in its last" + f" dimension, got strides {tuple(packed.stride())}." + ) + + def _packed_layout(fmt: str, num: int) -> str: + # bshd + 3 @ -3 -> bs3hd; bshd + 3 @ -2 -> bsh3d; thd + 2 @ -2 -> th2d + pos = len(fmt) + qkv_interleave_dim + 1 + return fmt[:pos] + str(num) + fmt[pos:] + + if qkv_layer is not None: + if any(x is not None for x in (query_layer, key_layer, value_layer)): + raise ValueError( + "qkv_layer already packs Q, K and V: query_layer, key_layer and" + " value_layer must be None when qkv_layer is provided." + ) + expected_rank = 4 if qkv_format == "thd" else 5 + if qkv_layer.dim() != expected_rank: + raise ValueError( + f"qkv_layer must be a {expected_rank}D tensor for" + f" qkv_format={qkv_format!r}, got {qkv_layer.dim()}D." + ) + if qkv_layer.shape[qkv_interleave_dim] != 3: + raise ValueError( + f"qkv_layer must have size 3 at dim {qkv_interleave_dim}" + f" (qkv_interleave_dim), got shape {tuple(qkv_layer.shape)}." + ) + query_layer, key_layer, value_layer = ( + qkv_layer.select(qkv_interleave_dim, i) for i in range(3) + ) + return query_layer, key_layer, value_layer, _packed_layout(qkv_format, 3) + + if query_layer is None: + raise ValueError( + "kv_layer packs only K and V: query_layer is required when kv_layer is provided." + ) + if key_layer is not None or value_layer is not None: + raise ValueError( + "kv_layer already packs K and V: key_layer and value_layer must be" + " None when kv_layer is provided." + ) + if kv_layer.dim() != query_layer.dim() + 1: + raise ValueError( + "kv_layer must have one more dimension than query_layer, got" + f" {kv_layer.dim()}D kv_layer and {query_layer.dim()}D query_layer." + ) + if kv_layer.shape[qkv_interleave_dim] != 2: + raise ValueError( + f"kv_layer must have size 2 at dim {qkv_interleave_dim}" + f" (qkv_interleave_dim), got shape {tuple(kv_layer.shape)}." + ) + key_layer, value_layer = (kv_layer.select(qkv_interleave_dim, i) for i in range(2)) + return ( + query_layer, + key_layer, + value_layer, + f"{qkv_format}_{_packed_layout(qkv_format, 2)}", + ) + + class DotProductAttention(TransformerEngineBaseModule): r"""Allows the model to jointly attend to information from different representation subspaces as described in the paper: @@ -988,9 +1093,9 @@ def get_quantizer_roles( @no_torch_dynamo(recursive=False) def forward( self, - query_layer: torch.Tensor, - key_layer: torch.Tensor, - value_layer: torch.Tensor, + query_layer: Optional[torch.Tensor] = None, + 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, cu_seqlens_q: torch.Tensor = None, @@ -1015,6 +1120,9 @@ def forward( score_mod_bprop: Optional[Callable] = None, score_mod_tensors: Optional[Dict[str, torch.Tensor]] = None, score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]] = None, + qkv_layer: Optional[torch.Tensor] = None, + kv_layer: Optional[torch.Tensor] = None, + qkv_interleave_dim: int = -3, ) -> torch.Tensor: r""" Dot Product Attention Layer. @@ -1226,8 +1334,44 @@ def forward( Runtime tensors exposed to score_mod_bprop as cuDNN graph tensors. Keys are user-defined string names consumed by the callback through ``tensors[name]``; there is no predefined set of accepted keys. + qkv_layer: Optional[torch.Tensor], default = None + Fully packed QKV tensor. When the QKV projection produces one packed buffer + (e.g. a fused QKV GEMM), it can be passed here directly instead of slicing + it into :attr:`query_layer`/:attr:`key_layer`/:attr:`value_layer` views. + For :attr:`qkv_format` = {"bshd", "sbhd"}, it must be a 5D tensor of shape + ``[b, s, 3, h, d]``/``[s, b, 3, h, d]`` (:attr:`qkv_interleave_dim` = -3) or + ``[b, s, h, 3, d]``/``[s, b, h, 3, d]`` (:attr:`qkv_interleave_dim` = -2); + for :attr:`qkv_format` = "thd", a 4D tensor of shape ``[t, 3, h, d]`` or + ``[t, h, 3, d]``. Q/K/V are derived as zero-copy views and the memory layout + (e.g. ``bs3hd``) is declared from the packing itself, so no pointer-based + layout detection runs on this path -- including for "thd" and FP8 attention. + Mutually exclusive with :attr:`query_layer`, :attr:`key_layer`, + :attr:`value_layer` and :attr:`kv_layer`. + kv_layer: Optional[torch.Tensor], default = None + Packed KV tensor, used together with :attr:`query_layer` + (e.g. ``[b, s, 2, hg, d]`` for :attr:`qkv_interleave_dim` = -3, or + ``[b, s, hg, 2, d]`` for :attr:`qkv_interleave_dim` = -2). K/V are derived + as zero-copy views and the layout (e.g. ``bshd_bs2hd``) is declared, not + detected. Mutually exclusive with :attr:`key_layer`, :attr:`value_layer` + and :attr:`qkv_layer`. + qkv_interleave_dim: int, default = -3 + Dimension of :attr:`qkv_layer`/:attr:`kv_layer` where the 3 (QKV) or 2 (KV) + interleave sits; must be -3 (e.g. ``bs3hd``) or -2 (e.g. ``bsh3d``, + Megatron-style). This is an explicit knob rather than shape inference, + since e.g. ``h == 3`` would make the shapes ambiguous. """ + query_layer, key_layer, value_layer, declared_qkv_layout = _unpack_packed_qkv( + qkv_layer, + kv_layer, + query_layer, + key_layer, + value_layer, + qkv_format if qkv_format is not None else self.qkv_format, + qkv_interleave_dim, + inference_params, + ) + with self.prepare_forward_ctx( query_layer, num_gemms=3, @@ -1413,7 +1557,15 @@ def forward( cu_seqlens_kv_padded = None # get qkv's memory layout - if all( + if declared_qkv_layout is not None: + # Packed inputs (qkv_layer/kv_layer) declare the layout: the enum is + # truthful by construction, so the pointer-based detection in + # get_qkv_layout is skipped entirely -- for dense, thd (t3hd/th3d) + # and FP8 DPA alike. + qkv_layout = declared_qkv_layout + q_format = qkv_format + kv_format = qkv_format + elif all( isinstance(x, Float8TensorStorage) for x in [query_layer, key_layer, value_layer] ): ( @@ -1801,6 +1953,8 @@ def forward( inference_params=inference_params, softmax_offset=softmax_offset, fp8_output=fp8_output, + packed_qkv=qkv_layer, + packed_kv=kv_layer, ) return self.fused_attention( query_layer, @@ -1836,6 +1990,8 @@ def forward( score_mod_bprop=score_mod_bprop, score_mod_tensors=score_mod_tensors, score_mod_bprop_tensors=score_mod_bprop_tensors, + packed_qkv=qkv_layer, + packed_kv=kv_layer, ) if use_unfused_attention: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 7be94a6fa1..4949641810 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -35,6 +35,7 @@ META_DP, ) from transformer_engine.pytorch.attention.inference import InferenceParams +from transformer_engine.pytorch.cpu_offload import is_cpu_offload_enabled from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage from transformer_engine.pytorch.tensor.float8_tensor import ( Float8Tensor, @@ -2432,6 +2433,22 @@ def run_iteratively(q, k, v): if qkv_layout == "not_supported": raise RuntimeError("The provided qkv memory layout is not supported!") + if len(qkv_layout.split("_")) < 3: + # q/k/v were recognized as views of a packed buffer only by inspecting + # their data pointers, strides and storage offsets. Skip the nudge while + # CPU offloading is enabled: offloading forces MultiheadAttention onto + # its sliced-views fallback, so packed views reaching detection are + # expected there and the caller has no migration option. + if not is_cpu_offload_enabled(): + warnings.warn( + "Relying on pointer-based detection of packed q/k/v layouts" + f" (detected {qkv_layout!r}) is deprecated: pass the packed buffer" + " explicitly via qkv_layer/kv_layer (with qkv_interleave_dim) to" + " DotProductAttention instead.", + DeprecationWarning, + stacklevel=2, + ) + if inference_params is not None and inference_params.is_paged: qkv_layout = "paged_kv_" + qkv_layout @@ -2787,8 +2804,18 @@ def combine_and_quantize( used_in_forward=True, used_in_backward=False, keep_same_data_and_scale_inv_format=False, + combined_qkv: Optional[torch.Tensor] = None, + combined_kv: Optional[torch.Tensor] = None, ): - """Combine Q, K, V tensors based on qkv_layout and quantize them together.""" + """Combine Q, K, V tensors based on qkv_layout and quantize them together. + + When ``combined_qkv`` (for ``3`` layouts such as ``bs3hd``) or ``combined_kv`` + (for ``2`` layouts such as ``bshd_bs2hd``) is provided, it must be the + caller's original packed buffer that q/k/v are views of. It is then quantized + directly instead of re-deriving the packed buffer from the q/k/v views via + ``combine_tensors`` (which rebuilds it with a raw ``set_`` under a silent + adjacency/interleave assumption). Ignored for MXFP8 quantization. + """ if isinstance(qkv_quantizer, MXFP8Quantizer): qkv_format, q_format, kv_format = get_qkv_format(qkv_layout) assert qkv_format in ("bshd", "sbhd"), ( @@ -2888,12 +2915,26 @@ def combine_and_quantize( match qkv_group: case 1: dim = qkv_layout.find("3") - qkv = combine_tensors([q, k, v], dim) + if combined_qkv is not None: + assert combined_qkv.shape[dim] == 3, ( + f"combined_qkv does not match qkv_layout {qkv_layout}: expected" + f" size 3 at dim {dim}, got shape {tuple(combined_qkv.shape)}." + ) + qkv = combined_qkv + else: + qkv = combine_tensors([q, k, v], dim) qkv_fp8 = qkv_quantizer(qkv) q_data, k_data, v_data = SplitAlongDim.apply(qkv_fp8._data, dim, [1, 1, 1], True) case 2: dim = qkv_layout.split("_")[1].find("2") - kv = combine_tensors([k, v], dim) + if combined_kv is not None: + assert combined_kv.shape[dim] == 2, ( + f"combined_kv does not match qkv_layout {qkv_layout}: expected" + f" size 2 at dim {dim}, got shape {tuple(combined_kv.shape)}." + ) + kv = combined_kv + else: + kv = combine_tensors([k, v], dim) tensors = [q, kv] num_tensors = len(tensors) shapes = [x.shape for x in tensors] diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index 70ae9dfc21..0b73daf667 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -9,6 +9,7 @@ import torch from transformer_engine.pytorch.quantization import FP8GlobalStateManager, QuantizerRole +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.module import LayerNormLinear, Linear, RMSNorm, LayerNorm @@ -903,6 +904,23 @@ def forward( self._update_output_quantizer_roles(qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) + # Packed pass-through to DotProductAttention: the fused QKV/KV projection + # already produces one packed buffer, which DPA accepts directly via its + # declarative qkv_layer/kv_layer arguments (deriving q/k/v as zero-copy + # views and skipping pointer-based layout detection). Only possible when + # no per-tensor operation (RoPE, QK normalization, KV caching, CPU + # offloading) needs the individual q/k/v slices. + packed_dpa_eligible = ( + rotary_pos_emb is None + and self.q_norm is None + and self.k_norm is None + and inference_params is None + and not is_cpu_offload_enabled() + ) + packed_qkv_layer = None + packed_kv_layer = None + packed_interleave_dim = -3 + layernorm_output = None if self.attention_type == "self": # Attention heads [sq, b, h] --> [sq, b, ng * (np/ng + 2) * hn] @@ -947,28 +965,43 @@ def forward( mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) - # qkv_weight_interleaved: - # [sq, b, ng, (np/ng + 2), hn] - # --> [sq, b, ng, np/ng, hn], [sq, b, ng, 1, hn], [sq, b, ng, 1, hn] - # not qkv_weight_interleaved: - # [sq, b, (np/ng + 2), ng, hn] - # --> [sq, b, np/ng, np, hn], [sq, b, 1, ng, hn], [sq, b, 1, ng, hn] - query_layer, key_layer, value_layer = SplitAlongDim.apply( - mixed_x_layer, split_dim, (num_queries_per_key_value, 1, 1) - ) - - if self.qkv_format == "thd": - query_layer, key_layer, value_layer = ( - x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) - for x in (query_layer, key_layer, value_layer) - ) + if ( + num_queries_per_key_value == 1 + and packed_dpa_eligible + and not isinstance(mixed_x_layer, QuantizedTensorStorage) + ): + # np == ng: the projection output is a uniform 3-interleave + # ([.., h, 3, d] interleaved / [.., 3, h, d] otherwise), which + # DotProductAttention accepts directly as a declared packed + # qkv_layer -- no slicing here, no layout detection there. + packed_qkv_layer = mixed_x_layer + packed_interleave_dim = split_dim + query_layer = None + key_layer = None + value_layer = None else: - # query: -> [sq, b, np, hn] - # key, value: -> [sq, b, ng, hn] - query_layer, key_layer, value_layer = ( - x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) - for x in (query_layer, key_layer, value_layer) + # qkv_weight_interleaved: + # [sq, b, ng, (np/ng + 2), hn] + # --> [sq, b, ng, np/ng, hn], [sq, b, ng, 1, hn], [sq, b, ng, 1, hn] + # not qkv_weight_interleaved: + # [sq, b, (np/ng + 2), ng, hn] + # --> [sq, b, np/ng, np, hn], [sq, b, 1, ng, hn], [sq, b, 1, ng, hn] + query_layer, key_layer, value_layer = SplitAlongDim.apply( + mixed_x_layer, split_dim, (num_queries_per_key_value, 1, 1) ) + + if self.qkv_format == "thd": + query_layer, key_layer, value_layer = ( + x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) + for x in (query_layer, key_layer, value_layer) + ) + else: + # query: -> [sq, b, np, hn] + # key, value: -> [sq, b, ng, hn] + query_layer, key_layer, value_layer = ( + x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) + for x in (query_layer, key_layer, value_layer) + ) elif self.attention_type == "cross": # Attention heads [sk, b, h] --> [sk, b, (ng * 2 * hn)] mixed_kv_layer = self.key_value( @@ -996,34 +1029,56 @@ def forward( mixed_kv_layer = mixed_kv_layer.view(*new_tensor_shape) - # mixed_kv_layer --> 2 [sk, b, ng, hn] - key_layer, value_layer = SplitAlongDim.apply( - mixed_kv_layer, - split_dim, - mixed_kv_layer.shape[split_dim] // 2, - ) - key_layer, value_layer = ( - x.reshape( - x.size(0), - x.size(1), - -1, - self.hidden_size_per_attention_head, - ) - for x in (key_layer, value_layer) - ) - - if self.qkv_format == "thd": - key_layer, value_layer = ( - x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) - for x in (key_layer, value_layer) - ) + if packed_dpa_eligible and not isinstance(mixed_kv_layer, QuantizedTensorStorage): + # Declare the packed KV to DotProductAttention instead of + # slicing it: expose the 2-interleave as its own dimension. + if self.qkv_weight_interleaved: + # [.., ng, 2 * hn] --> [.., ng, 2, hn] + packed_kv_shape = mixed_kv_layer.size()[:-1] + ( + 2, + self.hidden_size_per_attention_head, + ) + packed_interleave_dim = -2 + else: + # [.., 2 * ng, hn] --> [.., 2, ng, hn] + packed_kv_shape = mixed_kv_layer.size()[:-2] + ( + 2, + self.num_gqa_groups_per_partition, + self.hidden_size_per_attention_head, + ) + packed_interleave_dim = -3 + packed_kv_layer = mixed_kv_layer.view(*packed_kv_shape) + key_layer = None + value_layer = None else: - # key, value: -> [sq, b, ng, hn] + # mixed_kv_layer --> 2 [sk, b, ng, hn] + key_layer, value_layer = SplitAlongDim.apply( + mixed_kv_layer, + split_dim, + mixed_kv_layer.shape[split_dim] // 2, + ) key_layer, value_layer = ( - x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) + x.reshape( + x.size(0), + x.size(1), + -1, + self.hidden_size_per_attention_head, + ) for x in (key_layer, value_layer) ) + if self.qkv_format == "thd": + key_layer, value_layer = ( + x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) + for x in (key_layer, value_layer) + ) + else: + # key, value: -> [sq, b, ng, hn] + key_layer, value_layer = ( + x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) + for x in (key_layer, value_layer) + ) + # Attention head [sq, b, h] --> [sq, b, hp] if self.input_layernorm: layernorm_query_outputs = self.layernorm_query( @@ -1143,6 +1198,9 @@ def forward( inference_params=inference_params, pad_between_seqs=pad_between_seqs, fp8_output=dpa_fp8_output, + qkv_layer=packed_qkv_layer, + kv_layer=packed_kv_layer, + qkv_interleave_dim=packed_interleave_dim, ) # ===================