diff --git a/fla/ops/cp/KCP.md b/fla/ops/cp/KCP.md new file mode 100644 index 0000000000..4a28839fe7 --- /dev/null +++ b/fla/ops/cp/KCP.md @@ -0,0 +1,280 @@ +# KCP: Kimi Context Parallel + +Context Parallel for GDN (Gated Delta Rule) and KDA (Kimi Delta Attention). + +> CP was first introduced in [PR #691](https://github.com/fla-org/flash-linear-attention/pull/691), Special thanks to [mdy666](https://github.com/mdy666) + +## Core Recurrence + +Both GDN and KDA share the delta rule recurrence: + +``` +S_t = decay(g_t) * S_{t-1} + beta_t * k_t (x) (v_t - S_{t-1} @ k_t) +o_t = q_t^T @ S_t +``` + +Where `(x)` is outer product, `S` is `[K, V]` state matrix. FLA's kernels currently use `[K, V]` state matrix, but there are other backends that can store it by transposing `[V, K]`. + +In the chunk-parallel formulation, we first compute the WY representation +to get `w` and `u`, then the inter-chunk state recurrence becomes: + +``` +h_{c+1} = decay(g_last_c) * h_c + k_c^T @ (u_c - w_c @ h_c) +``` + +where `c` indexes chunks, `g_last_c` is the gate value at the last token of chunk `c`. + +--- + +## GDN vs KDA: Gate Handling + +### GDN: scalar per-head gate + +- `g` shape: `[B, T, H]` — one scalar per head per token +- `g` is chunk-local cumsum'd by `chunk_local_cumsum` +- All kernels receive **original** `k`, `q`, and **scalar** `g` +- Kernels internally apply gate via `USE_G=True`: + - Decay: `h *= exp(g_last)` + - Gate k: `k_gated = k * exp(g_last - g_token)` (done inside kernel) + - Gate q: `q_gated = q * exp(g_token)` (done inside kernel, backward only) + +### KDA: per-dim gate + +- `g` shape: `[B, T, H, K]` — one scalar per dimension per token +- `g` is chunk-local cumsum'd + scaled by `kda_gate_chunk_cumsum` (includes gate activation) + or `chunk_local_cumsum` (if gate pre-computed) +- The WY repr step (`chunk_kda_fwd_intra` / `recompute_w_u_fwd`) pre-computes gated tensors: + - `kg = k * exp2(gk_last_chunk - gk)` — relative gate, k aligned to chunk end + - `qg = q * exp2(gk)` — absolute gate on q (used in backward only) +- All kernels receive **pre-gated** `kg` (and `qg` in backward), plus `gk=g` for inter-chunk decay +- Kernels apply only the **chunk-level** decay via `USE_GK=True`: + - Decay: `h *= exp2(gk_last)` (diagonal, per-dim) + - No further gating on k/q — already done externally + +**Why the difference**: GDN's scalar gate is cheap to apply inside kernels. KDA's per-dim +gate `[H, K]` is more efficiently pre-applied during the WY representation step. + +--- + +## CP Architecture + +### Data Flow + +Each rank holds a local chunk of the sequence. CP computes cross-rank initial states +via an all-gather + merge pattern: + +``` +1. Each rank computes local (h_ext, M) from its chunk + - h_ext: accumulated state assuming h_0 = 0 + - M: transition matrix (product of per-chunk transition) + +2. All-gather [h_ext, M] across all ranks + +3. Rank r merges from ranks < r: + h_r = fold over [rank_{r-pre}, ..., rank_{r-1}]: + h = M_j @ h + h_ext_j +``` + +### Pre-Process Forward + +Computes `(h_ext, M)` for the first sequence of the local chunk. + +**Stage 1 — h_ext `[K, V]`:** + +``` +h = 0 +for each sub-chunk c: + h *= decay(g_last_c) # inter-chunk decay + v_new = u_c - w_c @ h # (computed via w @ h subtraction) + h += k_c^T @ v_new # accumulate +``` + +**Stage 2 — M `[K, K]` (transition matrix):** + +``` +M = I +for each sub-chunk c: + M_c = diag(decay(g_last_c)) - k_c^T @ w_c + M = M_c @ M # chain multiply +``` + +**Merge (forward direction):** + +For rank r with `pre_num_ranks` previous ranks: +``` +h = 0 +for j from (r - pre_num_ranks) to (r - 1): + h = M_j @ h + h_ext_j +``` + +### Pre-Process Backward + +Same structure but **reversed** direction — merges from ranks **after** current rank. + +**Stage 1 — dh_ext `[K, V]`:** + +``` +dh = 0 +for each sub-chunk c (reverse order): + dh *= decay(g_last_c) + dv = k_c @ dh + original_dv_c + dh += q_c^T @ do_c * scale - w_c^T @ dv +``` + +**Stage 2 — dM `[K, K]`:** + +``` +dM = I +for each sub-chunk c (reverse order): + dM_c = diag(decay(g_last_c)) - w_c^T @ k_c # NOTE: transposed vs forward + dM = dM_c @ dM +``` + +**Merge (backward direction):** + +For rank r with `post_num_ranks` following ranks: +``` +dh = 0 +for j from (r + post_num_ranks) down to (r + 1): + dh = dM_j @ dh + dh_ext_j +``` + +--- + +## Actual Code Flow + +### GDN Forward + +```python +g = chunk_local_cumsum(g, chunk_size=64) +w, u = recompute_w_u_fwd(k, v, beta, A, g=g) + +# CP pre-process: original k, scalar g +initial_state = chunk_gated_delta_rule_fwd_h_pre_process( + k=k, w=w, u=u, g=g, # USE_G=True, USE_GK=False + context=cp_context, +) + +# Main kernel: original k, scalar g +h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=k, w=w, u=u, g=g, + initial_state=initial_state, +) +``` + +### GDN Backward + +```python +w, u = recompute_w_u_fwd(k, v, beta, A, g=g) +h, v_new, _ = chunk_gated_delta_rule_fwd_h(k=k, w=w, u=u, g=g, ...) +dv = chunk_bwd_dv_local(q=q, k=k, g=g, do=do, ...) + +# CP pre-process: original q, k, scalar g +dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( + q=q, k=k, w=w, do=do, dv=dv, g=g, # USE_G=True, USE_GK=False + context=cp_context, +) + +# Main kernel: original q, k, scalar g +dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=q, k=k, w=w, g=g, + dht=dht, ... +) +``` + +### KDA Forward + +```python +# 1. Intra-chunk: compute WY repr + pre-gated tensors +w, u, qg, kg, Aqk, Akk = chunk_kda_fwd_intra(q, k, v, gk=g, beta, ...) +# kg = k * exp2(gk_last_chunk - gk) [relative gate to chunk end] +# qg = q * exp2(gk) [absolute gate, saved for backward] + +# 2. CP pre-process: pre-gated kg, per-dim gk=g +initial_state = chunk_gated_delta_rule_fwd_h_pre_process( + k=kg, w=w, u=u, gk=g, # USE_G=False, USE_GK=True, use_exp2=True + context=cp_context, +) + +# 3. Main kernel: pre-gated kg, per-dim gk=g +h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=kg, w=w, u=u, gk=g, + initial_state=initial_state, + use_exp2=True, +) +``` + +### KDA Backward + +```python +# 1. Recompute WY repr +w, u, qg, kg = recompute_w_u_fwd(q, k, v, beta, A=Akk, gk=g, ...) +# qg = q * exp2(gk) +# kg = k * exp2(gk_last_chunk - gk) + +# 2. Recompute h +h, v_new, _ = chunk_gated_delta_rule_fwd_h(k=kg, w=w, u=u, gk=g, ...) + +# 3. Compute local dv +dAqk, dv = chunk_kda_bwd_dAv(q, k, v=v_new, do, A=Aqk, ...) + +# 4. CP pre-process: pre-gated qg, kg, per-dim gk=g +dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( + q=qg, k=kg, w=w, do=do, dv=dv, gk=g, # USE_G=False, USE_GK=True, use_exp2=True + context=cp_context, +) + +# 5. Main kernel: pre-gated qg, kg +dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=qg, k=kg, w=w, gk=g, + dht=dht, ... + use_exp2=True, +) +``` + +--- + +## Input Tensor Summary + +| Function | GDN | KDA | Gate Path | +|----------|-----|-----|-----------| +| pre_process_fwd | k=k, g=g | k=kg, gk=g | GDN: USE_G, KDA: USE_GK | +| fwd_h | k=k, g=g | k=kg, gk=g | Same as pre_process | +| pre_process_bwd | q=q, k=k, g=g | q=qg, k=kg, gk=g | GDN: USE_G, KDA: USE_GK | +| bwd_dhu | q=q, k=k, g=g | q=qg, k=kg, gk=g | Same as pre_process | + +**Key consistency**: pre_process and main kernel always receive the **same** tensors. +For KDA, both receive pre-gated `kg`/`qg`. For GDN, both receive original tensors. + +--- + +## M (Transition Matrix) + +The transition matrix captures how the state transforms across a chunk: + +``` +M_c = diag(decay) - k_c^T @ w_c (forward) +dM_c = diag(decay) - w_c^T @ k_c (backward, transposed) +``` + +Where `decay` is: +- GDN: `exp(g_last)` scalar → `diag(exp(g_last)) = exp(g_last) * I` +- KDA: `exp2(gk_last)` per-dim → `diag(exp2(gk_last_0), exp2(gk_last_1), ...)` + +Cross-rank state is computed by chaining M matrices: +``` +h_r = M_{r-1} @ (M_{r-2} @ (... @ h_ext_0 + h_ext_1) + ...) + h_ext_{r-1} +``` + +**Precision note**: The M chain multiply `b_m = M_i @ b_m` must stay in fp32 to +avoid accumulated precision loss. In bf16, repeatedly casting fp32 accumulators back +to bf16 between iterations causes significant error growth over many chunks. + +--- + +## compress_h0 / expand_h0 + +Optimization for CP mode. Since only the first sequence in the local batch can be +a continuation from a previous rank, only its initial_state is non-zero. +`compress_h0` extracts just that one state to save memory during `save_for_backward`. +`expand_h0` restores the full `[N, H, K, V]` tensor in backward. diff --git a/fla/ops/cp/chunk_delta_h.py b/fla/ops/cp/chunk_delta_h.py index 6f97bc542d..c97b1a2a3c 100644 --- a/fla/ops/cp/chunk_delta_h.py +++ b/fla/ops/cp/chunk_delta_h.py @@ -505,7 +505,7 @@ def pre_process_fwd_kernel_merged( # Compute m update: m = (diag - k^T @ w) @ m b_kw = tl.dot(tl.trans(b_k.to(b_w.dtype)), b_w) b_m_i = b_diag - b_kw - b_m = tl.dot(b_m_i.to(b_w.dtype), b_m.to(b_w.dtype)) + b_m = tl.dot(b_m_i.to(tl.float32), b_m.to(tl.float32)) # Store m result p_m = tl.make_block_ptr(hm + V, (K, K), (K + V, 1), (0, i_k_col * BLOCK_SIZE), (BK1, BLOCK_SIZE), (1, 0)) @@ -551,7 +551,8 @@ def merge_fwd_bwd_kernel( b_ag_h = tl.load(p_ag_h, boundary_check=(0, 1)) p_ag_m = tl.make_block_ptr(ag_hm + cur_rank * stride + V, (K, K), (K + V, 1), (0, 0), (BK, BK), (1, 0)) b_ag_m = tl.load(p_ag_m, boundary_check=(0, 1)) - b_h = tl.dot(b_ag_m, b_h.to(b_ag_m.dtype)) + b_ag_h + # h = M @ h + h_ext + b_h = tl.dot(b_ag_m.to(tl.float32), b_h.to(tl.float32)) + b_ag_h.to(tl.float32) p_h = tl.make_block_ptr(h, (K, V), (V, 1), (0, i_v * BV), (BK, BV), (1, 0)) tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) @@ -1005,16 +1006,14 @@ def pre_process_bwd_kernel_merged( # Note: FORWARD=False uses tl.trans(b_w) @ b_k instead of tl.trans(b_k) @ b_w b_kw = tl.dot(tl.trans(b_w), b_k.to(b_w.dtype)) b_m_i = b_diag - b_kw - b_m = tl.dot(b_m_i.to(b_w.dtype), b_m.to(b_w.dtype)) + # Keep m chain in fp32 to avoid precision loss from repeated bf16 casting + b_m = tl.dot(b_m_i.to(tl.float32), b_m.to(tl.float32)) # Store dm result p_m = tl.make_block_ptr(dhm + V, (K, K), (V + K, 1), (0, i_k_col * BLOCK_SIZE), (BK1, BLOCK_SIZE), (1, 0)) tl.store(p_m, b_m.to(p_m.dtype.element_ty), boundary_check=(0, 1)) -DTYPE = torch.float32 - - def chunk_gated_delta_rule_fwd_h_pre_process( k: torch.Tensor, w: torch.Tensor, @@ -1043,8 +1042,8 @@ def chunk_gated_delta_rule_fwd_h_pre_process( N = len(cu_seqlens) - 1 assert K <= 256, "current kernel does not support head dimension larger than 256." - hm = k.new_zeros(H, K, (V + K), dtype=DTYPE) - initial_state = k.new_zeros(N, H, K, V, dtype=DTYPE) + hm = k.new_zeros(H, K, (V + K), dtype=torch.float32) + initial_state = k.new_zeros(N, H, K, V, dtype=torch.float32) if not context.is_last_rank: BLOCK_SIZE = 32 if K <= 64 else 64 grid = (triton.cdiv(V, BLOCK_SIZE) + triton.cdiv(K, BLOCK_SIZE), H) @@ -1112,8 +1111,9 @@ def chunk_gated_delta_rule_bwd_dhu_pre_process( else: N = len(cu_seqlens) - 1 - dhm = q.new_zeros(H, K, V + K, dtype=DTYPE) - dht = q.new_zeros(N, H, K, V, dtype=DTYPE) + dhm = q.new_zeros(H, K, V + K, dtype=torch.float32) + dht = q.new_zeros(N, H, K, V, dtype=torch.float32) + if not context.is_first_rank: BLOCK_SIZE = 32 if K <= 64 else 64 grid = (triton.cdiv(V, BLOCK_SIZE) + triton.cdiv(K, BLOCK_SIZE), H) @@ -1137,7 +1137,9 @@ def chunk_gated_delta_rule_bwd_dhu_pre_process( USE_EXP2=use_exp2, BLOCK_SIZE=BLOCK_SIZE, ) + ag_dhm, _ = all_gather_into_tensor(dhm, group=context.group) + if not context.is_last_rank: def grid(meta): return (triton.cdiv(V, meta['BV']), H) merge_fwd_bwd_kernel[grid]( @@ -1151,6 +1153,7 @@ def grid(meta): return (triton.cdiv(V, meta['BV']), H) BK=BK, FORWARD=False ) + # initial_state is None in the CP mode # We only need to compute dht of current rank and pass it to the backward kernel return dht, None diff --git a/tests/context_parallel/test_cp_conv.py b/tests/context_parallel/test_cp_conv.py index 27b80e2a1c..4d1008b150 100644 --- a/tests/context_parallel/test_cp_conv.py +++ b/tests/context_parallel/test_cp_conv.py @@ -41,6 +41,7 @@ 4. Single long sequence spanning all ranks """ +import logging import os import pytest @@ -50,10 +51,17 @@ from fla.modules.convolution import causal_conv1d from fla.ops.cp import build_cp_context +from fla.utils import assert_close + +# Configure logging to see assert_close messages +logging.basicConfig(level=logging.INFO, format='%(message)s') def init_distributed(rank, world_size): """Initialize distributed environment for a single process.""" + # Configure logging in worker process + logging.basicConfig(level=logging.INFO, format='%(message)s') + os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = '29500' os.environ['RANK'] = str(rank) @@ -70,13 +78,6 @@ def cleanup_distributed(): dist.destroy_process_group() -def assert_close(name, ref, tri, atol=1e-3, rtol=1e-3): - """Helper function for assertion with detailed diff output.""" - diff = (ref - tri).abs().max().item() - print(f" Checking {name:<10} | Max Diff: {diff:.6f}") - assert diff < atol, f"{name} mismatch. Max diff: {diff}" - - def run_cp_conv_test_worker( rank: int, world_size: int, @@ -109,7 +110,7 @@ def run_cp_conv_test_worker( torch.manual_seed(42) B = 1 - x_global = torch.randn(B, T, D, device=device, dtype=dtype) + x_global = torch.randn(B, T, D, device=device, dtype=dtype) * 100 dy_global = torch.randn(B, T, D, device=device, dtype=dtype) weight = torch.randn(D, W, device=device, dtype=dtype) @@ -197,10 +198,10 @@ def run_cp_conv_test_worker( if rank == 0: print(f"\n[{test_name}] Verification Results:") try: - assert_close("Output", ref_out, y_cp_global, atol=1e-3) - assert_close("dx", ref_dx, dx_cp_global, atol=1e-3) - assert_close("dw", ref_dw, dw_cp, atol=1e-3) - assert_close("db", ref_db, db_cp, atol=1e-3) + assert_close("Output", ref_out, y_cp_global, ratio=0.001) + assert_close("dx", ref_dx, dx_cp_global, ratio=0.001) + assert_close("dw", ref_dw, dw_cp, ratio=0.001) + assert_close("db", ref_db, db_cp, ratio=0.001) print(f"✅ [{test_name}] Test Passed!\n") except AssertionError as e: print(f"❌ [{test_name}] Test Failed: {e}\n") @@ -378,115 +379,3 @@ def setup_distributed_torchrun(): local_rank = int(os.environ["LOCAL_RANK"]) torch.cuda.set_device(local_rank) return True - - -if __name__ == "__main__": - # Check if running with torchrun - if 'RANK' in os.environ: - # Running with torchrun - if setup_distributed_torchrun(): - from fla.utils import device as fla_device - - world_size = dist.get_world_size() - rank = dist.get_rank() - - try: - if rank == 0: - print("=" * 60) - print("Running CP Conv1d Tests (torchrun mode)") - print("=" * 60) - - # Define test configs based on world_size - if world_size == 2: - test_configs = [ - ("CP2_SequenceCut", 1024, 128, 4, [300, 400, 324]), - ("CP2_BoundaryAligned", 1024, 128, 4, [512, 512]), - ("CP2_ManyShortSequences", 1024, 128, 4, [100, 150, 200, 250, 124, 100, 100]), - ] - elif world_size == 4: - test_configs = [ - ("CP4_Complex", 1024, 128, 4, [700, 324]), - ("CP4_SingleSequence", 1024, 128, 4, [1024]), - ] - else: - test_configs = [ - (f"CP{world_size}_SingleSequence", 1024, 128, 4, [1024]), - ] - - for test_name, T, D, W, lengths in test_configs: - # Run test inline (already in distributed context) - torch.manual_seed(42) - B = 1 - - x_global = torch.randn(B, T, D, device=fla_device, dtype=torch.float32) - dy_global = torch.randn(B, T, D, device=fla_device, dtype=torch.float32) - weight = torch.randn(D, W, device=fla_device, dtype=torch.float32) - bias = torch.randn(D, device=fla_device, dtype=torch.float32) - - dist.broadcast(weight, src=0) - dist.broadcast(bias, src=0) - - cu_seqlens_list = [0] + torch.cumsum(torch.tensor(lengths), 0).tolist() - cu_seqlens_global = torch.tensor(cu_seqlens_list, device=fla_device, dtype=torch.int32) - - # Reference - ref_out, ref_dx, ref_dw, ref_db = None, None, None, None - if rank == 0: - x_ref = x_global.clone().detach().requires_grad_(True) - weight_ref = weight.clone().detach().requires_grad_(True) - bias_ref = bias.clone().detach().requires_grad_(True) - y_ref, _ = causal_conv1d(x=x_ref, weight=weight_ref, bias=bias_ref, - activation='swish', backend='triton', cu_seqlens=cu_seqlens_global) - y_ref.backward(dy_global) - ref_out, ref_dx = y_ref.detach(), x_ref.grad.detach() - ref_dw, ref_db = weight_ref.grad.detach(), bias_ref.grad.detach() - - dist.barrier() - context = build_cp_context(cu_seqlens_global, group=dist.group.WORLD, conv1d_kernel_size=W) - - chunk_size = T // world_size - start_idx, end_idx = rank * chunk_size, (rank + 1) * chunk_size - x_local = x_global[:, start_idx:end_idx, :].clone().detach().requires_grad_(True) - dy_local = dy_global[:, start_idx:end_idx, :].clone() - weight_local = weight.clone().detach().requires_grad_(True) - bias_local = bias.clone().detach().requires_grad_(True) - - y_local, _ = causal_conv1d(x=x_local, weight=weight_local, bias=bias_local, - activation='swish', cp_context=context) - y_local.backward(dy_local) - - y_gathered = [torch.zeros_like(y_local) for _ in range(world_size)] - dist.all_gather(y_gathered, y_local) - y_cp_global = torch.cat(y_gathered, dim=1) - - dx_gathered = [torch.zeros_like(x_local.grad) for _ in range(world_size)] - dist.all_gather(dx_gathered, x_local.grad) - dx_cp_global = torch.cat(dx_gathered, dim=1) - - dw_cp, db_cp = weight_local.grad.clone(), bias_local.grad.clone() - dist.all_reduce(dw_cp, op=dist.ReduceOp.SUM) - dist.all_reduce(db_cp, op=dist.ReduceOp.SUM) - - if rank == 0: - print(f"\n[{test_name}] Verification:") - assert_close("Output", ref_out, y_cp_global, atol=1e-3) - assert_close("dx", ref_dx, dx_cp_global, atol=1e-3) - assert_close("dw", ref_dw, dw_cp, atol=1e-3) - assert_close("db", ref_db, db_cp, atol=1e-3) - print(f"✅ [{test_name}] Passed!") - - dist.barrier() - - if rank == 0: - print("\n" + "=" * 60) - print("All tests passed!") - print("=" * 60) - - finally: - cleanup_distributed() - else: - # Not running with torchrun, show usage - print("Run tests with pytest or torchrun:") - print(" pytest tests/context_parallel/test_cp_conv.py -v") - print(" torchrun --nproc_per_node=2 tests/context_parallel/test_cp_conv.py") - print(" torchrun --nproc_per_node=4 tests/context_parallel/test_cp_conv.py") diff --git a/tests/context_parallel/test_cp_gdn.py b/tests/context_parallel/test_cp_gdn.py index b4dab0fbc5..9c65f6a6eb 100644 --- a/tests/context_parallel/test_cp_gdn.py +++ b/tests/context_parallel/test_cp_gdn.py @@ -1,26 +1,49 @@ """ Test for Context Parallel (CP) Gated Delta Rule (GDN) -Context Parallel Principle for GDN: -=================================== - -GDN has a recurrent state dependency across tokens. The hidden state h evolves -as tokens are processed, creating dependencies that span the sequence. +Implementation Hierarchy and Relationships: +========================================== + +1. chunk_gated_delta_rule (fla/ops/gated_delta_rule/chunk.py): + - Production Triton kernel for GDN + - Input g is per-token log-space decay, shape [B, T, H] (scalar per head, NOT per-dim) + - Internally does chunk_local_cumsum on g + - Recurrence (delta rule with gating): + S_t = S_{t-1} * exp(g_t) + beta_t * k_t (x) (v_t - S_{t-1} @ k_t) + o_t = q_t^T @ S_t + where (x) denotes outer product, S is [K, V] state matrix + - No fused gate activation (unlike KDA), no lowerbound + - Supports variable-length sequences via cu_seqlens + - Supports context parallel via cp_context + +2. Context Parallel (CP) chunk_gated_delta_rule: + - Extension for multi-GPU distributed training + - Sequence is partitioned across ranks, with state communication between ranks + - Uses build_cp_context() to manage cross-rank dependencies + - Forward: Non-first ranks receive initial_state from previous rank + - Backward: Gradient dht flows back across rank boundaries + +Test Architecture Notes: +======================== +- Reference: single-GPU chunk_gated_delta_rule with cu_seqlens (varlen, no CP) +- CP path: same function with cp_context, sequence split across ranks +- Both should produce identical results + +Differences from KDA test: +- No use_gate_in_kernel / safe_gate / lower_bound / A_log / dt_bias +- No L2 normalization (q/k are pre-normalized before input) +- g shape is [B, T, H] not [B, T, H, D] + +Context Parallel Principle: +=========================== With Context Parallel: -1. Sequence Partitioning: The input sequence is split across ranks along the sequence dimension. - - Rank 0: tokens [0, T/N) - - Rank 1: tokens [T/N, 2T/N) - - ... +1. Sequence Partitioning: input sequence split across ranks along sequence dim + - Rank i: tokens [i*T/N, (i+1)*T/N) -2. Forward Pass: - - Each rank computes its local chunk - - Non-first ranks need the final state from previous rank as initial_state - - Communication: All-reduce style state passing between ranks +2. Forward: each rank computes local chunk; non-first ranks receive state from prev rank -3. Backward Pass: - - Gradients flow back through the recurrent state - - Communication: Gradient synchronization across ranks +3. Backward: gradients flow back through recurrent state across ranks Test Scenarios: =============== @@ -30,6 +53,7 @@ 4. CP4 with single long sequence """ +import logging import os import pytest @@ -40,12 +64,18 @@ from fla.ops.cp import build_cp_context from fla.ops.gated_delta_rule import chunk_gated_delta_rule +from fla.utils import assert_close + +# Configure logging to see assert_close messages +logging.basicConfig(level=logging.INFO, format='%(message)s') def init_distributed(rank, world_size): """Initialize distributed environment for a single process.""" + logging.basicConfig(level=logging.INFO, format='%(message)s') + os.environ['MASTER_ADDR'] = 'localhost' - os.environ['MASTER_PORT'] = '29502' # Different port from other tests + os.environ['MASTER_PORT'] = '29502' # Different port from KDA test os.environ['RANK'] = str(rank) os.environ['WORLD_SIZE'] = str(world_size) os.environ['LOCAL_RANK'] = str(rank) @@ -60,13 +90,6 @@ def cleanup_distributed(): dist.destroy_process_group() -def assert_close(name, ref, tri, atol=1e-3, rtol=1e-3): - """Helper function for assertion with detailed diff output.""" - diff = (ref - tri).abs().max().item() - print(f" Checking {name:<10} | Max Diff: {diff:.6f}") - assert diff < atol, f"{name} mismatch. Max diff: {diff}" - - def run_cp_gdn_test_worker( rank: int, world_size: int, @@ -95,16 +118,23 @@ def run_cp_gdn_test_worker( print(f"Sequence lengths: {lengths}") print(f"{'='*60}") - # Step 1: Prepare Global Data - torch.manual_seed(42) + # Step 1: Prepare Global Data (all generated on rank 0, broadcast to all) B = 1 + q_global = torch.empty(B, T, H, D, device=device, dtype=dtype) + k_global = torch.empty(B, T, H, D, device=device, dtype=dtype) + v_global = torch.empty(B, T, H, D, device=device, dtype=dtype) + g_global = torch.empty(B, T, H, device=device, dtype=dtype) + beta_global = torch.empty(B, T, H, device=device, dtype=torch.float32) + do_global = torch.empty(B, T, H, D, device=device, dtype=dtype) - # Generate inputs - note: g is [B, T, H] for gated_delta_rule, not [B, T, H, D] - q_global = torch.randn(B, T, H, D, device=device, dtype=dtype) - k_global = F.normalize(torch.randn(B, T, H, D, device=device, dtype=torch.float32), p=2, dim=-1).to(dtype) - v_global = torch.randn(B, T, H, D, device=device, dtype=dtype) - g_global = F.logsigmoid(torch.randn(B, T, H, device=device, dtype=dtype)) - beta_global = torch.randn(B, T, H, device=device, dtype=torch.float32).sigmoid() + if rank == 0: + torch.manual_seed(42) + q_global.copy_(F.normalize(torch.randn(B, T, H, D, device=device, dtype=torch.float32), p=2, dim=-1).to(dtype)) + k_global.copy_(F.normalize(torch.randn(B, T, H, D, device=device, dtype=torch.float32), p=2, dim=-1).to(dtype)) + v_global.copy_(torch.randn(B, T, H, D, device=device, dtype=dtype)) + g_global.copy_(F.logsigmoid(torch.randn(B, T, H, device=device, dtype=dtype))) + beta_global.copy_(torch.randn(B, T, H, device=device, dtype=torch.float32).sigmoid()) + do_global.copy_(torch.randn(B, T, H, D, device=device, dtype=dtype)) # Broadcast to ensure all ranks have same data dist.broadcast(q_global, src=0) @@ -112,16 +142,13 @@ def run_cp_gdn_test_worker( dist.broadcast(v_global, src=0) dist.broadcast(g_global, src=0) dist.broadcast(beta_global, src=0) + dist.broadcast(do_global, src=0) # Prepare cu_seqlens cu_seqlens_list = [0] + torch.cumsum(torch.tensor(lengths), 0).tolist() cu_seqlens_global = torch.tensor(cu_seqlens_list, device=device, dtype=torch.long) - # Prepare gradients - do_global = torch.randn(B, T, H, D, device=device, dtype=dtype) - dist.broadcast(do_global, src=0) - - # Step 2: Reference Run (single GPU, varlen) + # Step 2: Reference Run (single GPU, varlen, no CP) ref_out = None ref_dq, ref_dk, ref_dv, ref_dg, ref_db = None, None, None, None, None @@ -132,7 +159,6 @@ def run_cp_gdn_test_worker( g_ref = g_global.clone().detach().requires_grad_(True) beta_ref = beta_global.clone().detach().requires_grad_(True) - # Use chunk_gated_delta_rule with varlen (no CP) o_ref, _ = chunk_gated_delta_rule( q=q_ref, k=k_ref, @@ -154,14 +180,13 @@ def run_cp_gdn_test_worker( # Step 3: Context Parallel Run dist.barrier() - # Build CP context context = build_cp_context(cu_seqlens_global, group=dist.group.WORLD) chunk_size = T // world_size start_idx = rank * chunk_size end_idx = (rank + 1) * chunk_size - # Get local slices - note: g is [B, T, H] + # Get local slices - note: g is [B, T, H], beta is [B, T, H] q_local = q_global[:, start_idx:end_idx, :].clone().detach().requires_grad_(True) k_local = k_global[:, start_idx:end_idx, :].clone().detach().requires_grad_(True) v_local = v_global[:, start_idx:end_idx, :].clone().detach().requires_grad_(True) @@ -214,17 +239,23 @@ def run_cp_gdn_test_worker( test_passed = True if rank == 0: - print(f"\n[{test_name}] Verification Results:") + print(f"\n[{test_name}] Verifying results...") + + tensors_to_verify = [ + ("Output", ref_out, o_cp_global), + ("dq", ref_dq, dq_cp_global), + ("dk", ref_dk, dk_cp_global), + ("dv", ref_dv, dv_cp_global), + ("dg", ref_dg, dg_cp_global), + ("db", ref_db, db_cp_global), + ] + try: - assert_close("Output", ref_out, o_cp_global, atol=5e-3) - assert_close("dq", ref_dq, dq_cp_global, atol=8e-3) - assert_close("dk", ref_dk, dk_cp_global, atol=8e-3) - assert_close("dv", ref_dv, dv_cp_global, atol=8e-3) - assert_close("dg", ref_dg, dg_cp_global, atol=2e-2) - assert_close("db", ref_db, db_cp_global, atol=2e-2) - print(f"✅ [{test_name}] Test Passed!\n") + for name, ref, cp in tensors_to_verify: + assert_close(name, ref, cp, ratio=2e-3, warning=False) + print(f"[{test_name}] Test Passed!\n") except AssertionError as e: - print(f"❌ [{test_name}] Test Failed: {e}\n") + print(f"[{test_name}] Test Failed: {e}\n") test_passed = False dist.barrier() @@ -245,7 +276,7 @@ def run_cp_test_with_spawn( H: int, D: int, lengths: list[int], - dtype=torch.float16, + dtype=torch.bfloat16, ): """ Run CP test using torch.multiprocessing.spawn. @@ -265,116 +296,86 @@ def run_cp_test_with_spawn( # ============================================================ def test_cp2_sequence_cut(): - """ - Test Case 1: CP2 with sequences cut in the middle. - - Scenario: - - world_size=2, T=1024, chunk_size=512 - - lengths=[300, 400, 324] -> sequences span across rank boundary - - Rank 0: tokens [0, 512) contains seq0 (300) + part of seq1 (212) - - Rank 1: tokens [512, 1024) contains rest of seq1 (188) + seq2 (324) - """ + """CP2: sequences cut across rank boundary.""" if torch.cuda.device_count() < 2: pytest.skip("At least 2 GPUs required") run_cp_test_with_spawn( world_size=2, test_name="CP2_SequenceCut", - T=1024, - H=4, - D=64, - lengths=[300, 400, 324], - dtype=torch.float16, + T=10240, H=4, D=64, + lengths=[3000, 4000, 3240], + dtype=torch.bfloat16, ) def test_cp2_boundary_aligned(): - """ - Test Case 2: CP2 with sequence boundaries aligned with rank boundaries. - - Scenario: - - world_size=2, T=1024, chunk_size=512 - - lengths=[512, 512] -> sequence boundary exactly at rank boundary - - Rank 0: tokens [0, 512) contains exactly seq0 - - Rank 1: tokens [512, 1024) contains exactly seq1 - """ + """CP2: sequence boundaries aligned with rank boundaries.""" if torch.cuda.device_count() < 2: pytest.skip("At least 2 GPUs required") run_cp_test_with_spawn( world_size=2, test_name="CP2_BoundaryAligned", - T=1024, - H=4, - D=128, - lengths=[512, 512], - dtype=torch.float16, + T=10240, H=4, D=128, + lengths=[5120, 5120], + dtype=torch.bfloat16, ) def test_cp4_complex(): - """ - Test Case 3: CP4 with complex sequence distribution. - - Scenario: - - world_size=4, T=1024, chunk_size=256 - - lengths=[700, 324] -> first sequence spans 3 ranks - """ + """CP4: complex sequence distribution, first sequence spans 3 ranks.""" if torch.cuda.device_count() < 4: pytest.skip("At least 4 GPUs required") run_cp_test_with_spawn( world_size=4, test_name="CP4_Complex", - T=1024, - H=4, - D=128, - lengths=[700, 324], - dtype=torch.float16, + T=10240, H=4, D=128, + lengths=[7000, 3240], + dtype=torch.bfloat16, ) def test_cp4_single_sequence(): - """ - Test Case 4: CP4 with a single long sequence spanning all ranks. - - Scenario: - - world_size=4, T=1024, chunk_size=256 - - lengths=[1024] -> single sequence spans all 4 ranks - """ + """CP4: single long sequence spanning all ranks.""" if torch.cuda.device_count() < 4: pytest.skip("At least 4 GPUs required") run_cp_test_with_spawn( world_size=4, test_name="CP4_SingleSequence", - T=1024, - H=4, - D=64, - lengths=[1024], - dtype=torch.float16, + T=10240, H=4, D=128, + lengths=[10240], + dtype=torch.bfloat16, ) -def test_cp2_many_short_sequences(): - """ - Test Case 5: CP2 with many short sequences. +def test_cp8_single_sequence(): + """CP8: single long sequence spanning all ranks.""" + if torch.cuda.device_count() < 8: + pytest.skip("At least 8 GPUs required") - Scenario: - - world_size=2, T=1024, chunk_size=512 - - lengths=[100, 150, 200, 250, 124, 100, 100] -> many short sequences - """ + run_cp_test_with_spawn( + world_size=8, + test_name="CP8_SingleSequence", + T=65536, H=4, D=128, + lengths=[65536], + dtype=torch.bfloat16, + ) + + +def test_cp2_many_short_sequences(): + """CP2: many short sequences.""" if torch.cuda.device_count() < 2: pytest.skip("At least 2 GPUs required") run_cp_test_with_spawn( world_size=2, test_name="CP2_ManyShortSequences", - T=1024, - H=4, - D=64, - lengths=[100, 150, 200, 250, 124, 100, 100], - dtype=torch.float16, + T=10240, H=4, D=128, + lengths=[1000, 1500, 2000, 2500, 1240, 1000, 1000], + dtype=torch.bfloat16, ) @@ -391,163 +392,3 @@ def setup_distributed_torchrun(): local_rank = int(os.environ["LOCAL_RANK"]) torch.cuda.set_device(local_rank) return True - - -if __name__ == "__main__": - # Check if running with torchrun - if 'RANK' in os.environ: - # Running with torchrun - if setup_distributed_torchrun(): - from fla.utils import device as fla_device - - world_size = dist.get_world_size() - rank = dist.get_rank() - - try: - if rank == 0: - print("=" * 60) - print("Running CP GDN Tests (torchrun mode)") - print("=" * 60) - - # Define test configs based on world_size - if world_size == 2: - test_configs = [ - ("CP2_SequenceCut", 1024, 4, 64, [300, 400, 324]), - ("CP2_BoundaryAligned", 1024, 4, 64, [512, 512]), - ("CP2_ManyShortSequences", 1024, 4, 64, [100, 150, 200, 250, 124, 100, 100]), - ] - elif world_size == 4: - test_configs = [ - ("CP4_Complex", 1024, 4, 64, [700, 324]), - ("CP4_SingleSequence", 1024, 4, 64, [1024]), - ] - else: - test_configs = [ - (f"CP{world_size}_SingleSequence", 1024, 4, 64, [1024]), - ] - - for test_name, T, H, D, lengths in test_configs: - torch.manual_seed(42) - B = 1 - - # Generate global data - note: g is [B, T, H] for gated_delta_rule - q_global = torch.randn(B, T, H, D, device=fla_device, dtype=torch.float16) - k_global = F.normalize(torch.randn(B, T, H, D, device=fla_device, - dtype=torch.float32), p=2, dim=-1).to(torch.float16) - v_global = torch.randn(B, T, H, D, device=fla_device, dtype=torch.float16) - g_global = F.logsigmoid(torch.randn(B, T, H, device=fla_device, dtype=torch.float16)) - beta_global = torch.randn(B, T, H, device=fla_device, dtype=torch.float32).sigmoid() - do_global = torch.randn(B, T, H, D, device=fla_device, dtype=torch.float16) - - # Broadcast - dist.broadcast(q_global, src=0) - dist.broadcast(k_global, src=0) - dist.broadcast(v_global, src=0) - dist.broadcast(g_global, src=0) - dist.broadcast(beta_global, src=0) - dist.broadcast(do_global, src=0) - - cu_seqlens_list = [0] + torch.cumsum(torch.tensor(lengths), 0).tolist() - cu_seqlens_global = torch.tensor(cu_seqlens_list, device=fla_device, dtype=torch.long) - - # Reference - ref_out, ref_dq, ref_dk, ref_dv, ref_dg, ref_db = None, None, None, None, None, None - if rank == 0: - q_ref = q_global.clone().detach().requires_grad_(True) - k_ref = k_global.clone().detach().requires_grad_(True) - v_ref = v_global.clone().detach().requires_grad_(True) - g_ref = g_global.clone().detach().requires_grad_(True) - beta_ref = beta_global.clone().detach().requires_grad_(True) - - o_ref, _ = chunk_gated_delta_rule( - q=q_ref, - k=k_ref, - v=v_ref, - g=g_ref, - beta=beta_ref, - cu_seqlens=cu_seqlens_global, - ) - o_ref.backward(do_global) - ref_out = o_ref.detach() - ref_dq = q_ref.grad.detach() - ref_dk = k_ref.grad.detach() - ref_dv = v_ref.grad.detach() - ref_dg = g_ref.grad.detach() - ref_db = beta_ref.grad.detach() - - dist.barrier() - - # CP run - context = build_cp_context(cu_seqlens_global, group=dist.group.WORLD) - - chunk_size = T // world_size - start_idx, end_idx = rank * chunk_size, (rank + 1) * chunk_size - - # note: g is [B, T, H] - q_local = q_global[:, start_idx:end_idx, :].clone().detach().requires_grad_(True) - k_local = k_global[:, start_idx:end_idx, :].clone().detach().requires_grad_(True) - v_local = v_global[:, start_idx:end_idx, :].clone().detach().requires_grad_(True) - g_local = g_global[:, start_idx:end_idx].clone().detach().requires_grad_(True) - beta_local = beta_global[:, start_idx:end_idx].clone().detach().requires_grad_(True) - do_local = do_global[:, start_idx:end_idx, :].clone() - - o_local, _ = chunk_gated_delta_rule( - q=q_local, - k=k_local, - v=v_local, - g=g_local, - beta=beta_local, - cp_context=context, - ) - o_local.backward(do_local) - - # Gather output and gradients - o_gathered = [torch.zeros_like(o_local) for _ in range(world_size)] - dist.all_gather(o_gathered, o_local) - o_cp_global = torch.cat(o_gathered, dim=1) - - dq_gathered = [torch.zeros_like(q_local.grad) for _ in range(world_size)] - dist.all_gather(dq_gathered, q_local.grad) - dq_cp_global = torch.cat(dq_gathered, dim=1) - - dk_gathered = [torch.zeros_like(k_local.grad) for _ in range(world_size)] - dist.all_gather(dk_gathered, k_local.grad) - dk_cp_global = torch.cat(dk_gathered, dim=1) - - dv_gathered = [torch.zeros_like(v_local.grad) for _ in range(world_size)] - dist.all_gather(dv_gathered, v_local.grad) - dv_cp_global = torch.cat(dv_gathered, dim=1) - - dg_gathered = [torch.zeros_like(g_local.grad) for _ in range(world_size)] - dist.all_gather(dg_gathered, g_local.grad) - dg_cp_global = torch.cat(dg_gathered, dim=1) - - db_gathered = [torch.zeros_like(beta_local.grad) for _ in range(world_size)] - dist.all_gather(db_gathered, beta_local.grad) - db_cp_global = torch.cat(db_gathered, dim=1) - - if rank == 0: - print(f"\n[{test_name}] Verification:") - assert_close("Output", ref_out, o_cp_global, atol=5e-3) - assert_close("dq", ref_dq, dq_cp_global, atol=8e-3) - assert_close("dk", ref_dk, dk_cp_global, atol=8e-3) - assert_close("dv", ref_dv, dv_cp_global, atol=8e-3) - assert_close("dg", ref_dg, dg_cp_global, atol=2e-2) - assert_close("db", ref_db, db_cp_global, atol=2e-2) - print(f"✅ [{test_name}] Passed!") - - dist.barrier() - - if rank == 0: - print("\n" + "=" * 60) - print("All tests passed!") - print("=" * 60) - - finally: - cleanup_distributed() - else: - # Not running with torchrun, show usage - print("Run tests with pytest or torchrun:") - print(" pytest tests/context_parallel/test_cp_gdn.py -v") - print(" torchrun --nproc_per_node=2 tests/context_parallel/test_cp_gdn.py") - print(" torchrun --nproc_per_node=4 tests/context_parallel/test_cp_gdn.py") diff --git a/tests/context_parallel/test_cp_kda.py b/tests/context_parallel/test_cp_kda.py index a9a1f0fe94..5ef323783f 100644 --- a/tests/context_parallel/test_cp_kda.py +++ b/tests/context_parallel/test_cp_kda.py @@ -1,6 +1,80 @@ """ Test for Context Parallel (CP) KDA (Kimi Delta Attention) +Implementation Hierarchy and Relationships: +========================================== + +1. naive_recurrent_kda (fla/ops/kda/naive.py): + - The mathematical gold standard - sequential token-by-token computation + - Input g is per-token log-space decay (NOT cumulative) + - Each g_t is used independently as: S_t = S_{t-1} * exp(g_t) + - Recurrence: + S_t = S_{t-1} * exp(g_t) + beta_t * k_t (x) (v_t - S_{t-1} @ k_t) + o_t = q_t^T @ S_t + where (x) denotes outer product, S is [K, V] state matrix + - No internal gate computation or L2 normalization + - Used as the reference baseline for correctness verification + +2. naive_chunk_kda (fla/ops/kda/naive.py): + - Chunk-parallel version mathematically equivalent to naive_recurrent_kda + - Input g is also per-token (same as naive_recurrent_kda) + - Performs cumsum on g internally: g = g.cumsum(-2) at line 60 + - Still a PyTorch implementation, but exploits chunk-level parallelism + - The chunk algorithm uses the "w-y representation" for efficient computation + +3. chunk_kda (fla/ops/kda/chunk.py ChunkKDAFunction): + - Production Triton kernel implementation with optimizations: + a) Fused L2 normalization (when use_qk_l2norm_in_kernel=True) + b) Fused gate computation (when use_gate_in_kernel=True) + c) Variable-length sequence support (cu_seqlens) + - Gate processing flow when use_gate_in_kernel=True: + - Raw g -> kda_gate_chunk_cumsum() -> applies gate formula + cumsum + scale(RCP_LN2) + - For safe_gate=True: gate_t = lower_bound * sigmoid(exp(A_log) * (g_t + dt_bias)) + - For safe_gate=False: gate_t = -exp(A_log) * softplus(g_t + dt_bias) + - kda_gate_chunk_cumsum outputs CUMSUM'd + SCALED values (gate + cumsum in one fused op) + - L2 normalization is applied AFTER gate computation + +4. Context Parallel (CP) chunk_kda: + - Extension of chunk_kda for multi-GPU distributed training + - Sequence is partitioned across ranks, with state communication between ranks + - Uses build_cp_context() to manage cross-rank dependencies + - Forward: Non-first ranks receive initial_state from previous rank + - Backward: Gradient dht flows back across rank boundaries + +Gate Functions (fla/ops/kda/gate.py): +===================================== + +- naive_kda_gate / naive_kda_lowerbound_gate: + Pure PyTorch reference. Output is per-token gate values (NOT cumsum'd). + Used in this test's reference path to compute g for naive_recurrent_kda. + +- kda_gate_chunk_cumsum (Triton kernel): + Fused gate + chunk-local cumsum + optional scale. Output IS cumsum'd. + Used internally by chunk_kda when use_gate_in_kernel=True. + +Test Architecture Notes: +======================== +This test uses naive_recurrent_kda as the reference baseline because: +- It represents the exact mathematical definition without chunking approximations +- It expects per-token g (non-cumsum'd) and pre-normalized q/k + +When use_gate_in_kernel=True, the reference path must manually: +1. Apply L2 normalization to q/k before passing to naive_recurrent_kda +2. Apply the gate formula (naive_kda_lowerbound_gate or naive_kda_gate) +3. The output g from these naive gate functions is per-token (non-cumsum'd), + which is exactly what naive_recurrent_kda expects + +When use_gate_in_kernel=False: +- For chunk_kda: input g must be pre-processed (gate applied + cumsum'd) +- For naive_recurrent_kda reference: g should be gate-applied but NOT cumsum'd + +Variable-Length Sequence Handling: +================================== +- All sequences are flattened to batch size 1 with cu_seqlens markers +- Each sequence is processed independently with h0=0 (no cross-sequence state) +- The reference computation loops over sequences and concatenates results +- CP path uses the same cu_seqlens for correct chunk boundary handling + Context Parallel Principle for KDA: =================================== @@ -30,6 +104,7 @@ 4. CP4 with single long sequence """ +import logging import os import pytest @@ -40,10 +115,19 @@ from fla.ops.cp import build_cp_context from fla.ops.kda import chunk_kda +from fla.ops.kda.gate import naive_kda_lowerbound_gate +from fla.ops.kda.naive import naive_recurrent_kda +from fla.utils import assert_close + +# Configure logging to see assert_close messages +logging.basicConfig(level=logging.INFO, format='%(message)s') def init_distributed(rank, world_size): """Initialize distributed environment for a single process.""" + # Configure logging in worker process + logging.basicConfig(level=logging.INFO, format='%(message)s') + os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = '29501' # Different port from conv test os.environ['RANK'] = str(rank) @@ -60,13 +144,6 @@ def cleanup_distributed(): dist.destroy_process_group() -def assert_close(name, ref, tri, atol=1e-3, rtol=1e-3): - """Helper function for assertion with detailed diff output.""" - diff = (ref - tri).abs().max().item() - print(f" Checking {name:<10} | Max Diff: {diff:.6f}") - assert diff < atol, f"{name} mismatch. Max diff: {diff}" - - def run_cp_kda_test_worker( rank: int, world_size: int, @@ -77,6 +154,9 @@ def run_cp_kda_test_worker( lengths: list[int], dtype, disable_recompute: bool = False, + use_gate_in_kernel: bool = False, + safe_gate: bool = False, + lower_bound: float | None = None, ): """ Worker function for CP KDA test. @@ -93,19 +173,37 @@ def run_cp_kda_test_worker( print(f"\n{'='*60}") print(f"Test: {test_name}") print(f"Config: T={T}, H={H}, D={D}, world_size={world_size}, disable_recompute={disable_recompute}") + print(f"use_gate_in_kernel={use_gate_in_kernel}, safe_gate={safe_gate}, lower_bound={lower_bound}") print(f"Sequence lengths: {lengths}") print(f"{'='*60}") - # Step 1: Prepare Global Data - torch.manual_seed(42) + # Step 1: Prepare Global Data (all generated on rank 0, broadcast to other ranks) B = 1 + q_global = torch.empty(B, T, H, D, device=device, dtype=dtype) + k_global = torch.empty(B, T, H, D, device=device, dtype=dtype) + v_global = torch.empty(B, T, H, D, device=device, dtype=dtype) + g_global = torch.empty(B, T, H, D, device=device, dtype=dtype if use_gate_in_kernel else torch.float) + beta_global = torch.empty(B, T, H, device=device, dtype=dtype) + do_global = torch.empty(B, T, H, D, device=device, dtype=dtype) + A_log_global = None + dt_bias_global = None - # Generate inputs - q_global = torch.randn(B, T, H, D, device=device, dtype=dtype) - k_global = torch.randn(B, T, H, D, device=device, dtype=dtype) - v_global = torch.randn(B, T, H, D, device=device, dtype=dtype) - g_global = F.logsigmoid(torch.randn(B, T, H, D, device=device, dtype=torch.float)) - beta_global = torch.randn(B, T, H, device=device, dtype=dtype).sigmoid() + if rank == 0: + torch.manual_seed(42) + # Generate inputs + # Asymmetric initialization for q and k to test L2 norm + # q: small positive values around 1.0 + # k: large negative values around -50.0 + q_global.copy_(torch.randn(B, T, H, D, device=device, dtype=dtype) + 1.0) + k_global.copy_(torch.randn(B, T, H, D, device=device, dtype=dtype) * 10.0 - 50.0) + v_global.copy_(torch.randn(B, T, H, D, device=device, dtype=dtype)) + if use_gate_in_kernel: + g_global.copy_(torch.randn(B, T, H, D, device=device, dtype=dtype)) + else: + g_global.copy_(F.logsigmoid(torch.randn(B, T, H, D, device=device, dtype=torch.float))) + g_global.clamp_(min=-5.0) + beta_global.copy_(torch.randn(B, T, H, device=device, dtype=dtype).sigmoid()) + do_global.copy_(torch.randn(B, T, H, D, device=device, dtype=dtype)) # Broadcast to ensure all ranks have same data dist.broadcast(q_global, src=0) @@ -113,45 +211,146 @@ def run_cp_kda_test_worker( dist.broadcast(v_global, src=0) dist.broadcast(g_global, src=0) dist.broadcast(beta_global, src=0) + dist.broadcast(do_global, src=0) + + # Prepare and broadcast A_log and dt_bias for gate computation + # Always generate these for reference, even when use_gate_in_kernel=False + import triton + num_even_heads = triton.next_power_of_2(H) + projection_size = H * D + + # All ranks create tensors with same shape first + A_log_global = torch.empty(num_even_heads, dtype=torch.float32, device=device) + dt_bias_global = torch.empty(projection_size, dtype=torch.float32, device=device) + + # Rank 0 fills in the actual data + if rank == 0: + import math + if safe_gate and lower_bound is not None: + A_log_global.copy_(torch.log(torch.ones(num_even_heads, dtype=torch.float32, device=device))) + else: + A_log_global.copy_(torch.log(torch.empty(num_even_heads, dtype=torch.float32, device=device).uniform_(1, 16))) + # dt initialization from real training + dt = torch.exp(torch.rand(projection_size, device=device) * + (math.log(0.1) - math.log(0.001)) + math.log(0.001)).clamp_(min=1e-4) + dt_bias_global.copy_(dt + torch.log(-torch.expm1(-dt))) + + # Broadcast from rank 0 to all ranks + dist.broadcast(A_log_global, src=0) + dist.broadcast(dt_bias_global, src=0) # Prepare cu_seqlens cu_seqlens_list = [0] + torch.cumsum(torch.tensor(lengths), 0).tolist() cu_seqlens_global = torch.tensor(cu_seqlens_list, device=device, dtype=torch.long) - # Prepare gradients - do_global = torch.randn(B, T, H, D, device=device, dtype=dtype) - dist.broadcast(do_global, src=0) - - # Step 2: Reference Run (single GPU, varlen) + # Step 2: Reference Run using naive_recurrent_kda (ground truth) + # Each sequence is computed independently with h0=0 for simplicity ref_out = None ref_dq, ref_dk, ref_dv, ref_dg, ref_db = None, None, None, None, None if rank == 0: - q_ref = q_global.clone().detach().requires_grad_(True) - k_ref = k_global.clone().detach().requires_grad_(True) - v_ref = v_global.clone().detach().requires_grad_(True) - g_ref = g_global.clone().detach().requires_grad_(True) - beta_ref = beta_global.clone().detach().requires_grad_(True) - - # Use chunk_kda with varlen (no CP) - o_ref, _ = chunk_kda( - q=F.normalize(q_ref, p=2, dim=-1), - k=F.normalize(k_ref, p=2, dim=-1), - v=v_ref, - g=g_ref, - beta=beta_ref, - cu_seqlens=cu_seqlens_global, - disable_recompute=disable_recompute, - ) - - o_ref.backward(do_global) - - ref_out = o_ref.detach() - ref_dq = q_ref.grad.detach() - ref_dk = k_ref.grad.detach() - ref_dv = v_ref.grad.detach() - ref_dg = g_ref.grad.detach() - ref_db = beta_ref.grad.detach() + N = len(lengths) + ref_outputs = [] + + for i in range(N): + seq_start = cu_seqlens_list[i] + seq_end = cu_seqlens_list[i + 1] + + # Extract sequence data and create leaf tensors + q_seq = q_global[:, seq_start:seq_end].clone().detach().requires_grad_(True) + k_seq = k_global[:, seq_start:seq_end].clone().detach().requires_grad_(True) + v_seq = v_global[:, seq_start:seq_end].clone().detach().requires_grad_(True) + g_seq = g_global[:, seq_start:seq_end].clone().detach() + beta_seq = beta_global[:, seq_start:seq_end].clone().detach().requires_grad_(True) + do_seq = do_global[:, seq_start:seq_end].clone() + + # Apply L2 normalization to q (naive implementation expects normalized q) + q_seq_norm = F.normalize(q_seq, p=2, dim=-1) + k_seq_norm = F.normalize(k_seq, p=2, dim=-1) + + # Compute gate using naive implementation + # When use_gate_in_kernel=True, reference should compute g_raw's grad + # When use_gate_in_kernel=False, reference should compute g_processed's grad + if use_gate_in_kernel: + # For kernel mode: compare g_raw gradients + g_seq_input = g_seq.requires_grad_(True) + if safe_gate and lower_bound is not None: + g_processed = naive_kda_lowerbound_gate( + g_seq_input.to(torch.float), + A_log_global[:H].contiguous(), + dt_bias_global, + lower_bound=lower_bound + ) + else: + from fla.ops.kda.gate import naive_kda_gate + g_processed = naive_kda_gate( + g_seq_input.to(torch.float), + A_log_global[:H].contiguous(), + dt_bias_global + ) + g_for_grad = g_seq_input + else: + # For non-kernel mode: compare g_processed gradients + g_seq_input = g_seq # No grad needed for raw input + if safe_gate and lower_bound is not None: + g_processed = naive_kda_lowerbound_gate( + g_seq_input.to(torch.float), + A_log_global[:H].contiguous(), + dt_bias_global, + lower_bound=lower_bound + ).requires_grad_(True) + else: + from fla.ops.kda.gate import naive_kda_gate + g_processed = naive_kda_gate( + g_seq_input.to(torch.float), + A_log_global[:H].contiguous(), + dt_bias_global + ).requires_grad_(True) + g_for_grad = g_processed + + # Run naive forward with h0=0 (independent sequences) + o_seq, _ = naive_recurrent_kda( + q=q_seq_norm, + k=k_seq_norm, + v=v_seq, + g=g_processed, + beta=beta_seq, + initial_state=None, + output_final_state=False, + ) + + ref_outputs.append({ + 'o': o_seq, + 'q': q_seq, + 'k': k_seq, + 'v': v_seq, + 'g': g_for_grad, + 'beta': beta_seq, + 'do': do_seq, + }) + + # Backward pass for each sequence + all_dq = [] + all_dk = [] + all_dv = [] + all_dg = [] + all_db = [] + + for item in ref_outputs: + (item['o'] * item['do']).sum().backward() + all_dq.append(item['q'].grad.detach()) + all_dk.append(item['k'].grad.detach()) + all_dv.append(item['v'].grad.detach()) + all_dg.append(item['g'].grad.detach()) + all_db.append(item['beta'].grad.detach()) + + # Concatenate outputs and gradients + ref_out = torch.cat([item['o'].detach() for item in ref_outputs], dim=1) + ref_dq = torch.cat(all_dq, dim=1) + ref_dk = torch.cat(all_dk, dim=1) + ref_dv = torch.cat(all_dv, dim=1) + ref_dg = torch.cat(all_dg, dim=1) + ref_db = torch.cat(all_db, dim=1) # Step 3: Context Parallel Run dist.barrier() @@ -178,13 +377,19 @@ def run_cp_kda_test_worker( # CP Forward o_local, _ = chunk_kda( - q=F.normalize(q_local, p=2, dim=-1), - k=F.normalize(k_local, p=2, dim=-1), + q=q_local, + k=k_local, v=v_local, g=g_local, beta=beta_local, cp_context=context, disable_recompute=disable_recompute, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=use_gate_in_kernel, + safe_gate=safe_gate, + lower_bound=lower_bound, + A_log=A_log_global[:H].contiguous() if use_gate_in_kernel else None, + dt_bias=dt_bias_global if use_gate_in_kernel else None, ) # CP Backward @@ -217,14 +422,20 @@ def run_cp_kda_test_worker( test_passed = True if rank == 0: - print(f"\n[{test_name}] Verification Results:") + print(f"\n[{test_name}] Verifying results...") + + tensors_to_verify = [ + ("Output", ref_out, o_cp_global), + ("dq", ref_dq, dq_cp_global), + ("dk", ref_dk, dk_cp_global), + ("dv", ref_dv, dv_cp_global), + ("dg", ref_dg, dg_cp_global), + ("db", ref_db, db_cp_global), + ] + try: - assert_close("Output", ref_out, o_cp_global, atol=5e-3) - assert_close("dq", ref_dq, dq_cp_global, atol=8e-3) - assert_close("dk", ref_dk, dk_cp_global, atol=8e-3) - assert_close("dv", ref_dv, dv_cp_global, atol=8e-3) - assert_close("dg", ref_dg, dg_cp_global, atol=2e-2) - assert_close("db", ref_db, db_cp_global, atol=2e-2) + for name, ref, cp in tensors_to_verify: + assert_close(name, ref, cp, ratio=5e-2, warning=False) print(f"✅ [{test_name}] Test Passed!\n") except AssertionError as e: print(f"❌ [{test_name}] Test Failed: {e}\n") @@ -248,8 +459,11 @@ def run_cp_test_with_spawn( H: int, D: int, lengths: list[int], - dtype=torch.float16, + dtype=torch.bfloat16, disable_recompute: bool = False, + use_gate_in_kernel: bool = False, + safe_gate: bool = False, + lower_bound: float | None = None, ): """ Run CP test using torch.multiprocessing.spawn. @@ -257,7 +471,7 @@ def run_cp_test_with_spawn( """ mp.start_processes( run_cp_kda_test_worker, - args=(world_size, test_name, T, H, D, lengths, dtype, disable_recompute), + args=(world_size, test_name, T, H, D, lengths, dtype, disable_recompute, use_gate_in_kernel, safe_gate, lower_bound), nprocs=world_size, join=True, start_method='spawn', @@ -266,143 +480,115 @@ def run_cp_test_with_spawn( # ============================================================ # Test Scenario Definitions +# All tests use: use_gate_in_kernel=True, safe_gate=True, lower_bound=-5.0 # ============================================================ -def test_cp2_sequence_cut(): - """ - Test Case 1: CP2 with sequences cut in the middle. +GATE_KWARGS = dict(use_gate_in_kernel=True, safe_gate=True, lower_bound=-5.0) - Scenario: - - world_size=2, T=1024, chunk_size=512 - - lengths=[300, 400, 324] -> sequences span across rank boundary - - Rank 0: tokens [0, 512) contains seq0 (300) + part of seq1 (212) - - Rank 1: tokens [512, 1024) contains rest of seq1 (188) + seq2 (324) - """ + +def test_cp2_sequence_cut(): + """CP2: sequences cut across rank boundary.""" if torch.cuda.device_count() < 2: pytest.skip("At least 2 GPUs required") run_cp_test_with_spawn( world_size=2, test_name="CP2_SequenceCut", - T=1024, - H=4, - D=64, - lengths=[300, 400, 324], - dtype=torch.float16, + T=10240, H=4, D=128, + lengths=[3000, 4000, 3240], + dtype=torch.bfloat16, + **GATE_KWARGS, ) def test_cp2_boundary_aligned(): - """ - Test Case 2: CP2 with sequence boundaries aligned with rank boundaries. - - Scenario: - - world_size=2, T=1024, chunk_size=512 - - lengths=[512, 512] -> sequence boundary exactly at rank boundary - - Rank 0: tokens [0, 512) contains exactly seq0 - - Rank 1: tokens [512, 1024) contains exactly seq1 - """ + """CP2: sequence boundaries aligned with rank boundaries.""" if torch.cuda.device_count() < 2: pytest.skip("At least 2 GPUs required") run_cp_test_with_spawn( world_size=2, test_name="CP2_BoundaryAligned", - T=1024, - H=4, - D=128, - lengths=[512, 512], - dtype=torch.float16, + T=10240, H=4, D=128, + lengths=[5120, 5120], + dtype=torch.bfloat16, + **GATE_KWARGS, ) def test_cp4_complex(): - """ - Test Case 3: CP4 with complex sequence distribution. - - Scenario: - - world_size=4, T=1024, chunk_size=256 - - lengths=[700, 324] -> first sequence spans 3 ranks - """ + """CP4: complex sequence distribution, first sequence spans 3 ranks.""" if torch.cuda.device_count() < 4: pytest.skip("At least 4 GPUs required") run_cp_test_with_spawn( world_size=4, test_name="CP4_Complex", - T=1024, - H=4, - D=128, - lengths=[700, 324], - dtype=torch.float16, + T=10240, H=4, D=128, + lengths=[7000, 3240], + dtype=torch.bfloat16, + **GATE_KWARGS, ) def test_cp4_single_sequence(): - """ - Test Case 4: CP4 with a single long sequence spanning all ranks. - - Scenario: - - world_size=4, T=1024, chunk_size=256 - - lengths=[1024] -> single sequence spans all 4 ranks - """ + """CP4: single long sequence spanning all ranks.""" if torch.cuda.device_count() < 4: pytest.skip("At least 4 GPUs required") run_cp_test_with_spawn( world_size=4, test_name="CP4_SingleSequence", - T=1024, - H=4, - D=64, - lengths=[1024], - dtype=torch.float16, + T=10240, H=4, D=128, + lengths=[10240], + dtype=torch.bfloat16, + **GATE_KWARGS, ) -def test_cp2_many_short_sequences(): - """ - Test Case 5: CP2 with many short sequences. +def test_cp8_single_sequence(): + """CP8: single long sequence spanning all ranks.""" + if torch.cuda.device_count() < 8: + pytest.skip("At least 8 GPUs required") - Scenario: - - world_size=2, T=1024, chunk_size=512 - - lengths=[100, 150, 200, 250, 124, 100, 100] -> many short sequences - """ + run_cp_test_with_spawn( + world_size=8, + test_name="CP8_SingleSequence", + T=65536, H=4, D=128, + lengths=[65536], + dtype=torch.bfloat16, + **GATE_KWARGS, + ) + + +def test_cp2_many_short_sequences(): + """CP2: many short sequences.""" if torch.cuda.device_count() < 2: pytest.skip("At least 2 GPUs required") run_cp_test_with_spawn( world_size=2, test_name="CP2_ManyShortSequences", - T=1024, - H=4, - D=64, - lengths=[100, 150, 200, 250, 124, 100, 100], - dtype=torch.float16, + T=10240, H=4, D=128, + lengths=[1000, 1500, 2000, 2500, 1240, 1000, 1000], + dtype=torch.bfloat16, + **GATE_KWARGS, ) def test_cp2_disable_recompute(): - """ - Test Case 6: CP2 with disable_recompute=True. - - Scenario: - - world_size=2, T=1024, chunk_size=512 - - lengths=[300, 400, 324] -> sequences span across rank boundary - - disable_recompute=True - """ + """CP2: disable_recompute=True.""" if torch.cuda.device_count() < 2: pytest.skip("At least 2 GPUs required") run_cp_test_with_spawn( world_size=2, test_name="CP2_DisableRecompute", - T=1024, - H=4, - D=64, - lengths=[300, 400, 324], - dtype=torch.float16, + T=10240, H=4, D=128, + lengths=[3000, 4000, 3240], + dtype=torch.bfloat16, disable_recompute=True, + **GATE_KWARGS, ) @@ -419,161 +605,3 @@ def setup_distributed_torchrun(): local_rank = int(os.environ["LOCAL_RANK"]) torch.cuda.set_device(local_rank) return True - - -if __name__ == "__main__": - # Check if running with torchrun - if 'RANK' in os.environ: - # Running with torchrun - if setup_distributed_torchrun(): - from fla.utils import device as fla_device - - world_size = dist.get_world_size() - rank = dist.get_rank() - - try: - if rank == 0: - print("=" * 60) - print("Running CP KDA Tests (torchrun mode)") - print("=" * 60) - - # Define test configs based on world_size - if world_size == 2: - test_configs = [ - ("CP2_SequenceCut", 1024, 4, 64, [300, 400, 324]), - ("CP2_BoundaryAligned", 1024, 4, 64, [512, 512]), - ("CP2_ManyShortSequences", 1024, 4, 64, [100, 150, 200, 250, 124, 100, 100]), - ] - elif world_size == 4: - test_configs = [ - ("CP4_Complex", 1024, 4, 64, [700, 324]), - ("CP4_SingleSequence", 1024, 4, 64, [1024]), - ] - else: - test_configs = [ - (f"CP{world_size}_SingleSequence", 1024, 4, 64, [1024]), - ] - - for test_name, T, H, D, lengths in test_configs: - torch.manual_seed(42) - B = 1 - - # Generate global data - q_global = torch.randn(B, T, H, D, device=fla_device, dtype=torch.float16) - k_global = torch.randn(B, T, H, D, device=fla_device, dtype=torch.float16) - v_global = torch.randn(B, T, H, D, device=fla_device, dtype=torch.float16) - g_global = F.logsigmoid(torch.randn(B, T, H, D, device=fla_device, dtype=torch.float)) - beta_global = torch.randn(B, T, H, device=fla_device, dtype=torch.float16).sigmoid() - do_global = torch.randn(B, T, H, D, device=fla_device, dtype=torch.float16) - - # Broadcast - dist.broadcast(q_global, src=0) - dist.broadcast(k_global, src=0) - dist.broadcast(v_global, src=0) - dist.broadcast(g_global, src=0) - dist.broadcast(beta_global, src=0) - dist.broadcast(do_global, src=0) - - cu_seqlens_list = [0] + torch.cumsum(torch.tensor(lengths), 0).tolist() - cu_seqlens_global = torch.tensor(cu_seqlens_list, device=fla_device, dtype=torch.long) - - # Reference - ref_out, ref_dq, ref_dk, ref_dv, ref_dg, ref_db = None, None, None, None, None, None - if rank == 0: - q_ref = q_global.clone().detach().requires_grad_(True) - k_ref = k_global.clone().detach().requires_grad_(True) - v_ref = v_global.clone().detach().requires_grad_(True) - g_ref = g_global.clone().detach().requires_grad_(True) - beta_ref = beta_global.clone().detach().requires_grad_(True) - - o_ref, _ = chunk_kda( - q=F.normalize(q_ref, p=2, dim=-1), - k=F.normalize(k_ref, p=2, dim=-1), - v=v_ref, - g=g_ref, - beta=beta_ref, - cu_seqlens=cu_seqlens_global, - ) - o_ref.backward(do_global) - ref_out = o_ref.detach() - ref_dq = q_ref.grad.detach() - ref_dk = k_ref.grad.detach() - ref_dv = v_ref.grad.detach() - ref_dg = g_ref.grad.detach() - ref_db = beta_ref.grad.detach() - - dist.barrier() - - # CP run - context = build_cp_context(cu_seqlens_global, group=dist.group.WORLD) - - chunk_size = T // world_size - start_idx, end_idx = rank * chunk_size, (rank + 1) * chunk_size - - q_local = q_global[:, start_idx:end_idx, :].clone().detach().requires_grad_(True) - k_local = k_global[:, start_idx:end_idx, :].clone().detach().requires_grad_(True) - v_local = v_global[:, start_idx:end_idx, :].clone().detach().requires_grad_(True) - g_local = g_global[:, start_idx:end_idx, :].clone().detach().requires_grad_(True) - beta_local = beta_global[:, start_idx:end_idx].clone().detach().requires_grad_(True) - do_local = do_global[:, start_idx:end_idx, :].clone() - - o_local, _ = chunk_kda( - q=F.normalize(q_local, p=2, dim=-1), - k=F.normalize(k_local, p=2, dim=-1), - v=v_local, - g=g_local, - beta=beta_local, - cp_context=context, - ) - o_local.backward(do_local) - - # Gather output and gradients - o_gathered = [torch.zeros_like(o_local) for _ in range(world_size)] - dist.all_gather(o_gathered, o_local) - o_cp_global = torch.cat(o_gathered, dim=1) - - dq_gathered = [torch.zeros_like(q_local.grad) for _ in range(world_size)] - dist.all_gather(dq_gathered, q_local.grad) - dq_cp_global = torch.cat(dq_gathered, dim=1) - - dk_gathered = [torch.zeros_like(k_local.grad) for _ in range(world_size)] - dist.all_gather(dk_gathered, k_local.grad) - dk_cp_global = torch.cat(dk_gathered, dim=1) - - dv_gathered = [torch.zeros_like(v_local.grad) for _ in range(world_size)] - dist.all_gather(dv_gathered, v_local.grad) - dv_cp_global = torch.cat(dv_gathered, dim=1) - - dg_gathered = [torch.zeros_like(g_local.grad) for _ in range(world_size)] - dist.all_gather(dg_gathered, g_local.grad) - dg_cp_global = torch.cat(dg_gathered, dim=1) - - db_gathered = [torch.zeros_like(beta_local.grad) for _ in range(world_size)] - dist.all_gather(db_gathered, beta_local.grad) - db_cp_global = torch.cat(db_gathered, dim=1) - - if rank == 0: - print(f"\n[{test_name}] Verification:") - assert_close("Output", ref_out, o_cp_global, atol=5e-3) - assert_close("dq", ref_dq, dq_cp_global, atol=8e-3) - assert_close("dk", ref_dk, dk_cp_global, atol=8e-3) - assert_close("dv", ref_dv, dv_cp_global, atol=8e-3) - assert_close("dg", ref_dg, dg_cp_global, atol=2e-2) - assert_close("db", ref_db, db_cp_global, atol=2e-2) - print(f"✅ [{test_name}] Passed!") - - dist.barrier() - - if rank == 0: - print("\n" + "=" * 60) - print("All tests passed!") - print("=" * 60) - - finally: - cleanup_distributed() - else: - # Not running with torchrun, show usage - print("Run tests with pytest or torchrun:") - print(" pytest tests/context_parallel/test_cp_kda.py -v") - print(" torchrun --nproc_per_node=2 tests/context_parallel/test_cp_kda.py") - print(" torchrun --nproc_per_node=4 tests/context_parallel/test_cp_kda.py")