From f5023b95961005c02b1e380e4ec3ca68adb1cf2e Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Thu, 13 Aug 2026 10:23:20 +0000 Subject: [PATCH] distributed: admit exact hybrid GDN under Ulysses CP Add aligned packed-sequence cuts, a receipt-validated canonical scan-state handoff, serving-paired convolution halos, and fail-closed topology and toolchain admission. Cover collator attestation, topology guards, kernel-pin ownership, and composed byte-contract behavior with conventional tests. --- .../data/collators/sequence_shard_collator.py | 158 ++++++++++++- src/xorl/models/auto.py | 92 +++++++- src/xorl/ops/kernel_config_pin.py | 181 +++++++++++++++ .../linear_attention/layers/gated_deltanet.py | 19 +- .../linear_attention/modules/conv_contract.py | 161 +++++++++++++ src/xorl/ops/linear_attention/ops/cp/chain.py | 213 ++++++++++++++++++ .../ops/gated_delta_rule/chunk.py | 32 ++- .../test_ulysses_byte_alignment.py | 85 ++++++- .../test_qwen35_hybrid_ulysses_admission.py | 180 +++++++++++++++ tests/ops/test_gdn_conv_contract.py | 7 +- 10 files changed, 1105 insertions(+), 23 deletions(-) create mode 100644 src/xorl/ops/kernel_config_pin.py create mode 100644 src/xorl/ops/linear_attention/ops/cp/chain.py create mode 100644 tests/models/test_qwen35_hybrid_ulysses_admission.py diff --git a/src/xorl/data/collators/sequence_shard_collator.py b/src/xorl/data/collators/sequence_shard_collator.py index 1fc70d8c..f9c438a6 100644 --- a/src/xorl/data/collators/sequence_shard_collator.py +++ b/src/xorl/data/collators/sequence_shard_collator.py @@ -1,5 +1,6 @@ +import logging from dataclasses import dataclass -from typing import Dict +from typing import Dict, List, Tuple import torch @@ -10,6 +11,84 @@ from .packing_concat_collator import add_flash_attention_kwargs_from_position_ids +logger = logging.getLogger("xorl.gdn_cp_collator") +_alignment_engagement_logged = False + +# The GDN chunk kernels operate on a 64-token grid per document. The exact-CP +# collator contract keeps every shard cut on every crossing document's own +# grid by 64-aligning every document start via inter-document pad-docs and +# making the per-rank shard length a multiple of 64. +GDN_CP_CHUNK = 64 + + +def find_document_boundaries(position_ids_row: torch.Tensor) -> List[int]: + """Document boundaries of a packed stream from position-id resets.""" + pos = position_ids_row.reshape(-1) + starts = (pos == 0).nonzero(as_tuple=False).reshape(-1).tolist() + if not starts or starts[0] != 0: + starts = [0, *starts] + return [*starts, pos.numel()] + + +def gdn_cp_alignment_segments(doc_bounds: List[int], chunk: int = GDN_CP_CHUNK) -> List[Tuple[str, int, int]]: + """C2 splice plan: ('pad', 0, gap) and ('doc', start, end) segments such + that every document START in the spliced stream is a multiple of `chunk`, + and the stream END is too (so C1 tail pad-docs also start on the grid). + + Pure function of the document layout so admission checks and the collator + use the same splice plan. + """ + segments: List[Tuple[str, int, int]] = [] + out_len = 0 + for start, end in zip(doc_bounds[:-1], doc_bounds[1:]): + gap = (-out_len) % chunk + if gap: + segments.append(("pad", 0, gap)) + out_len += gap + segments.append(("doc", start, end)) + out_len += end - start + tail_gap = (-out_len) % chunk + if tail_gap: + segments.append(("pad", 0, tail_gap)) + return segments + + +def gdn_cp_spliced_length(segments: List[Tuple[str, int, int]]) -> int: + return sum(end - start if kind == "doc" else end for kind, start, end in segments) + + +def apply_alignment_segments( + tensor: torch.Tensor, + segments: List[Tuple[str, int, int]], + dim: int, + pad_value: float | int = 0, + sequential_positions: bool = False, +) -> torch.Tensor: + """Splice a token-aligned tensor along `dim` per the C2 plan. Pad segments + are filled with `pad_value`, or with `arange(gap)` when + `sequential_positions` (each inter-document pad is its own document).""" + if dim < 0: + dim = tensor.ndim + dim + pieces: List[torch.Tensor] = [] + for kind, start, end in segments: + if kind == "doc": + pieces.append(tensor.narrow(dim, start, end - start)) + else: + gap = end + shape = list(tensor.shape) + shape[dim] = gap + if sequential_positions: + seq = torch.arange(gap, device=tensor.device, dtype=tensor.dtype) + view = [1] * tensor.ndim + view[dim] = gap + pieces.append(seq.view(view).expand(shape).contiguous()) + else: + pieces.append( + torch.full(shape, fill_value=pad_value, dtype=tensor.dtype, device=tensor.device) + ) + return torch.cat(pieces, dim=dim) + + def zigzag_reorder_packed_sequence( tensor: torch.Tensor, position_ids: torch.Tensor, @@ -92,15 +171,27 @@ class TextSequenceShardCollator(DataCollator): fa_max_length_bucket: If > 0, round the flash-attn max_length up to a multiple of this value (upper bound only; correctness via cu_seqlens) to avoid torch.compile recompiling on ragged packs. 0 = off. + gdn_exact_cp_align: Engage the exact-GDN CP alignment contract: + 64-align every document start via inter-document pad-docs and pad + the packed total to a multiple of 64 * cp_size, so every shard cut + lands on each crossing document's 64-token chunk grid. + Incompatible with ring attention. """ pad_token_id: int = 0 fa_max_length_bucket: int = 0 + gdn_exact_cp_align: bool = False def __post_init__(self): self.cp_size = get_parallel_state().cp_size self.cp_rank = get_parallel_state().cp_rank self.ringattn_size = get_parallel_state().ringattn_size + if self.gdn_exact_cp_align and self.ringattn_size > 1: + raise ValueError( + "gdn_exact_cp_align is incompatible with ring attention " + "because zigzag reorder breaks the per-document 64-token " + "chunk grid. Fail closed.", + ) def sp_slice(self, tensor: "torch.Tensor", dim: int = -1) -> "torch.Tensor": """ @@ -193,6 +284,37 @@ def __call__(self, batch: Dict[str, "torch.Tensor"]) -> Dict[str, "torch.Tensor" f"This suggests data is not properly shifted." ) + # C2 (exact-GDN CP alignment): splice inter-document pad-docs so every + # document start (and the stream end) sits on the 64-token chunk grid. + # Runs BEFORE the _original_position_ids capture: after splicing, the + # spliced stream IS the stream the model consumes (minus C1 tail pads), + # so token-aligned consumers must see the spliced layout. + alignment_segments = None + if self.gdn_exact_cp_align: + doc_bounds = find_document_boundaries(position_ids[0]) + alignment_segments = gdn_cp_alignment_segments(doc_bounds) + inserted = gdn_cp_spliced_length(alignment_segments) - input_ids.size(-1) + input_ids = apply_alignment_segments(input_ids, alignment_segments, -1, self.pad_token_id) + labels = apply_alignment_segments(labels, alignment_segments, -1, IGNORE_INDEX) + position_ids = apply_alignment_segments( + position_ids, alignment_segments, -1, sequential_positions=True + ) + if "attention_mask" in batch: + batch["attention_mask"] = apply_alignment_segments( + batch["attention_mask"], alignment_segments, -1, 1 + ) + # The pre-splice layout no longer matches the model's stream. + batch["_original_position_ids"] = position_ids.clone() + global _alignment_engagement_logged + if not _alignment_engagement_logged: + logger.info( + "gdn-cp collator alignment engaged: +%d pad tokens on the first batch " + "(%.3f%% of the spliced stream)", + inserted, + 100.0 * inserted / max(1, position_ids.size(-1)), + ) + _alignment_engagement_logged = True + # Store original position_ids before padding for unpacking per-token outputs later if "_original_position_ids" not in batch: batch["_original_position_ids"] = position_ids.clone() @@ -201,7 +323,11 @@ def __call__(self, batch: Dict[str, "torch.Tensor"]) -> Dict[str, "torch.Tensor" # With zigzag, each doc must be divisible by 2*ringattn_size sub-chunks, # and each sub-chunk must be divisible by ulysses_size. So total # sequence must be divisible by 2*ringattn_size*ulysses_size = 2*cp_size. + # C1 (exact-GDN CP alignment): the per-rank shard length must also be a + # multiple of the 64-token chunk grid. pad_multiple = 2 * self.cp_size if self.ringattn_size > 1 else self.cp_size + if self.gdn_exact_cp_align: + pad_multiple = GDN_CP_CHUNK * self.cp_size seq_length = input_ids.size(-1) cp_chunk_size = (seq_length + pad_multiple - 1) // pad_multiple * pad_multiple // self.cp_size pad_length = cp_chunk_size * self.cp_size - seq_length @@ -289,6 +415,10 @@ def __call__(self, batch: Dict[str, "torch.Tensor"]) -> Dict[str, "torch.Tensor" # Determine pad value: IGNORE_INDEX for target_tokens, 0 for others pad_value = IGNORE_INDEX if field == "target_tokens" else 0.0 + if alignment_segments is not None: + field_tensor = apply_alignment_segments( + field_tensor, alignment_segments, seq_dim, pad_value + ) field_tensor = self.sp_padding(field_tensor, dim=seq_dim, pad_value=pad_value, pad_length=pad_length) if self.ringattn_size > 1: field_tensor = zigzag_reorder_packed_sequence( @@ -316,4 +446,30 @@ def __call__(self, batch: Dict[str, "torch.Tensor"]) -> Dict[str, "torch.Tensor" else: add_flash_attention_kwargs_from_position_ids(batch, max_length_bucket=self.fa_max_length_bucket) + # Defense in depth for the exact-GDN CP contract: every shard cut must + # land on the 64-token chunk grid of whichever document it crosses + # (real doc or pad-doc — the GDN kernels cannot tell them apart). + # NOTE: pad-doc STARTS are unaligned by construction (they fill the + # gap up to the next 64-multiple) — that is fine precisely because a + # sub-64 gap pad can never contain a 64-multiple cut in its interior. + # Fail closed — a violated grid means byte-broken GDN handoffs. + if self.gdn_exact_cp_align: + if cp_chunk_size % GDN_CP_CHUNK != 0: + raise AssertionError( + f"gdn_exact_cp_align postcondition violated: shard length " + f"{cp_chunk_size} is not a multiple of {GDN_CP_CHUNK}", + ) + bounds = find_document_boundaries(position_ids[0]) + cuts = [cp_chunk_size * r for r in range(1, self.cp_size)] + bad = [] + for cut in cuts: + doc_start = max(s for s in bounds[:-1] if s <= cut) + if cut != doc_start and (cut - doc_start) % GDN_CP_CHUNK != 0: + bad.append((cut, doc_start)) + if bad: + raise AssertionError( + "gdn_exact_cp_align postcondition violated: shard cuts off " + f"the crossing document's 64-token grid: {bad[:4]}", + ) + return batch diff --git a/src/xorl/models/auto.py b/src/xorl/models/auto.py index 1ea08412..b482d731 100644 --- a/src/xorl/models/auto.py +++ b/src/xorl/models/auto.py @@ -386,12 +386,96 @@ def _validate_exact_qwen35_topology(config: PretrainedConfig, parallel_state: An if _is_qwen35_moe(config) else ((1, 1, 1, 1, 1, 1, 1, 1, 1, 1),) ) - if topology not in admitted: + if topology in admitted: + return + if not _is_qwen35_moe(config) and _admit_qwen35_hybrid_ulysses(config, parallel_state): + return + raise ValueError( + "The Qwen3.5-family exact server-training path is admitted only for " + "WORLD/DP/DP-replicate/DP-shard/TP/PP/EP/CP/Ring/Ulysses=" + f"{admitted}; got {topology}" + ) + + +# Hybrid (GDN + full-attention) dense Qwen3.5 under Ulysses: degrees covered +# by the composed-program and FA4 head-bucket byte-contract tests. +_QWEN35_HYBRID_ULYSSES_DEGREES = (2, 4, 8) +_GDN_CP_COLLATOR_ATTESTATION_ENV = "XORL_GDN_CP_ALIGN_COLLATOR" + + +def _admit_qwen35_hybrid_ulysses(config: PretrainedConfig, parallel_state: Any) -> bool: + """Conditional Ulysses admission for the exact hybrid dense program. + + Returns False when the topology is not the pure-Ulysses hybrid shape (the + caller then raises its generic refusal). When the shape matches, every + requirement below must hold or this RAISES with an actionable message: + + 1. hybrid layer_types (a GDN-free dense config has no qualified U>1 + program — its contract remains single-rank); + 2. the Ulysses degree checks (heads divisible by the degree; GQA + replication degree divisible by kv-heads) — admission-time versions + of the pre-collective raises in UlyssesSyncStrategy; + 3. the C1/C2 aligned-collator attestation + (``XORL_GDN_CP_ALIGN_COLLATOR=1``): admission cannot see the data + pipeline, so it demands an explicit attestation that + ``gdn_exact_cp_align`` is engaged; the GDN chain receipts and the + collator postcondition remain the runtime enforcement (fail-closed on + any misaligned shard cut); + 4. the kernel/toolchain pin (first-class): a seeded pin directory whose + toolchain fingerprint matches this runtime, installed as this rank's + per-rank Triton cache before any kernel compiles. + """ + import os # noqa: PLC0415 + + u = parallel_state.ulysses_size + shape_matches = ( + u in _QWEN35_HYBRID_ULYSSES_DEGREES + and parallel_state.world_size == u + and parallel_state.cp_size == u + and parallel_state.ringattn_size == 1 + and parallel_state.dp_size == 1 + and parallel_state.dp_replicate_size == 1 + and parallel_state.dp_shard_size == 1 + and parallel_state.tp_size == 1 + and parallel_state.pp_size == 1 + and parallel_state.ep_size == 1 + ) + if not shape_matches: + return False + + layer_types = getattr(config, "layer_types", None) or [] + if "linear_attention" not in layer_types: + return False # GDN-free dense: keep the single-rank refusal. + + num_heads = getattr(config, "num_attention_heads", None) + num_kv = getattr(config, "num_key_value_heads", None) + if not num_heads or num_heads % u != 0: + raise ValueError( + f"Exact hybrid Qwen3.5 at Ulysses {u}: num_attention_heads ({num_heads}) must be " + f"divisible by the Ulysses degree; an uneven head split cannot be scattered " + "byte-safely", + ) + if not num_kv or (u > num_kv and u % num_kv != 0): + raise ValueError( + f"Exact hybrid Qwen3.5 at Ulysses {u}: num_key_value_heads ({num_kv}) must divide " + "the Ulysses degree for GQA replication", + ) + if os.environ.get(_GDN_CP_COLLATOR_ATTESTATION_ENV) != "1": raise ValueError( - "The Qwen3.5-family exact server-training path is admitted only for " - "WORLD/DP/DP-replicate/DP-shard/TP/PP/EP/CP/Ring/Ulysses=" - f"{admitted}; got {topology}" + f"Exact hybrid Qwen3.5 at Ulysses {u} requires the C1/C2 aligned collator " + f"(TextSequenceShardCollator(gdn_exact_cp_align=True)); set " + f"{_GDN_CP_COLLATOR_ATTESTATION_ENV}=1 to attest it. Without 64-aligned shard " + "cuts the GDN chain receipts fail closed at the first crossing document.", ) + from xorl.ops.kernel_config_pin import pin_exact_kernel_configs # noqa: PLC0415 + + pin_clone = pin_exact_kernel_configs() + logger.info( + "exact hybrid Qwen3.5 Ulysses-%d admission engaged: heads %d/%d kv, collator " + "attested, kernel pin -> %s", + u, num_heads, num_kv, pin_clone, + ) + return True @dataclass(frozen=True) diff --git a/src/xorl/ops/kernel_config_pin.py b/src/xorl/ops/kernel_config_pin.py new file mode 100644 index 00000000..344c464c --- /dev/null +++ b/src/xorl/ops/kernel_config_pin.py @@ -0,0 +1,181 @@ +"""First-class kernel/toolchain pinning for exact contracts. + +Byte-exact programs are TOOLCHAIN-SCOPED claims: different Triton or FA4 +builds can compile the same source into different arithmetic. Autotune configs +are likewise per-process unless pinned: Triton's ``cache_results`` replays +tuned configs from the cache directory, so all ranks and the qualification +oracle must share one seeded cache. Each rank receives a separate clone to +avoid concurrent writes to the shared seed. + +The admission contract is mechanical: + +- A qualification run SEEDS a pin directory: its triton cache plus a + toolchain manifest (torch/triton/flash-attn versions). +- Admission (``pin_exact_kernel_configs``) FAILS CLOSED unless + ``XORL_EXACT_KERNEL_CONFIG_DIR`` names a seeded pin directory whose + manifest matches the running toolchain, then points this process's + ``TRITON_CACHE_DIR`` at a per-rank clone of the seeded cache. It must run + BEFORE the first kernel compilation (admission time satisfies this). +""" + +from __future__ import annotations + +import json +import logging +import os +import shutil + +import torch +import triton + + +logger = logging.getLogger("xorl.kernel_config_pin") + +PIN_DIR_ENV = "XORL_EXACT_KERNEL_CONFIG_DIR" +MANIFEST_NAME = "toolchain_manifest.json" +CACHE_SUBDIR = "triton-cache" + + +class KernelConfigPinError(RuntimeError): + """The kernel/toolchain pin is missing or violated. Fail closed.""" + + +# Ownership sentinel: this module only ever deletes directories it created +# itself (the marker IS the authorization). Env-var-fed paths never reach +# rmtree without it. +OWNED_SENTINEL = ".xorl-kernel-pin-owned" + + +def _mark_owned(path: str) -> None: + with open(os.path.join(path, OWNED_SENTINEL), "w") as f: + f.write("created by xorl.ops.kernel_config_pin; safe for it to replace\n") + + +def _rmtree_owned(path: str) -> None: + """Delete `path` only if this module created it (sentinel present).""" + if not os.path.isdir(path): + return + if not os.path.isfile(os.path.join(path, OWNED_SENTINEL)): + raise KernelConfigPinError( + f"refusing to delete {path!r}: it lacks the ownership sentinel " + f"{OWNED_SENTINEL!r} and was not created by this module. Remove or " + "relocate it manually if it is stale.", + ) + shutil.rmtree(path) + + +def _runtime_fingerprint() -> dict: + # Use the distribution version rather than only the module attribute: + # distinct flash-attn wheel builds can share a torch/triton fingerprint, + # while some builds do not expose a useful ``flash_attn.__version__``. + from importlib import metadata # noqa: PLC0415 + + fa = "unavailable" + for dist in ("flash-attn-4", "flash_attn_4", "flash-attn", "flash_attn"): + try: + fa = metadata.version(dist) + break + except metadata.PackageNotFoundError: + continue + if fa == "unavailable": + try: + import flash_attn # noqa: PLC0415 + + fa = getattr(flash_attn, "__version__", "unavailable") + except Exception: # pragma: no cover - build dependent + pass + return { + "torch": torch.__version__, + "triton": triton.__version__, + "flash_attn": fa, + "cuda": torch.version.cuda or "none", + } + + +def seed_exact_kernel_config_pin(pin_dir: str, *, source_cache: str | None = None) -> dict: + """Create/refresh a pin directory from the CURRENT runtime. + + Called by qualification runs (e.g. the fixture oracle phase) after their + kernels have been tuned. Copies `source_cache` (default: the active + TRITON_CACHE_DIR or ~/.triton/cache) into the pin and writes the + toolchain manifest. + """ + fingerprint = _runtime_fingerprint() + pin_dir = os.path.realpath(pin_dir) + parent = os.path.dirname(pin_dir) + if not os.path.isdir(parent): + raise KernelConfigPinError( + f"pin directory parent {parent!r} does not exist; refusing to create a pin " + "at an implausible location", + ) + os.makedirs(pin_dir, exist_ok=True) + cache_src = os.path.realpath( + source_cache + or os.environ.get("TRITON_CACHE_DIR", os.path.expanduser("~/.triton/cache")) + ) + cache_dst = os.path.join(pin_dir, CACHE_SUBDIR) + if os.path.commonpath([cache_src, cache_dst]) == cache_src: + raise KernelConfigPinError( + f"seed source cache {cache_src!r} contains the pin destination {cache_dst!r}; " + "copying would recurse into itself", + ) + _rmtree_owned(cache_dst) + if os.path.isdir(cache_src): + shutil.copytree(cache_src, cache_dst) + else: + os.makedirs(cache_dst, exist_ok=True) + _mark_owned(cache_dst) + with open(os.path.join(pin_dir, MANIFEST_NAME), "w") as f: + json.dump(fingerprint, f, indent=2) + logger.info("exact kernel-config pin seeded at %s: %s", pin_dir, fingerprint) + return fingerprint + + +def pin_exact_kernel_configs(*, rank: int | None = None) -> str: + """Admission-time pin. Returns the per-rank TRITON_CACHE_DIR it installed. + + Fail-closed on: env unset, pin dir or manifest missing, toolchain + fingerprint mismatch. Engagement-logged once per process. + """ + pin_dir = os.environ.get(PIN_DIR_ENV) + if not pin_dir: + raise KernelConfigPinError( + f"{PIN_DIR_ENV} is not set. Exact hybrid-Ulysses admission requires a seeded " + "kernel/toolchain pin directory (seed_exact_kernel_config_pin from the " + "qualification run). Byte claims are toolchain-scoped; refusing to run unpinned.", + ) + # The env var feeds filesystem mutations below: resolve it and refuse + # anything that is not an existing, seeded pin directory. + pin_dir = os.path.realpath(pin_dir) + if not os.path.isdir(pin_dir): + raise KernelConfigPinError( + f"{PIN_DIR_ENV}={pin_dir!r} is not an existing directory. Fail closed.", + ) + manifest_path = os.path.join(pin_dir, MANIFEST_NAME) + if not os.path.isfile(manifest_path): + raise KernelConfigPinError( + f"{PIN_DIR_ENV}={pin_dir!r} has no {MANIFEST_NAME}; the pin directory was never " + "seeded by a qualification run. Fail closed.", + ) + with open(manifest_path) as f: + pinned = json.load(f) + running = _runtime_fingerprint() + mismatches = {k: (pinned.get(k), running[k]) for k in running if pinned.get(k) != running[k]} + if mismatches: + raise KernelConfigPinError( + "Toolchain fingerprint mismatch against the kernel-config pin — the byte " + f"qualification does not transfer: {mismatches}. Re-qualify or restore the " + "pinned environment.", + ) + if rank is None: + rank = int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", "0"))) + clone = os.path.join(pin_dir, "clones", f"rank{rank}") + _rmtree_owned(clone) + shutil.copytree(os.path.join(pin_dir, CACHE_SUBDIR), clone) + _mark_owned(clone) + os.environ["TRITON_CACHE_DIR"] = clone + logger.info( + "exact kernel-config pin engaged: %s (rank %d clone %s, toolchain %s)", + pin_dir, rank, clone, running, + ) + return clone diff --git a/src/xorl/ops/linear_attention/layers/gated_deltanet.py b/src/xorl/ops/linear_attention/layers/gated_deltanet.py index 7edff547..5201a803 100644 --- a/src/xorl/ops/linear_attention/layers/gated_deltanet.py +++ b/src/xorl/ops/linear_attention/layers/gated_deltanet.py @@ -324,12 +324,26 @@ def _forward_impl( raise RuntimeError("Exact Qwen3.5 GDN requires short convolution") if self.use_short_conv and _is_gdn_contract_enabled(): - if cp_context is not None: - raise RuntimeError("Exact Qwen3.5 GDN does not support CP yet (conv prefix exchange)") if use_cache or last_state is not None: raise RuntimeError( "Exact Qwen3.5 trainer GDN supports packed prefill only, not recurrent cache updates" ) + if cp_context is None: + # Fail-closed backstop against the ring/silent-skip fail-open: + # build_linear_attention_cp_context returns None for ring>1 + # and for missing metadata — an exact GDN layer must never + # run U1 math on sequence-sharded rows. + from xorl.distributed.parallel_state import get_parallel_state # noqa: PLC0415 + + ps = get_parallel_state() + if ps.ulysses_size > 1 or ps.ringattn_size > 1: + raise RuntimeError( + "Exact Qwen3.5 GDN: sequence parallelism is active " + f"(ulysses={ps.ulysses_size}, ring={ps.ringattn_size}) but no " + "cp_context reached the layer. Ring is unsupported (contract C4) " + "and a missing Ulysses context would silently break bytes. " + "Fail closed.", + ) q, k, v = causal_conv1d_qkv_contract( q_input, k_input, @@ -338,6 +352,7 @@ def _forward_impl( self.k_conv1d, self.v_conv1d, cu_seqlens=cu_seqlens, + cp_context=cp_context, ) conv_state_q = conv_state_k = conv_state_v = None elif self.use_short_conv: diff --git a/src/xorl/ops/linear_attention/modules/conv_contract.py b/src/xorl/ops/linear_attention/modules/conv_contract.py index 7eb8c974..5ec73b82 100644 --- a/src/xorl/ops/linear_attention/modules/conv_contract.py +++ b/src/xorl/ops/linear_attention/modules/conv_contract.py @@ -26,10 +26,17 @@ from __future__ import annotations +import logging + import torch import torch.nn.functional as F from xorl.ops.linear_attention.ops.causal_conv1d_triton import causal_conv1d_fn +from xorl.ops.linear_attention.ops.cp.comm import conv_cp_send_recv_bwd, conv_cp_send_recv_fwd + + +logger = logging.getLogger("xorl.gdn_cp_conv") +_cp_engagement_logged = False def _pack_conv_weight(*weights: torch.Tensor) -> torch.Tensor: @@ -114,6 +121,106 @@ def backward(ctx, grad_output: torch.Tensor): return grads[0], grads[1], grads[2], grads[3], None, None, None +class _CausalConv1dContractCP(torch.autograd.Function): + """CP variant of the conv contract: SAME serving kernel, with the first + local sequence's window seeded by the (width-1)-token RAW halo from the + previous rank. Operands cross the wire, never results: the halo rows are + conv INPUTS; every output byte is computed by this rank inside the + contracted kernel. + + Backward recomputes the torch depthwise composition WITH the prefix and + ships the prefix gradient back to the previous rank's tail rows (the + reverse halo), so trainability is preserved without a handwritten kernel. + """ + + @staticmethod + def forward( + ctx, + x_packed: torch.Tensor, + weight_q: torch.Tensor, + weight_k: torch.Tensor, + weight_v: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: tuple[int, ...], + activation: str | None, + halo: torch.Tensor, # [width-1, dim] raw rows from prev rank (zeros if fresh) + first_seq_continues: bool, + cp_group, + ) -> torch.Tensor: + weight_packed = _pack_conv_weight(weight_q, weight_k, weight_v) + width = weight_packed.shape[-1] + num_seqs = len(seq_lens) + device = x_packed.device + conv_states = torch.zeros( + num_seqs, x_packed.shape[1], width - 1, device=device, dtype=x_packed.dtype + ) + has_init = torch.zeros(num_seqs, device=device, dtype=torch.bool) + if first_seq_continues: + # halo rows are chronological (oldest first); the kernel window + # axis is likewise oldest-first (causal_conv1d_triton.py:134-140). + conv_states[0] = halo.transpose(0, 1) + has_init[0] = True + out = causal_conv1d_fn( + x_packed.transpose(0, 1), + weight_packed, + None, + conv_states=conv_states, + query_start_loc=query_start_loc, + seq_lens_cpu=list(seq_lens), + cache_indices=torch.arange(num_seqs, device=device, dtype=torch.int32), + has_initial_state=has_init, + activation=activation, + ).transpose(0, 1) + ctx.save_for_backward(x_packed, weight_q, weight_k, weight_v, halo) + ctx.seq_lens = seq_lens + ctx.activation = activation + ctx.first_seq_continues = first_seq_continues + ctx.cp_group = cp_group + return out + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + x_packed, weight_q, weight_k, weight_v, halo = ctx.saved_tensors + prefix_width = halo.shape[0] + with torch.enable_grad(): + x_leaf = x_packed.detach().requires_grad_(True) + w_leaves = [w.detach().requires_grad_(True) for w in (weight_q, weight_k, weight_v)] + halo_leaf = halo.detach().requires_grad_(True) + if ctx.first_seq_continues: + # Recompute the ACTUAL forward composition: first sequence + # convolved with the prefix prepended, prefix outputs dropped. + first_len = ctx.seq_lens[0] + x_ext = torch.cat([halo_leaf, x_leaf[:first_len]], dim=0) + y_first = _depthwise_recompute( + x_ext, _pack_conv_weight(*w_leaves), ctx.activation, [prefix_width + first_len] + )[prefix_width:] + y_rest = ( + _depthwise_recompute( + x_leaf[first_len:], _pack_conv_weight(*w_leaves), ctx.activation, + list(ctx.seq_lens[1:]), + ) + if len(ctx.seq_lens) > 1 + else x_leaf.new_zeros(0, x_leaf.shape[1]) + ) + y = torch.cat([y_first, y_rest], dim=0) + else: + y = _depthwise_recompute( + x_leaf, _pack_conv_weight(*w_leaves), ctx.activation, list(ctx.seq_lens) + ) + grads = torch.autograd.grad(y, [x_leaf, *w_leaves, halo_leaf], grad_output, + allow_unused=True) + dx, dwq, dwk, dwv, dhalo = grads + # Reverse halo: my prefix gradient belongs to the PREVIOUS rank's last + # rows; symmetrically I receive my successor's and add it to my tail. + # Every rank participates in the collective (zeros when unused). + send = dhalo if dhalo is not None else torch.zeros_like(halo) + recv = conv_cp_send_recv_bwd(send.to(x_packed.dtype).contiguous(), ctx.cp_group) + tail = min(dx.shape[0], recv.shape[0]) + if tail > 0: + dx[-tail:] = dx[-tail:] + recv[-tail:] + return dx, dwq, dwk, dwv, None, None, None, None, None, None + + def causal_conv1d_qkv_contract( q_input: torch.Tensor, k_input: torch.Tensor, @@ -122,6 +229,7 @@ def causal_conv1d_qkv_contract( k_conv: torch.nn.Module, v_conv: torch.nn.Module, cu_seqlens: torch.Tensor | None = None, + cp_context=None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Serving-bit short conv over the packed (q|k|v) projection outputs. @@ -157,6 +265,59 @@ def causal_conv1d_qkv_contract( seq_lens = tuple(end - start for start, end in zip(starts[:-1], starts[1:], strict=False)) boundaries = cu_seqlens.to(device=q_input.device, dtype=torch.int32) + if cp_context is not None: + from xorl.ops.linear_attention.ops.cp.context import FLACPContext # noqa: PLC0415 + + if not isinstance(cp_context, FLACPContext): + raise TypeError( + f"Exact GDN CP conv requires a FLACPContext, got {type(cp_context).__name__}. " + "Fail closed: partial/duck-typed contexts cannot carry the alignment metadata.", + ) + if cu_seqlens is None: + raise ValueError("Exact GDN CP conv requires the LOCAL cu_seqlens from cp_context.") + if cp_context.group is None: + raise ValueError("Exact GDN CP conv requires an initialized cp_context group.") + width = q_conv.weight.shape[-1] + prefix_width = width - 1 + x_flat = x_packed.reshape(-1, x_packed.shape[-1]) + if x_flat.shape[0] < prefix_width: + raise RuntimeError( + "Exact GDN CP conv: local shard shorter than the conv window — " + "the C1 collator contract (shard length multiple of 64) is violated", + ) + pre_tokens = int(cp_context.pre_num_conv_tokens or 0) + first_seq_continues = pre_tokens > 0 + if 0 < pre_tokens < prefix_width: + raise RuntimeError( + f"Exact GDN CP conv: crossing document has only {pre_tokens} upstream " + f"tokens (< {prefix_width}) — impossible under the C1/C2 alignment " + "contract; the collator is misconfigured. Fail closed.", + ) + # Halo exchange of RAW operand rows (all ranks participate; the wire + # never carries conv outputs). + tails = x_flat[-prefix_width:].detach().contiguous() + halo = conv_cp_send_recv_fwd(tails, cp_context.group) + global _cp_engagement_logged + if not _cp_engagement_logged: + logger.info( + "gdn-cp conv halo engaged: width %d, first_seq_continues=%s, dim %d", + width, first_seq_continues, x_flat.shape[-1], + ) + _cp_engagement_logged = True + out = _CausalConv1dContractCP.apply( + x_flat, + q_conv.weight, + k_conv.weight, + v_conv.weight, + boundaries, + seq_lens, + activation, + halo, + first_seq_continues, + cp_context.group, + ).view_as(x_packed) + return out.split((q_input.shape[-1], k_input.shape[-1], v_input.shape[-1]), dim=-1) + out = _CausalConv1dContract.apply( x_packed.reshape(-1, x_packed.shape[-1]), q_conv.weight, diff --git a/src/xorl/ops/linear_attention/ops/cp/chain.py b/src/xorl/ops/linear_attention/ops/cp/chain.py new file mode 100644 index 00000000..573a4b95 --- /dev/null +++ b/src/xorl/ops/linear_attention/ops/cp/chain.py @@ -0,0 +1,213 @@ +"""Canonical scan chain for exact GDN under Ulysses context parallelism. + +Cross-rank recurrent state moves as a sequential fp32 final-state handoff in +declared canonical order (ascending rank equals document order). Each hop +transports operands, never results: the fp32 boundary state is captured +bit-exactly by the unmodified U1 scan kernel and re-injected bit-exactly +downstream. No arithmetic happens on the wire and no summary or fold +arithmetic exists outside the canonical U1 composition. + +A chain hop is admitted only if its receipt validates before any dependent +compute; validation failures raise (fail closed), and engagement is logged +once per (layer, direction) so a silent no-op is impossible. +""" + +from __future__ import annotations + +import hashlib +import logging +from dataclasses import dataclass + +import torch +import torch.distributed as dist + + +logger = logging.getLogger("xorl.gdn_cp_chain") + +_DIGEST_BYTES = 32 +_CHUNK = 64 + +# Engagement bookkeeping: {tag: hop_count}. Logged once per tag; counters +# always increment so gates can assert coverage. +_ENGAGEMENTS: dict[str, int] = {} + + +class ChainReceiptError(RuntimeError): + """A chain hop failed admission. Fail closed: never fall back.""" + + +@dataclass(frozen=True) +class ChainReceipt: + """Byte-relevant facts about one state handoff, derived independently on + both sides from shared collator/cp metadata (everything except + `state_digest`, which travels with the payload and certifies transport). + + in_doc_offset is the receiving shard's start measured from the crossing + document's start (== cp_context.pre_num_conv_tokens for the first local + sequence): the canonical grid demands it be a multiple of 64. + """ + + sender_rank: int + receiver_rank: int + in_doc_offset: int + state_shape: tuple[int, ...] + state_dtype: torch.dtype + state_digest: bytes | None = None + + +def state_digest(state: torch.Tensor) -> bytes: + return hashlib.sha256(state.detach().contiguous().cpu().numpy().tobytes()).digest() + + +def _record_engagement(tag: str, message: str) -> None: + if tag not in _ENGAGEMENTS: + _ENGAGEMENTS[tag] = 0 + logger.info("gdn-cp chain engaged: %s", message) + _ENGAGEMENTS[tag] += 1 + + +def engagement_count(tag: str) -> int: + return _ENGAGEMENTS.get(tag, 0) + + +def validate_chain_receipt(receipt: ChainReceipt, *, state: torch.Tensor) -> None: + """Admission check, run BEFORE any compute consumes the received state.""" + if receipt.sender_rank != receipt.receiver_rank - 1: + raise ChainReceiptError( + f"non-canonical hop order: sender {receipt.sender_rank} -> " + f"receiver {receipt.receiver_rank}; the declared order is ascending rank", + ) + if receipt.in_doc_offset <= 0 or receipt.in_doc_offset % _CHUNK != 0: + raise ChainReceiptError( + f"shard boundary at in-document offset {receipt.in_doc_offset} is not on the " + f"document's {_CHUNK}-token chunk grid; the collator alignment contract " + "is violated", + ) + if receipt.state_dtype != torch.float32 or state.dtype != torch.float32: + raise ChainReceiptError( + f"chain state must be fp32 end-to-end, got receipt={receipt.state_dtype}, " + f"tensor={state.dtype}", + ) + if tuple(state.shape) != tuple(receipt.state_shape): + raise ChainReceiptError( + f"state shape {tuple(state.shape)} != receipt shape {tuple(receipt.state_shape)}", + ) + if receipt.state_digest is None: + raise ChainReceiptError("receipt carries no transport digest") + actual = state_digest(state) + if actual != receipt.state_digest: + raise ChainReceiptError( + "state digest mismatch after transport: the wire did not preserve the " + f"fp32 bytes (expected {receipt.state_digest.hex()[:16]}..., " + f"got {actual.hex()[:16]}...)", + ) + + +def _wire_device(group: dist.ProcessGroup | None) -> str: + return "cpu" if dist.get_backend(group) == "gloo" else "cuda" + + +def chain_send_state( + state: torch.Tensor, + dst: int, + group: dist.ProcessGroup | None = None, + *, + tag_prefix: str = "gdn", +) -> None: + """Send one fp32 boundary state + its digest to the next rank (P2P). + + The payload is the state's raw bytes (contiguous fp32) followed by a + 32-byte SHA-256 digest tensor. Cast-free: fp32 in, fp32 on the wire. + """ + if state.dtype != torch.float32: + raise ChainReceiptError(f"refusing to send non-fp32 chain state ({state.dtype})") + wire = _wire_device(group) + payload = state.detach().contiguous().to(wire) + digest = torch.frombuffer(bytearray(state_digest(state)), dtype=torch.uint8).to(wire) + dist.send(payload, dst=dst, group=group) + dist.send(digest, dst=dst, group=group) + _record_engagement( + f"{tag_prefix}:send", + f"rank {dist.get_rank(group)} -> {dst}, state {tuple(state.shape)} fp32", + ) + + +def chain_pre_scan_initial_state( + *, + num_seqs: int, + heads: int, + key_dim: int, + value_dim: int, + context, + device: torch.device | str, + tag_prefix: str = "gdn", +) -> torch.Tensor: + """Receive-side of the canonical scan chain for one GDN forward. + + Returns the `[N, H, K, V]` fp32 initial-state tensor for the local + fragment scan: row 0 carries the received boundary state when the first + local sequence continues a document from the previous rank; all other + rows (locally-starting sequences) stay zero. Receipt validation happens + inside chain_recv_state BEFORE the state is returned (fail closed). + """ + initial_state = torch.zeros(num_seqs, heads, key_dim, value_dim, + dtype=torch.float32, device=device) + if not context.is_first_rank: + in_doc_offset = int(context.pre_num_conv_tokens or 0) + rank = dist.get_rank(context.group) + initial_state[0] = chain_recv_state( + (heads, key_dim, value_dim), + src=rank - 1, + group=context.group, + in_doc_offset=in_doc_offset, + device=device, + tag_prefix=tag_prefix, + ) + return initial_state + + +def chain_post_scan_send( + final_state: torch.Tensor, + *, + context, + tag_prefix: str = "gdn", +) -> None: + """Send-side: ship the LAST local sequence's fp32 final state onward when + that document continues into the next rank.""" + if context.is_last_rank: + return + rank = dist.get_rank(context.group) + chain_send_state(final_state[-1], dst=rank + 1, group=context.group, tag_prefix=tag_prefix) + + +def chain_recv_state( + shape: tuple[int, ...], + src: int, + group: dist.ProcessGroup | None = None, + *, + in_doc_offset: int, + device: torch.device | str = "cuda", + tag_prefix: str = "gdn", +) -> torch.Tensor: + """Receive one fp32 boundary state, VALIDATE its receipt, return it on + `device`. Raises ChainReceiptError before returning anything on failure.""" + wire = _wire_device(group) + payload = torch.empty(shape, dtype=torch.float32, device=wire) + digest = torch.empty(_DIGEST_BYTES, dtype=torch.uint8, device=wire) + dist.recv(payload, src=src, group=group) + dist.recv(digest, src=src, group=group) + receipt = ChainReceipt( + sender_rank=src, + receiver_rank=dist.get_rank(group), + in_doc_offset=in_doc_offset, + state_shape=tuple(shape), + state_dtype=torch.float32, + state_digest=bytes(digest.cpu().numpy().tobytes()), + ) + validate_chain_receipt(receipt, state=payload) + _record_engagement( + f"{tag_prefix}:recv", + f"rank {receipt.receiver_rank} <- {src}, state {tuple(shape)} fp32, " + f"in-doc offset {in_doc_offset}", + ) + return payload.to(device) diff --git a/src/xorl/ops/linear_attention/ops/gated_delta_rule/chunk.py b/src/xorl/ops/linear_attention/ops/gated_delta_rule/chunk.py index 747758d7..d47f7466 100644 --- a/src/xorl/ops/linear_attention/ops/gated_delta_rule/chunk.py +++ b/src/xorl/ops/linear_attention/ops/gated_delta_rule/chunk.py @@ -5,6 +5,7 @@ import torch +from xorl.ops.linear_attention.modules.bi_contract import _is_gdn_contract_enabled from xorl.ops.linear_attention.modules.l2norm import l2norm_bwd, l2norm_fwd from xorl.ops.linear_attention.ops.common.chunk_delta_h import ( chunk_gated_delta_rule_bwd_dhu, @@ -13,6 +14,10 @@ from xorl.ops.linear_attention.ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o from xorl.ops.linear_attention.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd from xorl.ops.linear_attention.ops.cp import FLACPContext +from xorl.ops.linear_attention.ops.cp.chain import ( + chain_post_scan_send, + chain_pre_scan_initial_state, +) from xorl.ops.linear_attention.ops.cp.chunk_delta_h import ( chunk_gated_delta_rule_bwd_dhu_pre_process, chunk_gated_delta_rule_fwd_h_pre_process, @@ -59,7 +64,22 @@ def chunk_gated_delta_rule_fwd( cu_seqlens=cu_seqlens, ) - if cp_context is not None: + exact_chain = cp_context is not None and _is_gdn_contract_enabled() + if exact_chain: + # The initial state is the previous rank's fp32 boundary state, + # received and receipt-validated by the canonical scan chain. It is + # not the (h_e, M) summary fold, which executes a different floating- + # point program from the U1 scan. + assert initial_state is None, "When enable CP, the provided initial_state must be None." + initial_state = chain_pre_scan_initial_state( + num_seqs=len(cu_seqlens) - 1 if cu_seqlens is not None else k.shape[0], + heads=k.shape[2], + key_dim=k.shape[3], + value_dim=u.shape[-1], + context=cp_context, + device=k.device, + ) + elif cp_context is not None: initial_state = chunk_gated_delta_rule_fwd_h_pre_process( k=k, w=w, @@ -76,10 +96,18 @@ def chunk_gated_delta_rule_fwd( u=u, g=g, initial_state=initial_state, - output_final_state=output_final_state, + output_final_state=output_final_state or (exact_chain and not cp_context.is_last_rank), cu_seqlens=cu_seqlens, ) + if exact_chain: + if not cp_context.is_last_rank: + # Ship the LAST local sequence's fp32 final state onward + # (bit-exact capture of the running scan registers). + chain_post_scan_send(final_state, context=cp_context) + if not output_final_state: + final_state = None + if cp_context is not None: initial_state = compress_h0(initial_state, context=cp_context) diff --git a/tests/distributed/test_ulysses_byte_alignment.py b/tests/distributed/test_ulysses_byte_alignment.py index 9ad26b8c..90daf84e 100644 --- a/tests/distributed/test_ulysses_byte_alignment.py +++ b/tests/distributed/test_ulysses_byte_alignment.py @@ -14,8 +14,17 @@ sequence shards in rank order and byte-compares hidden + logprobs against the reference npz. Includes the collator cp-multiple padding case (pad tokens appended as their own documents; real-token bytes must match - the unpadded U1 reference) and the hybrid negative: a GDN layer under - Ulysses must RAISE the exact-contract CP refusal, not compute. + the unpadded U1 reference) and hybrid cells under the composed GDN-CP + contract. + +The hybrid cells cover the composed program (canonical scan chain, conv halo, +C1/C2 collator, conditional admission, and kernel/toolchain pin): +`hybrid_admits_and_matches` (full contract engaged -> ADMITS and +byte-matches the U1 hybrid reference; this fixture's pack is 64-aligned +in-document at every tested degree, so the chain receipts admit it) and +`hybrid_raises_without_attestation` (the fail-closed negative: +admission REFUSES a hybrid Ulysses topology when the C1/C2 collator +attestation is absent). """ @@ -173,11 +182,18 @@ def _run_reference(out_path: str) -> None: # would itself be an unproven assumption). short = _short_batch(batch, device) hidden_short = _forward_hidden(model, short) + # Hybrid U1 reference for the composed-contract positive: the same batch + # through the exact hybrid program at Ulysses 1. + hybrid_model, _ = _build_model( + ["linear_attention", "full_attention", "linear_attention", "full_attention"], device + ) + hidden_hybrid = _forward_hidden(hybrid_model, batch) np.savez( out_path, hidden_bf16=hidden.view(torch.int16).cpu().numpy(), logprobs=logprobs.view(torch.int32).cpu().numpy(), hidden_short_bf16=hidden_short.view(torch.int16).cpu().numpy(), + hidden_hybrid_bf16=hidden_hybrid.view(torch.int16).cpu().numpy(), ) print(f"[ulysses-gate] reference written: {out_path}", flush=True) @@ -185,7 +201,7 @@ def _run_reference(out_path: str) -> None: def _run_sharded() -> None: import torch.distributed as dist - from xorl.distributed.parallel_state import init_parallel_state + from xorl.distributed.parallel_state import get_parallel_state, init_parallel_state from xorl.utils.device import get_nccl_backend degree = int(os.environ["ULYSSES_GATE_DEGREE"]) @@ -206,6 +222,7 @@ def _run_sharded() -> None: device_type="cuda", cp_fsdp_mode="none", ) + ps = get_parallel_state() device = torch.device("cuda", local_rank) rank = dist.get_rank() @@ -216,14 +233,14 @@ def _run_sharded() -> None: reference = np.load(ref_path) verdicts = {} - def _sharded_forward(batch): + def _sharded_forward(batch, forward_model=None): seq_len = batch["input_ids"].shape[-1] assert seq_len % degree == 0 shard = seq_len // degree local_ids = batch["input_ids"][:, rank * shard : (rank + 1) * shard].contiguous() local_batch = dict(batch) local_batch["input_ids"] = local_ids - hidden_local = _forward_hidden(model, local_batch) + hidden_local = _forward_hidden(forward_model if forward_model is not None else model, local_batch) gathered = [torch.empty_like(hidden_local) for _ in range(degree)] dist.all_gather(gathered, hidden_local.contiguous()) return torch.cat(gathered, dim=1) @@ -253,18 +270,62 @@ def _sharded_forward(batch): torch.equal(hidden_padded[:, :SHORT_LEN].contiguous().view(torch.int16).cpu(), ref_short) ) - # --- hybrid negative: GDN under Ulysses must RAISE the contract floor -- + # --- hybrid cells under the composed contract --------------------------- + # Positive: full contract engaged -> ADMIT + byte-match the U1 hybrid + # reference (this fixture's pack is 64-aligned in-document at every + # tested degree: doc starts 0 and 192, shard = 512/degree, all multiples + # of 64, so the chain receipts admit every cut). Negative: admission + # REFUSES the hybrid Ulysses topology when the C1/C2 collator + # attestation is absent (representative missing-piece). + import tempfile as _tempfile + + from xorl.models.auto import _validate_exact_qwen35_topology + from xorl.ops.kernel_config_pin import seed_exact_kernel_config_pin + hybrid_model, hybrid_config = _build_model( ["linear_attention", "full_attention", "linear_attention", "full_attention"], device ) + for tensor in list(hybrid_model.parameters()) + list(hybrid_model.buffers()): + dist.broadcast(tensor.data, src=0) hybrid_batch = _make_batch(SEQ_LEN, hybrid_config.vocab_size, device) - shard = SEQ_LEN // degree - hybrid_batch["input_ids"] = hybrid_batch["input_ids"][:, rank * shard : (rank + 1) * shard].contiguous() + + saved_env = { + key: os.environ.get(key) + for key in ("XORL_GDN_CP_ALIGN_COLLATOR", "XORL_EXACT_KERNEL_CONFIG_DIR", "TRITON_CACHE_DIR") + } + pin_dir = _tempfile.mkdtemp(prefix=f"ulysses-gate-pin-r{rank}-") + empty_cache_src = _tempfile.mkdtemp(prefix=f"ulysses-gate-cache-r{rank}-") try: - _forward_hidden(hybrid_model, hybrid_batch) - verdicts["hybrid_raises"] = False - except RuntimeError as exc: - verdicts["hybrid_raises"] = "does not support CP yet" in str(exc) + # Negative first (attestation absent; pin present so the refusal is + # attributable to exactly the missing attestation). Seed from an + # empty cache dir: this cell asserts admission semantics, not + # cross-process config replay, and avoids copying the ambient default + # Triton cache. + seed_exact_kernel_config_pin(pin_dir, source_cache=empty_cache_src) + os.environ["XORL_EXACT_KERNEL_CONFIG_DIR"] = pin_dir + os.environ.pop("XORL_GDN_CP_ALIGN_COLLATOR", None) + try: + _validate_exact_qwen35_topology(hybrid_config, ps) + verdicts["hybrid_raises_without_attestation"] = False + except ValueError as exc: + verdicts["hybrid_raises_without_attestation"] = "aligned collator" in str(exc) + + # Positive: full composed contract engaged. + os.environ["XORL_GDN_CP_ALIGN_COLLATOR"] = "1" + _validate_exact_qwen35_topology(hybrid_config, ps) # must ADMIT (no raise) + hidden_hybrid = _sharded_forward(hybrid_batch, forward_model=hybrid_model) + verdicts["hybrid_admits_and_matches"] = bool( + torch.equal( + hidden_hybrid.view(torch.int16).cpu(), + torch.from_numpy(reference["hidden_hybrid_bf16"]), + ) + ) + finally: + for key, value in saved_env.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value gathered_verdicts: list = [None] * dist.get_world_size() dist.all_gather_object(gathered_verdicts, verdicts) diff --git a/tests/models/test_qwen35_hybrid_ulysses_admission.py b/tests/models/test_qwen35_hybrid_ulysses_admission.py new file mode 100644 index 00000000..93fc9119 --- /dev/null +++ b/tests/models/test_qwen35_hybrid_ulysses_admission.py @@ -0,0 +1,180 @@ +"""Fail-closed admission for the exact hybrid Qwen3.5 program at Ulysses > 1. + +Covers Ulysses degree checks, the aligned-collator attestation, and the +first-class kernel/toolchain pin. CPU-only. +""" + +from __future__ import annotations + +import json +import os +from types import SimpleNamespace + +import pytest +import torch +import triton + +from xorl.models.auto import _validate_exact_qwen35_topology +from xorl.ops.kernel_config_pin import ( + MANIFEST_NAME, + KernelConfigPinError, + pin_exact_kernel_configs, + seed_exact_kernel_config_pin, +) + + +HYBRID_LAYERS = ["linear_attention", "linear_attention", "linear_attention", "full_attention"] + + +def _ps(u=8, **overrides): + fields = dict( + world_size=u, dp_size=1, dp_replicate_size=1, dp_shard_size=1, + tp_size=1, pp_size=1, ep_size=1, cp_size=u, ringattn_size=1, ulysses_size=u, + ) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def _config(layer_types=None, heads=16, kv=2): + return SimpleNamespace( + _qwen35_exact_contract=True, + model_type="qwen3_5", + layer_types=layer_types, + num_attention_heads=heads, + num_key_value_heads=kv, + ) + + +@pytest.fixture() +def seeded_pin(tmp_path, monkeypatch): + pin = tmp_path / "pin" + cache = tmp_path / "cache-src" + cache.mkdir() + (cache / "dummy.json").write_text("{}") + seed_exact_kernel_config_pin(str(pin), source_cache=str(cache)) + monkeypatch.setenv("XORL_EXACT_KERNEL_CONFIG_DIR", str(pin)) + monkeypatch.setenv("XORL_GDN_CP_ALIGN_COLLATOR", "1") + monkeypatch.setenv("RANK", "3") + return pin + + +class TestHybridUlyssesAdmission: + def test_admitted_with_full_attestations(self, seeded_pin, monkeypatch): + monkeypatch.delenv("TRITON_CACHE_DIR", raising=False) + _validate_exact_qwen35_topology(_config(HYBRID_LAYERS), _ps(8)) + clone = os.environ["TRITON_CACHE_DIR"] + assert clone == os.path.realpath(str(seeded_pin / "clones" / "rank3")) + assert os.path.isfile(os.path.join(clone, "dummy.json")) + + def test_gdn_free_dense_keeps_single_rank_refusal(self, seeded_pin): + with pytest.raises(ValueError, match="admitted only for"): + _validate_exact_qwen35_topology(_config(["full_attention"] * 4), _ps(8)) + + def test_missing_collator_attestation_raises(self, seeded_pin, monkeypatch): + monkeypatch.delenv("XORL_GDN_CP_ALIGN_COLLATOR", raising=False) + with pytest.raises(ValueError, match="aligned collator"): + _validate_exact_qwen35_topology(_config(HYBRID_LAYERS), _ps(8)) + + def test_missing_kernel_pin_raises(self, seeded_pin, monkeypatch): + monkeypatch.delenv("XORL_EXACT_KERNEL_CONFIG_DIR", raising=False) + with pytest.raises(KernelConfigPinError, match="toolchain pin"): + _validate_exact_qwen35_topology(_config(HYBRID_LAYERS), _ps(8)) + + def test_toolchain_mismatch_raises(self, seeded_pin): + manifest = seeded_pin / MANIFEST_NAME + doctored = json.loads(manifest.read_text()) + doctored["torch"] = "0.0.0+nope" + manifest.write_text(json.dumps(doctored)) + with pytest.raises(KernelConfigPinError, match="fingerprint mismatch"): + _validate_exact_qwen35_topology(_config(HYBRID_LAYERS), _ps(8)) + + def test_uneven_head_split_raises(self, seeded_pin): + with pytest.raises(ValueError, match="divisible by the Ulysses degree"): + _validate_exact_qwen35_topology(_config(HYBRID_LAYERS, heads=6), _ps(4)) + + def test_kv_non_divisor_raises(self, seeded_pin): + with pytest.raises(ValueError, match="GQA replication"): + _validate_exact_qwen35_topology(_config(HYBRID_LAYERS, kv=3), _ps(8)) + + def test_ring_shape_keeps_generic_refusal(self, seeded_pin): + ps = _ps(8, ringattn_size=2, cp_size=16, world_size=16) + with pytest.raises(ValueError, match="admitted only for"): + _validate_exact_qwen35_topology(_config(HYBRID_LAYERS), ps) + + def test_unlisted_degree_keeps_generic_refusal(self, seeded_pin): + with pytest.raises(ValueError, match="admitted only for"): + _validate_exact_qwen35_topology(_config(HYBRID_LAYERS), _ps(16)) + + def test_single_rank_still_admitted_without_envs(self, monkeypatch): + monkeypatch.delenv("XORL_EXACT_KERNEL_CONFIG_DIR", raising=False) + monkeypatch.delenv("XORL_GDN_CP_ALIGN_COLLATOR", raising=False) + _validate_exact_qwen35_topology(_config(HYBRID_LAYERS), _ps(1)) + + +class TestKernelConfigPin: + def test_seed_manifest_matches_runtime(self, tmp_path): + src = tmp_path / "src" + src.mkdir() + fp = seed_exact_kernel_config_pin(str(tmp_path / "p"), source_cache=str(src)) + assert fp["torch"] == torch.__version__ + assert fp["triton"] == triton.__version__ + + def test_seed_refuses_self_recursive_copy(self, tmp_path): + with pytest.raises(KernelConfigPinError, match="recurse"): + seed_exact_kernel_config_pin(str(tmp_path / "p"), source_cache=str(tmp_path)) + + def test_unseeded_dir_raises(self, tmp_path, monkeypatch): + empty = tmp_path / "empty" + empty.mkdir() + monkeypatch.setenv("XORL_EXACT_KERNEL_CONFIG_DIR", str(empty)) + with pytest.raises(KernelConfigPinError, match="never .*seeded|no toolchain"): + pin_exact_kernel_configs(rank=0) + + def test_pin_env_not_a_directory_raises(self, tmp_path, monkeypatch): + monkeypatch.setenv("XORL_EXACT_KERNEL_CONFIG_DIR", str(tmp_path / "missing")) + with pytest.raises(KernelConfigPinError, match="not an existing directory"): + pin_exact_kernel_configs(rank=0) + + def _seeded(self, tmp_path): + src = tmp_path / "src" + src.mkdir() + (src / "cfg.json").write_text("{}") + pin = tmp_path / "pin" + seed_exact_kernel_config_pin(str(pin), source_cache=str(src)) + return pin + + def test_pin_refuses_to_delete_unowned_clone(self, tmp_path, monkeypatch): + """The ownership sentinel IS the deletion authorization: a clone-path + directory this module did not create must never reach rmtree.""" + pin = self._seeded(tmp_path) + unowned = pin / "clones" / "rank0" + unowned.mkdir(parents=True) + (unowned / "precious.txt").write_text("not ours to delete") + monkeypatch.setenv("XORL_EXACT_KERNEL_CONFIG_DIR", str(pin)) + with pytest.raises(KernelConfigPinError, match="ownership sentinel"): + pin_exact_kernel_configs(rank=0) + assert (unowned / "precious.txt").exists() + + def test_pin_replaces_its_own_clone(self, tmp_path, monkeypatch): + pin = self._seeded(tmp_path) + monkeypatch.setenv("XORL_EXACT_KERNEL_CONFIG_DIR", str(pin)) + first = pin_exact_kernel_configs(rank=1) + second = pin_exact_kernel_configs(rank=1) # replaces the owned clone + assert first == second and os.path.isdir(second) + + def test_seed_refuses_to_delete_unowned_cache(self, tmp_path): + from xorl.ops.kernel_config_pin import CACHE_SUBDIR, OWNED_SENTINEL + + pin = self._seeded(tmp_path) + (pin / CACHE_SUBDIR / OWNED_SENTINEL).unlink() # simulate a foreign dir + src2 = tmp_path / "src2" + src2.mkdir() + with pytest.raises(KernelConfigPinError, match="ownership sentinel"): + seed_exact_kernel_config_pin(str(pin), source_cache=str(src2)) + + def test_seed_refuses_implausible_parent(self, tmp_path): + with pytest.raises(KernelConfigPinError, match="parent .*does not exist"): + seed_exact_kernel_config_pin( + str(tmp_path / "no" / "such" / "parent" / "pin"), + source_cache=str(tmp_path), + ) diff --git a/tests/ops/test_gdn_conv_contract.py b/tests/ops/test_gdn_conv_contract.py index 796096b6..721ef432 100644 --- a/tests/ops/test_gdn_conv_contract.py +++ b/tests/ops/test_gdn_conv_contract.py @@ -259,7 +259,7 @@ def test_weight_pack_routing_admission_and_state_lifecycle_policy(self, monkeypa def _assert_forward_routes_through_contract_when_armed(self, monkeypatch): calls = [] - def fake_contract(q_in, k_in, v_in, *convs, cu_seqlens=None): + def fake_contract(q_in, k_in, v_in, *convs, cu_seqlens=None, cp_context=None): calls.append(cu_seqlens) return F.silu(q_in), F.silu(k_in), F.silu(v_in) @@ -298,7 +298,10 @@ def conv_bias(): cases = [ ("decode cache", use_cache, RuntimeError, "prefill only"), - ("context parallelism", cp_context, RuntimeError, "does not support CP"), + # The blanket "does not support CP" raise is gone: the exact CP + # program runs under a real FLACPContext. Malformed or duck-typed + # contexts stay fail-closed at the convolution boundary. + ("context parallelism, malformed context", cp_context, TypeError, "FLACPContext"), ("missing short convolution", no_short_conv, RuntimeError, "requires short convolution"), ("convolution bias", conv_bias, NotImplementedError, "bias"), ]