diff --git a/autoparallel/ops.py b/autoparallel/ops.py index c2c1e2ad..36b7fc45 100644 --- a/autoparallel/ops.py +++ b/autoparallel/ops.py @@ -3,8 +3,13 @@ # This source code is licensed under the BSD license found in the # LICENSE file in the root directory of this source tree. +from typing import Any + import torch +# Importing triggers registration of torch_attn::_varlen_attn{,_backward}. +import torch.nn.attention.varlen as _varlen # noqa: F401 + def permutation(x: torch.Tensor, axis: int = 0, independent: bool = False): """Randomly permute elements of a tensor along an axis. @@ -39,3 +44,296 @@ def permutation(x: torch.Tensor, axis: int = 0, independent: bool = False): # generate random permutation matrix which is independent per axis idxs = torch.rand_like(x, dtype=torch.float32).argsort(axis) return x.gather(axis, idxs) + + +# --------------------------------------------------------------------------- +# doc_packed_attn: document-packed variable-length attention with a shape that +# AutoParallel can shard. +# +# The user-facing op takes q/k/v in [B, S, H, D] layout and a cu_seq_q tensor +# in [B, MAX_DOCS+1] layout. Sharding ``Shard(0)`` on the leading dim of all +# four inputs gives clean DP across the batch. +# +# Internally we reshape q/k/v to THD layout, build a flat cu_seq_q with +# per-batch-element offsets of ``b*S``, and dispatch to +# ``torch_attn::_varlen_attn``. ``n_docs`` is uniform across batch elements; +# batch elements with fewer real documents pad their row with repeated ``S`` +# so the kernel sees zero-length trailing docs and skips them. +# +# The custom op returns ``(out, lse, rng_state)`` (mirroring the upstream +# ``torch_attn::_varlen_attn`` shape) so ``setup_context`` can save the +# residuals for backward without re-running the forward kernel. End users who +# want just the output should call :func:`doc_packed_attn`, the thin Python +# helper below. +# --------------------------------------------------------------------------- + + +def _build_flat_cu_seq_q(cu_seq_q: torch.Tensor, n_docs: int, S: int) -> torch.Tensor: + """Slice ``cu_seq_q`` to ``n_docs+1`` entries per row, then offset and flatten. + + For ``B == 1`` returns a view (no copy, no sync). + For ``B > 1`` returns ``[0, d0_1, ..., S, S+d1_1, ..., 2S, ...]`` of length + ``B*n_docs + 1`` — sync-free arithmetic on a tiny tensor. + """ + B = cu_seq_q.shape[0] + sliced = cu_seq_q[:, : n_docs + 1] # [B, n_docs+1] + if B == 1: + return sliced[0] + # Offset each row by b*S so doc boundaries are absolute positions in the + # flattened T=B*S dimension. + offsets = torch.arange(B, dtype=cu_seq_q.dtype, device=cu_seq_q.device) * S + shifted = sliced + offsets.unsqueeze(1) # [B, n_docs+1] + # Each row's leading entry equals b*S (== prior row's trailing entry), so + # drop it from rows 1..B-1 to avoid duplicating boundary positions. + return torch.cat([shifted[0], shifted[1:, 1:].flatten()]) + + +@torch.library.custom_op("autoparallel::doc_packed_attn_op", mutates_args=()) +def _doc_packed_attn_op( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + cu_seq_q: torch.Tensor, + n_docs: int, + scale: float | None = None, + window_size: list[int] | None = None, + enable_gqa: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Custom op: returns ``(out, lse, rng_state)``. + + ``out`` has the same shape as ``query``. ``lse`` and ``rng_state`` come + straight from the underlying flash-attention kernel and are passed back + into ``torch_attn::_varlen_attn_backward`` during backward. + """ + B, S = query.shape[:2] + cu = _build_flat_cu_seq_q(cu_seq_q, n_docs, S) + + ws = list(window_size) if window_size is not None else [-1, -1] + is_causal = ws == [-1, 0] + out_thd, lse, rng_state = torch.ops.torch_attn._varlen_attn( + query.reshape(B * S, *query.shape[2:]), + key.reshape(B * S, *key.shape[2:]), + value.reshape(B * S, *value.shape[2:]), + cu, + cu, + S, + S, + is_causal, + scale, + ws, + enable_gqa, + None, # seqused_k — inference-only per upstream + None, # block_table — inference-only per upstream + None, # num_splits — perf-tuning knob, not exposed + ) + return out_thd.view_as(query), lse, rng_state + + +@_doc_packed_attn_op.register_fake +def _doc_packed_attn_op_fake( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + cu_seq_q: torch.Tensor, + n_docs: int, + scale: float | None = None, + window_size: list[int] | None = None, + enable_gqa: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, S, Hq = query.shape[:3] + out = torch.empty_like(query) + # Matches torch_attn::_varlen_attn fake shapes (after THD flatten): + # lse: (num_heads, total_q) float32 + # rng_state: (2,) int64 placeholder + lse = torch.empty((Hq, B * S), dtype=torch.float32, device=query.device) + rng_state = torch.empty((2,), dtype=torch.int64, device=query.device) + return out, lse, rng_state + + +def _doc_packed_attn_setup_context(ctx: Any, inputs: tuple, output: Any) -> None: + ( + query, + key, + value, + cu_seq_q, + n_docs, + scale, + window_size, + enable_gqa, + ) = inputs + out, lse, rng_state = output + ctx.save_for_backward(query, key, value, cu_seq_q, out, lse, rng_state) + ctx.n_docs = n_docs + ctx.scale = scale + ctx.window_size = window_size + + +@torch.library.custom_op("autoparallel::doc_packed_attn_backward_op", mutates_args=()) +def _doc_packed_attn_backward_op( + grad_out: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + out: torch.Tensor, + lse: torch.Tensor, + cu_seq_q: torch.Tensor, + rng_state: torch.Tensor, + n_docs: int, + scale: float | None = None, + window_size: list[int] | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Backward custom op for ``doc_packed_attn``. + + Takes the same ``[B, S, H, D]`` / ``[B, MAX_DOCS+1]`` shapes as the forward. + The THD reshape and the flat ``cu_seq_q`` construction happen inside the + op body and are invisible to AutoParallel — so the AP sharding strategy + can reason about the natural ``[B, S, H, D]`` layout symmetrically with + the forward, and the divisibility check uses B directly. + """ + B, S = query.shape[:2] + ws = list(window_size) if window_size is not None else [-1, -1] + is_causal = ws == [-1, 0] + cu = _build_flat_cu_seq_q(cu_seq_q, n_docs, S) + + dq_thd, dk_thd, dv_thd = torch.ops.torch_attn._varlen_attn_backward( + grad_out.reshape(B * S, *grad_out.shape[2:]), + query.reshape(B * S, *query.shape[2:]), + key.reshape(B * S, *key.shape[2:]), + value.reshape(B * S, *value.shape[2:]), + out.reshape(B * S, *out.shape[2:]), + lse, + cu, + cu, + S, + S, + is_causal, + rng_state, + scale, + ws, + ) + return ( + dq_thd.view_as(query), + dk_thd.view_as(key), + dv_thd.view_as(value), + ) + + +@_doc_packed_attn_backward_op.register_fake +def _doc_packed_attn_backward_op_fake( + grad_out: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + out: torch.Tensor, + lse: torch.Tensor, + cu_seq_q: torch.Tensor, + rng_state: torch.Tensor, + n_docs: int, + scale: float | None = None, + window_size: list[int] | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return torch.empty_like(query), torch.empty_like(key), torch.empty_like(value) + + +def _doc_packed_attn_backward( + ctx: Any, + grad_out: torch.Tensor, + grad_lse: torch.Tensor, + grad_rng: torch.Tensor, +) -> tuple[torch.Tensor | None, ...]: + """Backward callback for doc_packed_attn. + + Routes to the ``doc_packed_attn_backward_op`` custom op so that the + THD reshape stays hidden from AutoParallel. + """ + query, key, value, cu_seq_q, out, lse, rng_state = ctx.saved_tensors + dq, dk, dv = torch.ops.autoparallel.doc_packed_attn_backward_op( + grad_out, + query, + key, + value, + out, + lse, + cu_seq_q, + rng_state, + ctx.n_docs, + ctx.scale, + ctx.window_size, + ) + # Match the forward signature: query, key, value, cu_seq_q, n_docs, + # scale, window_size, enable_gqa. Only q/k/v get gradients. + return (dq, dk, dv, None, None, None, None, None) + + +torch.library.register_autograd( + "autoparallel::doc_packed_attn_op", + _doc_packed_attn_backward, + setup_context=_doc_packed_attn_setup_context, +) + + +def doc_packed_attn( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + cu_seq_q: torch.Tensor, + n_docs: int, + *, + scale: float | None = None, + window_size: tuple[int, int] = (-1, -1), + enable_gqa: bool = False, +) -> torch.Tensor: + """Document-packed variable-length attention. + + Mirrors :func:`torch.nn.attention.varlen.varlen_attn`'s API, but takes + ``query``/``key``/``value`` in ``[B, S, H, D]`` layout and ``cu_seq_q`` as + a padded ``[B, MAX_DOCS+1]`` tensor so the leading dim can be sharded as + ``Shard(0)`` by AutoParallel. ``cu_seq_k`` is assumed to equal ``cu_seq_q`` + (self-attention) and ``max_q == max_k == S``. + + Args: + query: ``[B, S, H_q, D]`` query tensor. + key: ``[B, S, H_kv, D]`` key tensor. + value: ``[B, S, H_kv, D]`` value tensor. + cu_seq_q: ``[B, MAX_DOCS+1]`` int32 padded cumulative sequence positions. + ``cu_seq_q[b, 0] == 0`` and ``cu_seq_q[b, n_docs] == S``. Real + document boundaries occupy entries ``1..n_real_docs[b]``; entries + ``n_real_docs[b]..n_docs`` should be padded with ``S`` so the + trailing slots appear as zero-length documents (which the kernel + skips). ``n_docs`` is uniform across batch elements. + n_docs: Document count per batch element (SymInt under ``dynamic=True``). + Batches with fewer real documents pad as described above. + scale: Softmax scale. ``None`` uses ``1 / sqrt(D)``. + window_size: ``(left, right)`` sliding-window sizes per token. + ``(-1, -1)`` is full attention (default), ``(-1, 0)`` is causal, + ``(W, 0)`` is causal sliding window of size ``W``. + enable_gqa: Allow ``H_kv != H_q`` via grouped-query attention. + + Returns: + Output tensor of the same shape as ``query``. + """ + num_heads_q = query.size(2) + num_heads_k = key.size(2) + if not enable_gqa and num_heads_q != num_heads_k: + raise ValueError( + f"Expect query and key/value to have the same number of heads " + f"but got Hq={num_heads_q} and Hkv={num_heads_k}. " + f"Try setting enable_gqa=True for GQA." + ) + if enable_gqa and num_heads_q % num_heads_k != 0: + raise ValueError( + f"Expect number of query heads to be a multiple of kv heads for GQA " + f"but got Hq={num_heads_q} and Hkv={num_heads_k}." + ) + + out, _lse, _rng = torch.ops.autoparallel.doc_packed_attn_op( + query, + key, + value, + cu_seq_q, + n_docs, + scale, + list(window_size), + enable_gqa, + ) + return out diff --git a/autoparallel/shardings/propagation_rules.py b/autoparallel/shardings/propagation_rules.py index 0a3eaf2a..7b686e04 100644 --- a/autoparallel/shardings/propagation_rules.py +++ b/autoparallel/shardings/propagation_rules.py @@ -42,6 +42,9 @@ # need to import this to have the dtype_cast registered from ..cast_parametrization import dtype_cast # noqa + +# need to import this to have doc_packed_attn_op registered +from ..ops import doc_packed_attn # noqa: F401 from .dtensor_sharding_helpers import _try_single_dim_strategy, get_op_strategy _op_rules = {} @@ -670,3 +673,187 @@ def uniform_strategy(mesh, op_schema: OpSchema): return expand_to_full_mesh_op_strategy( mesh, op_schema, single_mesh_dim_strategies, input_index=1 ) + + +@register_rule(torch.ops.autoparallel.doc_packed_attn_op.default) +def doc_packed_attn_op_rule(mesh, op_schema): + """Sharding strategy for autoparallel::doc_packed_attn_op. + + Schema: + (query, key, value, cu_seq_q, n_docs, scale, window_size, enable_gqa) + -> (out, lse, rng_state) + + Tensor layouts: + query : [B, S, H_q, D] -- batch dim 0, heads dim 2 + key : [B, S, H_kv, D] + value : [B, S, H_kv, D] + cu_seq_q: [B, MAX_DOCS+1] -- batch dim 0; MAX_DOCS dim never sharded + out : same as query + lse : [H_q, B*S] -- heads dim 0, T=B*S dim 1 + rng_state: [2] -- always Replicate + + Per-mesh-dim options: + - "R" Replicate everywhere. + - "B" Shard batch: q/k/v/cu/out dim 0, lse dim 1. + - "H" Shard heads: q/k/v/out dim 2, lse dim 0. Requires H_q and H_kv + both divisible by the mesh-dim size so GQA broadcasting keeps + the H_q/H_kv ratio uniform per rank. + + Not shardable in Phase 1: + - S dim (sequence) -- would require context parallelism. + - cu_seq_q's MAX_DOCS+1 dim -- doc-boundary descriptors are not a + parallel axis. + """ + from torch.distributed.tensor._ops.utils import expand_to_full_mesh_op_strategy + + q_strat = op_schema.args_schema[0] + k_strat = op_schema.args_schema[1] + q_meta = q_strat.strategies[0].output_spec.tensor_meta + k_meta = k_strat.strategies[0].output_spec.tensor_meta + B, S, H_q, _D = q_meta.shape + H_kv = k_meta.shape[2] + + # Layout per role: [out, lse, rng | q, k, v, cu_seq_q, n_docs, scale, window_size, enable_gqa] + # input_index = 3 (three outputs). + R = Replicate() + # Build PlacementLists for each per-mesh-dim option. + # "B" sharding — q/k/v/cu/out shard dim 0, lse shards dim 1, rng replicated. + b_strat = [ + Shard(0), + Shard(1), + R, # outputs + Shard(0), + Shard(0), + Shard(0), + Shard(0), # tensor inputs + None, + None, + None, + None, + ] # n_docs, scale, window_size, enable_gqa + # "H" sharding — q/k/v/out shard dim 2, lse shards dim 0, cu/rng replicated. + h_strat = [ + Shard(2), + Shard(0), + R, + Shard(2), + Shard(2), + Shard(2), + R, + None, + None, + None, + None, + ] + # "R" replicate-everywhere. + r_strat = [R, R, R, R, R, R, R, None, None, None, None] + + single_mesh_dim_strategies = [r_strat] + if B > 1: + single_mesh_dim_strategies.append(b_strat) + if H_q > 1 and H_kv > 1: + single_mesh_dim_strategies.append(h_strat) + + lse_meta = _gen_tensor_meta( + torch.empty((H_q, B * S), dtype=torch.float32, device="meta") + ) + rng_meta = _gen_tensor_meta(torch.empty((2,), dtype=torch.int64, device="meta")) + + return expand_to_full_mesh_op_strategy( + mesh, + op_schema, + single_mesh_dim_strategies, + output_tensor_meta=(q_meta, lse_meta, rng_meta), + input_index=3, + ) + + +@register_rule(torch.ops.autoparallel.doc_packed_attn_backward_op.default) +def doc_packed_attn_backward_op_rule(mesh, op_schema): + """Sharding strategy for autoparallel::doc_packed_attn_backward_op. + + Schema: + (grad_out, query, key, value, out, lse, cu_seq_q, rng_state, + n_docs, scale, window_size) + -> (dq, dk, dv) + + Tensor layouts (same as the forward — the THD reshape is hidden inside + the op body): + grad_out, q, out : [B, S, H_q, D] + k, v : [B, S, H_kv, D] + lse : [H_q, B*S] + cu_seq_q : [B, MAX_DOCS+1] + rng_state : [2] + + Per-mesh-dim options match the forward: + - "R" Replicate everywhere. + - "B" Shard batch: q/k/v/grad_out/out/cu_seq_q/dq/dk/dv dim 0, lse dim 1. + - "H" Shard heads: q/k/v/grad_out/out/dq/dk/dv dim 2, lse dim 0. + """ + from torch.distributed.tensor._ops.utils import expand_to_full_mesh_op_strategy + + q_strat = op_schema.args_schema[1] + k_strat = op_schema.args_schema[2] + q_meta = q_strat.strategies[0].output_spec.tensor_meta + k_meta = k_strat.strategies[0].output_spec.tensor_meta + B, S, H_q, _D = q_meta.shape + H_kv = k_meta.shape[2] + + # Layout per role: [dq, dk, dv | grad_out, q, k, v, out, lse, cu_seq_q, + # rng_state, n_docs, scale, window_size] + # input_index = 3. + R = Replicate() + # "B" — batch sharding: q/k/v/grad_out/out/cu_seq_q/dq/dk/dv shard dim 0, + # lse shards dim 1, rng_state replicated. + b_strat = [ + Shard(0), + Shard(0), + Shard(0), # dq, dk, dv + Shard(0), + Shard(0), + Shard(0), + Shard(0), + Shard(0), + Shard(1), + # grad_out, q, k, v, out, lse + Shard(0), # cu_seq_q + R, # rng_state + None, + None, + None, + ] # n_docs, scale, window_size + # "H" — heads sharding: q/k/v/grad_out/out/dq/dk/dv shard dim 2, + # lse shards dim 0, cu/rng replicated. + h_strat = [ + Shard(2), + Shard(2), + Shard(2), + Shard(2), + Shard(2), + Shard(2), + Shard(2), + Shard(2), + Shard(0), + R, + R, + None, + None, + None, + ] + # "R" — replicate everywhere. + r_strat = [R, R, R, R, R, R, R, R, R, R, R, None, None, None] + + single_mesh_dim_strategies = [r_strat] + if B > 1: + single_mesh_dim_strategies.append(b_strat) + if H_q > 1 and H_kv > 1: + single_mesh_dim_strategies.append(h_strat) + + return expand_to_full_mesh_op_strategy( + mesh, + op_schema, + single_mesh_dim_strategies, + # Outputs are (dq, dk, dv) — same shapes as q, k, v (v has k's shape). + output_tensor_meta=(q_meta, k_meta, k_meta), + input_index=3, + ) diff --git a/tests/test_ops.py b/tests/test_ops.py index 904d1a12..99f0c52b 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -10,7 +10,7 @@ from autoparallel import AutoParallel, with_sharding_constraint from autoparallel.collectives import local_map -from autoparallel.ops import permutation +from autoparallel.ops import doc_packed_attn, permutation def get_local_map_nodes(graph, is_backward=False): @@ -396,3 +396,495 @@ def test_cuda_tensor(self): result = permutation(x, axis=0) assert result.device == x.device assert result.shape == x.shape + + +# --------------------------------------------------------------------------- +# doc_packed_attn +# --------------------------------------------------------------------------- + + +def _build_padded_cu_seq_q( + per_batch_doc_lens: list[list[int]], + n_docs: int, + seq_len: int, + device: str | torch.device = "cuda", +) -> torch.Tensor: + """Build a ``[B, n_docs+1]`` int32 cu_seq_q tensor from per-batch doc lengths. + + Rows with fewer than ``n_docs`` real documents are padded by repeating + ``seq_len`` at the trailing positions, so those slots represent zero-length + documents that the kernel skips. + """ + B = len(per_batch_doc_lens) + cu = torch.full((B, n_docs + 1), seq_len, dtype=torch.int32, device=device) + for b, lens in enumerate(per_batch_doc_lens): + assert ( + sum(lens) == seq_len + ), f"batch {b}: doc lengths {lens} sum to {sum(lens)}, expected {seq_len}" + assert ( + len(lens) <= n_docs + ), f"batch {b}: {len(lens)} real docs > n_docs={n_docs}" + cumsum = torch.tensor( + [0] + list(torch.tensor(lens).cumsum(0).tolist()), + dtype=torch.int32, + device=device, + ) + cu[b, : len(cumsum)] = cumsum + return cu + + +def _per_doc_sdpa_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + per_batch_doc_lens: list[list[int]], + window_size: tuple[int, int], +) -> torch.Tensor: + """Reference implementation: SDPA per document, then concat. + + Builds an explicit mask for sliding-window cases so it exactly matches + flash-attention's per-doc causal+window semantics. + """ + import torch.nn.functional as F + + is_full = window_size == (-1, -1) + is_causal = window_size == (-1, 0) + out = torch.zeros_like(q) + for b in range(q.shape[0]): + offset = 0 + for L in per_batch_doc_lens[b]: + q_doc = q[b, offset : offset + L].transpose(0, 1).unsqueeze(0) + k_doc = k[b, offset : offset + L].transpose(0, 1).unsqueeze(0) + v_doc = v[b, offset : offset + L].transpose(0, 1).unsqueeze(0) + if is_full: + o = F.scaled_dot_product_attention( + q_doc, k_doc, v_doc, is_causal=False, enable_gqa=True + ) + elif is_causal: + o = F.scaled_dot_product_attention( + q_doc, k_doc, v_doc, is_causal=True, enable_gqa=True + ) + else: + window_left = window_size[0] + 1 # inclusive of self + i_idx = torch.arange(L, device=q.device).unsqueeze(1) + j_idx = torch.arange(L, device=q.device).unsqueeze(0) + mask = (j_idx <= i_idx) & (i_idx - j_idx < window_left) + o = F.scaled_dot_product_attention( + q_doc, k_doc, v_doc, attn_mask=mask, enable_gqa=True + ) + out[b, offset : offset + L] = o.squeeze(0).transpose(0, 1) + offset += L + return out + + +class TestDocPackedAttn: + """Tests for autoparallel.ops.doc_packed_attn. + + Phase-1 contract: + - q/k/v: ``[B, S, H, D]`` (THD reshaped internally). + - cu_seq_q: ``[B, n_docs+1]`` int32; padding is repeated ``S`` past real docs. + - n_docs is uniform across batch elements. + - Returns a single ``[B, S, H, D]`` tensor (the underlying op also yields + ``lse``/``rng_state`` for backward but the public helper hides them). + """ + + # -- schema / fake mode ---------------------------------------------------- + + def test_schema(self): + """Op schema mirrors varlen_attn's argument style and returns 3 tensors.""" + schema = torch.ops.autoparallel.doc_packed_attn_op.default._schema + s = str(schema) + assert "Tensor query" in s + assert "Tensor cu_seq_q" in s + assert "SymInt n_docs" in s + assert "SymInt[]? window_size" in s + assert "bool enable_gqa" in s + # 3-tuple return: (out, lse, rng_state) + assert "-> (Tensor, Tensor, Tensor)" in s + + def test_fake_tensor_trace(self): + """FakeTensorMode tracing produces a tensor with an autograd grad_fn.""" + from torch._subclasses import FakeTensorMode + + with FakeTensorMode(): + B, S, Hq, Hkv, D = 1, 128, 8, 4, 64 + q = torch.empty( + B, S, Hq, D, dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + k = torch.empty( + B, S, Hkv, D, dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + v = torch.empty( + B, S, Hkv, D, dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + cu = torch.empty(B, 5, dtype=torch.int32, device="cuda") + out = doc_packed_attn( + q, k, v, cu, n_docs=3, window_size=(-1, 0), enable_gqa=True + ) + assert out.shape == q.shape + assert out.dtype == q.dtype + assert out.grad_fn is not None + + # -- input validation ------------------------------------------------------ + + def test_gqa_disabled_requires_matching_heads(self): + """Without enable_gqa, Hq must equal Hkv.""" + from torch._subclasses import FakeTensorMode + + with FakeTensorMode(): + q = torch.empty(1, 128, 8, 64, dtype=torch.bfloat16, device="cuda") + k = torch.empty(1, 128, 4, 64, dtype=torch.bfloat16, device="cuda") + v = torch.empty(1, 128, 4, 64, dtype=torch.bfloat16, device="cuda") + cu = torch.empty(1, 5, dtype=torch.int32, device="cuda") + with pytest.raises(ValueError, match="enable_gqa=True"): + doc_packed_attn(q, k, v, cu, n_docs=3) + + def test_gqa_requires_divisible_heads(self): + """With enable_gqa, Hq must be a multiple of Hkv.""" + from torch._subclasses import FakeTensorMode + + with FakeTensorMode(): + q = torch.empty(1, 128, 9, 64, dtype=torch.bfloat16, device="cuda") + k = torch.empty(1, 128, 4, 64, dtype=torch.bfloat16, device="cuda") + v = torch.empty(1, 128, 4, 64, dtype=torch.bfloat16, device="cuda") + cu = torch.empty(1, 5, dtype=torch.int32, device="cuda") + with pytest.raises(ValueError, match="multiple of kv heads"): + doc_packed_attn(q, k, v, cu, n_docs=3, enable_gqa=True) + + # -- parity (B=1, B>1, uniform & mixed doc counts) ------------------------- + + # Curated cases: (name, per_batch_doc_lens, n_docs, window_size) + _PARITY_CASES = [ + # B=1 + ("B1_causal", [[100, 80, 76]], 3, (-1, 0)), + ("B1_full", [[100, 80, 76]], 3, (-1, -1)), + ("B1_swa_W32", [[100, 80, 76]], 3, (31, 0)), + ("B1_single_doc", [[256]], 1, (-1, 0)), + # B>1, uniform doc counts (no padding) + ("B2_uniform", [[100, 80, 76], [128, 64, 64]], 3, (-1, 0)), + ("B2_uniform_swa", [[100, 80, 76], [128, 64, 64]], 3, (31, 0)), + ( + "B4_uniform", + [ + [64, 64, 64, 64], + [128, 64, 32, 32], + [200, 24, 24, 8], + [80, 80, 48, 48], + ], + 4, + (-1, 0), + ), + # B>1, mixed doc counts (zero-length-doc padding) + ( + "B2_mixed_3v5", + [[100, 80, 76], [60, 60, 60, 60, 16]], + 5, + (-1, 0), + ), + ( + "B3_mixed_2v4v6", + [ + [200, 56], + [100, 60, 50, 46], + [60, 50, 50, 40, 30, 26], + ], + 6, + (-1, 0), + ), + ] + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize( + "name, per_batch_doc_lens, n_docs, window_size", + _PARITY_CASES, + ids=[c[0] for c in _PARITY_CASES], + ) + def test_parity_vs_per_doc_sdpa( + self, name, per_batch_doc_lens, n_docs, window_size + ): + """Forward + backward match per-document SDPA reference at bf16 precision.""" + torch.manual_seed(42) + S = 256 + B = len(per_batch_doc_lens) + Hq, Hkv, D = 8, 4, 64 + cu_seq_q = _build_padded_cu_seq_q(per_batch_doc_lens, n_docs, S) + + q = torch.randn( + B, S, Hq, D, dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + k = torch.randn( + B, S, Hkv, D, dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + v = torch.randn( + B, S, Hkv, D, dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + + out = doc_packed_attn( + q, k, v, cu_seq_q, n_docs, window_size=window_size, enable_gqa=True + ) + grad = torch.randn_like(out) + out.backward(grad) + + q_ref = q.detach().clone().requires_grad_(True) + k_ref = k.detach().clone().requires_grad_(True) + v_ref = v.detach().clone().requires_grad_(True) + out_ref = _per_doc_sdpa_reference( + q_ref, k_ref, v_ref, per_batch_doc_lens, window_size + ) + out_ref.backward(grad) + + # bf16 epsilon ~ 8e-3; allow ~4x headroom because flash and SDPA + # accumulate dk across the sequence dim in different orders, and + # GQA broadcasting compounds rounding for the longest single-doc cases. + tol = 3e-2 + assert torch.isfinite(out).all() + assert (out - out_ref).abs().max().item() < tol, ( + f"forward parity failure: max_abs_diff=" + f"{(out - out_ref).abs().max().item():.2e}" + ) + for name_g, g, g_ref in [ + ("q", q.grad, q_ref.grad), + ("k", k.grad, k_ref.grad), + ("v", v.grad, v_ref.grad), + ]: + diff = (g - g_ref).abs().max().item() + assert diff < tol, f"{name_g}.grad parity failure: max_abs_diff={diff:.2e}" + + # -- padding semantics ----------------------------------------------------- + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_zero_length_padding_does_not_affect_real_docs(self): + """Padding past n_docs (zero-length trailing docs) must not change output. + + Two cu_seq_q tensors with identical real-doc boundaries but different + padding (n_docs=3 tight vs n_docs=8 padded) should produce bit-identical + results because the trailing entries describe zero-length documents. + """ + torch.manual_seed(0) + B, S, Hq, Hkv, D = 1, 256, 8, 4, 64 + per_batch = [[100, 80, 76]] + + q = torch.randn(B, S, Hq, D, dtype=torch.bfloat16, device="cuda") + k = torch.randn(B, S, Hkv, D, dtype=torch.bfloat16, device="cuda") + v = torch.randn(B, S, Hkv, D, dtype=torch.bfloat16, device="cuda") + + cu_tight = _build_padded_cu_seq_q(per_batch, n_docs=3, seq_len=S) + cu_padded = _build_padded_cu_seq_q(per_batch, n_docs=8, seq_len=S) + + out_tight = doc_packed_attn( + q, k, v, cu_tight, n_docs=3, window_size=(-1, 0), enable_gqa=True + ) + out_padded = doc_packed_attn( + q, k, v, cu_padded, n_docs=8, window_size=(-1, 0), enable_gqa=True + ) + assert torch.equal(out_tight, out_padded), ( + "zero-length padding changed the output — kernel may not handle " + "trailing empty segments correctly" + ) + + # -- no recompute in backward --------------------------------------------- + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_backward_does_not_recompute_forward(self): + """Profile N iterations: flash-fwd count must equal iteration count. + + If the backward recomputed the forward, fwd_count would be 2*N. We + save (out, lse, rng_state) from forward and dispatch directly to + _flash_attention_backward, so fwd_count == bwd_count == N. + """ + from torch.profiler import ProfilerActivity, profile + + torch.manual_seed(0) + B, S, Hq, Hkv, D = 1, 256, 8, 4, 64 + per_batch = [[100, 80, 76]] + cu_seq_q = _build_padded_cu_seq_q(per_batch, n_docs=3, seq_len=S) + + q = torch.randn( + B, S, Hq, D, dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + k = torch.randn( + B, S, Hkv, D, dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + v = torch.randn( + B, S, Hkv, D, dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + + n_iters = 3 + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for _ in range(n_iters): + out = doc_packed_attn( + q, k, v, cu_seq_q, 3, window_size=(-1, 0), enable_gqa=True + ) + out.backward(torch.randn_like(out)) + q.grad = k.grad = v.grad = None + + fwd_count = bwd_count = 0 + for row in prof.key_averages(): + if "flash_attention_forward" in row.key: + fwd_count = row.count + elif "flash_attention_backward" in row.key: + bwd_count = row.count + assert fwd_count == n_iters, ( + f"flash_attention_forward called {fwd_count} times for {n_iters} " + f"iterations — backward is recomputing the forward" + ) + assert bwd_count == n_iters, ( + f"flash_attention_backward called {bwd_count} times for {n_iters} " + f"iterations" + ) + + +class TestDocPackedAttnShardingStrategy: + """Tests for AP's sharding strategy on autoparallel::doc_packed_attn_op. + + These verify the strategy enumerates the right placement options and that + the solver picks the expected DP / DP+TP shardings when constraints pin + the inputs. They use ``JustAttn`` (no surrounding linear layers) so that + the test is independent of how the solver scores larger graphs. + """ + + class _JustAttn(nn.Module): + """Minimal model: a single doc_packed_attn call, no linears. + + Used to isolate AP's choice for the attention op from upstream/down- + stream sharding decisions (which a real transformer would dominate). + """ + + def __init__(self, n_docs: int): + super().__init__() + self.n_docs = n_docs + + def forward(self, q, k, v, cu_seq_q): + return doc_packed_attn( + q, + k, + v, + cu_seq_q, + n_docs=self.n_docs, + window_size=(-1, 0), + enable_gqa=True, + ) + + @staticmethod + def _input_fn_factory(B, S, Hq, Hkv, D, n_docs): + """Returns an input_fn that builds q/k/v plus uniform-doc cu_seq_q.""" + + def input_fn(): + q = torch.randn(B, S, Hq, D, device="cuda", dtype=torch.bfloat16) + k = torch.randn(B, S, Hkv, D, device="cuda", dtype=torch.bfloat16) + v = torch.randn(B, S, Hkv, D, device="cuda", dtype=torch.bfloat16) + per_doc = S // n_docs + base = torch.arange(n_docs + 1, dtype=torch.int32, device="cuda") * per_doc + cu = base.unsqueeze(0).expand(B, -1).contiguous() + return q, k, v, cu + + return input_fn + + @staticmethod + def _find_doc_packed_node(gm): + for node in gm.graph.nodes: + if ( + node.op == "call_function" + and getattr(node.target, "name", lambda: "")() + == "autoparallel::doc_packed_attn_op" + ): + return node + raise AssertionError("autoparallel::doc_packed_attn_op not found in graph") + + def test_dp_only_pins_shard0(self, device_mesh_1d): + """1D dp mesh, all inputs pinned Shard(0): op picks Shard(0) end-to-end.""" + # B must be divisible by mesh size (256 in the fixture). + B, S, Hq, Hkv, D, n_docs = 256, 256, 8, 4, 64, 4 + model = self._JustAttn(n_docs).to(dtype=torch.bfloat16) + input_fn = self._input_fn_factory(B, S, Hq, Hkv, D, n_docs) + + with AutoParallel( + model, input_fn, device_mesh_1d, repeated_subgraphs=False + ) as autop: + autop.add_input_constraints( + [ + (Shard(0),), # q + (Shard(0),), # k + (Shard(0),), # v + (Shard(0),), # cu_seq_q + ] + ) + autop.add_output_constraints([(Shard(0),)]) + sol = autop.optimize_placement() + + node = self._find_doc_packed_node(autop.gm) + spec = sol[node] + in_placements = [s.placements for s in spec.input_specs] + out_placements = [s.placements for s in spec.output_specs] + + # All four inputs must be Shard(0). + for i, p in enumerate(in_placements): + assert p == (Shard(0),), f"input {i} placement {p} != (Shard(0),)" + + # out: Shard(0); lse [Hq, B*S] sharded on dim 1; rng replicated. + assert out_placements[0] == (Shard(0),), out_placements[0] + assert out_placements[1] == (Shard(1),), out_placements[1] + assert out_placements[2] == (Replicate(),), out_placements[2] + + def test_dp_tp_pins_shard0_shard2(self, device_mesh_2d): + """2D dp×tp mesh, q/k/v pinned (Shard(0), Shard(2)). + + cu_seq_q has no head dim so it's (Shard(0), Replicate()). Output + inherits q's placement. lse rearranges to (Shard(1), Shard(0)) + because T == B*S occupies dim 1 and Hq occupies dim 0. + """ + # Mesh is (32, 8): dp=32 must divide B; tp=8 must divide Hq and Hkv. + dp_size, tp_size = device_mesh_2d.shape + B, S, Hq, Hkv, D, n_docs = 32, 256, 8, 8, 64, 4 # Hq=Hkv=tp_size + if Hq % tp_size != 0 or Hkv % tp_size != 0: + pytest.skip(f"tp_size={tp_size} doesn't divide Hq={Hq} or Hkv={Hkv}") + if B % dp_size != 0: + pytest.skip(f"dp_size={dp_size} doesn't divide B={B}") + + model = self._JustAttn(n_docs).to(dtype=torch.bfloat16) + input_fn = self._input_fn_factory(B, S, Hq, Hkv, D, n_docs) + + with AutoParallel( + model, input_fn, device_mesh_2d, repeated_subgraphs=False + ) as autop: + autop.add_input_constraints( + [ + (Shard(0), Shard(2)), + (Shard(0), Shard(2)), + (Shard(0), Shard(2)), + (Shard(0), Replicate()), + ] + ) + autop.add_output_constraints([(Shard(0), Shard(2))]) + sol = autop.optimize_placement() + + node = self._find_doc_packed_node(autop.gm) + spec = sol[node] + in_placements = [s.placements for s in spec.input_specs] + out_placements = [s.placements for s in spec.output_specs] + + assert in_placements[0] == (Shard(0), Shard(2)), in_placements[0] # q + assert in_placements[1] == (Shard(0), Shard(2)), in_placements[1] # k + assert in_placements[2] == (Shard(0), Shard(2)), in_placements[2] # v + assert in_placements[3] == (Shard(0), Replicate()), in_placements[3] # cu + + assert out_placements[0] == (Shard(0), Shard(2)), out_placements[0] # out + assert out_placements[1] == (Shard(1), Shard(0)), out_placements[1] # lse + assert out_placements[2] == (Replicate(), Replicate()), out_placements[2] + + def test_apply_placement_succeeds(self, device_mesh_1d): + """End-to-end smoke test that apply_placement produces a runnable module.""" + B, S, Hq, Hkv, D, n_docs = 256, 256, 8, 4, 64, 4 + model = self._JustAttn(n_docs).to(dtype=torch.bfloat16) + input_fn = self._input_fn_factory(B, S, Hq, Hkv, D, n_docs) + + with AutoParallel( + model, input_fn, device_mesh_1d, repeated_subgraphs=False + ) as autop: + autop.add_input_constraints( + [(Shard(0),), (Shard(0),), (Shard(0),), (Shard(0),)] + ) + autop.add_output_constraints([(Shard(0),)]) + sol = autop.optimize_placement() + parallel_mod = autop.apply_placement(sol) + assert parallel_mod is not None