From 630c49a5dea19ee4b42934a693ff7be2cad85835 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Thu, 13 Aug 2026 10:37:40 +0000 Subject: [PATCH] distributed: make Ulysses strategy selection fail closed Validate head divisibility before collectives, make sync-versus-async selection explicit, and reject unsupported GDN fallback. Add CPU admission, distributed byte-alignment, and FA4 head-bucket invariance tests. --- .../distributed/sequence_parallel/strategy.py | 46 ++- .../transformers/qwen3_5/modeling_qwen3_5.py | 11 +- src/xorl/ops/linear_attention/backend.py | 15 +- .../test_ulysses_byte_alignment.py | 329 ++++++++++++++++++ .../distributed/test_ulysses_byte_contract.py | 152 ++++++++ tests/ops/test_fa4_head_bucket_invariance.py | 118 +++++++ 6 files changed, 662 insertions(+), 9 deletions(-) create mode 100644 tests/distributed/test_ulysses_byte_alignment.py create mode 100644 tests/distributed/test_ulysses_byte_contract.py create mode 100644 tests/ops/test_fa4_head_bucket_invariance.py diff --git a/src/xorl/distributed/sequence_parallel/strategy.py b/src/xorl/distributed/sequence_parallel/strategy.py index 5608abf5..335b3efc 100644 --- a/src/xorl/distributed/sequence_parallel/strategy.py +++ b/src/xorl/distributed/sequence_parallel/strategy.py @@ -179,12 +179,26 @@ def project_qkv(self, module, hidden_states, position_embeddings): # Model-specific QKV projection (MHA, MLA, etc.) q, k, v = module._project_qkv(hidden_states, position_embeddings) + # Fail closed BEFORE any collective: an uneven Q-head split does not + # reliably error downstream (the 4-D all_to_all path allocates every + # receive buffer from the first split's shape and would silently + # mis-size them; the 3-D path fails with an opaque reshape error). + q_head_num = q.shape[2] + if q_head_num % self.ulysses_size != 0: + raise ValueError( + f"Ulysses requires num_attention_heads ({q_head_num}) to be divisible by " + f"ulysses_size ({self.ulysses_size}); an uneven head split cannot be " + f"scattered byte-safely" + ) + # GQA expand if ulysses_size > num_kv_heads kv_head_num = k.shape[2] if self.ulysses_size > kv_head_num: - assert self.ulysses_size % kv_head_num == 0, ( - f"ulysses_size ({self.ulysses_size}) must be divisible by num_key_value_heads ({kv_head_num})" - ) + if self.ulysses_size % kv_head_num != 0: + raise ValueError( + f"ulysses_size ({self.ulysses_size}) must be divisible by " + f"num_key_value_heads ({kv_head_num}) for GQA replication" + ) n_repeat = self.ulysses_size // kv_head_num # repeat_kv expects [batch, num_heads, seq, head_dim] k = k.transpose(1, 2) @@ -482,7 +496,7 @@ def prepare_position_embeddings(self, position_embeddings, dim, sp_group, **kwar _NOOP = NoopStrategy() -def get_cp_strategy(num_kv_heads: Optional[int] = None) -> CPStrategy: +def get_cp_strategy(num_kv_heads: Optional[int] = None, variant: str = "auto") -> CPStrategy: """Resolve the SP strategy from the current ParallelState. Returns a singleton NoopStrategy when SP is disabled, or the @@ -495,16 +509,32 @@ def get_cp_strategy(num_kv_heads: Optional[int] = None) -> CPStrategy: 3. Ring only (ringattn_size > 1) Args: - num_kv_heads: Number of key-value heads in the model. Required when - Ulysses SP is enabled to choose between sync and async variants. + num_kv_heads: Number of key-value heads in the model. Used by the + ``"auto"`` variant when Ulysses SP is enabled to choose between + sync and async variants. + variant: ``"auto"`` keeps the historical heuristic; ``"sync"`` / + ``"async"`` select the Ulysses variant EXPLICITLY. The choice is + bit-relevant: the sync variant applies RoPE BEFORE the + head-scattering all-to-all on sequence-sliced tables, the async + variant AFTER it on full-length tables — flipping the variant + silently relocates RoPE relative to the exchange. Exact lanes + pin the variant instead of relying on whether a call site + happens to pass ``num_kv_heads``. """ from ...distributed.parallel_state import get_parallel_state # noqa: PLC0415 + if variant not in ("auto", "sync", "async"): + raise ValueError(f"Unknown CP strategy variant {variant!r}; expected 'auto', 'sync', or 'async'") + ps = get_parallel_state() if not ps.cp_enabled: return _NOOP if ps.ulysses_enabled and ps.ringattn_enabled: + if variant != "auto": + raise NotImplementedError( + f"Explicit Ulysses variant {variant!r} is not supported with hybrid Ulysses+Ring" + ) # Hybrid Ulysses + Ring return HybridUlyssesRingStrategy( ulysses_group=ps.ulysses_group, @@ -513,6 +543,10 @@ def get_cp_strategy(num_kv_heads: Optional[int] = None) -> CPStrategy: ) if ps.ulysses_enabled: + if variant == "sync": + return UlyssesSyncStrategy(group=ps.ulysses_group, ulysses_size=ps.ulysses_size) + if variant == "async": + return UlyssesAsyncStrategy(group=ps.ulysses_group, ulysses_size=ps.ulysses_size) if num_kv_heads is not None and ps.ulysses_size <= num_kv_heads: return UlyssesAsyncStrategy(group=ps.ulysses_group, ulysses_size=ps.ulysses_size) else: diff --git a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py index 501a8531..2cf90883 100644 --- a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py +++ b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py @@ -309,7 +309,12 @@ def forward( **kwargs: Unpack[AttentionKwargs], ) -> tuple[torch.Tensor, torch.Tensor | None]: del position_ids, past_key_values - attn_strategy = get_cp_strategy() + # Qwen3.5 PINS the sync Ulysses variant: this call site and the + # prepare_position_embeddings site must agree (RoPE is applied before + # the head-scattering all-to-all on sequence-sliced tables). The + # historical "auto" heuristic flips the variant — and with it the + # RoPE placement — based on whether num_kv_heads is passed. + attn_strategy = get_cp_strategy(variant="sync") query_states, key_states, value_states = attn_strategy.project_qkv(self, hidden_states, position_embeddings) attn_output = attn_strategy.compute_attention( self, query_states, key_states, value_states, attention_mask, **kwargs @@ -552,7 +557,9 @@ def forward( linear_attn_mask = None position_embeddings = self.rotary_emb(hidden_states, position_ids) - position_embeddings = get_cp_strategy().prepare_position_embeddings( + # Same explicit variant as the attention call site: sequence-slice the + # cos/sin tables because RoPE runs before the sync all-to-all. + position_embeddings = get_cp_strategy(variant="sync").prepare_position_embeddings( position_embeddings, dim=1, sp_group=ps.sp_group, diff --git a/src/xorl/ops/linear_attention/backend.py b/src/xorl/ops/linear_attention/backend.py index c0473301..b2789bca 100644 --- a/src/xorl/ops/linear_attention/backend.py +++ b/src/xorl/ops/linear_attention/backend.py @@ -70,7 +70,20 @@ def resolve_flashqla_auto_cp(auto_cp: bool | None) -> bool: def warn_cp_fallback_once() -> None: - """Warn (once) that a FlashQLA request fell back to the FLA Triton GDN kernel.""" + """Warn (once) that a FlashQLA request fell back to the FLA Triton GDN kernel. + + Under the exact GDN contract a silent backend swap is a byte hazard, not + a performance note: the contract RAISES instead. (Today the contract pins + the backend to ``fla`` before any FlashQLA request can be made, so this + is defense-in-depth against a future reordering of backend resolution.) + """ + from xorl.ops.linear_attention.modules.bi_contract import _is_gdn_contract_enabled # noqa: PLC0415 + + if _is_gdn_contract_enabled(): + raise RuntimeError( + "Exact Qwen3.5 GDN: a FlashQLA->FLA backend fallback was requested while the GDN " + "contract is active; silently swapping the kernel program is not admitted" + ) global _warned_cp_fallback if not _warned_cp_fallback: warnings.warn( diff --git a/tests/distributed/test_ulysses_byte_alignment.py b/tests/distributed/test_ulysses_byte_alignment.py new file mode 100644 index 00000000..9ad26b8c --- /dev/null +++ b/tests/distributed/test_ulysses_byte_alignment.py @@ -0,0 +1,329 @@ +"""Ulysses U1-vs-UN byte-equality gate for the exact logit path. + +Two-phase structure (the Ulysses strategy reads GLOBAL parallel state, so the +U1 reference must run in its own process): + +- phase "ref" (1 GPU, no torch.distributed): the exact-contract program on a + FULL-ATTENTION-ONLY tiny model with the production head geometry + (8 Q-heads / 2 KV-heads / head_dim 256, bf16, FA4, sglang_fused BI RMSNorm, + bi_fused head) over packed varlen inputs; writes last-hidden and per-token + logprob bytes to an npz. +- phase "shard" (torchrun, ULYSSES_GATE_DEGREE ranks): the same model and + tokens through the production Ulysses path (sequence-sharded input_ids, + FULL position_ids and cu_seqlens per the collator convention); gathers the + 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. + +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import tempfile +from pathlib import Path + +import numpy as np +import pytest +import torch + +from xorl.models.layers.normalization import set_rmsnorm_mode +from xorl.models.transformers.qwen3_5.configuration_qwen3_5 import Qwen3_5Config +from xorl.models.transformers.qwen3_5.modeling_qwen3_5 import Qwen3_5ForCausalLM +from xorl.ops.loss.causallm_loss import causallm_loss_function + + +THIS_DIR = Path(__file__).resolve().parent +if str(THIS_DIR) not in sys.path: + sys.path.insert(0, str(THIS_DIR)) + +from distributed_utils import run_distributed_script, skip_if_gpu_count_less_than # noqa: E402 + + +pytestmark = [pytest.mark.distributed, pytest.mark.gpu] + +SEQ_LEN = 512 # divisible by every tested degree +# Padding case: 505 is odd, so EVERY tested degree needs collator padding +# (degree 2 -> pad 1, degree 8 -> pad 7). +SHORT_LEN = SEQ_LEN - 7 +DOC_BOUNDARY = 192 # two packed documents +IGNORE_INDEX = -100 + + +def _build_config(layer_types) -> Qwen3_5Config: + set_rmsnorm_mode("sglang_fused") + config = Qwen3_5Config( + vocab_size=512, + hidden_size=256, + intermediate_size=512, + num_hidden_layers=len(layer_types), + num_attention_heads=8, + num_key_value_heads=2, + head_dim=256, + linear_num_key_heads=8, + linear_num_value_heads=8, + linear_key_head_dim=32, + linear_value_head_dim=32, + layer_types=list(layer_types), + max_position_embeddings=2048, + use_cache=False, + tie_word_embeddings=False, + ) + config._attn_implementation = "flash_attention_4" + config._qwen35_exact_contract = True + config._qwen35_rmsnorm_family = "v1" + config.dtype = torch.bfloat16 + return config + + +def _build_model(layer_types, device): + torch.manual_seed(1729) + config = _build_config(layer_types) + return Qwen3_5ForCausalLM(config).to(torch.bfloat16).to(device).eval(), config + + +def _make_batch(seq_len: int, vocab_size: int, device): + generator = torch.Generator().manual_seed(777) + input_ids = torch.randint(0, vocab_size, (1, seq_len), generator=generator).to(device) + labels = torch.roll(input_ids, -1, dims=-1).clone() + labels[:, -1] = IGNORE_INDEX + cu = torch.tensor([0, DOC_BOUNDARY, seq_len], dtype=torch.int32, device=device) + return { + "input_ids": input_ids, + "labels": labels, + "position_ids": torch.arange(seq_len, device=device).unsqueeze(0), + "cu_seq_lens_q": cu, + "cu_seq_lens_k": cu, + "max_length_q": seq_len - DOC_BOUNDARY, + "max_length_k": seq_len - DOC_BOUNDARY, + } + + +def _pad_like_collator(batch: dict, cp_size: int, device): + """Mirror sequence_shard_collator: pad the packed length to a multiple of + cp_size; pad tokens get sequential position ids and their own document.""" + seq_len = batch["input_ids"].shape[-1] + target = (seq_len + cp_size - 1) // cp_size * cp_size + pad = target - seq_len + if pad == 0: + return batch + out = dict(batch) + out["input_ids"] = torch.nn.functional.pad(batch["input_ids"], (0, pad), value=0) + out["labels"] = torch.nn.functional.pad(batch["labels"], (0, pad), value=IGNORE_INDEX) + out["position_ids"] = torch.cat( + [batch["position_ids"], torch.arange(pad, device=device).unsqueeze(0)], dim=-1 + ) + cu = batch["cu_seq_lens_q"].tolist() + [target] + out["cu_seq_lens_q"] = torch.tensor(cu, dtype=torch.int32, device=device) + out["cu_seq_lens_k"] = out["cu_seq_lens_q"].clone() + return out + + +def _forward_hidden(model, batch): + with torch.no_grad(): + outputs = model( + input_ids=batch["input_ids"], + position_ids=batch["position_ids"], + use_cache=False, + output_hidden_states=False, + cu_seq_lens_q=batch["cu_seq_lens_q"], + cu_seq_lens_k=batch["cu_seq_lens_k"], + max_length_q=batch["max_length_q"], + max_length_k=batch["max_length_k"], + ) + return outputs.last_hidden_state + + +def _logprobs(model, hidden, labels): + with torch.no_grad(): + result = causallm_loss_function( + hidden_states=hidden, + weight=model.lm_head.weight, + labels=labels, + return_per_token=True, + ce_mode="bi_fused", + lm_head_fp32=True, + ) + return result.per_token_logprobs + + +def _short_batch(batch: dict, device): + short = dict(batch) + for key in ("input_ids", "labels", "position_ids"): + short[key] = batch[key][:, :SHORT_LEN].contiguous() + short["cu_seq_lens_q"] = torch.tensor([0, DOC_BOUNDARY, SHORT_LEN], dtype=torch.int32, device=device) + short["cu_seq_lens_k"] = short["cu_seq_lens_q"].clone() + short["max_length_q"] = SHORT_LEN - DOC_BOUNDARY + short["max_length_k"] = SHORT_LEN - DOC_BOUNDARY + return short + + +def _run_reference(out_path: str) -> None: + device = torch.device("cuda") + model, config = _build_model(["full_attention"] * 4, device) + batch = _make_batch(SEQ_LEN, config.vocab_size, device) + hidden = _forward_hidden(model, batch) + logprobs = _logprobs(model, hidden, batch["labels"]) + # The padding case compares against a REAL unpadded short-sequence run, + # not a truncated slice of the long run (truncation-bitwise-equivalence + # would itself be an unproven assumption). + short = _short_batch(batch, device) + hidden_short = _forward_hidden(model, short) + 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(), + ) + print(f"[ulysses-gate] reference written: {out_path}", flush=True) + + +def _run_sharded() -> None: + import torch.distributed as dist + + from xorl.distributed.parallel_state import init_parallel_state + from xorl.utils.device import get_nccl_backend + + degree = int(os.environ["ULYSSES_GATE_DEGREE"]) + ref_path = os.environ["ULYSSES_GATE_REF"] + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend=get_nccl_backend()) + init_parallel_state( + dp_size=1, + dp_replicate_size=1, + dp_shard_size=1, + tp_size=1, + ep_size=1, + pp_size=1, + ulysses_size=degree, + ringattn_size=1, + dp_mode="none", + device_type="cuda", + cp_fsdp_mode="none", + ) + device = torch.device("cuda", local_rank) + rank = dist.get_rank() + + model, config = _build_model(["full_attention"] * 4, device) + for tensor in list(model.parameters()) + list(model.buffers()): + dist.broadcast(tensor.data, src=0) + + reference = np.load(ref_path) + verdicts = {} + + def _sharded_forward(batch): + 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) + gathered = [torch.empty_like(hidden_local) for _ in range(degree)] + dist.all_gather(gathered, hidden_local.contiguous()) + return torch.cat(gathered, dim=1) + + # --- core: clean divisible packed batch -------------------------------- + batch = _make_batch(SEQ_LEN, config.vocab_size, device) + hidden_full = _sharded_forward(batch) + hidden_ok = bool( + torch.equal( + hidden_full.view(torch.int16).cpu(), torch.from_numpy(reference["hidden_bf16"]) + ) + ) + logprobs = _logprobs(model, hidden_full, batch["labels"]) + logprob_ok = bool( + torch.equal(logprobs.view(torch.int32).cpu(), torch.from_numpy(reference["logprobs"])) + ) + verdicts["core_hidden"] = hidden_ok + verdicts["core_logprobs"] = logprob_ok + + # --- padding case: collator cp-multiple padding, real rows must match -- + short = _short_batch(_make_batch(SEQ_LEN, config.vocab_size, device), device) + padded = _pad_like_collator(short, degree, device) + assert padded["input_ids"].shape[-1] > SHORT_LEN, "padding case degenerated (no pad added)" + hidden_padded = _sharded_forward(padded) + ref_short = torch.from_numpy(reference["hidden_short_bf16"]) + verdicts["padded_hidden"] = bool( + torch.equal(hidden_padded[:, :SHORT_LEN].contiguous().view(torch.int16).cpu(), ref_short) + ) + + # --- hybrid negative: GDN under Ulysses must RAISE the contract floor -- + hybrid_model, hybrid_config = _build_model( + ["linear_attention", "full_attention", "linear_attention", "full_attention"], device + ) + 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() + 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) + + gathered_verdicts: list = [None] * dist.get_world_size() + dist.all_gather_object(gathered_verdicts, verdicts) + merged = {} + for rank_verdicts in gathered_verdicts: + for name, ok in (rank_verdicts or {}).items(): + merged[name] = merged.get(name, True) and bool(ok) + if rank == 0: + for name in sorted(merged): + print(f"[{'PASS' if merged[name] else 'FAIL'}] ulysses{degree}_{name}", flush=True) + failed = [name for name, ok in merged.items() if not ok] + assert not failed, f"Ulysses{degree} byte gate failed: {failed}" + if rank == 0: + print(f"Ulysses byte gate passed (degree={degree})", flush=True) + + +def _pytest_run(degree: int, num_gpus: int) -> None: + with tempfile.TemporaryDirectory() as tmp: + ref_path = os.path.join(tmp, "reference.npz") + env = dict(os.environ) + env["ULYSSES_GATE_PHASE"] = "ref" + env["ULYSSES_GATE_REF"] = ref_path + result = subprocess.run( + [sys.executable, __file__], env=env, capture_output=True, text=True, timeout=900 + ) + assert result.returncode == 0, f"reference phase failed:\n{result.stderr[-2000:]}" + dist_result = run_distributed_script( + __file__, + num_gpus=num_gpus, + timeout=900, + extra_env={ + "ULYSSES_GATE_PHASE": "shard", + "ULYSSES_GATE_DEGREE": str(degree), + "ULYSSES_GATE_REF": ref_path, + }, + ) + dist_result.assert_success(f"Ulysses{degree} must be byte-identical to Ulysses1") + + +if __name__ != "__main__": + + @skip_if_gpu_count_less_than(2) + def test_ulysses2_byte_alignment_exact_program(): + _pytest_run(degree=2, num_gpus=2) + + @skip_if_gpu_count_less_than(4) + def test_ulysses4_byte_alignment_exact_program(): + # Degree 4 > kv_heads=2: the smallest degree that exercises the GQA + # KV-replication branch end-to-end (degree 2 == kv_heads skips it). + _pytest_run(degree=4, num_gpus=4) + + @skip_if_gpu_count_less_than(8) + def test_ulysses8_byte_alignment_exact_program(): + _pytest_run(degree=8, num_gpus=8) + + +if __name__ == "__main__": + phase = os.environ.get("ULYSSES_GATE_PHASE", "ref") + if phase == "ref": + _run_reference(os.environ["ULYSSES_GATE_REF"]) + else: + _run_sharded() diff --git a/tests/distributed/test_ulysses_byte_contract.py b/tests/distributed/test_ulysses_byte_contract.py new file mode 100644 index 00000000..e8fe66d7 --- /dev/null +++ b/tests/distributed/test_ulysses_byte_contract.py @@ -0,0 +1,152 @@ +"""Fail-closed admission tests for the Ulysses byte-contract surface (CPU-only). + +Covers missing Q-head divisibility, the implicit sync-versus-async strategy +choice that can relocate RoPE across the all-to-all, and the GDN backend +fallback under the exact contract. +The bitwise gates themselves are GPU tests (kernel head-bucket gate and the +U1-vs-U8 fixture gate). +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +import xorl.distributed.parallel_state as parallel_state_module +import xorl.ops.linear_attention.backend as gdn_backend +from xorl.distributed.sequence_parallel.strategy import ( + NoopStrategy, + UlyssesAsyncStrategy, + UlyssesSyncStrategy, + get_cp_strategy, +) +from xorl.models.auto import _validate_exact_qwen35_topology +from xorl.ops.linear_attention.modules.bi_contract import gdn_contract + + +def _fake_state(*, ulysses: bool = True, ring: bool = False, size: int = 8) -> SimpleNamespace: + return SimpleNamespace( + cp_enabled=ulysses or ring, + ulysses_enabled=ulysses, + ringattn_enabled=ring, + ulysses_size=size, + ulysses_group=None, + ringattn_group=None, + ) + + +@pytest.fixture +def ulysses8_state(monkeypatch): + monkeypatch.setattr(parallel_state_module, "get_parallel_state", lambda: _fake_state()) + + +# --------------------------------------------------------------------------- +# Explicit strategy selection (RoPE placement is bit-relevant) +# --------------------------------------------------------------------------- + + +def test_explicit_variant_overrides_the_kv_heads_heuristic(ulysses8_state): + # The historical heuristic sends num_kv_heads >= ulysses_size to ASYNC — + # which applies RoPE after the all-to-all. An exact contract pinning "sync" + # must win regardless of num_kv_heads. + assert isinstance(get_cp_strategy(num_kv_heads=64, variant="sync"), UlyssesSyncStrategy) + assert isinstance(get_cp_strategy(variant="async"), UlyssesAsyncStrategy) + + +def test_auto_variant_keeps_the_historical_heuristic(ulysses8_state): + assert isinstance(get_cp_strategy(num_kv_heads=8), UlyssesAsyncStrategy) + assert isinstance(get_cp_strategy(num_kv_heads=2), UlyssesSyncStrategy) + assert isinstance(get_cp_strategy(), UlyssesSyncStrategy) + + +def test_unknown_variant_raises(ulysses8_state): + with pytest.raises(ValueError, match="Unknown CP strategy variant"): + get_cp_strategy(variant="bogus") + + +def test_explicit_variant_rejects_hybrid_ring(monkeypatch): + monkeypatch.setattr( + parallel_state_module, "get_parallel_state", lambda: _fake_state(ring=True) + ) + with pytest.raises(NotImplementedError, match="hybrid Ulysses\\+Ring"): + get_cp_strategy(variant="sync") + + +def test_variant_is_moot_when_cp_disabled(monkeypatch): + monkeypatch.setattr( + parallel_state_module, + "get_parallel_state", + lambda: _fake_state(ulysses=False, ring=False), + ) + assert isinstance(get_cp_strategy(variant="sync"), NoopStrategy) + + +# --------------------------------------------------------------------------- +# Head divisibility fails closed BEFORE any collective +# --------------------------------------------------------------------------- + + +class _StubAttention: + def __init__(self, q_heads: int, kv_heads: int, seq_len: int = 16, head_dim: int = 4): + self._q = torch.zeros(1, seq_len, q_heads, head_dim, dtype=torch.bfloat16) + self._k = torch.zeros(1, seq_len, kv_heads, head_dim, dtype=torch.bfloat16) + self._v = torch.zeros(1, seq_len, kv_heads, head_dim, dtype=torch.bfloat16) + + def _project_qkv(self, hidden_states, position_embeddings): + return self._q, self._k, self._v + + +def test_uneven_q_head_split_raises_before_comm(): + strategy = UlyssesSyncStrategy(group=None, ulysses_size=8) + with pytest.raises(ValueError, match="num_attention_heads \\(6\\) to be divisible"): + strategy.project_qkv(_StubAttention(q_heads=6, kv_heads=8), None, None) + + +def test_non_divisor_kv_heads_raise_before_comm(): + # Promoted from a bare assert (which vanishes under -O) to a ValueError; + # ordering: the Q-head check passes first, then GQA replication refuses. + strategy = UlyssesSyncStrategy(group=None, ulysses_size=8) + with pytest.raises(ValueError, match="num_key_value_heads \\(3\\) for GQA replication"): + strategy.project_qkv(_StubAttention(q_heads=8, kv_heads=3), None, None) + + +# --------------------------------------------------------------------------- +# GDN backend fallback under the exact contract +# --------------------------------------------------------------------------- + + +def test_gdn_cp_fallback_raises_under_exact_contract(): + with gdn_contract(True): + with pytest.raises(RuntimeError, match="silently swapping the kernel program is not admitted"): + gdn_backend.warn_cp_fallback_once() + + +def test_gdn_cp_fallback_still_warns_outside_the_contract(monkeypatch): + monkeypatch.setattr(gdn_backend, "_warned_cp_fallback", False) + with pytest.warns(UserWarning, match="FlashQLA requires 128-dim heads"): + gdn_backend.warn_cp_fallback_once() + + +# --------------------------------------------------------------------------- +# Topology admission stays fail-closed at Ulysses > 1 for the exact contract +# --------------------------------------------------------------------------- + + +def test_exact_dense_topology_rejects_ulysses8(): + config = SimpleNamespace(_qwen35_exact_contract=True, model_type="qwen3_5") + ps = SimpleNamespace( + world_size=8, + dp_size=1, + dp_replicate_size=1, + dp_shard_size=1, + tp_size=1, + pp_size=1, + ep_size=1, + cp_size=8, + ringattn_size=1, + ulysses_size=8, + ) + with pytest.raises(ValueError, match="admitted only for"): + _validate_exact_qwen35_topology(config, ps) diff --git a/tests/ops/test_fa4_head_bucket_invariance.py b/tests/ops/test_fa4_head_bucket_invariance.py new file mode 100644 index 00000000..bbf65abf --- /dev/null +++ b/tests/ops/test_fa4_head_bucket_invariance.py @@ -0,0 +1,118 @@ +"""FA4 per-head head-batch invariance gate for Ulysses. + +Ulysses degree d hands each rank the FULL sequence with 1/d of the Q heads +(and replicated KV: one KV head per rank at d > kv_heads). The exact +contract requires that a given head's attention output is byte-identical no +matter which head-batch it is computed in — the analogue of the BI-GEMM +row-bucket gate, for the head axis. + +This gate runs the SAME FA4 varlen entry the production backend calls +(``fa4_flash_attn_varlen_func`` with ``num_splits=1``, mirroring +src/xorl/models/layers/attention/backend/flash_attention.py) on the +Qwen3.5-0.8B exact attention geometry (8 Q-heads, 2 KV-heads, head_dim 256, +bf16, packed varlen), slicing the head axis exactly as Ulysses degrees +{1,2,4,8} would (including the GQA-ratio changes 4 -> 4 -> 2 -> 1 and the +fresh contiguous allocations the all-to-all produces), and asserts BYTE +equality per head against the full-batch reference. + +Requires one GPU and FA4. +""" + +from __future__ import annotations + +import pytest +import torch + + +pytestmark = [pytest.mark.gpu] + +if not torch.cuda.is_available(): + pytest.skip("FA4 head-bucket invariance requires CUDA", allow_module_level=True) + +fa4 = pytest.importorskip("flash_attn.cute", reason="FA4 (flash_attn.cute) not installed") + +NUM_Q_HEADS = 8 +NUM_KV_HEADS = 2 +HEAD_DIM = 256 +# Packed varlen: uneven document lengths, not multiples of typical tiles. +CU_SEQLENS = [0, 384, 1000] +DEGREES = (1, 2, 4, 8) + + +def _make_qkv(device): + generator = torch.Generator(device="cpu").manual_seed(1729) + total = CU_SEQLENS[-1] + q = torch.randn(total, NUM_Q_HEADS, HEAD_DIM, generator=generator, dtype=torch.float32) + k = torch.randn(total, NUM_KV_HEADS, HEAD_DIM, generator=generator, dtype=torch.float32) + v = torch.randn(total, NUM_KV_HEADS, HEAD_DIM, generator=generator, dtype=torch.float32) + return ( + q.to(torch.bfloat16).to(device), + k.to(torch.bfloat16).to(device), + v.to(torch.bfloat16).to(device), + ) + + +def _run_fa4(q, k, v, device): + from flash_attn.cute import flash_attn_varlen_func # noqa: PLC0415 + + cu = torch.tensor(CU_SEQLENS, dtype=torch.int32, device=device) + max_len = max(b - a for a, b in zip(CU_SEQLENS, CU_SEQLENS[1:])) + out = flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q=cu, + cu_seqlens_k=cu, + max_seqlen_q=max_len, + max_seqlen_k=max_len, + softmax_scale=HEAD_DIM**-0.5, + causal=True, + num_splits=1, + ) + if isinstance(out, tuple): + out = out[0] + return out + + +@pytest.mark.parametrize("degree", DEGREES) +def test_fa4_head_bucket_bytes_match_full_batch(degree): + """Every Ulysses-degree head grouping must reproduce the full-batch bytes.""" + device = torch.device("cuda") + q, k, v = _make_qkv(device) + reference = _run_fa4(q, k, v, device) + + q_per_rank = NUM_Q_HEADS // degree + q_heads_per_kv = NUM_Q_HEADS // NUM_KV_HEADS + mismatches = [] + for rank in range(degree): + q_lo, q_hi = rank * q_per_rank, (rank + 1) * q_per_rank + # KV heads this rank's Q-group attends to (replication at high degree + # is a pure copy, so slicing the ORIGINAL kv head is byte-equivalent). + kv_heads = sorted({h // q_heads_per_kv for h in range(q_lo, q_hi)}) + # .contiguous(): the all-to-all hands the kernel fresh allocations. + q_group = q[:, q_lo:q_hi, :].contiguous() + k_group = k[:, kv_heads, :].contiguous() + v_group = v[:, kv_heads, :].contiguous() + out_group = _run_fa4(q_group, k_group, v_group, device) + ref_group = reference[:, q_lo:q_hi, :] + if not torch.equal(out_group.view(torch.int16), ref_group.contiguous().view(torch.int16)): + diff = (out_group.float() - ref_group.float()).abs() + mismatched = int((out_group.view(torch.int16) != ref_group.contiguous().view(torch.int16)).sum()) + mismatches.append( + f"degree={degree} rank={rank} heads[{q_lo}:{q_hi}] kv={kv_heads}: " + f"{mismatched} mismatched int16 lanes, max|diff|={diff.max().item():.3e}" + ) + assert not mismatches, ( + "FA4 head-batch invariance failed: per-head bytes depend on the " + "head-batch composition:\n" + "\n".join(mismatches) + ) + + +def test_fa4_single_head_repeatability(): + """Control: the same call twice must be byte-stable (rules out + run-to-run nondeterminism masquerading as head-batch variance).""" + device = torch.device("cuda") + q, k, v = _make_qkv(device) + first = _run_fa4(q, k, v, device) + second = _run_fa4(q, k, v, device) + assert torch.equal(first.view(torch.int16), second.view(torch.int16))