diff --git a/.gitignore b/.gitignore index b16e6212199..a102c77894a 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ -__pycache__ +__pycache__/ +*.pyc *.so build .coverage_* @@ -17,6 +18,7 @@ onelogger.err runs/ /test_cases/ **/dist/ +AGENTS.md # Sphinx documentation docs/_build diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 00000000000..b7f96a4c60b --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "third_party/mcore_gdn_opt"] + path = third_party/mcore_gdn_opt + url = ssh://git@gitlab-master.nvidia.com:12051/bhsueh/mcore_gdn_opt.git diff --git a/docs/gdn_cuda_optimization_reproduction.md b/docs/gdn_cuda_optimization_reproduction.md new file mode 100644 index 00000000000..90721e0da6f --- /dev/null +++ b/docs/gdn_cuda_optimization_reproduction.md @@ -0,0 +1,144 @@ +# GDN CUDA Optimization Reproduction + +This note covers the current GatedDeltaNet CUDA optimization test flow for +Megatron-LM on B200/H100. The optimized kernels are provided by +`third_party/mcore_gdn_opt`; FLA routes its gated-delta-rule calls through that +package. + +## Install + +Run these from the Megatron-LM repository root inside the GPU container. + +```bash +git submodule update --init --recursive third_party/mcore_gdn_opt +pip install -e . --user --no-build-isolation + +cd third_party/mcore_gdn_opt +./install_gdn_opt.sh +cd ../.. + +# FLA must contain the mcore_gdn_opt routing patch. +cd third_party/flash-linear-attention +pip install -e . --user --no-build-isolation +cd ../.. +``` + +Do not use `PYTHONPATH` or ad-hoc `sys.modules` injection for these tests. The +submodules should be installed in editable mode. + +## Runtime Flags + +| Case | Flags | +|---|---| +| Triton baseline | unset all `FLA_CUTE_*` flags | +| `wy_bwd` only | `FLA_CUTE_WY_BWD=1` | +| `dhu` only | `FLA_CUTE_BWD_DHU=1` | +| `dqkwg` only | `FLA_CUTE_BWD_DQKWG=1` | +| fused backward | `FLA_CUTE_WY_BWD=1 FLA_CUTE_BWD_DHU_DQKWG=1` | +| all three separate | `FLA_CUTE_WY_BWD=1 FLA_CUTE_BWD_DHU=1 FLA_CUTE_BWD_DQKWG=1` | +| all four | `FLA_CUTE_FWD_H=1 CHUNK_DELTA_FWD_USE_BWD_PORT=1 FLA_CUTE_WY_BWD=1 FLA_CUTE_BWD_DHU=1 FLA_CUTE_BWD_DQKWG=1` | + +## GDN-Only Direct Test + +This bypasses the full GPT layer and measures a direct `GatedDeltaNet` +forward/backward. It checks output, input grad, and parameter grads against the +Triton baseline. + +```bash +python -m tests.unit_tests.ssm.bench_gdn_cuda_opt \ + --dtype bf16 \ + --loss sum \ + --scenarios baseline,fused,separate,all_four \ + --warmup 5 --repeats 20 --rounds 3 +``` + +Use `--loss square_mean` to reproduce the earlier loss used during debugging, +and add `--fail-on-accuracy` when the command should return non-zero on any +accuracy mismatch. + +The submodule is currently pinned at `mcore_gdn_opt@12605c5` (its `main`/HEAD). +This is the first training-safe commit of the kernel lineage: the backward +kernels (`wy_bwd`, `dhu`, `dqkwg`) produce correct, NaN-free gradients at +`DV_DHU=0`, so the optimized path is safe for both forward and backward. It +supersedes the earlier pin `9121702`, whose backward kernels produced wrong +gradients (forward/inference-only). + +> **Re-measure pending.** The `loss=sum` spot-check table below was captured on +> the now-superseded pin `9121702`, where the optimized scenarios *failed* the +> strict gradient comparison. After the bump to `12605c5` (which fixes those +> gradients) the table must be regenerated on B200; the numbers below are kept +> only as the historical record for the old pin and no longer reflect the +> pinned kernels. + +Historical B200 spot check (superseded pin) for +`B=2,T=8192,H=64,D=128,bf16,loss=sum,warmup=3,repeats=10,rounds=3` +on Megatron-LM `c42dc298a`, `mcore_gdn_opt@9121702`, and +`gated_delta_rule_bwd@949c959`: + +| Scenario | Accuracy vs Triton | Mean us | Speedup | +|---|---:|---:|---:| +| Triton baseline | PASS | 15220.135 | 1.000x | +| CUDA `wy+dhu+dqkwg fused` | FAIL | 13470.359 | 1.130x | +| CUDA all three separate | FAIL | 13420.346 | 1.134x | +| CUDA all four | FAIL | 12875.528 | 1.182x | + +On the old pin this direct `loss=sum` GDN-only check failed the strict gradient +comparison, so the production workflow was validated with `loss=square_mean` +instead. That `loss=square_mean` full workflow validation (logged at +`mcore_gdn_opt@cb51345`) passed all requested scenarios and measured +`CUDA all four` at `12895.830 us` (`1.182x`) and `CUDA fwd_h+wy+dv_dhu+dqkwg` +at `12732.651 us` (`1.197x`). Logs: +`third_party/gdn_doc_loss_sum_20260528_205336.log` and +`third_party/gdn_full_validation_cb51345_20260528_204219.log`. With the pin now +at `12605c5`, re-run the `loss=sum` check above — it is expected to PASS. + +## E2E Pytest + +This runs the focused GDN CUDA optimization pytest path. It checks correctness +by default and can print the benchmark table when `MCORE_GDN_UNIT_TEST_PERF=1`. + +```bash +MCORE_GDN_UNIT_TEST_SCENARIOS=baseline,all_four_dv_dhu \ +pytest -s tests/unit_tests/ssm/test_gated_delta_net_cuda_opt.py::test_gated_delta_net_cuda_opt_correctness_and_optional_perf -k bf16 +``` + +To generate the E2E benchmark table with NVTX labels: + +```bash +MCORE_GDN_UNIT_TEST_SCENARIOS=baseline,wy,dhu,dqkwg,fused,separate,all_four,all_four_dv_dhu \ +MCORE_GDN_UNIT_TEST_PERF=1 \ +MCORE_GDN_UNIT_TEST_WARMUP=5 \ +MCORE_GDN_UNIT_TEST_REPEATS=20 \ +MCORE_GDN_UNIT_TEST_ROUNDS=3 \ +pytest -s tests/unit_tests/ssm/test_gated_delta_net_cuda_opt.py::test_gated_delta_net_cuda_opt_correctness_and_optional_perf -k bf16 +``` + +Latest B200 full workflow validation for `loss=square_mean` passed correctness +for wrapper forced FLA, wrapper auto, wrapper forced CUDA, `CUDA all four`, and +`CUDA fwd_h+wy+dv_dhu+dqkwg`. Observed speedups were `1.198x` for wrapper auto, +`1.182x` for `CUDA all four`, and `1.197x` for +`CUDA fwd_h+wy+dv_dhu+dqkwg` versus the Triton baseline. + +## Nsight Systems + +Use the E2E pytest command above under `nsys profile`. The benchmark emits NVTX +labels in this format: + +```text +gdn_only/_/round_/iter_ +``` + +Example: + +```bash +MCORE_GDN_UNIT_TEST_SCENARIOS=baseline,all_four_dv_dhu \ +MCORE_GDN_UNIT_TEST_PERF=1 \ +MCORE_GDN_UNIT_TEST_WARMUP=5 \ +MCORE_GDN_UNIT_TEST_REPEATS=20 \ +MCORE_GDN_UNIT_TEST_ROUNDS=3 \ +nsys profile -f true -o gdn_e2e_b200 \ + pytest -s tests/unit_tests/ssm/test_gated_delta_net_cuda_opt.py::test_gated_delta_net_cuda_opt_correctness_and_optional_perf -k bf16 +``` + +Profiler outputs (`*.nsys-rep`, `*.sqlite`, `*.qdrep`) and local run directories +are ignored by `.gitignore`. diff --git a/megatron/core/fusions/fused_mega_pre_gated_delta_rule.py b/megatron/core/fusions/fused_mega_pre_gated_delta_rule.py new file mode 100644 index 00000000000..66ce10a99ea --- /dev/null +++ b/megatron/core/fusions/fused_mega_pre_gated_delta_rule.py @@ -0,0 +1,1109 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Mega fused pre-gated-delta-rule kernels. + +This is the "mega" sibling of :mod:`fused_pre_gated_delta_rule`. The streamed +path splits the pre-gated-delta-rule front-end into four separate Triton launch +scopes (QK / V / Z / g-beta) plus an external conv backward, optimized for +per-kernel quality and overlap under CUDA-graph capture. The mega path instead +folds **all forward tasks into a single Triton launch** using a flat logical +task space, trading a little per-kernel efficiency for far fewer host-side +launches. It is the right choice for non-CUDA-graph recipes where launch +overhead is visible in the trace. + +Public contract is identical across unfused / streamed / mega: +``(query, key, value, gate, beta, g)``. + +Design notes: + +* The forward kernel maps ``program_id(0)`` onto a flat row space partitioned + into QK, V, Z, and g/beta ranges; ``program_id(1)`` tiles the sequence axis. + Each program inspects its row id and runs exactly one task body. This keeps + every sub-computation in one launch while still letting Triton schedule + memory-bound (Z copy) and compute-bound (QK/V conv) tiles concurrently on the + SMs. +* Numerics mirror the streamed/unfused reference **bit-for-bit within the unit + test tolerance**: the conv accumulator is rounded through the activation + dtype before SiLU; the SiLU output is rounded again before the L2-norm + reduction; ``g`` uses an fp32 ``log(1+exp(...))`` softplus; ``beta`` uses an + fp32 sigmoid. The QK ``silu(conv(x))`` intermediate is persisted channel-last + exactly as the streamed path saves it, so the backward can be shared. +* The kernel assumes ``key_head_dim == value_head_dim`` so a single + ``HEAD_DIM`` constexpr drives the QK/V/Z channel tiles (true for the GDN + production shapes and the unit tests). This is asserted at the Python entry. + +The backward mirrors the forward: a single fused Triton kernel folds the four +streamed branch backward scopes (QK l2norm/repeat, V layout, Z layout, g/beta +chain rule) into one flat-task launch, then the depthwise conv input/weight +gradients are delegated to the same external ``causal_conv1d_bwd_function`` the +streamed path uses (its hand-tuned C++ remains the conv-backward anchor). That +is two launches total (one fused branch kernel + one external conv backward), +down from the streamed path's five, while staying numerically bit-identical to +the streamed branch kernels. +""" + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl +from torch import Tensor + +# Reuse the streamed module's validated constants and the external conv backward +# binding. Importing is not a modification of that module. +from megatron.core.fusions.fused_pre_gated_delta_rule import ( + _L2NORM_EPS, + _causal_conv1d_bwd_function, + _is_power_of_two, + _resolve_packed_seq_idx, +) + + +# --------------------------------------------------------------------------- +# Forward kernel +# --------------------------------------------------------------------------- + + +def _mega_autotune_configs(): + return [ + triton.Config({"BLOCK_S": 32}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_S": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=8, num_stages=3), + triton.Config({"BLOCK_S": 256}, num_warps=8, num_stages=2), + ] + + +@triton.jit +def _mega_seq_bounds(cu_seqlens_ptr, token_offsets, total_tokens, num_packed_seqs): + """Lane-wise packed-sequence [start, end) bounds for flattened THD tokens. + + Local copy of the streamed helper so this kernel never depends on + cross-module ``@triton.jit`` symbol resolution. + """ + + safe_tokens = tl.minimum(token_offsets, total_tokens - 1) + seq_start = token_offsets * 0 + seq_end = token_offsets * 0 + total_tokens + + seq_id = 0 + while seq_id < num_packed_seqs: + start = tl.load(cu_seqlens_ptr + seq_id) + end = tl.load(cu_seqlens_ptr + seq_id + 1) + in_seq = (safe_tokens >= start) & (safe_tokens < end) + seq_start = tl.where(in_seq, start, seq_start) + seq_end = tl.where(in_seq, end, seq_end) + seq_id += 1 + + return seq_start, seq_end + + +@triton.autotune( + configs=_mega_autotune_configs(), + key=["seq_len", "HEAD_DIM", "K_W", "num_key_heads", "num_value_heads", "REPEAT", "HAS_THD"], +) +@triton.jit +def _mega_forward_kernel( + qkvzba_ptr, + weight_ptr, + A_log_ptr, + dt_bias_ptr, + qk_out_ptr, + value_ptr, + gate_ptr, + g_ptr, + beta_ptr, + silu_save_ptr, + cu_seqlens_ptr, + seq_len, + num_packed_seqs, + num_key_heads, + num_value_heads, + qk_channels, + v_channels, + R_qk, + R_v, + R_z, + qkvzba_s_stride, + qkvzba_b_stride, + qkvzba_c_stride, + weight_c_stride, + weight_w_stride, + qk_g_stride, + qk_b_stride, + qk_s_stride, + qk_h_stride, + v_b_stride, + v_s_stride, + v_h_stride, + z_b_stride, + z_s_stride, + z_h_stride, + g_b_stride, + g_s_stride, + g_h_stride, + beta_b_stride, + beta_s_stride, + beta_h_stride, + silu_b_stride, + silu_c_stride, + silu_s_stride, + eps, + HEAD_DIM: tl.constexpr, + K_W: tl.constexpr, + REPEAT: tl.constexpr, + HAS_THD: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """All-in-one forward for the pre-gated-delta-rule front-end. + + Flat task space on ``program_id(0)``: + rows [0, R_qk) -> QK conv+silu+l2norm+repeat + rows [R_qk, R_qk+R_v) -> V conv+silu + rows [R_qk+R_v, +R_z) -> Z copy + rows [.., end) -> g/beta + ``program_id(1)`` tiles the (flattened, for THD) sequence axis. + """ + + pid_row = tl.program_id(0) + pid_s = tl.program_id(1) + + # Common rounding dtype (activation dtype, e.g. bf16). All q/k/v/gate/beta + # outputs share this; g is fp32. + out_ty = qk_out_ptr.dtype.element_ty + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + chan_off = tl.arange(0, HEAD_DIM) + + R_qkv = R_qk + R_v + R_qkvz = R_qkv + R_z + + if pid_row < R_qk: + # ---- QK: depthwise causal conv + silu + l2norm + head repeat ---- + local = pid_row + heads_per_batch = 2 * num_key_heads + batch_id = local // heads_per_batch + lb = local - batch_id * heads_per_batch + group_id = lb // num_key_heads # 0 -> Q, 1 -> K + head_id = lb - group_id * num_key_heads + chan = group_id * qk_channels + head_id * HEAD_DIM + chan_off + + if HAS_THD: + seq_start, seq_end = _mega_seq_bounds( + cu_seqlens_ptr, s_offs, seq_len, num_packed_seqs + ) + + acc = tl.zeros([BLOCK_S, HEAD_DIM], dtype=tl.float32) + for i in tl.static_range(K_W): + x_s = s_offs - (K_W - 1) + i + if HAS_THD: + x_mask = s_mask & (x_s >= seq_start) & (x_s < seq_end) + safe_x_s = tl.minimum(tl.maximum(x_s, 0), seq_len - 1) + else: + x_mask = (x_s >= 0) & (x_s < seq_len) + safe_x_s = x_s + x_ptr = ( + qkvzba_ptr + + safe_x_s[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + chan[None, :] * qkvzba_c_stride + ) + x_val = tl.load(x_ptr, mask=x_mask[:, None], other=0.0).to(tl.float32) + w_tap = tl.load( + weight_ptr + chan * weight_c_stride + i * weight_w_stride + ).to(tl.float32) + acc += w_tap[None, :] * x_val + + acc = acc.to(out_ty).to(tl.float32) # F.conv1d rounding + silu_out = acc * tl.sigmoid(acc) + silu_out = silu_out.to(out_ty).to(tl.float32) # round before l2norm + + # Persist silu(conv(x)) for the QK channels, channel-last (b, 2*qk, s). + silu_chan = group_id * qk_channels + head_id * HEAD_DIM + chan_off + silu_ptrs = ( + silu_save_ptr + + batch_id * silu_b_stride + + silu_chan[None, :] * silu_c_stride + + s_offs[:, None] * silu_s_stride + ) + tl.store( + silu_ptrs, + silu_out.to(silu_save_ptr.dtype.element_ty), + mask=s_mask[:, None], + ) + + norm_sq = tl.sum(silu_out * silu_out, axis=1) + rstd = 1.0 / tl.sqrt(norm_sq + eps) + out_typed = (silu_out * rstd[:, None]).to(out_ty) + + for r in tl.static_range(REPEAT): + v_head = head_id * REPEAT + r + write_ptr = ( + qk_out_ptr + + group_id * qk_g_stride + + batch_id * qk_b_stride + + s_offs[:, None] * qk_s_stride + + v_head * qk_h_stride + + chan_off[None, :] + ) + tl.store(write_ptr, out_typed, mask=s_mask[:, None]) + + elif pid_row < R_qkv: + # ---- V: depthwise causal conv + silu (no l2norm, no repeat) ---- + local = pid_row - R_qk + batch_id = local // num_value_heads + head_id = local - batch_id * num_value_heads + chan = 2 * qk_channels + head_id * HEAD_DIM + chan_off + + if HAS_THD: + seq_start, seq_end = _mega_seq_bounds( + cu_seqlens_ptr, s_offs, seq_len, num_packed_seqs + ) + + acc = tl.zeros([BLOCK_S, HEAD_DIM], dtype=tl.float32) + for i in tl.static_range(K_W): + x_s = s_offs - (K_W - 1) + i + if HAS_THD: + x_mask = s_mask & (x_s >= seq_start) & (x_s < seq_end) + safe_x_s = tl.minimum(tl.maximum(x_s, 0), seq_len - 1) + else: + x_mask = (x_s >= 0) & (x_s < seq_len) + safe_x_s = x_s + x_ptr = ( + qkvzba_ptr + + safe_x_s[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + chan[None, :] * qkvzba_c_stride + ) + x_val = tl.load(x_ptr, mask=x_mask[:, None], other=0.0).to(tl.float32) + w_tap = tl.load( + weight_ptr + chan * weight_c_stride + i * weight_w_stride + ).to(tl.float32) + acc += w_tap[None, :] * x_val + + acc = acc.to(out_ty).to(tl.float32) + silu_out = acc * tl.sigmoid(acc) + out_typed = silu_out.to(out_ty) + write_ptr = ( + value_ptr + + batch_id * v_b_stride + + s_offs[:, None] * v_s_stride + + head_id * v_h_stride + + chan_off[None, :] + ) + tl.store(write_ptr, out_typed, mask=s_mask[:, None]) + + elif pid_row < R_qkvz: + # ---- Z: copy qkvzba z slice into the final gate layout ---- + local = pid_row - R_qkv + batch_id = local // num_value_heads + head_id = local - batch_id * num_value_heads + z_chan = 2 * qk_channels + v_channels + head_id * HEAD_DIM + chan_off + src_ptr = ( + qkvzba_ptr + + s_offs[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + z_chan[None, :] * qkvzba_c_stride + ) + z_val = tl.load(src_ptr, mask=s_mask[:, None]) + write_ptr = ( + gate_ptr + + batch_id * z_b_stride + + s_offs[:, None] * z_s_stride + + head_id * z_h_stride + + chan_off[None, :] + ) + tl.store(write_ptr, z_val, mask=s_mask[:, None]) + + else: + # ---- g/beta: -exp(A_log)*softplus(alpha+dt_bias) and sigmoid(beta) ---- + local = pid_row - R_qkvz + batch_id = local // num_value_heads + head_id = local - batch_id * num_value_heads + beta_chan = 2 * qk_channels + 2 * v_channels + head_id + alpha_chan = beta_chan + num_value_heads + + alpha_ptr = ( + qkvzba_ptr + + s_offs * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + alpha_chan * qkvzba_c_stride + ) + beta_raw_ptr = ( + qkvzba_ptr + + s_offs * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + beta_chan * qkvzba_c_stride + ) + alpha = tl.load(alpha_ptr, mask=s_mask, other=0.0).to(tl.float32) + beta_raw = tl.load(beta_raw_ptr, mask=s_mask, other=0.0).to(tl.float32) + A_log = tl.load(A_log_ptr + head_id).to(tl.float32) + dt_bias = tl.load(dt_bias_ptr + head_id).to(tl.float32) + + pre = alpha + dt_bias + softplus_val = tl.log(1.0 + tl.exp(pre)) + g = -tl.exp(A_log) * softplus_val + beta_sig = tl.sigmoid(beta_raw) + + g_store_ptr = ( + g_ptr + batch_id * g_b_stride + s_offs * g_s_stride + head_id * g_h_stride + ) + beta_store_ptr = ( + beta_ptr + + batch_id * beta_b_stride + + s_offs * beta_s_stride + + head_id * beta_h_stride + ) + tl.store(g_store_ptr, g.to(g_ptr.dtype.element_ty), mask=s_mask) + tl.store(beta_store_ptr, beta_sig.to(beta_ptr.dtype.element_ty), mask=s_mask) + + +# --------------------------------------------------------------------------- +# Forward orchestration +# --------------------------------------------------------------------------- + + +def _mega_pre_gated_delta_rule_forward( + qkvzba: Tensor, + conv1d_weight: Tensor, + A_log: Tensor, + dt_bias: Tensor, + *, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, + cu_seqlens: Optional[Tensor] = None, +) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + """Single-launch mega forward. + + Returns ``(query, key, value, gate, beta, g, silu_qk_save)``; the last + element is the bf16-rounded QK ``silu(conv(x))`` laid out channel-last, + matching the streamed forward so the shared backward can consume it. + """ + + seq_len, batch, total_channels = qkvzba.shape + is_packed_thd = cu_seqlens is not None + num_packed_seqs = (cu_seqlens.shape[0] - 1) if is_packed_thd else 0 + + assert key_head_dim == value_head_dim, ( + "fused_mega_pre_gated_delta_rule currently requires " + f"key_head_dim == value_head_dim; got {key_head_dim=} {value_head_dim=}." + ) + assert _is_power_of_two(key_head_dim), ( + f"Mega kernel expects key_head_dim to be a power of two; got {key_head_dim=}." + ) + head_dim = key_head_dim + + qk_channels = num_key_heads * key_head_dim + v_channels = num_value_heads * value_head_dim + repeat_factor = num_value_heads // num_key_heads + k_w = conv1d_weight.shape[-1] + + expected_channels = 2 * qk_channels + 2 * v_channels + 2 * num_value_heads + assert total_channels == expected_channels, ( + f"qkvzba last-dim mismatch: got {total_channels}, expected {expected_channels}." + ) + + out_dtype = qkvzba.dtype + device = qkvzba.device + + # Output buffers (identical layouts to the streamed path). + qk_out = torch.empty( + 2, batch, seq_len, num_value_heads, key_head_dim, dtype=out_dtype, device=device + ) + query = qk_out[0] + key = qk_out[1] + value = torch.empty( + batch, seq_len, num_value_heads, value_head_dim, dtype=out_dtype, device=device + ) + gate = torch.empty( + batch, seq_len, num_value_heads, value_head_dim, dtype=out_dtype, device=device + ) + g = torch.empty(batch, seq_len, num_value_heads, dtype=torch.float32, device=device) + beta = torch.empty(batch, seq_len, num_value_heads, dtype=out_dtype, device=device) + + # QK silu(conv(x)) persisted channel-last: (b, 2*qk_channels, s), stride(1)==1. + silu_qk_save = torch.empty( + (batch, seq_len, 2 * qk_channels), dtype=out_dtype, device=device + ).permute(0, 2, 1) + + weight_2d = conv1d_weight.view(conv1d_weight.shape[0], k_w) + + # Flat task-space row partition. + R_qk = batch * 2 * num_key_heads + R_v = batch * num_value_heads + R_z = batch * num_value_heads + R_gb = batch * num_value_heads + num_rows = R_qk + R_v + R_z + R_gb + + cu_seqlens_arg = cu_seqlens if is_packed_thd else qkvzba # dummy when dense + + grid = lambda meta: (num_rows, triton.cdiv(seq_len, meta["BLOCK_S"])) + _mega_forward_kernel[grid]( + qkvzba, + weight_2d, + A_log, + dt_bias, + qk_out, + value, + gate, + g, + beta, + silu_qk_save, + cu_seqlens_arg, + seq_len, + num_packed_seqs, + num_key_heads, + num_value_heads, + qk_channels, + v_channels, + R_qk, + R_v, + R_z, + qkvzba.stride(0), + qkvzba.stride(1), + qkvzba.stride(2), + weight_2d.stride(0), + weight_2d.stride(1), + qk_out.stride(0), + qk_out.stride(1), + qk_out.stride(2), + qk_out.stride(3), + value.stride(0), + value.stride(1), + value.stride(2), + gate.stride(0), + gate.stride(1), + gate.stride(2), + g.stride(0), + g.stride(1), + g.stride(2), + beta.stride(0), + beta.stride(1), + beta.stride(2), + silu_qk_save.stride(0), + silu_qk_save.stride(1), + silu_qk_save.stride(2), + _L2NORM_EPS, + HEAD_DIM=head_dim, + K_W=k_w, + REPEAT=repeat_factor, + HAS_THD=is_packed_thd, + ) + + return query, key, value, gate, beta, g, silu_qk_save + + +# --------------------------------------------------------------------------- +# Backward kernel +# --------------------------------------------------------------------------- + + +def _mega_backward_autotune_configs(): + return [ + triton.Config({"BLOCK_S": 32}, num_warps=2, num_stages=2), + triton.Config({"BLOCK_S": 32}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_S": 64}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 256}, num_warps=8, num_stages=2), + ] + + +@triton.autotune( + configs=_mega_backward_autotune_configs(), + key=["seq_len", "HEAD_DIM", "REPEAT", "num_key_heads", "num_value_heads"], + # The g/beta task atomic-adds per-head partials into these accumulators. + # reset_to_zero clears them before each autotune trial so trials don't stack. + reset_to_zero=["d_A_log_ptr", "d_dt_bias_ptr"], +) +@triton.jit +def _mega_backward_kernel( + # inputs + dq_ptr, + dk_ptr, + dv_ptr, + dgate_ptr, + dg_ptr, + dbeta_ptr, + silu_save_ptr, + qkvzba_ptr, + A_log_ptr, + dt_bias_ptr, + # outputs + d_silu_conv_ptr, + d_qkvzba_ptr, + d_A_log_ptr, + d_dt_bias_ptr, + # sizes / layout + seq_len, + num_key_heads, + num_value_heads, + qk_channels, + v_channels, + R_qk, + R_v, + R_z, + eps, + # dq / dk strides (b, s, h, d) + dq_b_stride, + dq_s_stride, + dq_h_stride, + dk_b_stride, + dk_s_stride, + dk_h_stride, + # dv strides (b, s, h, d) + dv_b_stride, + dv_s_stride, + dv_h_stride, + # dgate strides (b, s, h, d) + dgate_b_stride, + dgate_s_stride, + dgate_h_stride, + # dg / dbeta strides (b, s, h) + dg_b_stride, + dg_s_stride, + dg_h_stride, + dbeta_b_stride, + dbeta_s_stride, + dbeta_h_stride, + # silu_save strides (b, 2*qk, s) + silu_b_stride, + silu_c_stride, + silu_s_stride, + # qkvzba / d_qkvzba strides (s, b, C) + qkvzba_s_stride, + qkvzba_b_stride, + qkvzba_c_stride, + # d_silu_conv strides (b, conv_dim, s) + dsc_b_stride, + dsc_c_stride, + dsc_s_stride, + HEAD_DIM: tl.constexpr, + REPEAT: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """All-in-one backward for the QK / V / Z / g-beta branches. + + Mirrors the four streamed branch kernels in one flat task space. Conv input + and weight gradients are NOT produced here; the caller feeds ``d_silu_conv`` + into the external ``causal_conv1d_bwd_function`` exactly as the streamed + path does. Flat task space on ``program_id(0)``: + rows [0, R_qk) -> QK l2norm + repeat backward -> d_silu_conv[Q/K] + rows [R_qk, +R_v) -> V layout copy -> d_silu_conv[V] + rows [.., +R_z) -> Z layout copy -> d_qkvzba[z] + rows [.., end) -> g/beta chain rule -> d_qkvzba[alpha,beta], + atomic d_A_log/d_dt_bias + """ + + pid_row = tl.program_id(0) + pid_s = tl.program_id(1) + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + chan_off = tl.arange(0, HEAD_DIM) + + R_qkv = R_qk + R_v + R_qkvz = R_qkv + R_z + + if pid_row < R_qk: + # ---- QK: repeat-reduce + l2norm backward -> d_silu_conv[Q/K] ---- + local = pid_row + heads_per_batch = 2 * num_key_heads + batch_id = local // heads_per_batch + lb = local - batch_id * heads_per_batch + group_id = lb // num_key_heads + head_id = lb - group_id * num_key_heads + is_query = group_id == 0 + is_key = group_id == 1 + chan = group_id * qk_channels + head_id * HEAD_DIM + chan_off + + d_normed = tl.zeros([BLOCK_S, HEAD_DIM], dtype=tl.float32) + for r in tl.static_range(REPEAT): + v_head = head_id * REPEAT + r + dq_ptrs = ( + dq_ptr + + batch_id * dq_b_stride + + s_offs[:, None] * dq_s_stride + + v_head * dq_h_stride + + chan_off[None, :] + ) + dk_ptrs = ( + dk_ptr + + batch_id * dk_b_stride + + s_offs[:, None] * dk_s_stride + + v_head * dk_h_stride + + chan_off[None, :] + ) + d_normed += tl.load(dq_ptrs, mask=s_mask[:, None] & is_query, other=0.0).to(tl.float32) + d_normed += tl.load(dk_ptrs, mask=s_mask[:, None] & is_key, other=0.0).to(tl.float32) + + silu_ptrs = ( + silu_save_ptr + + batch_id * silu_b_stride + + chan[None, :] * silu_c_stride + + s_offs[:, None] * silu_s_stride + ) + silu_bf16 = tl.load(silu_ptrs, mask=s_mask[:, None], other=0.0).to(tl.float32) + + norm_sq = tl.sum(silu_bf16 * silu_bf16, axis=1) + rstd = 1.0 / tl.sqrt(norm_sq + eps) + s_row = tl.sum(d_normed * silu_bf16, axis=1) + rstd3 = rstd * rstd * rstd + d_silu = rstd[:, None] * d_normed - rstd3[:, None] * silu_bf16 * s_row[:, None] + + dsc_ptrs = ( + d_silu_conv_ptr + + batch_id * dsc_b_stride + + chan[None, :] * dsc_c_stride + + s_offs[:, None] * dsc_s_stride + ) + tl.store(dsc_ptrs, d_silu.to(d_silu_conv_ptr.dtype.element_ty), mask=s_mask[:, None]) + + elif pid_row < R_qkv: + # ---- V: relayout dv -> d_silu_conv[V] ---- + local = pid_row - R_qk + batch_id = local // num_value_heads + head_id = local - batch_id * num_value_heads + dv_ptrs = ( + dv_ptr + + batch_id * dv_b_stride + + s_offs[:, None] * dv_s_stride + + head_id * dv_h_stride + + chan_off[None, :] + ) + dv_val = tl.load(dv_ptrs, mask=s_mask[:, None], other=0.0) + dsc_chan = 2 * qk_channels + head_id * HEAD_DIM + chan_off + dsc_ptrs = ( + d_silu_conv_ptr + + batch_id * dsc_b_stride + + dsc_chan[None, :] * dsc_c_stride + + s_offs[:, None] * dsc_s_stride + ) + tl.store(dsc_ptrs, dv_val, mask=s_mask[:, None]) + + elif pid_row < R_qkvz: + # ---- Z: relayout dgate -> d_qkvzba[z] ---- + local = pid_row - R_qkv + batch_id = local // num_value_heads + head_id = local - batch_id * num_value_heads + dgate_ptrs = ( + dgate_ptr + + batch_id * dgate_b_stride + + s_offs[:, None] * dgate_s_stride + + head_id * dgate_h_stride + + chan_off[None, :] + ) + dgate_val = tl.load(dgate_ptrs, mask=s_mask[:, None], other=0.0) + dz_chan = 2 * qk_channels + v_channels + head_id * HEAD_DIM + chan_off + dz_ptrs = ( + d_qkvzba_ptr + + s_offs[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + dz_chan[None, :] * qkvzba_c_stride + ) + tl.store(dz_ptrs, dgate_val, mask=s_mask[:, None]) + + else: + # ---- g/beta: chain rule -> d_qkvzba[alpha,beta] + atomic d_A_log/d_dt_bias ---- + local = pid_row - R_qkvz + batch_id = local // num_value_heads + head_id = local - batch_id * num_value_heads + beta_chan = 2 * qk_channels + 2 * v_channels + head_id + alpha_chan = beta_chan + num_value_heads + + alpha_ptrs = ( + qkvzba_ptr + + s_offs * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + alpha_chan * qkvzba_c_stride + ) + beta_ptrs = ( + qkvzba_ptr + + s_offs * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + beta_chan * qkvzba_c_stride + ) + alpha = tl.load(alpha_ptrs, mask=s_mask, other=0.0).to(tl.float32) + beta_raw = tl.load(beta_ptrs, mask=s_mask, other=0.0).to(tl.float32) + A_log = tl.load(A_log_ptr + head_id).to(tl.float32) + dt_bias = tl.load(dt_bias_ptr + head_id).to(tl.float32) + + pre = alpha + dt_bias + sigmoid_pre = tl.sigmoid(pre) + softplus_pre = tl.log(1.0 + tl.exp(pre)) + exp_A = tl.exp(A_log) + g = -exp_A * softplus_pre + beta_sig = tl.sigmoid(beta_raw) + + dg_ptrs = ( + dg_ptr + batch_id * dg_b_stride + s_offs * dg_s_stride + head_id * dg_h_stride + ) + dbeta_ptrs = ( + dbeta_ptr + + batch_id * dbeta_b_stride + + s_offs * dbeta_s_stride + + head_id * dbeta_h_stride + ) + d_g = tl.load(dg_ptrs, mask=s_mask, other=0.0).to(tl.float32) + d_beta_out = tl.load(dbeta_ptrs, mask=s_mask, other=0.0).to(tl.float32) + + d_alpha = d_g * (-exp_A * sigmoid_pre) + d_beta_raw = d_beta_out * beta_sig * (1.0 - beta_sig) + + d_alpha_ptrs = ( + d_qkvzba_ptr + + s_offs * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + alpha_chan * qkvzba_c_stride + ) + d_beta_ptrs = ( + d_qkvzba_ptr + + s_offs * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + beta_chan * qkvzba_c_stride + ) + tl.store(d_alpha_ptrs, d_alpha.to(d_qkvzba_ptr.dtype.element_ty), mask=s_mask) + tl.store(d_beta_ptrs, d_beta_raw.to(d_qkvzba_ptr.dtype.element_ty), mask=s_mask) + + d_g_masked = tl.where(s_mask, d_g, 0.0) + d_alpha_masked = tl.where(s_mask, d_alpha, 0.0) + d_A_log_partial = tl.sum(d_g_masked * g) + d_dt_bias_partial = tl.sum(d_alpha_masked) + tl.atomic_add(d_A_log_ptr + head_id, d_A_log_partial) + tl.atomic_add(d_dt_bias_ptr + head_id, d_dt_bias_partial) + + +# --------------------------------------------------------------------------- +# Backward orchestration +# --------------------------------------------------------------------------- + + +def _mega_pre_gated_delta_rule_backward( + qkvzba: Tensor, + conv1d_weight: Tensor, + silu_qk_save: Tensor, + dq: Tensor, + dk: Tensor, + dv: Tensor, + dgate: Tensor, + dbeta: Tensor, + dg: Tensor, + A_log: Tensor, + dt_bias: Tensor, + *, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, + seq_idx: Optional[Tensor] = None, +) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + """Two-launch mega backward: one fused branch kernel + external conv bwd. + + Collapses the four streamed branch kernels (QK l2norm/repeat, V layout, Z + layout, g/beta chain rule) into a single Triton launch, then delegates the + depthwise conv input/weight gradients to ``causal_conv1d_bwd_function`` as + the streamed path does. Returns ``(d_qkvzba, d_weight, d_A_log, d_dt_bias)``. + """ + + seq_len, batch, _ = qkvzba.shape + qk_channels = num_key_heads * key_head_dim + v_channels = num_value_heads * value_head_dim + conv_dim = 2 * qk_channels + v_channels + k_w = conv1d_weight.shape[-1] + device = qkvzba.device + head_dim = key_head_dim + + weight_2d = conv1d_weight.view(conv1d_weight.shape[0], k_w) + # Channel-last conv input view (stride(1)==1) — no copy. + qkvzba_conv = qkvzba[:, :, :conv_dim].permute(1, 2, 0) + + # d_silu_conv channel-last (b, conv_dim, s), stride(1)==1. + d_silu_conv = torch.empty( + (batch, seq_len, conv_dim), dtype=qkvzba.dtype, device=device + ).permute(0, 2, 1) + d_qkvzba = torch.empty_like(qkvzba) + d_A_log_fp32 = torch.zeros(num_value_heads, dtype=torch.float32, device=device) + d_dt_bias_fp32 = torch.zeros(num_value_heads, dtype=torch.float32, device=device) + + R_qk = batch * 2 * num_key_heads + R_v = batch * num_value_heads + R_z = batch * num_value_heads + R_gb = batch * num_value_heads + num_rows = R_qk + R_v + R_z + R_gb + + grid = lambda meta: (num_rows, triton.cdiv(seq_len, meta["BLOCK_S"])) + _mega_backward_kernel[grid]( + dq, + dk, + dv, + dgate, + dg, + dbeta, + silu_qk_save, + qkvzba, + A_log, + dt_bias, + d_silu_conv, + d_qkvzba, + d_A_log_fp32, + d_dt_bias_fp32, + seq_len, + num_key_heads, + num_value_heads, + qk_channels, + v_channels, + R_qk, + R_v, + R_z, + _L2NORM_EPS, + dq.stride(0), + dq.stride(1), + dq.stride(2), + dk.stride(0), + dk.stride(1), + dk.stride(2), + dv.stride(0), + dv.stride(1), + dv.stride(2), + dgate.stride(0), + dgate.stride(1), + dgate.stride(2), + dg.stride(0), + dg.stride(1), + dg.stride(2), + dbeta.stride(0), + dbeta.stride(1), + dbeta.stride(2), + silu_qk_save.stride(0), + silu_qk_save.stride(1), + silu_qk_save.stride(2), + qkvzba.stride(0), + qkvzba.stride(1), + qkvzba.stride(2), + d_silu_conv.stride(0), + d_silu_conv.stride(1), + d_silu_conv.stride(2), + HEAD_DIM=head_dim, + REPEAT=num_value_heads // num_key_heads, + ) + + # External conv backward: writes d_x into d_qkvzba's conv slice (strided + # view, no copy) and returns d_weight. Same call shape as the streamed path. + seq_stride = qkvzba.stride(0) + batch_stride = qkvzba.stride(1) + d_x_conv_view = d_qkvzba.as_strided( + (batch, conv_dim, seq_len), + (batch_stride, 1, seq_stride), + ) + if _causal_conv1d_bwd_function is None: + raise RuntimeError( + "Fused pre-gated-delta-rule backward requires the 'causal_conv1d' package. " + "Install it, or use pre_gated_delta_rule_impl='unfused'." + ) + _, d_weight_fp32, _, _ = _causal_conv1d_bwd_function( + qkvzba_conv, + weight_2d, + None, # no bias + d_silu_conv, + seq_idx, + None, # initial_states + None, # dfinal_states + d_x_conv_view, # dx pre-allocated into d_qkvzba's conv slice + False, # return_dinitial_states + True, # activation (silu folded into conv bwd) + ) + + d_weight = d_weight_fp32.view(*conv1d_weight.shape).to(conv1d_weight.dtype) + d_A_log = d_A_log_fp32.to(A_log.dtype) + d_dt_bias = d_dt_bias_fp32.to(dt_bias.dtype) + return d_qkvzba, d_weight, d_A_log, d_dt_bias + + +# --------------------------------------------------------------------------- +# Autograd wiring +# --------------------------------------------------------------------------- + + +class _FusedMegaPreGatedDeltaRuleFunction(torch.autograd.Function): + """Autograd entry point for the mega path. + + Forward dispatches to the single-launch mega forward. Backward currently + reuses the streamed conv-backend-delegated backward (which consumes the + same saved ``silu_qk_save`` layout); a dedicated mega backward is layered + in behind this same entry point. + """ + + @staticmethod + def forward( + ctx, + qkvzba, + conv1d_weight, + A_log, + dt_bias, + cu_seqlens, + seq_idx, + num_key_heads, + num_value_heads, + key_head_dim, + value_head_dim, + ): + ctx.num_key_heads = num_key_heads + ctx.num_value_heads = num_value_heads + ctx.key_head_dim = key_head_dim + ctx.value_head_dim = value_head_dim + query, key, value, gate, beta, g, silu_qk_save = ( + _mega_pre_gated_delta_rule_forward( + qkvzba, + conv1d_weight, + A_log, + dt_bias, + num_key_heads=num_key_heads, + num_value_heads=num_value_heads, + key_head_dim=key_head_dim, + value_head_dim=value_head_dim, + cu_seqlens=cu_seqlens, + ) + ) + ctx.has_seq_idx = seq_idx is not None + if ctx.has_seq_idx: + ctx.save_for_backward(qkvzba, conv1d_weight, A_log, dt_bias, silu_qk_save, seq_idx) + else: + ctx.save_for_backward(qkvzba, conv1d_weight, A_log, dt_bias, silu_qk_save) + return query, key, value, gate, beta, g + + @staticmethod + def backward(ctx, dq, dk, dv, dgate, dbeta, dg): + if ctx.has_seq_idx: + qkvzba, conv1d_weight, A_log, dt_bias, silu_qk_save, seq_idx = ctx.saved_tensors + else: + qkvzba, conv1d_weight, A_log, dt_bias, silu_qk_save = ctx.saved_tensors + seq_idx = None + d_qkvzba, d_weight, d_A_log, d_dt_bias = _mega_pre_gated_delta_rule_backward( + qkvzba, + conv1d_weight, + silu_qk_save, + dq, + dk, + dv, + dgate, + dbeta, + dg, + A_log, + dt_bias, + num_key_heads=ctx.num_key_heads, + num_value_heads=ctx.num_value_heads, + key_head_dim=ctx.key_head_dim, + value_head_dim=ctx.value_head_dim, + seq_idx=seq_idx, + ) + return (d_qkvzba, d_weight, d_A_log, d_dt_bias, None, None, None, None, None, None) + + +def fused_mega_pre_gated_delta_rule( + qkvzba: Tensor, + conv1d_weight: Tensor, + conv1d_bias: Optional[Tensor], + A_log: Tensor, + dt_bias: Tensor, + *, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, + use_qk_l2norm: bool = True, + cu_seqlens: Optional[Tensor] = None, + seq_idx: Optional[Tensor] = None, +) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + """Mega fused pre-gated-delta-rule entry point. + + Args: + qkvzba: ``[seq_len, batch, in_proj_dim]`` projection output. + conv1d_weight: ``[conv_dim, 1, k_w]`` depthwise conv weight. + conv1d_bias: Must be ``None`` in the mega path. + A_log: ``[num_value_heads]`` raw decay parameter. + dt_bias: ``[num_value_heads]`` time-step bias. + num_key_heads / num_value_heads / key_head_dim / value_head_dim: GDN + architecture parameters. ``num_value_heads`` must be a multiple of + ``num_key_heads`` and ``key_head_dim == value_head_dim``. + use_qk_l2norm: Must be ``True`` for parity with the streamed path. + cu_seqlens: Optional packed THD cumulative sequence lengths. + seq_idx: Optional precomputed token-to-sequence map for packed THD mode. + + Returns: + ``(query, key, value, gate, beta, g)`` matching the unfused and streamed + fused pre-GDR APIs. + """ + + assert qkvzba.is_cuda, ( + "fused_mega_pre_gated_delta_rule requires CUDA inputs; " + f"got qkvzba.device={qkvzba.device}." + ) + assert conv1d_bias is None, ( + "Conv bias is not supported by fused_mega_pre_gated_delta_rule " + "(production GDN config has none)." + ) + assert use_qk_l2norm, ( + "use_qk_l2norm=False is not supported by fused_mega_pre_gated_delta_rule " + "(the backward closes over the l2norm path)." + ) + assert num_value_heads % num_key_heads == 0, ( + f"{num_value_heads=} must be a multiple of {num_key_heads=}." + ) + assert key_head_dim == value_head_dim, ( + "fused_mega_pre_gated_delta_rule currently requires " + f"key_head_dim == value_head_dim; got {key_head_dim=} {value_head_dim=}." + ) + if cu_seqlens is not None: + assert cu_seqlens.is_cuda, ( + "Packed fused_mega_pre_gated_delta_rule requires CUDA cu_seqlens; " + f"got cu_seqlens.device={cu_seqlens.device}." + ) + assert cu_seqlens.dtype == torch.int32, ( + "Packed fused_mega_pre_gated_delta_rule requires int32 cu_seqlens; " + f"got {cu_seqlens.dtype=}." + ) + assert cu_seqlens.dim() == 1, ( + "Packed fused_mega_pre_gated_delta_rule expects 1-D cu_seqlens; " + f"got {cu_seqlens.shape=}." + ) + assert qkvzba.shape[1] == 1, ( + "Packed THD fused_mega_pre_gated_delta_rule expects batch dimension 1; " + f"got qkvzba.shape={qkvzba.shape}." + ) + assert cu_seqlens.shape[0] >= 2, ( + "Packed fused_mega_pre_gated_delta_rule requires at least one packed sequence; " + f"got {cu_seqlens.shape=}." + ) + assert cu_seqlens[0].item() == 0, ( + "Packed fused_mega_pre_gated_delta_rule requires cu_seqlens[0] == 0, " + f"got {cu_seqlens[0].item()}." + ) + assert cu_seqlens[-1].item() == qkvzba.shape[0], ( + "Packed fused_mega_pre_gated_delta_rule requires cu_seqlens[-1] to match " + f"seq_len, got {cu_seqlens[-1].item()} vs {qkvzba.shape[0]}." + ) + cu_seqlens = cu_seqlens.contiguous() + seq_idx = _resolve_packed_seq_idx(cu_seqlens, seq_idx, qkvzba.shape[0]) + else: + assert seq_idx is None, "seq_idx requires cu_seqlens for packed THD mode." + + return _FusedMegaPreGatedDeltaRuleFunction.apply( + qkvzba, + conv1d_weight, + A_log, + dt_bias, + cu_seqlens, + seq_idx, + num_key_heads, + num_value_heads, + key_head_dim, + value_head_dim, + ) diff --git a/megatron/core/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py index 6eed7581d03..4319af230ff 100644 --- a/megatron/core/fusions/fused_mla_yarn_rope_apply.py +++ b/megatron/core/fusions/fused_mla_yarn_rope_apply.py @@ -41,12 +41,17 @@ def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size): last_cum_seqlen = cur_cum_seqlen seq_idx += 1 if cp_size > 1: - if token_idx < this_seq_len // 2: - token_idx = token_idx + cp_rank * this_seq_len // 2 + first_cp_seg = (this_seq_len + 1) // 2 + second_cp_seg = this_seq_len // 2 + if token_idx < first_cp_seg: + token_idx = token_idx + cp_rank * first_cp_seg else: - token_idx = (token_idx - this_seq_len // 2) + ( - 2 * cp_size - cp_rank - 1 - ) * this_seq_len // 2 + token_idx = ( + token_idx + - first_cp_seg + + cp_size * first_cp_seg + + (cp_size - cp_rank - 1) * second_cp_seg + ) return token_idx diff --git a/megatron/core/fusions/fused_mrope.py b/megatron/core/fusions/fused_mrope.py new file mode 100644 index 00000000000..6ebad4df933 --- /dev/null +++ b/megatron/core/fusions/fused_mrope.py @@ -0,0 +1,871 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Triton fused multimodal RoPE apply. + +The fused path consumes the raw three-axis mRoPE frequencies with shape +``[3, batch, seq, rotary_dim / 2]`` and applies the rotation directly to a BSHD +tensor. It supports both Qwen2-VL section-based mRoPE and Qwen3.5-VL +stride-3 interleaved mRoPE layouts. +""" + +from __future__ import annotations + +from typing import List, Optional +from unittest.mock import MagicMock + +import torch + +from megatron.core.utils import null_decorator + +try: + import triton + import triton.language as tl + + HAVE_TRITON = True +except ImportError: + HAVE_TRITON = False + +if not HAVE_TRITON: + triton = MagicMock() + triton.jit = null_decorator + tl = MagicMock() + + +def _smallest_power_of_2_at_least(x: int) -> int: + block = 1 + while block < x: + block *= 2 + return block + + +def _expected_interleaved_mrope_section(half_rotary_dim: int) -> tuple[int, int, int]: + return ((half_rotary_dim + 2) // 3, (half_rotary_dim + 1) // 3, half_rotary_dim // 3) + + +def _validate_mrope_section( + mrope_section: List[int], half_rotary_dim: int, interleaved_mrope: bool +) -> tuple[int, int, int]: + assert len(mrope_section) == 3, f"mrope_section must have length 3, got {mrope_section}" + + sec_t, sec_h, sec_w = (int(section) for section in mrope_section) + assert ( + min(sec_t, sec_h, sec_w) >= 0 + ), f"mrope_section values must be non-negative, got {mrope_section}" + assert half_rotary_dim > 0, "raw mRoPE rotary dim must be greater than 0" + assert ( + sec_t + sec_h + sec_w == half_rotary_dim + ), f"mrope_section {mrope_section} must sum to rotary_dim / 2 = {half_rotary_dim}" + if interleaved_mrope: + expected = _expected_interleaved_mrope_section(half_rotary_dim) + assert (sec_t, sec_h, sec_w) == expected, ( + f"interleaved mRoPE with rotary_dim / 2 = {half_rotary_dim} requires " + f"mrope_section {list(expected)}, got {mrope_section}" + ) + return sec_t, sec_h, sec_w + + +def _validate_mrope_inputs( + t: torch.Tensor, freqs: torch.Tensor, mrope_section: List[int], interleaved_mrope: bool +) -> tuple[int, int, int, int, int, int, int, int]: + assert t.dim() == 4, f"t must have shape [seq, batch, heads, head_dim], got {t.shape}" + assert freqs.dim() == 4, ( + "raw mRoPE freqs must have shape [3, batch, seq, rotary_dim / 2], " f"got {freqs.shape}" + ) + + seq, batch, heads, head_dim = t.shape + axes, freq_batch, freq_seq, half_rotary_dim = freqs.shape + assert axes == 3, f"raw mRoPE freqs first dimension must be 3, got {axes}" + assert ( + freq_batch == batch and freq_seq == seq + ), f"freqs shape {tuple(freqs.shape)} is incompatible with t shape {tuple(t.shape)}" + + sec_t, sec_h, sec_w = _validate_mrope_section(mrope_section, half_rotary_dim, interleaved_mrope) + + rotary_dim = half_rotary_dim * 2 + assert ( + rotary_dim <= head_dim + ), f"raw mRoPE rotary dim {rotary_dim} exceeds input head dim {head_dim}" + return seq, batch, heads, head_dim, half_rotary_dim, sec_t, sec_h, sec_w + + +def _validate_mrope_thd_inputs( + t: torch.Tensor, + cu_seqlens: torch.Tensor, + freqs: torch.Tensor, + mrope_section: List[int], + interleaved_mrope: bool, + cp_size: int, +) -> tuple[int, int, int, int, int, int, int]: + assert t.dim() == 3, f"t must have shape [tokens, heads, head_dim], got {t.shape}" + assert freqs.dim() == 4, ( + "raw mRoPE freqs must have shape [3, 1, total_seqlen, rotary_dim / 2], " + f"got {freqs.shape}" + ) + assert cu_seqlens.dim() == 1, f"cu_seqlens must be 1D, got {cu_seqlens.shape}" + + tokens, heads, head_dim = t.shape + axes, freq_batch, freq_seq, half_rotary_dim = freqs.shape + assert axes == 3, f"raw mRoPE freqs first dimension must be 3, got {axes}" + assert freq_batch == 1, ( + "raw mRoPE THD freqs must have singleton batch dimension, " f"got {freqs.shape}" + ) + assert freq_seq == tokens * cp_size, ( + "raw mRoPE THD freqs sequence length must match local tokens times cp_size, " + f"got freqs.shape[2]={freq_seq}, tokens={tokens}, cp_size={cp_size}" + ) + + sec_t, sec_h, sec_w = _validate_mrope_section(mrope_section, half_rotary_dim, interleaved_mrope) + rotary_dim = half_rotary_dim * 2 + assert ( + rotary_dim <= head_dim + ), f"raw mRoPE rotary dim {rotary_dim} exceeds input head dim {head_dim}" + return tokens, heads, head_dim, half_rotary_dim, sec_t, sec_h, sec_w + + +def get_fused_mrope_unavailable_reason( + t: Optional[torch.Tensor] = None, + freqs: Optional[torch.Tensor] = None, + rotary_interleaved: bool = False, +) -> Optional[str]: + """Return why fused mRoPE cannot run, or None when it is launchable.""" + if not HAVE_TRITON: + return "Triton is not available" + if rotary_interleaved: + return "rotary_interleaved=True is not supported" + if t is None or freqs is None: + return None + if not t.is_cuda or not freqs.is_cuda: + return "Triton fused mRoPE requires CUDA tensors" + if t.device != freqs.device: + return ( + "Triton fused mRoPE requires t and freqs on the same device, " + f"got {t.device} and {freqs.device}" + ) + if freqs.dtype != torch.float32: + return f"raw mRoPE freqs must be float32, got {freqs.dtype}" + if t.dtype not in (torch.float16, torch.bfloat16, torch.float32): + return f"input dtype {t.dtype} is not supported" + if t.stride(-1) != 1: + return f"input head dimension must be contiguous, got stride {t.stride()}" + try: + capability = torch.cuda.get_device_capability(t.device) + except RuntimeError as exc: + return f"could not query CUDA device capability: {exc}" + if capability < (7, 0): + return f"requires CUDA compute capability >= 7.0, got {capability[0]}.{capability[1]}" + if t.dtype == torch.bfloat16 and capability < (8, 0): + return ( + "requires CUDA compute capability >= 8.0 for bfloat16 inputs, " + f"got {capability[0]}.{capability[1]}" + ) + return None + + +def get_fused_mrope_thd_unavailable_reason( + t: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.Tensor] = None, + freqs: Optional[torch.Tensor] = None, + rotary_interleaved: bool = False, + cp_size: int = 1, + cp_rank: int = 0, +) -> Optional[str]: + """Return why fused THD mRoPE cannot run, or None when it is launchable.""" + if not HAVE_TRITON: + return "Triton is not available" + if rotary_interleaved: + return "rotary_interleaved=True is not supported" + if cp_size < 1: + return f"cp_size must be positive, got {cp_size}" + if cp_rank < 0 or cp_rank >= cp_size: + return f"cp_rank must be in [0, {cp_size}), got {cp_rank}" + if t is None or cu_seqlens is None or freqs is None: + return None + if t.dim() != 3: + return ( + f"THD fused mRoPE expects t with shape [tokens, heads, head_dim], got {tuple(t.shape)}" + ) + if freqs.dim() != 4: + return ( + "raw mRoPE THD freqs must have shape [3, 1, total_seqlen, rotary_dim / 2], " + f"got {tuple(freqs.shape)}" + ) + if cu_seqlens.dim() != 1: + return f"cu_seqlens must be 1D, got {tuple(cu_seqlens.shape)}" + if not t.is_cuda or not freqs.is_cuda or not cu_seqlens.is_cuda: + return "Triton fused THD mRoPE requires CUDA tensors" + if t.device != freqs.device or t.device != cu_seqlens.device: + return ( + "Triton fused THD mRoPE requires t, freqs, and cu_seqlens on the same device, " + f"got {t.device}, {freqs.device}, and {cu_seqlens.device}" + ) + if freqs.dtype != torch.float32: + return f"raw mRoPE freqs must be float32, got {freqs.dtype}" + if t.dtype not in (torch.float16, torch.bfloat16, torch.float32): + return f"input dtype {t.dtype} is not supported" + if cu_seqlens.dtype not in (torch.int32, torch.int64): + return f"cu_seqlens dtype {cu_seqlens.dtype} is not supported" + if t.stride(-1) != 1: + return f"input head dimension must be contiguous, got stride {t.stride()}" + if freqs.shape[0] != 3 or freqs.shape[1] != 1: + return ( + "raw mRoPE THD freqs must have shape [3, 1, total_seqlen, rotary_dim / 2], " + f"got {tuple(freqs.shape)}" + ) + if cp_size > 1 and freqs.shape[2] % cp_size != 0: + return ( + "raw mRoPE THD freqs sequence length must be divisible by context parallel size, " + f"got freqs.shape[2]={freqs.shape[2]}, cp_size={cp_size}" + ) + if cp_size > 1: + # Guard: each packed sub-sequence length must satisfy seqlen % cp_size == 0. + seq_bounds = cu_seqlens.tolist() + for seq_start, seq_end in zip(seq_bounds[:-1], seq_bounds[1:]): + if (seq_end - seq_start) % cp_size != 0: + return ( + "each packed THD sub-sequence length must be divisible by context " + f"parallel size, got sub-sequence length {seq_end - seq_start} " + f"with cp_size={cp_size}" + ) + if freqs.shape[2] != t.shape[0] * cp_size: + return ( + "raw mRoPE THD freqs sequence length must match local tokens times cp_size, " + f"got freqs.shape[2]={freqs.shape[2]}, tokens={t.shape[0]}, cp_size={cp_size}" + ) + try: + capability = torch.cuda.get_device_capability(t.device) + except RuntimeError as exc: + return f"could not query CUDA device capability: {exc}" + if capability < (7, 0): + return f"requires CUDA compute capability >= 7.0, got {capability[0]}.{capability[1]}" + if t.dtype == torch.bfloat16 and capability < (8, 0): + return ( + "requires CUDA compute capability >= 8.0 for bfloat16 inputs, " + f"got {capability[0]}.{capability[1]}" + ) + return None + + +def can_launch_fused_mrope( + t: Optional[torch.Tensor] = None, + freqs: Optional[torch.Tensor] = None, + rotary_interleaved: bool = False, +) -> bool: + """Return whether the Triton fused mRoPE kernel can be launched.""" + return get_fused_mrope_unavailable_reason(t, freqs, rotary_interleaved) is None + + +def can_launch_fused_mrope_thd( + t: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.Tensor] = None, + freqs: Optional[torch.Tensor] = None, + rotary_interleaved: bool = False, + cp_size: int = 1, + cp_rank: int = 0, +) -> bool: + """Return whether the Triton fused THD mRoPE kernel can be launched.""" + return ( + get_fused_mrope_thd_unavailable_reason( + t, + cu_seqlens, + freqs, + rotary_interleaved=rotary_interleaved, + cp_size=cp_size, + cp_rank=cp_rank, + ) + is None + ) + + +def mrope_freqs_to_rotary_emb( + freqs: torch.Tensor, + mrope_section: List[int], + interleaved_mrope: bool = False, + rotary_interleaved: bool = False, +) -> torch.Tensor: + """Convert raw mRoPE freqs to the unfused RoPE embedding layout. + + Args: + freqs: Raw mRoPE frequencies with shape ``[3, batch, seq, rotary_dim / 2]``. + mrope_section: Temporal, height, and width channel sections. + interleaved_mrope: Use Qwen3.5-VL stride-3 T/H/W layout when True. Use + Qwen2-VL section layout when False. + rotary_interleaved: Use adjacent-pair RoPE layout when True. This is + available for reference conversion; fused Triton currently supports + split-half layout only. + + Returns: + Tensor with shape ``[seq, batch, 1, rotary_dim]``. + """ + assert freqs.dim() == 4, ( + "raw mRoPE freqs must have shape [3, batch, seq, rotary_dim / 2], " f"got {freqs.shape}" + ) + assert freqs.size(0) == 3, f"raw mRoPE freqs first dimension must be 3, got {freqs.size(0)}" + assert len(mrope_section) == 3, f"mrope_section must have length 3, got {mrope_section}" + + half_rotary_dim = freqs.size(-1) + sec_t, sec_h, sec_w = _validate_mrope_section(mrope_section, half_rotary_dim, interleaved_mrope) + + if interleaved_mrope: + freqs_out = freqs[0].clone() + for dim_idx, offset in enumerate((1, 2), start=1): + length = int(mrope_section[dim_idx]) * 3 + idx = slice(offset, length, 3) + freqs_out[..., idx] = freqs[dim_idx, ..., idx] + if rotary_interleaved: + batch = freqs_out.shape[0] + emb = torch.stack( + (freqs_out.reshape(batch, -1, 1), freqs_out.reshape(batch, -1, 1)), dim=-1 + ) + emb = emb.view(batch, freqs_out.shape[1], -1) + else: + emb = torch.cat((freqs_out, freqs_out), dim=-1) + else: + if rotary_interleaved: + batch = freqs.shape[1] + emb = torch.stack( + (freqs.reshape(3, batch, -1, 1), freqs.reshape(3, batch, -1, 1)), dim=-1 + ).view(3, batch, freqs.shape[2], -1) + mrope_section_doubled = list(mrope_section) * 2 + emb = torch.cat( + [chunk[i % 3] for i, chunk in enumerate(emb.split(mrope_section_doubled, dim=-1))], + dim=-1, + ) + else: + freqs_out = torch.empty_like(freqs[0]) + freqs_out[..., :sec_t] = freqs[0, ..., :sec_t] + freqs_out[..., sec_t : sec_t + sec_h] = freqs[1, ..., sec_t : sec_t + sec_h] + freqs_out[..., sec_t + sec_h :] = freqs[2, ..., sec_t + sec_h :] + emb = torch.cat((freqs_out, freqs_out), dim=-1) + return emb[..., None, :].transpose(0, 1).contiguous() + + +@triton.jit +def _mrope_axis( + k, + SEC_T: tl.constexpr, + SEC_H: tl.constexpr, + SEC_W: tl.constexpr, + INTERLEAVED_MROPE: tl.constexpr, +): + if INTERLEAVED_MROPE: + rem = k % 3 + section_idx = k // 3 + is_h = (rem == 1) & (section_idx < SEC_H) + is_w = (rem == 2) & (section_idx < SEC_W) + return tl.where(is_h, 1, tl.where(is_w, 2, 0)) + + is_h = (k >= SEC_T) & (k < (SEC_T + SEC_H)) + is_w = k >= (SEC_T + SEC_H) + return tl.where(is_h, 1, tl.where(is_w, 2, 0)) + + +@triton.jit +def _fused_mrope_kernel( + T, + FREQS, + OUT, + t_s_seq, + t_s_batch, + t_s_head, + t_s_dim, + f_s_axis, + f_s_batch, + f_s_seq, + f_s_dim, + o_s_seq, + o_s_batch, + o_s_head, + o_s_dim, + HEAD_DIM: tl.constexpr, + HALF_ROTARY_DIM: tl.constexpr, + PASS_DIM: tl.constexpr, + SEC_T: tl.constexpr, + SEC_H: tl.constexpr, + SEC_W: tl.constexpr, + INTERLEAVED_MROPE: tl.constexpr, + ROTARY_INTERLEAVED: tl.constexpr, + INVERSE: tl.constexpr, + BLOCK_HALF: tl.constexpr, + BLOCK_PASS: tl.constexpr, +): + seq_idx = tl.program_id(0) + batch_idx = tl.program_id(1) + head_idx = tl.program_id(2) + + k = tl.arange(0, BLOCK_HALF) + mask = k < HALF_ROTARY_DIM + + axis = _mrope_axis(k, SEC_T, SEC_H, SEC_W, INTERLEAVED_MROPE) + + freqs_offset = axis * f_s_axis + batch_idx * f_s_batch + seq_idx * f_s_seq + k * f_s_dim + freqs = tl.load(FREQS + freqs_offset, mask=mask, other=0.0) + # Match PyTorch pointwise dtype semantics: cast cos/sin before the multiply. + cos_v = tl.cos(freqs).to(OUT.dtype.element_ty) + sin_v = tl.sin(freqs).to(OUT.dtype.element_ty) + if INVERSE: + sin_v = -sin_v + + t_base = T + seq_idx * t_s_seq + batch_idx * t_s_batch + head_idx * t_s_head + out_base = OUT + seq_idx * o_s_seq + batch_idx * o_s_batch + head_idx * o_s_head + + if ROTARY_INTERLEAVED: + lo_offset = (2 * k) * t_s_dim + hi_offset = (2 * k + 1) * t_s_dim + out_lo_offset = (2 * k) * o_s_dim + out_hi_offset = (2 * k + 1) * o_s_dim + else: + lo_offset = k * t_s_dim + hi_offset = (k + HALF_ROTARY_DIM) * t_s_dim + out_lo_offset = k * o_s_dim + out_hi_offset = (k + HALF_ROTARY_DIM) * o_s_dim + + t_lo = tl.load(t_base + lo_offset, mask=mask, other=0.0).to(OUT.dtype.element_ty) + t_hi = tl.load(t_base + hi_offset, mask=mask, other=0.0).to(OUT.dtype.element_ty) + + lo_cos = (t_lo * cos_v).to(OUT.dtype.element_ty) + hi_sin = (t_hi * sin_v).to(OUT.dtype.element_ty) + hi_cos = (t_hi * cos_v).to(OUT.dtype.element_ty) + lo_sin = (t_lo * sin_v).to(OUT.dtype.element_ty) + + out_lo = (lo_cos - hi_sin).to(OUT.dtype.element_ty) + out_hi = (hi_cos + lo_sin).to(OUT.dtype.element_ty) + + tl.store(out_base + out_lo_offset, out_lo, mask=mask) + tl.store(out_base + out_hi_offset, out_hi, mask=mask) + + if PASS_DIM > 0: + pass_idx = tl.arange(0, BLOCK_PASS) + pass_mask = pass_idx < PASS_DIM + src_dim = 2 * HALF_ROTARY_DIM + pass_idx + pass_values = tl.load(t_base + src_dim * t_s_dim, mask=pass_mask, other=0.0) + tl.store(out_base + src_dim * o_s_dim, pass_values, mask=pass_mask) + + +@triton.jit +def _fused_mrope_thd_kernel( + T, + CU_SEQLENS, + FREQS, + OUT, + t_s_token, + t_s_head, + t_s_dim, + cu_s_idx, + f_s_axis, + f_s_seq, + f_s_dim, + o_s_token, + o_s_head, + o_s_dim, + NUM_SEQS, + HEAD_DIM: tl.constexpr, + HALF_ROTARY_DIM: tl.constexpr, + PASS_DIM: tl.constexpr, + SEC_T: tl.constexpr, + SEC_H: tl.constexpr, + SEC_W: tl.constexpr, + INTERLEAVED_MROPE: tl.constexpr, + ROTARY_INTERLEAVED: tl.constexpr, + INVERSE: tl.constexpr, + CP_SIZE: tl.constexpr, + CP_RANK: tl.constexpr, + FP32_COMPUTE: tl.constexpr, + BLOCK_HALF: tl.constexpr, + BLOCK_PASS: tl.constexpr, +): + token_idx = tl.program_id(0) + head_idx = tl.program_id(1) + + freq_seq_idx = token_idx + seq_i = 0 + while seq_i < NUM_SEQS: + global_start = tl.load(CU_SEQLENS + seq_i * cu_s_idx) + global_end = tl.load(CU_SEQLENS + (seq_i + 1) * cu_s_idx) + local_start = global_start // CP_SIZE + local_end = global_end // CP_SIZE + in_seq = (token_idx >= local_start) & (token_idx < local_end) + local_offset = token_idx - local_start + + if CP_SIZE > 1: + local_seq_len = local_end - local_start + first_cp_seg = (local_seq_len + 1) // 2 + second_cp_seg = local_seq_len // 2 + first_freq_idx = global_start + CP_RANK * first_cp_seg + local_offset + second_freq_idx = ( + global_end - (CP_RANK + 1) * second_cp_seg + (local_offset - first_cp_seg) + ) + seq_freq_idx = tl.where(local_offset < first_cp_seg, first_freq_idx, second_freq_idx) + else: + seq_freq_idx = global_start + local_offset + + freq_seq_idx = tl.where(in_seq, seq_freq_idx, freq_seq_idx) + seq_i += 1 + + k = tl.arange(0, BLOCK_HALF) + mask = k < HALF_ROTARY_DIM + axis = _mrope_axis(k, SEC_T, SEC_H, SEC_W, INTERLEAVED_MROPE) + + freqs_offset = axis * f_s_axis + freq_seq_idx * f_s_seq + k * f_s_dim + freqs = tl.load(FREQS + freqs_offset, mask=mask, other=0.0) + if FP32_COMPUTE: + cos_v = tl.cos(freqs) + sin_v = tl.sin(freqs) + else: + cos_v = tl.cos(freqs).to(OUT.dtype.element_ty) + sin_v = tl.sin(freqs).to(OUT.dtype.element_ty) + if INVERSE: + sin_v = -sin_v + + t_base = T + token_idx * t_s_token + head_idx * t_s_head + out_base = OUT + token_idx * o_s_token + head_idx * o_s_head + + if ROTARY_INTERLEAVED: + lo_offset = (2 * k) * t_s_dim + hi_offset = (2 * k + 1) * t_s_dim + out_lo_offset = (2 * k) * o_s_dim + out_hi_offset = (2 * k + 1) * o_s_dim + else: + lo_offset = k * t_s_dim + hi_offset = (k + HALF_ROTARY_DIM) * t_s_dim + out_lo_offset = k * o_s_dim + out_hi_offset = (k + HALF_ROTARY_DIM) * o_s_dim + + if FP32_COMPUTE: + t_lo = tl.load(t_base + lo_offset, mask=mask, other=0.0).to(tl.float32) + t_hi = tl.load(t_base + hi_offset, mask=mask, other=0.0).to(tl.float32) + else: + t_lo = tl.load(t_base + lo_offset, mask=mask, other=0.0).to(OUT.dtype.element_ty) + t_hi = tl.load(t_base + hi_offset, mask=mask, other=0.0).to(OUT.dtype.element_ty) + + if FP32_COMPUTE: + lo_cos = t_lo * cos_v + hi_sin = t_hi * sin_v + hi_cos = t_hi * cos_v + lo_sin = t_lo * sin_v + else: + lo_cos = (t_lo * cos_v).to(OUT.dtype.element_ty) + hi_sin = (t_hi * sin_v).to(OUT.dtype.element_ty) + hi_cos = (t_hi * cos_v).to(OUT.dtype.element_ty) + lo_sin = (t_lo * sin_v).to(OUT.dtype.element_ty) + + if FP32_COMPUTE: + out_lo = lo_cos - hi_sin + out_hi = hi_cos + lo_sin + else: + out_lo = (lo_cos - hi_sin).to(OUT.dtype.element_ty) + out_hi = (hi_cos + lo_sin).to(OUT.dtype.element_ty) + + tl.store(out_base + out_lo_offset, out_lo, mask=mask) + tl.store(out_base + out_hi_offset, out_hi, mask=mask) + + if PASS_DIM > 0: + pass_idx = tl.arange(0, BLOCK_PASS) + pass_mask = pass_idx < PASS_DIM + src_dim = 2 * HALF_ROTARY_DIM + pass_idx + pass_values = tl.load(t_base + src_dim * t_s_dim, mask=pass_mask, other=0.0) + tl.store(out_base + src_dim * o_s_dim, pass_values, mask=pass_mask) + + +def _launch_fused_mrope( + t: torch.Tensor, + freqs: torch.Tensor, + mrope_section: List[int], + interleaved_mrope: bool, + rotary_interleaved: bool, + inverse: bool, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + unavailable_reason = get_fused_mrope_unavailable_reason(t, freqs, rotary_interleaved) + assert unavailable_reason is None, unavailable_reason + + seq, batch, heads, head_dim, half_rotary_dim, sec_t, sec_h, sec_w = _validate_mrope_inputs( + t, freqs, mrope_section, interleaved_mrope + ) + + if out is None: + out = torch.empty_like(t) + else: + assert out.shape == t.shape and out.dtype == t.dtype + assert ( + out.stride(-1) == 1 + ), f"fused mRoPE requires output contiguous head dimension, got {out.stride()}" + + block_half = _smallest_power_of_2_at_least(half_rotary_dim) + pass_dim = head_dim - (2 * half_rotary_dim) + block_pass = _smallest_power_of_2_at_least(max(pass_dim, 1)) + + grid = (seq, batch, heads) + _fused_mrope_kernel[grid]( + t, + freqs, + out, + t.stride(0), + t.stride(1), + t.stride(2), + t.stride(3), + freqs.stride(0), + freqs.stride(1), + freqs.stride(2), + freqs.stride(3), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + HEAD_DIM=head_dim, + HALF_ROTARY_DIM=half_rotary_dim, + PASS_DIM=pass_dim, + SEC_T=sec_t, + SEC_H=sec_h, + SEC_W=sec_w, + INTERLEAVED_MROPE=interleaved_mrope, + ROTARY_INTERLEAVED=rotary_interleaved, + INVERSE=inverse, + BLOCK_HALF=block_half, + BLOCK_PASS=block_pass, + num_warps=4, + ) + return out + + +def _launch_fused_mrope_thd( + t: torch.Tensor, + cu_seqlens: torch.Tensor, + freqs: torch.Tensor, + mrope_section: List[int], + interleaved_mrope: bool, + rotary_interleaved: bool, + inverse: bool, + cp_size: int, + cp_rank: int, + fp32_compute: bool = False, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + unavailable_reason = get_fused_mrope_thd_unavailable_reason( + t, + cu_seqlens, + freqs, + rotary_interleaved=rotary_interleaved, + cp_size=cp_size, + cp_rank=cp_rank, + ) + assert unavailable_reason is None, unavailable_reason + + tokens, heads, head_dim, half_rotary_dim, sec_t, sec_h, sec_w = _validate_mrope_thd_inputs( + t, cu_seqlens, freqs, mrope_section, interleaved_mrope, cp_size + ) + + if out is None: + out = torch.empty_like(t) + else: + assert out.shape == t.shape and out.dtype == t.dtype + assert ( + out.stride(-1) == 1 + ), f"fused THD mRoPE requires output contiguous head dimension, got {out.stride()}" + + block_half = _smallest_power_of_2_at_least(half_rotary_dim) + pass_dim = head_dim - (2 * half_rotary_dim) + block_pass = _smallest_power_of_2_at_least(max(pass_dim, 1)) + num_seqs = cu_seqlens.numel() - 1 + + grid = (tokens, heads) + _fused_mrope_thd_kernel[grid]( + t, + cu_seqlens, + freqs, + out, + t.stride(0), + t.stride(1), + t.stride(2), + cu_seqlens.stride(0), + freqs.stride(0), + freqs.stride(2), + freqs.stride(3), + out.stride(0), + out.stride(1), + out.stride(2), + num_seqs, + HEAD_DIM=head_dim, + HALF_ROTARY_DIM=half_rotary_dim, + PASS_DIM=pass_dim, + SEC_T=sec_t, + SEC_H=sec_h, + SEC_W=sec_w, + INTERLEAVED_MROPE=interleaved_mrope, + ROTARY_INTERLEAVED=rotary_interleaved, + INVERSE=inverse, + CP_SIZE=cp_size, + CP_RANK=cp_rank, + FP32_COMPUTE=fp32_compute, + BLOCK_HALF=block_half, + BLOCK_PASS=block_pass, + num_warps=4, + ) + return out + + +class _FusedMRoPE(torch.autograd.Function): + """Autograd wrapper for fused mRoPE. + + The raw frequency table is generated from position IDs and inverse frequencies, + so gradients are only propagated to the rotated tensor. + """ + + @staticmethod + def forward(ctx, t, freqs, mrope_section, interleaved_mrope, rotary_interleaved): + assert not freqs.requires_grad, "fused mRoPE expects non-gradient raw frequency tensors" + ctx.mrope_section = tuple(int(section) for section in mrope_section) + ctx.interleaved_mrope = bool(interleaved_mrope) + ctx.rotary_interleaved = bool(rotary_interleaved) + ctx.save_for_backward(freqs) + return _launch_fused_mrope( + t, + freqs, + ctx.mrope_section, + ctx.interleaved_mrope, + ctx.rotary_interleaved, + inverse=False, + ) + + @staticmethod + def backward(ctx, grad_output): + (freqs,) = ctx.saved_tensors + grad_input = _launch_fused_mrope( + grad_output.contiguous(), + freqs, + ctx.mrope_section, + ctx.interleaved_mrope, + ctx.rotary_interleaved, + inverse=True, + ) + return grad_input, None, None, None, None + + +class _FusedMRoPETHD(torch.autograd.Function): + """Autograd wrapper for fused THD mRoPE.""" + + @staticmethod + def forward( + ctx, + t, + cu_seqlens, + freqs, + mrope_section, + interleaved_mrope, + rotary_interleaved, + cp_size, + cp_rank, + fp32_compute, + ): + assert not freqs.requires_grad, "fused THD mRoPE expects non-gradient raw frequency tensors" + ctx.mrope_section = tuple(int(section) for section in mrope_section) + ctx.interleaved_mrope = bool(interleaved_mrope) + ctx.rotary_interleaved = bool(rotary_interleaved) + ctx.cp_size = int(cp_size) + ctx.cp_rank = int(cp_rank) + ctx.fp32_compute = bool(fp32_compute) + ctx.save_for_backward(cu_seqlens, freqs) + return _launch_fused_mrope_thd( + t, + cu_seqlens, + freqs, + ctx.mrope_section, + ctx.interleaved_mrope, + ctx.rotary_interleaved, + inverse=False, + cp_size=ctx.cp_size, + cp_rank=ctx.cp_rank, + fp32_compute=ctx.fp32_compute, + ) + + @staticmethod + def backward(ctx, grad_output): + cu_seqlens, freqs = ctx.saved_tensors + grad_input = _launch_fused_mrope_thd( + grad_output.contiguous(), + cu_seqlens, + freqs, + ctx.mrope_section, + ctx.interleaved_mrope, + ctx.rotary_interleaved, + inverse=True, + cp_size=ctx.cp_size, + cp_rank=ctx.cp_rank, + fp32_compute=ctx.fp32_compute, + ) + return grad_input, None, None, None, None, None, None, None, None + + +def fused_apply_mrope( + t: torch.Tensor, + freqs: torch.Tensor, + mrope_section: List[int], + interleaved_mrope: bool = False, + rotary_interleaved: bool = False, +) -> torch.Tensor: + """Apply multimodal RoPE with a fused Triton kernel. + + Args: + t: Input tensor with shape ``[seq, batch, heads, head_dim]``. + freqs: Raw mRoPE frequencies with shape ``[3, batch, seq, rotary_dim / 2]``. + mrope_section: Temporal, height, and width channel sections. + interleaved_mrope: Use Qwen3.5-VL stride-3 T/H/W layout when True. Use + Qwen2-VL section layout when False. + rotary_interleaved: Must be False. The integrated fused mRoPE path + currently supports split-half RoPE layout. + + Returns: + Rotated tensor with the same shape and dtype as ``t``. + """ + return _FusedMRoPE.apply(t, freqs, mrope_section, interleaved_mrope, rotary_interleaved) + + +def fused_apply_mrope_thd( + t: torch.Tensor, + cu_seqlens: torch.Tensor, + freqs: torch.Tensor, + mrope_section: List[int], + interleaved_mrope: bool = False, + rotary_interleaved: bool = False, + cp_size: int = 1, + cp_rank: int = 0, + fp32_compute: bool = False, +) -> torch.Tensor: + """Apply multimodal RoPE to THD-packed tensors with a fused Triton kernel. + + Args: + t: Input tensor with shape ``[total_tokens, heads, head_dim]``. + cu_seqlens: Global cumulative sequence lengths for the packed batch. + freqs: Raw mRoPE frequencies with shape ``[3, 1, total_seqlen, rotary_dim / 2]``. + mrope_section: Temporal, height, and width channel sections. + interleaved_mrope: Use Qwen3.5-VL stride-3 T/H/W layout when True. + rotary_interleaved: Must be False. + cp_size: Context parallel world size for THD token mapping. + cp_rank: Context parallel rank for THD token mapping. + fp32_compute: Apply the rotary math in fp32 and cast directly to output dtype. + + Returns: + Rotated tensor with the same shape and dtype as ``t``. + """ + return _FusedMRoPETHD.apply( + t, + cu_seqlens, + freqs, + mrope_section, + interleaved_mrope, + rotary_interleaved, + cp_size, + cp_rank, + fp32_compute, + ) + + +def is_fused_mrope_available() -> bool: + """Return whether the Triton mRoPE fusion can be used on this host. + + This does not check tensor device, dtype, stride, or CUDA capability. Use + ``can_launch_fused_mrope`` or ``get_fused_mrope_unavailable_reason`` with + tensors before dispatching to the fused kernel. + """ + if not torch.cuda.is_available(): + return False + return can_launch_fused_mrope() diff --git a/megatron/core/fusions/fused_pre_gated_delta_rule.py b/megatron/core/fusions/fused_pre_gated_delta_rule.py new file mode 100644 index 00000000000..72c702ecbb4 --- /dev/null +++ b/megatron/core/fusions/fused_pre_gated_delta_rule.py @@ -0,0 +1,2191 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Fused pre-gated-delta-rule projection kernels. + +The public entry point consumes the dense ``qkvzba`` projection and returns +``query``, ``key``, ``value``, ``gate``, ``beta``, and ``g`` in the layouts +expected by the gated delta rule. The forward path keeps QK, V, Z, and +G/Beta as separate streamed scopes. The backward mirrors those scopes for +layout/l2norm/g-beta work, then delegates depthwise conv gradients to the +``causal_conv1d`` backend. + +Unsupported cases are rejected at the Python entry point: CPU tensors, +conv bias, and ``use_qk_l2norm=False``. Packed THD sequences use separate +QK/V causal-conv kernels so the dense BSHD kernels stay free of packed +metadata and runtime branches. +""" + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + +# The 1.6.1+ ``causal_conv1d`` package exposes the lower-level binding via +# ``causal_conv1d.cpp_functions.causal_conv1d_bwd_function``; older builds +# (still common in some older environments) expose the same +# function under ``causal_conv1d_cuda.causal_conv1d_bwd``. Try both so the +# fast path is taken everywhere the package is installed. +from torch import Tensor + +try: + from causal_conv1d.cpp_functions import ( + causal_conv1d_bwd_function as _causal_conv1d_bwd_function, + ) +except ImportError: + try: + import causal_conv1d_cuda as _causal_conv1d_cuda + + _causal_conv1d_bwd_function = _causal_conv1d_cuda.causal_conv1d_bwd + except ImportError: + # The external causal_conv1d package is optional: only the fused pre-GDR + # backward needs it. Importing this module (and hence GatedDeltaNet) must + # not fail when it is absent; raise a clear error only if the fused path + # is actually exercised. + _causal_conv1d_bwd_function = None + + +_L2NORM_EPS = 1e-6 + +_QK_STREAM_SLOT = 0 +_V_STREAM_SLOT = 2 +_G_BETA_STREAM_SLOT = 3 +_Z_STREAM_SLOT = 4 + +_LAYOUT_BLOCK_S = 64 + + +# --------------------------------------------------------------------------- +# Forward kernels +# --------------------------------------------------------------------------- + + +def _conv_autotune_configs(): + return [ + triton.Config({"BLOCK_S": 16}, num_warps=2, num_stages=2), + triton.Config({"BLOCK_S": 32}, num_warps=2, num_stages=2), + triton.Config({"BLOCK_S": 32}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 32}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_S": 64}, num_warps=2, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=4), + triton.Config({"BLOCK_S": 64}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_S": 128}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=8, num_stages=3), + triton.Config({"BLOCK_S": 256}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 256}, num_warps=8, num_stages=2), + ] + + +def _g_beta_autotune_configs(): + return [ + triton.Config({"BLOCK_S": 32, "BLOCK_H": 16}, num_warps=2, num_stages=2), + triton.Config({"BLOCK_S": 64, "BLOCK_H": 16}, num_warps=2, num_stages=2), + triton.Config({"BLOCK_S": 64, "BLOCK_H": 32}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 128, "BLOCK_H": 16}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 128, "BLOCK_H": 32}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 128, "BLOCK_H": 64}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 256, "BLOCK_H": 16}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 256, "BLOCK_H": 32}, num_warps=8, num_stages=2), + ] + + +@triton.autotune( + configs=_conv_autotune_configs(), + key=["seq_len", "HEAD_DIM", "K_W", "APPLY_L2", "REPEAT", "NUM_GROUPS"], +) +@triton.jit +def _conv_silu_project_kernel( + qkvzba_ptr, + weight_ptr, + bias_ptr, + out_ptr, + silu_save_ptr, + seq_len, + num_in_heads, + in_channel_offset, + in_group_stride, + silu_save_chan_offset, + silu_save_group_stride, + qkvzba_s_stride, + qkvzba_b_stride, + qkvzba_c_stride, + weight_c_stride, + weight_w_stride, + bias_stride, + out_group_dim_stride, + out_b_stride, + out_s_stride, + out_h_stride, + silu_save_b_stride, + silu_save_c_stride, + silu_save_s_stride, + eps, + HEAD_DIM: tl.constexpr, + K_W: tl.constexpr, + REPEAT: tl.constexpr, + NUM_GROUPS: tl.constexpr, + HAS_BIAS: tl.constexpr, + APPLY_L2: tl.constexpr, + SAVE_SILU: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """Depthwise conv1d + silu + (optional l2norm) + (optional head repeat). + + Grid layout (program_id): + 0: batch * NUM_GROUPS * num_in_heads (flat) + 1: num_seq_blocks + + Args: + in_channel_offset: starting channel index of the first group inside + ``qkvzba``. 0 for QK, ``v_channel_offset`` for V. + in_group_stride: channel distance between logical groups. For QK this + is ``qk_channels`` so group 0 is Q and group 1 is K. For V this is + 0 because ``NUM_GROUPS == 1``. + out_group_dim_stride: output-storage distance between logical groups. QK + passes a grouped output buffer and V passes 0. + """ + + pid_bgh = tl.program_id(0) + pid_s = tl.program_id(1) + + heads_per_batch = num_in_heads * NUM_GROUPS + batch_id = pid_bgh // heads_per_batch + local_bgh = pid_bgh - batch_id * heads_per_batch + group_id = local_bgh // num_in_heads + head_id = local_bgh - group_id * num_in_heads + + chan_off = tl.arange(0, HEAD_DIM) + group_channel_offset = in_channel_offset + group_id * in_group_stride + chan = group_channel_offset + head_id * HEAD_DIM + chan_off + + if HAS_BIAS: + bias = tl.load(bias_ptr + chan * bias_stride).to(tl.float32) + else: + bias = tl.zeros([HEAD_DIM], dtype=tl.float32) + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + + acc = tl.zeros([BLOCK_S, HEAD_DIM], dtype=tl.float32) + for i in tl.static_range(K_W): + x_s = s_offs - (K_W - 1) + i + x_mask = (x_s >= 0) & (x_s < seq_len) + x_ptr = ( + qkvzba_ptr + + x_s[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + chan[None, :] * qkvzba_c_stride + ) + x_val = tl.load(x_ptr, mask=x_mask[:, None], other=0.0).to(tl.float32) + w_tap = tl.load( + weight_ptr + chan * weight_c_stride + i * weight_w_stride + ).to(tl.float32) + acc += w_tap[None, :] * x_val + + acc += bias[None, :] + # Mimic the unfused F.conv1d rounding: the reference path stores the conv + # output in the input dtype (bf16) before silu, so do the same here. This + # keeps the fused output bit-aligned with the reference within one ULP. + acc = acc.to(out_ptr.dtype.element_ty).to(tl.float32) + silu_out = acc * tl.sigmoid(acc) + + if APPLY_L2: + # F.silu rounds to the input dtype before l2norm reads it. Round-trip + # via bf16 to match that precision. + silu_out = silu_out.to(out_ptr.dtype.element_ty).to(tl.float32) + if SAVE_SILU: + # Persist only the QK silu output in the channel-last layout + # consumed by the QK l2norm backward. + silu_save_chan = ( + silu_save_chan_offset + + group_id * silu_save_group_stride + + head_id * HEAD_DIM + + chan_off + ) + silu_save_ptrs = ( + silu_save_ptr + + batch_id * silu_save_b_stride + + silu_save_chan[None, :] * silu_save_c_stride + + s_offs[:, None] * silu_save_s_stride + ) + tl.store( + silu_save_ptrs, + silu_out.to(silu_save_ptr.dtype.element_ty), + mask=s_mask[:, None], + ) + norm_sq = tl.sum(silu_out * silu_out, axis=1) + rstd = 1.0 / tl.sqrt(norm_sq + eps) + out = silu_out * rstd[:, None] + else: + # No l2norm follows. The final store→bf16 already does the rounding; + # an intermediate bf16 round-trip would be redundant. + out = silu_out + + out_typed = out.to(out_ptr.dtype.element_ty) + + # Write the same data to ``REPEAT`` adjacent value heads. ``REPEAT == 1`` + # is the no-repeat case (V branch is handled by a separate kernel that + # always has REPEAT == 1, but using the same code here is convenient). + for r in tl.static_range(REPEAT): + v_head = head_id * REPEAT + r + write_ptr = ( + out_ptr + + group_id * out_group_dim_stride + + batch_id * out_b_stride + + s_offs[:, None] * out_s_stride + + v_head * out_h_stride + + chan_off[None, :] + ) + tl.store(write_ptr, out_typed, mask=s_mask[:, None]) + + +@triton.jit +def _thd_seq_bounds(cu_seqlens_ptr, token_offsets, total_tokens, num_packed_seqs): + """Return lane-wise packed sequence bounds for flattened THD tokens.""" + + safe_tokens = tl.minimum(token_offsets, total_tokens - 1) + seq_start = token_offsets * 0 + seq_end = token_offsets * 0 + total_tokens + + seq_id = 0 + while seq_id < num_packed_seqs: + start = tl.load(cu_seqlens_ptr + seq_id) + end = tl.load(cu_seqlens_ptr + seq_id + 1) + in_seq = (safe_tokens >= start) & (safe_tokens < end) + seq_start = tl.where(in_seq, start, seq_start) + seq_end = tl.where(in_seq, end, seq_end) + seq_id += 1 + + return seq_start, seq_end + + +@triton.autotune( + configs=_conv_autotune_configs(), + key=["seq_len", "HEAD_DIM", "K_W", "APPLY_L2", "REPEAT", "NUM_GROUPS"], +) +@triton.jit +def _conv_silu_project_thd_kernel( + qkvzba_ptr, + weight_ptr, + bias_ptr, + out_ptr, + silu_save_ptr, + cu_seqlens_ptr, + seq_len, + num_packed_seqs, + num_in_heads, + in_channel_offset, + in_group_stride, + silu_save_chan_offset, + silu_save_group_stride, + qkvzba_s_stride, + qkvzba_b_stride, + qkvzba_c_stride, + weight_c_stride, + weight_w_stride, + bias_stride, + out_group_dim_stride, + out_b_stride, + out_s_stride, + out_h_stride, + silu_save_b_stride, + silu_save_c_stride, + silu_save_s_stride, + eps, + HEAD_DIM: tl.constexpr, + K_W: tl.constexpr, + REPEAT: tl.constexpr, + NUM_GROUPS: tl.constexpr, + HAS_BIAS: tl.constexpr, + APPLY_L2: tl.constexpr, + SAVE_SILU: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """THD depthwise conv1d + silu + optional l2norm/repeat. + + This is intentionally separate from ``_conv_silu_project_kernel`` so + packed sequence boundary metadata never enters the dense BSHD hot path. + Only the causal-conv loads use ``cu_seqlens``; the following per-token + transforms and stores are identical to the dense path. + """ + + pid_bgh = tl.program_id(0) + pid_s = tl.program_id(1) + + heads_per_batch = num_in_heads * NUM_GROUPS + batch_id = pid_bgh // heads_per_batch + local_bgh = pid_bgh - batch_id * heads_per_batch + group_id = local_bgh // num_in_heads + head_id = local_bgh - group_id * num_in_heads + + chan_off = tl.arange(0, HEAD_DIM) + group_channel_offset = in_channel_offset + group_id * in_group_stride + chan = group_channel_offset + head_id * HEAD_DIM + chan_off + + if HAS_BIAS: + bias = tl.load(bias_ptr + chan * bias_stride).to(tl.float32) + else: + bias = tl.zeros([HEAD_DIM], dtype=tl.float32) + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + seq_start, seq_end = _thd_seq_bounds(cu_seqlens_ptr, s_offs, seq_len, num_packed_seqs) + + acc = tl.zeros([BLOCK_S, HEAD_DIM], dtype=tl.float32) + for i in tl.static_range(K_W): + x_s = s_offs - (K_W - 1) + i + x_mask = s_mask & (x_s >= seq_start) & (x_s < seq_end) + safe_x_s = tl.minimum(tl.maximum(x_s, 0), seq_len - 1) + x_ptr = ( + qkvzba_ptr + + safe_x_s[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + chan[None, :] * qkvzba_c_stride + ) + x_val = tl.load(x_ptr, mask=x_mask[:, None], other=0.0).to(tl.float32) + w_tap = tl.load( + weight_ptr + chan * weight_c_stride + i * weight_w_stride + ).to(tl.float32) + acc += w_tap[None, :] * x_val + + acc += bias[None, :] + acc = acc.to(out_ptr.dtype.element_ty).to(tl.float32) + silu_out = acc * tl.sigmoid(acc) + + if APPLY_L2: + silu_out = silu_out.to(out_ptr.dtype.element_ty).to(tl.float32) + if SAVE_SILU: + silu_save_chan = ( + silu_save_chan_offset + + group_id * silu_save_group_stride + + head_id * HEAD_DIM + + chan_off + ) + silu_save_ptrs = ( + silu_save_ptr + + batch_id * silu_save_b_stride + + silu_save_chan[None, :] * silu_save_c_stride + + s_offs[:, None] * silu_save_s_stride + ) + tl.store( + silu_save_ptrs, + silu_out.to(silu_save_ptr.dtype.element_ty), + mask=s_mask[:, None], + ) + norm_sq = tl.sum(silu_out * silu_out, axis=1) + rstd = 1.0 / tl.sqrt(norm_sq + eps) + out = silu_out * rstd[:, None] + else: + out = silu_out + + out_typed = out.to(out_ptr.dtype.element_ty) + + for r in tl.static_range(REPEAT): + v_head = head_id * REPEAT + r + write_ptr = ( + out_ptr + + group_id * out_group_dim_stride + + batch_id * out_b_stride + + s_offs[:, None] * out_s_stride + + v_head * out_h_stride + + chan_off[None, :] + ) + tl.store(write_ptr, out_typed, mask=s_mask[:, None]) + + +@triton.jit +def _copy_z_kernel( + qkvzba_ptr, + gate_ptr, + seq_len, + num_v_heads, + z_channel_offset, + qkvzba_s_stride, + qkvzba_b_stride, + qkvzba_c_stride, + gate_b_stride, + gate_s_stride, + gate_h_stride, + HEAD_DIM: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """Copy the z slice from qkvzba into the final gate layout.""" + + pid_bh = tl.program_id(0) + pid_s = tl.program_id(1) + + batch_id = pid_bh // num_v_heads + head_id = pid_bh - batch_id * num_v_heads + + chan_off = tl.arange(0, HEAD_DIM) + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + + z_chan = z_channel_offset + head_id * HEAD_DIM + chan_off + z_src_ptr = ( + qkvzba_ptr + + s_offs[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + z_chan[None, :] * qkvzba_c_stride + ) + z_val = tl.load(z_src_ptr, mask=s_mask[:, None]) + z_write_ptr = ( + gate_ptr + + batch_id * gate_b_stride + + s_offs[:, None] * gate_s_stride + + head_id * gate_h_stride + + chan_off[None, :] + ) + tl.store(z_write_ptr, z_val, mask=s_mask[:, None]) + + +@triton.autotune(configs=_g_beta_autotune_configs(), key=["seq_len", "num_v_heads"]) +@triton.jit +def _compute_g_and_beta_kernel( + qkvzba_ptr, + A_log_ptr, + dt_bias_ptr, + g_out_ptr, + beta_out_ptr, + seq_len, + num_v_heads, + beta_channel_offset, + alpha_channel_offset, + qkvzba_s_stride, + qkvzba_b_stride, + qkvzba_c_stride, + g_b_stride, + g_s_stride, + g_h_stride, + beta_b_stride, + beta_s_stride, + beta_h_stride, + BLOCK_S: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """Compute ``g = -exp(A_log) * softplus(alpha + dt_bias)`` and ``sigmoid(beta)``.""" + + pid_b = tl.program_id(0) + pid_s = tl.program_id(1) + pid_h = tl.program_id(2) + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + s_mask = s_offs < seq_len + h_mask = h_offs < num_v_heads + mask = s_mask[:, None] & h_mask[None, :] + + alpha_ptr = ( + qkvzba_ptr + + s_offs[:, None] * qkvzba_s_stride + + pid_b * qkvzba_b_stride + + (alpha_channel_offset + h_offs[None, :]) * qkvzba_c_stride + ) + beta_ptr = ( + qkvzba_ptr + + s_offs[:, None] * qkvzba_s_stride + + pid_b * qkvzba_b_stride + + (beta_channel_offset + h_offs[None, :]) * qkvzba_c_stride + ) + + alpha = tl.load(alpha_ptr, mask=mask, other=0.0).to(tl.float32) + beta = tl.load(beta_ptr, mask=mask, other=0.0).to(tl.float32) + + A_log = tl.load(A_log_ptr + h_offs, mask=h_mask, other=0.0).to(tl.float32) + dt_bias = tl.load(dt_bias_ptr + h_offs, mask=h_mask, other=0.0).to(tl.float32) + + pre = alpha + dt_bias[None, :] + # softplus(x) = log(1 + exp(x)); torch's softplus thresholds at x>20 but we + # rely on fp32 evaluation here, which stays well within range for typical + # GDN inputs (the unfused path computes the same expression). + softplus_val = tl.log(1.0 + tl.exp(pre)) + g = -tl.exp(A_log)[None, :] * softplus_val + beta_sig = tl.sigmoid(beta) + + g_ptr = ( + g_out_ptr + + pid_b * g_b_stride + + s_offs[:, None] * g_s_stride + + h_offs[None, :] * g_h_stride + ) + beta_out_ptr_calc = ( + beta_out_ptr + + pid_b * beta_b_stride + + s_offs[:, None] * beta_s_stride + + h_offs[None, :] * beta_h_stride + ) + tl.store(g_ptr, g.to(g_out_ptr.dtype.element_ty), mask=mask) + tl.store(beta_out_ptr_calc, beta_sig.to(beta_out_ptr.dtype.element_ty), mask=mask) + + +# --------------------------------------------------------------------------- +# Backward kernels +# --------------------------------------------------------------------------- + + +@triton.jit +def _conv_silu_l2norm_backward_kernel( + qkvzba_ptr, + weight_ptr, + d_out_ptr, + d_qkvzba_ptr, + d_w_partial_ptr, + seq_len, + num_qk_heads, + in_channel_offset, + eps, + d_out_scale, + qkvzba_s_stride, + qkvzba_b_stride, + qkvzba_c_stride, + weight_c_stride, + weight_w_stride, + d_out_b_stride, + d_out_s_stride, + d_out_h_stride, + d_wp_b_stride, + d_wp_h_stride, + d_wp_s_stride, + d_wp_c_stride, + d_wp_w_stride, + HEAD_DIM: tl.constexpr, + K_W: tl.constexpr, + REPEAT: tl.constexpr, + BLOCK_S: tl.constexpr, + USE_L2NORM: tl.constexpr, + V_HEAD_SHARED: tl.constexpr, +): + """Backward for the Q / K / V branches of ``_conv_silu_project_kernel``. + + ``USE_L2NORM`` is a constexpr branch: ``True`` for the QK branches (with + l2norm) and ``False`` for the V branch. The V case skips the l2norm + intermediates entirely — Triton DCE drops them at compile time. The + ``REPEAT=2`` workaround for the channel-collapse codegen bug still + applies in both branches. + + Forward (no bias, with l2norm, with REPEAT-way head broadcast): + acc = depthwise_conv(qkvzba_qk_slice, weight_qk_slice) + acc_bf16 = acc.to(bf16).to(fp32) # F.conv1d rounding + silu_out = acc_bf16 * sigmoid(acc_bf16) + silu_bf16 = silu_out.to(bf16).to(fp32) # round before l2norm + norm_sq = sum_c silu_bf16^2 + rstd = 1 / sqrt(norm_sq + eps) + out = silu_bf16 * rstd + # out is stored identically to REPEAT adjacent value heads. + + Backward (given d_out for each v_head): + d_qk_out = Σ_{r in REPEAT} d_v_out[head_id * REPEAT + r] + S = Σ_c d_qk_out_c * silu_bf16_c + d_silu_c = rstd * d_qk_out_c - rstd^3 * silu_bf16_c * S + d_acc_c = d_silu_c * silu'(acc_bf16) + d_w[c, i] = Σ_{b, t} d_acc[t, c] * x[t + i - (K_W - 1), c] + d_x[u, c] += Σ_i d_acc[u + (K_W - 1) - i, c] * w[c, i] + + ``d_w`` uses the per-program partial-buffer pattern (see the V backward + kernel). ``d_qkvzba`` uses bf16 atomic_add because the K_W − 1 boundary + input rows cross seq-block programs. + """ + + pid_bh = tl.program_id(0) + pid_s = tl.program_id(1) + + batch_id = pid_bh // num_qk_heads + head_id = pid_bh - batch_id * num_qk_heads + + chan_off = tl.arange(0, HEAD_DIM) + chan = in_channel_offset + head_id * HEAD_DIM + chan_off + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + + # ----- Forward recompute (conv + silu + l2norm) ----- + acc = tl.zeros([BLOCK_S, HEAD_DIM], dtype=tl.float32) + for i in tl.static_range(K_W): + x_s = s_offs - (K_W - 1) + i + x_mask = (x_s >= 0) & (x_s < seq_len) + x_ptr = ( + qkvzba_ptr + + x_s[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + chan[None, :] * qkvzba_c_stride + ) + x_val = tl.load(x_ptr, mask=x_mask[:, None], other=0.0).to(tl.float32) + w_tap = tl.load( + weight_ptr + chan * weight_c_stride + i * weight_w_stride + ).to(tl.float32) + acc += w_tap[None, :] * x_val + acc = acc.to(d_qkvzba_ptr.dtype.element_ty).to(tl.float32) + + # ----- Sum d_out across REPEAT v_heads ----- + # For QK: v_head = head_id*REPEAT + r, summing REPEAT distinct heads. + # For V (V_HEAD_SHARED=True): both r iterations load the SAME v_head + # (head_id), so d_qk_out = REPEAT * d_value[head_id]; the host passes + # d_out_scale = 1/REPEAT to recover d_value[head_id]. The duplicate + # load goes through L2, so the kernel-side cost is roughly one load; + # the trick was needed to avoid the REPEAT=1 codegen bug without + # having to allocate a doubled d_value tensor on the host. + d_qk_out = tl.zeros([BLOCK_S, HEAD_DIM], dtype=tl.float32) + for r in tl.static_range(REPEAT): + if V_HEAD_SHARED: + v_head = head_id + else: + v_head = head_id * REPEAT + r + d_out_ptrs = ( + d_out_ptr + + batch_id * d_out_b_stride + + s_offs[:, None] * d_out_s_stride + + v_head * d_out_h_stride + + chan_off[None, :] + ) + d_qk_out += tl.load(d_out_ptrs, mask=s_mask[:, None], other=0.0).to(tl.float32) + d_qk_out = d_qk_out * d_out_scale + + # ----- l2norm backward gated by USE_L2NORM constexpr. ----- + if USE_L2NORM: + silu_out = acc * tl.sigmoid(acc) + silu_bf16 = silu_out.to(d_qkvzba_ptr.dtype.element_ty).to(tl.float32) + norm_sq = tl.sum(silu_bf16 * silu_bf16, axis=1) + rstd = 1.0 / tl.sqrt(norm_sq + eps) + s_row = tl.sum(d_qk_out * silu_bf16, axis=1) + rstd3 = rstd * rstd * rstd + d_silu = rstd[:, None] * d_qk_out - rstd3[:, None] * silu_bf16 * s_row[:, None] + else: + d_silu = d_qk_out + + # ----- silu backward ----- + sig_acc = tl.sigmoid(acc) + silu_prime = sig_acc + acc * sig_acc * (1.0 - sig_acc) + d_acc = d_silu * silu_prime + d_acc = tl.where(s_mask[:, None], d_acc, 0.0) + + # ----- d_w via per-program partial, d_x via atomic_add ----- + partial_base = ( + d_w_partial_ptr + + batch_id * d_wp_b_stride + + head_id * d_wp_h_stride + + pid_s * d_wp_s_stride + ) + + for i in tl.static_range(K_W): + x_s = s_offs - (K_W - 1) + i + x_mask_inner = (x_s >= 0) & (x_s < seq_len) + x_ptr = ( + qkvzba_ptr + + x_s[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + chan[None, :] * qkvzba_c_stride + ) + x_val = tl.load(x_ptr, mask=x_mask_inner[:, None], other=0.0).to(tl.float32) + d_w_partial = tl.sum(d_acc * x_val, axis=0) + tl.store( + partial_base + chan_off * d_wp_c_stride + i * d_wp_w_stride, + d_w_partial, + ) + + w_tap = tl.load( + weight_ptr + chan * weight_c_stride + i * weight_w_stride + ).to(tl.float32) + contribution = d_acc * w_tap[None, :] + d_qkvzba_target = ( + d_qkvzba_ptr + + x_s[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + chan[None, :] * qkvzba_c_stride + ) + tl.atomic_add( + d_qkvzba_target, + contribution.to(d_qkvzba_ptr.dtype.element_ty), + mask=x_mask_inner[:, None], + ) + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_S": 32}, num_warps=2, num_stages=2), + triton.Config({"BLOCK_S": 32}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_S": 64}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=8, num_stages=3), + triton.Config({"BLOCK_S": 256}, num_warps=8, num_stages=2), + ], + key=["seq_len", "HEAD_DIM", "REPEAT"], +) +@triton.jit +def _l2norm_repeat_backward_kernel( + d_qk_out_ptr, # (b, s, num_v_heads, head_dim) — gradient from downstream + silu_bf16_ptr, # (b, conv_dim, s) — silu(conv(x)) recomputed for QK channels + d_silu_bf16_ptr, # (b, conv_dim, s) — output gradient w.r.t. silu(conv(x)) + seq_len, + num_qk_heads, + channel_offset, # 0 for Q, qk_channels for K — indexes into conv_dim + eps, + d_qk_b_stride, + d_qk_s_stride, + d_qk_h_stride, + silu_b_stride, + silu_c_stride, + silu_s_stride, + d_silu_b_stride, + d_silu_c_stride, + d_silu_s_stride, + HEAD_DIM: tl.constexpr, + REPEAT: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """l2norm + REPEAT-way head broadcast backward. + + Forward (per QK head): + silu_bf16 ∈ R^{HEAD_DIM} # silu(conv(x)) rounded to bf16 + norm_sq = Σ_c silu_bf16_c^2 + rstd = 1 / sqrt(norm_sq + eps) + out = silu_bf16 * rstd + # out is broadcast identically to REPEAT adjacent v_heads. + + Backward (given d_qk_out for each v_head): + d_normed = Σ_{r in REPEAT} d_qk_out[head_id * REPEAT + r] + S = Σ_c d_normed_c * silu_bf16_c + d_silu_c = rstd * d_normed_c - rstd^3 * silu_bf16_c * S + + The output ``d_silu_bf16`` is the gradient w.r.t. ``silu(conv(x))`` — + exactly what ``causal_conv1d_bwd_function`` consumes as its ``dout`` + argument when ``activation="silu"`` is in effect on the forward. + """ + + pid_bh = tl.program_id(0) + pid_s = tl.program_id(1) + + batch_id = pid_bh // num_qk_heads + head_id = pid_bh - batch_id * num_qk_heads + + chan_off = tl.arange(0, HEAD_DIM) + chan = channel_offset + head_id * HEAD_DIM + chan_off + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + + # ----- Sum d_qk_out across REPEAT v_heads ----- + d_normed = tl.zeros([BLOCK_S, HEAD_DIM], dtype=tl.float32) + for r in tl.static_range(REPEAT): + v_head = head_id * REPEAT + r + d_out_ptrs = ( + d_qk_out_ptr + + batch_id * d_qk_b_stride + + s_offs[:, None] * d_qk_s_stride + + v_head * d_qk_h_stride + + chan_off[None, :] + ) + d_normed += tl.load(d_out_ptrs, mask=s_mask[:, None], other=0.0).to(tl.float32) + + # ----- Load silu_bf16 ----- + silu_ptrs = ( + silu_bf16_ptr + + batch_id * silu_b_stride + + chan[None, :] * silu_c_stride + + s_offs[:, None] * silu_s_stride + ) + silu_bf16 = tl.load(silu_ptrs, mask=s_mask[:, None], other=0.0).to(tl.float32) + + # ----- l2norm backward ----- + norm_sq = tl.sum(silu_bf16 * silu_bf16, axis=1) + rstd = 1.0 / tl.sqrt(norm_sq + eps) + s_row = tl.sum(d_normed * silu_bf16, axis=1) + rstd3 = rstd * rstd * rstd + d_silu = rstd[:, None] * d_normed - rstd3[:, None] * silu_bf16 * s_row[:, None] + + # ----- Store d_silu (same (b, conv_dim, s) layout as silu_bf16_ptr) ----- + d_silu_ptrs = ( + d_silu_bf16_ptr + + batch_id * d_silu_b_stride + + chan[None, :] * d_silu_c_stride + + s_offs[:, None] * d_silu_s_stride + ) + tl.store(d_silu_ptrs, d_silu.to(d_silu_bf16_ptr.dtype.element_ty), mask=s_mask[:, None]) + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_S": 32}, num_warps=2, num_stages=2), + triton.Config({"BLOCK_S": 32}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_S": 64}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=8, num_stages=3), + triton.Config({"BLOCK_S": 256}, num_warps=8, num_stages=2), + ], + key=["seq_len", "HEAD_DIM", "REPEAT"], +) +@triton.jit +def _qk_l2norm_repeat_backward_kernel( + dq_ptr, + dk_ptr, + silu_bf16_ptr, + d_silu_bf16_ptr, + seq_len, + num_qk_heads, + qk_channels, + eps, + dq_b_stride, + dq_s_stride, + dq_h_stride, + dk_b_stride, + dk_s_stride, + dk_h_stride, + silu_b_stride, + silu_c_stride, + silu_s_stride, + d_silu_b_stride, + d_silu_c_stride, + d_silu_s_stride, + HEAD_DIM: tl.constexpr, + REPEAT: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """Merged Q/K l2norm + REPEAT-way head broadcast backward.""" + + pid_bgh = tl.program_id(0) + pid_s = tl.program_id(1) + + heads_per_batch = num_qk_heads * 2 + batch_id = pid_bgh // heads_per_batch + local_bgh = pid_bgh - batch_id * heads_per_batch + group_id = local_bgh // num_qk_heads + head_id = local_bgh - group_id * num_qk_heads + is_query = group_id == 0 + is_key = group_id == 1 + + chan_off = tl.arange(0, HEAD_DIM) + chan = group_id * qk_channels + head_id * HEAD_DIM + chan_off + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + + d_normed = tl.zeros([BLOCK_S, HEAD_DIM], dtype=tl.float32) + for r in tl.static_range(REPEAT): + v_head = head_id * REPEAT + r + dq_ptrs = ( + dq_ptr + + batch_id * dq_b_stride + + s_offs[:, None] * dq_s_stride + + v_head * dq_h_stride + + chan_off[None, :] + ) + dk_ptrs = ( + dk_ptr + + batch_id * dk_b_stride + + s_offs[:, None] * dk_s_stride + + v_head * dk_h_stride + + chan_off[None, :] + ) + d_normed += tl.load( + dq_ptrs, mask=s_mask[:, None] & is_query, other=0.0 + ).to(tl.float32) + d_normed += tl.load( + dk_ptrs, mask=s_mask[:, None] & is_key, other=0.0 + ).to(tl.float32) + + silu_ptrs = ( + silu_bf16_ptr + + batch_id * silu_b_stride + + chan[None, :] * silu_c_stride + + s_offs[:, None] * silu_s_stride + ) + silu_bf16 = tl.load(silu_ptrs, mask=s_mask[:, None], other=0.0).to(tl.float32) + + norm_sq = tl.sum(silu_bf16 * silu_bf16, axis=1) + rstd = 1.0 / tl.sqrt(norm_sq + eps) + s_row = tl.sum(d_normed * silu_bf16, axis=1) + rstd3 = rstd * rstd * rstd + d_silu = rstd[:, None] * d_normed - rstd3[:, None] * silu_bf16 * s_row[:, None] + + d_silu_ptrs = ( + d_silu_bf16_ptr + + batch_id * d_silu_b_stride + + chan[None, :] * d_silu_c_stride + + s_offs[:, None] * d_silu_s_stride + ) + tl.store(d_silu_ptrs, d_silu.to(d_silu_bf16_ptr.dtype.element_ty), mask=s_mask[:, None]) + + +@triton.jit +def _v_layout_to_conv_kernel( + dv_ptr, # (b, s, num_v_heads, value_head_dim) + d_silu_conv_ptr, # (b, conv_dim, s) — write into V channel slice + seq_len, + num_v_heads, + v_channel_offset, # = 2 * qk_channels + dv_b_stride, + dv_s_stride, + dv_h_stride, + d_silu_b_stride, + d_silu_c_stride, + d_silu_s_stride, + HEAD_DIM: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """Write V-branch gradients into the conv-backward layout. + + ``dv`` is the gradient of ``value`` (forward layout + ``(b, s, num_v_heads, value_head_dim)``). The conv backward needs + ``d_silu_conv`` in layout ``(b, conv_dim, s)`` for the V channel + slice. + """ + + pid_bh = tl.program_id(0) + pid_s = tl.program_id(1) + + batch_id = pid_bh // num_v_heads + head_id = pid_bh - batch_id * num_v_heads + + chan_off = tl.arange(0, HEAD_DIM) + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + + # Read dv at (batch, s, head, chan). + dv_ptrs = ( + dv_ptr + + batch_id * dv_b_stride + + s_offs[:, None] * dv_s_stride + + head_id * dv_h_stride + + chan_off[None, :] + ) + dv_val = tl.load(dv_ptrs, mask=s_mask[:, None], other=0.0) + + # Write to d_silu_conv at (batch, v_channel_offset + head*HEAD_DIM + chan, s). + d_silu_chan = v_channel_offset + head_id * HEAD_DIM + chan_off + d_silu_ptrs = ( + d_silu_conv_ptr + + batch_id * d_silu_b_stride + + d_silu_chan[None, :] * d_silu_c_stride + + s_offs[:, None] * d_silu_s_stride + ) + tl.store(d_silu_ptrs, dv_val, mask=s_mask[:, None]) + + +@triton.jit +def _z_layout_to_qkvzba_kernel( + dgate_ptr, # (b, s, num_v_heads, value_head_dim) + d_qkvzba_ptr, # (s, b, total_channels) — write into z channel slice + seq_len, + num_v_heads, + z_channel_offset, # = 2 * qk_channels + v_channels + dgate_b_stride, + dgate_s_stride, + dgate_h_stride, + d_qkvzba_s_stride, + d_qkvzba_b_stride, + d_qkvzba_c_stride, + HEAD_DIM: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """Write gate gradients into the z slice of ``d_qkvzba``. + + ``dgate`` is the autograd-supplied gradient of ``gate`` (= the z + slice of qkvzba in forward) with layout + ``(b, s, num_v_heads, value_head_dim)``. We need to write it into + ``d_qkvzba``'s z slice — layout ``(s, b, total_channels)`` with + channels in ``[z_channel_offset, z_channel_offset + v_channels)``. + """ + + pid_bh = tl.program_id(0) + pid_s = tl.program_id(1) + + batch_id = pid_bh // num_v_heads + head_id = pid_bh - batch_id * num_v_heads + + chan_off = tl.arange(0, HEAD_DIM) + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + + # Read dgate at (batch, s, head, chan). + dgate_ptrs = ( + dgate_ptr + + batch_id * dgate_b_stride + + s_offs[:, None] * dgate_s_stride + + head_id * dgate_h_stride + + chan_off[None, :] + ) + dgate_val = tl.load(dgate_ptrs, mask=s_mask[:, None], other=0.0) + + # Write to d_qkvzba at (s, batch, z_channel_offset + head*HEAD_DIM + chan). + d_qkvzba_chan = z_channel_offset + head_id * HEAD_DIM + chan_off + d_qkvzba_ptrs = ( + d_qkvzba_ptr + + s_offs[:, None] * d_qkvzba_s_stride + + batch_id * d_qkvzba_b_stride + + d_qkvzba_chan[None, :] * d_qkvzba_c_stride + ) + tl.store(d_qkvzba_ptrs, dgate_val, mask=s_mask[:, None]) + + +@triton.autotune( + configs=_g_beta_autotune_configs(), + key=["seq_len", "num_v_heads"], + # Each autotune trial atomic-adds partial sums into these accumulators. + # Without reset_to_zero the trials would stack on top of one another and + # produce values that are ``num_trials`` × the correct result. + reset_to_zero=["d_A_log_ptr", "d_dt_bias_ptr"], +) +@triton.jit +def _g_beta_backward_kernel( + qkvzba_ptr, + A_log_ptr, + dt_bias_ptr, + d_g_ptr, + d_beta_out_ptr, + d_qkvzba_ptr, + d_A_log_ptr, + d_dt_bias_ptr, + seq_len, + num_v_heads, + beta_channel_offset, + alpha_channel_offset, + qkvzba_s_stride, + qkvzba_b_stride, + qkvzba_c_stride, + d_g_b_stride, + d_g_s_stride, + d_g_h_stride, + d_beta_b_stride, + d_beta_s_stride, + d_beta_h_stride, + BLOCK_S: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """Backward for ``_compute_g_and_beta_kernel``. + + Forward: + pre = alpha + dt_bias # fp32 + softplus_pre = log(1 + exp(pre)) + g = -exp(A_log) * softplus_pre + beta_sig = sigmoid(beta_raw) + + Backward (given d_g and d_beta_out): + d_alpha = d_g * (-exp(A_log) * sigmoid(pre)) + d_beta_raw = d_beta_out * beta_sig * (1 - beta_sig) + d_dt_bias[h] = Σ_{b,s} d_alpha[b,s,h] + d_A_log[h] = Σ_{b,s} d_g[b,s,h] * g[b,s,h] + + ``d_alpha`` and ``d_beta_raw`` are written into the matching channel slices + of ``d_qkvzba``. ``d_A_log`` and ``d_dt_bias`` are reduced via per-element + atomic_add to fp32 buffers; the caller casts those to the parameter dtype. + """ + + pid_b = tl.program_id(0) + pid_s = tl.program_id(1) + pid_h = tl.program_id(2) + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + s_mask = s_offs < seq_len + h_mask = h_offs < num_v_heads + mask = s_mask[:, None] & h_mask[None, :] + + # ----- Forward recompute ----- + alpha_ptr = ( + qkvzba_ptr + + s_offs[:, None] * qkvzba_s_stride + + pid_b * qkvzba_b_stride + + (alpha_channel_offset + h_offs[None, :]) * qkvzba_c_stride + ) + beta_ptr = ( + qkvzba_ptr + + s_offs[:, None] * qkvzba_s_stride + + pid_b * qkvzba_b_stride + + (beta_channel_offset + h_offs[None, :]) * qkvzba_c_stride + ) + alpha = tl.load(alpha_ptr, mask=mask, other=0.0).to(tl.float32) + beta_raw = tl.load(beta_ptr, mask=mask, other=0.0).to(tl.float32) + A_log = tl.load(A_log_ptr + h_offs, mask=h_mask, other=0.0).to(tl.float32) + dt_bias = tl.load(dt_bias_ptr + h_offs, mask=h_mask, other=0.0).to(tl.float32) + + pre = alpha + dt_bias[None, :] + sigmoid_pre = tl.sigmoid(pre) + softplus_pre = tl.log(1.0 + tl.exp(pre)) + exp_A = tl.exp(A_log)[None, :] + g = -exp_A * softplus_pre + beta_sig = tl.sigmoid(beta_raw) + + # ----- Load upstream gradients ----- + d_g_ptrs = ( + d_g_ptr + + pid_b * d_g_b_stride + + s_offs[:, None] * d_g_s_stride + + h_offs[None, :] * d_g_h_stride + ) + d_beta_out_ptrs = ( + d_beta_out_ptr + + pid_b * d_beta_b_stride + + s_offs[:, None] * d_beta_s_stride + + h_offs[None, :] * d_beta_h_stride + ) + d_g = tl.load(d_g_ptrs, mask=mask, other=0.0).to(tl.float32) + d_beta_out = tl.load(d_beta_out_ptrs, mask=mask, other=0.0).to(tl.float32) + + # ----- Per-element gradients ----- + d_alpha = d_g * (-exp_A * sigmoid_pre) + d_beta_raw = d_beta_out * beta_sig * (1.0 - beta_sig) + + # ----- (b, s) → h reductions ----- + d_g_masked = tl.where(mask, d_g, 0.0) + d_alpha_masked = tl.where(mask, d_alpha, 0.0) + d_A_log_partial = tl.sum(d_g_masked * g, axis=0) + d_dt_bias_partial = tl.sum(d_alpha_masked, axis=0) + + # ----- Store per-element grads back to d_qkvzba ----- + d_alpha_ptrs = ( + d_qkvzba_ptr + + s_offs[:, None] * qkvzba_s_stride + + pid_b * qkvzba_b_stride + + (alpha_channel_offset + h_offs[None, :]) * qkvzba_c_stride + ) + d_beta_ptrs = ( + d_qkvzba_ptr + + s_offs[:, None] * qkvzba_s_stride + + pid_b * qkvzba_b_stride + + (beta_channel_offset + h_offs[None, :]) * qkvzba_c_stride + ) + tl.store( + d_alpha_ptrs, d_alpha.to(d_qkvzba_ptr.dtype.element_ty), mask=mask + ) + tl.store( + d_beta_ptrs, d_beta_raw.to(d_qkvzba_ptr.dtype.element_ty), mask=mask + ) + + # ----- Atomic-add (b, s) partials into per-head accumulators ----- + tl.atomic_add(d_A_log_ptr + h_offs, d_A_log_partial, mask=h_mask) + tl.atomic_add(d_dt_bias_ptr + h_offs, d_dt_bias_partial, mask=h_mask) + + +# --------------------------------------------------------------------------- +# Python entry points +# --------------------------------------------------------------------------- + + + + +def _is_power_of_two(value: int) -> bool: + return value > 0 and (value & (value - 1)) == 0 + + +_SIDE_STREAMS: dict = {} + + +def _get_side_stream(device: torch.device, slot: int) -> "torch.cuda.Stream": + """Lazily allocate and cache CUDA streams keyed by ``(device, slot)``. + + Reusing streams across calls keeps launches free of stream-creation + overhead, which would otherwise dominate the small kernels. + """ + + key = (device.index if device.index is not None else torch.cuda.current_device(), slot) + stream = _SIDE_STREAMS.get(key) + if stream is None: + stream = torch.cuda.Stream(device=device) + _SIDE_STREAMS[key] = stream + return stream + + +def _triton_l2norm_repeat_backward( + d_qk_out: Tensor, + silu_bf16: Tensor, + d_silu_bf16: Tensor, + *, + is_query: bool, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + eps: float = 1e-6, + stream: Optional["torch.cuda.Stream"] = None, +) -> Tensor: + """l2norm + REPEAT backward. + + ``silu_bf16`` is the (b, conv_dim, s) bf16 tensor produced by re-running + causal_conv1d_fn (forward, no-grad). Output ``d_silu_bf16`` is written + in place; only the matching channel slice (Q or K) is filled in. + """ + + batch = d_qk_out.shape[0] + seq_len = d_qk_out.shape[1] + qk_channels = num_key_heads * key_head_dim + repeat = num_value_heads // num_key_heads + channel_offset = 0 if is_query else qk_channels + + device = d_qk_out.device + + grid = lambda meta: ( + batch * num_key_heads, + triton.cdiv(seq_len, meta["BLOCK_S"]), + ) + + with _launch_context(device, stream): + _l2norm_repeat_backward_kernel[grid]( + d_qk_out, + silu_bf16, + d_silu_bf16, + seq_len, + num_key_heads, + channel_offset, + eps, + d_qk_out.stride(0), + d_qk_out.stride(1), + d_qk_out.stride(2), + silu_bf16.stride(0), + silu_bf16.stride(1), + silu_bf16.stride(2), + d_silu_bf16.stride(0), + d_silu_bf16.stride(1), + d_silu_bf16.stride(2), + HEAD_DIM=key_head_dim, + REPEAT=repeat, + ) + + return d_silu_bf16 + + +def _triton_qk_l2norm_repeat_backward( + dq: Tensor, + dk: Tensor, + silu_bf16: Tensor, + d_silu_bf16: Tensor, + *, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + eps: float = 1e-6, + stream: Optional["torch.cuda.Stream"] = None, +) -> Tensor: + """Merged Q/K l2norm + REPEAT backward launch.""" + + batch = dq.shape[0] + seq_len = dq.shape[1] + qk_channels = num_key_heads * key_head_dim + repeat = num_value_heads // num_key_heads + device = dq.device + + grid = lambda meta: ( + batch * 2 * num_key_heads, + triton.cdiv(seq_len, meta["BLOCK_S"]), + ) + + with _launch_context(device, stream): + _qk_l2norm_repeat_backward_kernel[grid]( + dq, + dk, + silu_bf16, + d_silu_bf16, + seq_len, + num_key_heads, + qk_channels, + eps, + dq.stride(0), + dq.stride(1), + dq.stride(2), + dk.stride(0), + dk.stride(1), + dk.stride(2), + silu_bf16.stride(0), + silu_bf16.stride(1), + silu_bf16.stride(2), + d_silu_bf16.stride(0), + d_silu_bf16.stride(1), + d_silu_bf16.stride(2), + HEAD_DIM=key_head_dim, + REPEAT=repeat, + ) + + return d_silu_bf16 + + +def _triton_v_layout_to_conv( + dv: Tensor, + d_silu_conv: Tensor, + *, + v_channel_offset: int, + num_value_heads: int, + value_head_dim: int, + stream: Optional["torch.cuda.Stream"] = None, +) -> None: + """Write ``dv`` into ``d_silu_conv``'s V channel slice.""" + + batch, seq_len, _, _ = dv.shape + device = dv.device + + BLOCK_S = _LAYOUT_BLOCK_S + num_seq_blocks = triton.cdiv(seq_len, BLOCK_S) + grid = (batch * num_value_heads, num_seq_blocks) + + with _launch_context(device, stream): + _v_layout_to_conv_kernel[grid]( + dv, + d_silu_conv, + seq_len, + num_value_heads, + v_channel_offset, + dv.stride(0), + dv.stride(1), + dv.stride(2), + d_silu_conv.stride(0), + d_silu_conv.stride(1), + d_silu_conv.stride(2), + HEAD_DIM=value_head_dim, + BLOCK_S=BLOCK_S, + num_warps=4, + num_stages=2, + ) + + +def _triton_z_layout_to_qkvzba( + dgate: Tensor, + d_qkvzba: Tensor, + *, + z_channel_offset: int, + num_value_heads: int, + value_head_dim: int, + stream: Optional["torch.cuda.Stream"] = None, +) -> None: + """Write ``dgate`` into ``d_qkvzba``'s z channel slice.""" + + batch, seq_len, _, _ = dgate.shape + device = dgate.device + + BLOCK_S = _LAYOUT_BLOCK_S + num_seq_blocks = triton.cdiv(seq_len, BLOCK_S) + grid = (batch * num_value_heads, num_seq_blocks) + + with _launch_context(device, stream): + _z_layout_to_qkvzba_kernel[grid]( + dgate, + d_qkvzba, + seq_len, + num_value_heads, + z_channel_offset, + dgate.stride(0), + dgate.stride(1), + dgate.stride(2), + d_qkvzba.stride(0), + d_qkvzba.stride(1), + d_qkvzba.stride(2), + HEAD_DIM=value_head_dim, + BLOCK_S=BLOCK_S, + num_warps=4, + num_stages=2, + ) + + +def _triton_g_beta_backward( + qkvzba: Tensor, + A_log: Tensor, + dt_bias: Tensor, + d_g: Tensor, + d_beta_out: Tensor, + *, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, + num_key_heads: int, + d_qkvzba_out: Optional[Tensor] = None, + stream: Optional["torch.cuda.Stream"] = None, +) -> Tuple[Tensor, Tensor, Tensor]: + """Launch ``_g_beta_backward_kernel`` and return its outputs. + + Returns: + ``(d_qkvzba_out, d_A_log, d_dt_bias)``. ``d_qkvzba_out`` only has its + alpha and beta slices filled in; the caller is expected to allocate + the buffer while the other backward kernels fill the rest. + ``d_A_log`` and ``d_dt_bias`` are fp32 and need to be cast back to + the parameter dtype by the caller. + """ + + seq_len, batch, total_channels = qkvzba.shape + qk_channels = num_key_heads * key_head_dim + v_channels = num_value_heads * value_head_dim + beta_channel_offset = 2 * qk_channels + 2 * v_channels + alpha_channel_offset = beta_channel_offset + num_value_heads + + if d_qkvzba_out is None: + d_qkvzba_out = torch.zeros_like(qkvzba) + + device = qkvzba.device + + g_beta_grid = lambda meta: ( + batch, + triton.cdiv(seq_len, meta["BLOCK_S"]), + triton.cdiv(num_value_heads, meta["BLOCK_H"]), + ) + with _launch_context(device, stream): + d_param_grads = torch.empty((2, num_value_heads), dtype=torch.float32, device=device) + d_param_grads.zero_() + d_A_log = d_param_grads[0] + d_dt_bias = d_param_grads[1] + _g_beta_backward_kernel[g_beta_grid]( + qkvzba, + A_log, + dt_bias, + d_g, + d_beta_out, + d_qkvzba_out, + d_A_log, + d_dt_bias, + seq_len, + num_value_heads, + beta_channel_offset, + alpha_channel_offset, + qkvzba.stride(0), + qkvzba.stride(1), + qkvzba.stride(2), + d_g.stride(0), + d_g.stride(1), + d_g.stride(2), + d_beta_out.stride(0), + d_beta_out.stride(1), + d_beta_out.stride(2), + ) + return d_qkvzba_out, d_A_log, d_dt_bias + + +class _NullContext: + def __enter__(self): + return None + + def __exit__(self, exc_type, exc_val, exc_tb): + return False + + +def _launch_context( + device: torch.device, + stream: Optional["torch.cuda.Stream"], +): + """Return a CUDA launch context after wiring the optional side stream.""" + + if stream is None: + return _NullContext() + stream.wait_stream(torch.cuda.current_stream(device)) + return torch.cuda.stream(stream) + + +def _wait_for_streams( + dst_stream: "torch.cuda.Stream", + *src_streams: "torch.cuda.Stream", +) -> None: + for stream in src_streams: + dst_stream.wait_stream(stream) + + +def _resolve_packed_seq_idx( + cu_seqlens: Optional[Tensor], + seq_idx: Optional[Tensor], + total_tokens: int, +) -> Optional[Tensor]: + """Return the token-level sequence-id buffer for causal-conv backward.""" + + if cu_seqlens is None: + assert seq_idx is None, "seq_idx requires cu_seqlens for packed THD mode." + return None + + if seq_idx is None: + seq_lengths = cu_seqlens[1:] - cu_seqlens[:-1] + seq_idx = torch.repeat_interleave( + torch.arange(seq_lengths.numel(), device=cu_seqlens.device, dtype=torch.int32), + seq_lengths, + ) + seq_idx = seq_idx.unsqueeze(0) + elif seq_idx.dim() == 1: + seq_idx = seq_idx.unsqueeze(0) + + assert seq_idx.is_cuda, f"Packed seq_idx must be CUDA, got {seq_idx.device}." + assert seq_idx.dtype == torch.int32, f"Packed seq_idx must be int32, got {seq_idx.dtype}." + assert seq_idx.shape == (1, total_tokens), ( + "Packed seq_idx must have shape [1, total_tokens], " + f"got {seq_idx.shape=} and {total_tokens=}." + ) + return seq_idx.contiguous() + + +def _triton_pre_gated_delta_rule_forward( + qkvzba: Tensor, + conv1d_weight: Tensor, + A_log: Tensor, + dt_bias: Tensor, + *, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, + cu_seqlens: Optional[Tensor] = None, +) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + """Triton-backed forward for the pre-gated-delta-rule front-end. + + Returns ``(query, key, value, gate, beta, g, silu_qk_save)``. The last + element is the bf16-rounded ``silu(conv(x))`` for the QK channel range + laid out channel-last so the backward can feed it straight into + ``causal_conv1d_bwd_function`` — see module docstring. + """ + + seq_len, batch, total_channels = qkvzba.shape + is_packed_thd = cu_seqlens is not None + if is_packed_thd: + assert batch == 1, ( + "Packed THD fused_pre_gated_delta_rule expects batch dimension 1; " + f"got {batch=}." + ) + num_packed_seqs = cu_seqlens.shape[0] - 1 + else: + num_packed_seqs = 0 + + qk_channels = num_key_heads * key_head_dim + v_channels = num_value_heads * value_head_dim + repeat_factor = num_value_heads // num_key_heads + k_w = conv1d_weight.shape[-1] + assert _is_power_of_two(key_head_dim), ( + "Triton kernel currently expects key_head_dim to be a power of two; " + f"got {key_head_dim=}." + ) + assert _is_power_of_two(value_head_dim), ( + "Triton kernel currently expects value_head_dim to be a power of two; " + f"got {value_head_dim=}." + ) + + expected_channels = 2 * qk_channels + 2 * v_channels + 2 * num_value_heads + assert total_channels == expected_channels, ( + f"qkvzba last-dim mismatch: got {total_channels}, expected {expected_channels}." + ) + + out_dtype = qkvzba.dtype + device = qkvzba.device + + # Output buffers: contiguous (b, s, h, d) for q/k/v and (b, s, h) for g/beta. + # Q and K share one allocation so the fused-streamed QK kernel can select + # the logical group by pointer stride instead of branching between two + # unrelated base pointers inside Triton. + qk_out = torch.empty( + 2, batch, seq_len, num_value_heads, key_head_dim, dtype=out_dtype, device=device + ) + query = qk_out[0] + key = qk_out[1] + value = torch.empty( + batch, seq_len, num_value_heads, value_head_dim, dtype=out_dtype, device=device + ) + g = torch.empty(batch, seq_len, num_value_heads, dtype=torch.float32, device=device) + beta = torch.empty(batch, seq_len, num_value_heads, dtype=out_dtype, device=device) + + # Conv weight is (conv_dim, 1, K_W); we treat it as (conv_dim, K_W). + weight_2d = conv1d_weight.view(conv1d_weight.shape[0], k_w) + + # No conv bias support: the entry point asserts this. We still pass a + # dummy ``bias_tensor`` to the kernel so the launch signature stays + # stable; ``HAS_BIAS=False`` ensures the kernel never reads it. + bias_tensor = qkvzba + bias_stride = 0 + + # Allocate the gate (z) output buffer that the independent Z kernel will + # populate. Keeping Z separate makes the forward scopes QK / V / Z / + # G-Beta explicit. + gate = torch.empty( + batch, seq_len, num_value_heads, value_head_dim, dtype=out_dtype, device=device + ) + + # Persist the QK silu(conv(x)) intermediate in channel-last layout so the + # backward can feed it directly into the l2norm backward. + silu_qk_save = torch.empty( + (batch, seq_len, 2 * qk_channels), dtype=out_dtype, device=device + ).permute(0, 2, 1) # → (b, 2*qk_c, s) with stride(1)==1 + silu_save_b_stride = silu_qk_save.stride(0) + silu_save_c_stride = silu_qk_save.stride(1) + silu_save_s_stride = silu_qk_save.stride(2) + + # Stream setup. Each side stream handles one of the four sub-computations + # (QK conv+l2norm, V conv, Z copy, g/beta). + main_stream = torch.cuda.current_stream(device=device) + qk_stream = _get_side_stream(device, slot=_QK_STREAM_SLOT) + v_stream = _get_side_stream(device, slot=_V_STREAM_SLOT) + g_beta_stream = _get_side_stream(device, slot=_G_BETA_STREAM_SLOT) + z_stream = _get_side_stream(device, slot=_Z_STREAM_SLOT) + for stream in (qk_stream, v_stream, g_beta_stream, z_stream): + stream.wait_stream(main_stream) + + # --- QK conv + silu + l2norm + repeat --- + qk_grid = lambda meta: ( + batch * 2 * num_key_heads, + triton.cdiv(seq_len, meta["BLOCK_S"]), + ) + with torch.cuda.stream(qk_stream): + if is_packed_thd: + _conv_silu_project_thd_kernel[qk_grid]( + qkvzba, + weight_2d, + bias_tensor, + qk_out, + silu_qk_save, + cu_seqlens, + seq_len, + num_packed_seqs, + num_key_heads, + 0, # QK starts at channel 0; group 1 starts at +qk_channels. + qk_channels, + 0, # silu_save_chan_offset + qk_channels, + qkvzba.stride(0), + qkvzba.stride(1), + qkvzba.stride(2), + weight_2d.stride(0), + weight_2d.stride(1), + bias_stride, + qk_out.stride(0), + qk_out.stride(1), + qk_out.stride(2), + qk_out.stride(3), + silu_save_b_stride, + silu_save_c_stride, + silu_save_s_stride, + _L2NORM_EPS, + HEAD_DIM=key_head_dim, + K_W=k_w, + REPEAT=repeat_factor, + NUM_GROUPS=2, + HAS_BIAS=False, + SAVE_SILU=True, + APPLY_L2=True, + ) + else: + _conv_silu_project_kernel[qk_grid]( + qkvzba, + weight_2d, + bias_tensor, + qk_out, + silu_qk_save, + seq_len, + num_key_heads, + 0, # QK starts at channel 0; group 1 starts at +qk_channels. + qk_channels, + 0, # silu_save_chan_offset + qk_channels, + qkvzba.stride(0), + qkvzba.stride(1), + qkvzba.stride(2), + weight_2d.stride(0), + weight_2d.stride(1), + bias_stride, + qk_out.stride(0), + qk_out.stride(1), + qk_out.stride(2), + qk_out.stride(3), + silu_save_b_stride, + silu_save_c_stride, + silu_save_s_stride, + _L2NORM_EPS, + HEAD_DIM=key_head_dim, + K_W=k_w, + REPEAT=repeat_factor, + NUM_GROUPS=2, + HAS_BIAS=False, + SAVE_SILU=True, + APPLY_L2=True, + ) + + # --- V conv + silu (no l2norm, no repeat) --- + v_channel_offset = 2 * qk_channels + z_channel_offset = 2 * qk_channels + v_channels + v_grid = lambda meta: (batch * num_value_heads, triton.cdiv(seq_len, meta["BLOCK_S"])) + with torch.cuda.stream(v_stream): + if is_packed_thd: + _conv_silu_project_thd_kernel[v_grid]( + qkvzba, + weight_2d, + bias_tensor, + value, + qkvzba, # silu_save unused (SAVE_SILU=False) + cu_seqlens, + seq_len, + num_packed_seqs, + num_value_heads, + v_channel_offset, + 0, # in_group_stride unused for NUM_GROUPS=1 + 0, # silu_save_chan_offset unused + 0, # silu_save_group_stride unused + qkvzba.stride(0), + qkvzba.stride(1), + qkvzba.stride(2), + weight_2d.stride(0), + weight_2d.stride(1), + bias_stride, + 0, # out_group_dim_stride unused for NUM_GROUPS=1 + value.stride(0), + value.stride(1), + value.stride(2), + 0, # silu_save strides unused + 0, + 0, + _L2NORM_EPS, + HEAD_DIM=value_head_dim, + K_W=k_w, + REPEAT=1, + NUM_GROUPS=1, + HAS_BIAS=False, + SAVE_SILU=False, + APPLY_L2=False, + ) + else: + _conv_silu_project_kernel[v_grid]( + qkvzba, + weight_2d, + bias_tensor, + value, + qkvzba, # silu_save unused (SAVE_SILU=False) + seq_len, + num_value_heads, + v_channel_offset, + 0, # in_group_stride unused for NUM_GROUPS=1 + 0, # silu_save_chan_offset unused + 0, # silu_save_group_stride unused + qkvzba.stride(0), + qkvzba.stride(1), + qkvzba.stride(2), + weight_2d.stride(0), + weight_2d.stride(1), + bias_stride, + 0, # out_group_dim_stride unused for NUM_GROUPS=1 + value.stride(0), + value.stride(1), + value.stride(2), + 0, # silu_save strides unused + 0, + 0, + _L2NORM_EPS, + HEAD_DIM=value_head_dim, + K_W=k_w, + REPEAT=1, + NUM_GROUPS=1, + HAS_BIAS=False, + SAVE_SILU=False, + APPLY_L2=False, + ) + + # --- Z copy --- + BLOCK_Z_S = _LAYOUT_BLOCK_S + z_grid = (batch * num_value_heads, triton.cdiv(seq_len, BLOCK_Z_S)) + with torch.cuda.stream(z_stream): + _copy_z_kernel[z_grid]( + qkvzba, + gate, + seq_len, + num_value_heads, + z_channel_offset, + qkvzba.stride(0), + qkvzba.stride(1), + qkvzba.stride(2), + gate.stride(0), + gate.stride(1), + gate.stride(2), + HEAD_DIM=value_head_dim, + BLOCK_S=BLOCK_Z_S, + num_warps=4, + num_stages=2, + ) + + # --- g and beta --- + beta_channel_offset = 2 * qk_channels + 2 * v_channels + alpha_channel_offset = beta_channel_offset + num_value_heads + g_beta_grid = lambda meta: ( + batch, + triton.cdiv(seq_len, meta["BLOCK_S"]), + triton.cdiv(num_value_heads, meta["BLOCK_H"]), + ) + with torch.cuda.stream(g_beta_stream): + _compute_g_and_beta_kernel[g_beta_grid]( + qkvzba, + A_log, + dt_bias, + g, + beta, + seq_len, + num_value_heads, + beta_channel_offset, + alpha_channel_offset, + qkvzba.stride(0), + qkvzba.stride(1), + qkvzba.stride(2), + g.stride(0), + g.stride(1), + g.stride(2), + beta.stride(0), + beta.stride(1), + beta.stride(2), + ) + + # Re-join the side streams so the caller's stream observes the writes. + _wait_for_streams(main_stream, qk_stream, v_stream, z_stream, g_beta_stream) + + return query, key, value, gate, beta, g, silu_qk_save + + +def _triton_pre_gated_delta_rule_backward( + qkvzba: Tensor, + conv1d_weight: Tensor, + silu_qk_save: Tensor, + dq: Tensor, + dk: Tensor, + dv: Tensor, + dgate: Tensor, + dbeta: Tensor, + dg: Tensor, + A_log: Tensor, + dt_bias: Tensor, + *, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, + seq_idx: Optional[Tensor] = None, +) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + """Triton-backed backward for the pre-gated-delta-rule front-end. + + Mirror of :func:`_triton_pre_gated_delta_rule_forward`. Takes upstream + gradients (``dq``/``dk``/``dv``/``dgate``/``dbeta``/``dg``) plus the + saved forward intermediates and returns input/parameter gradients + ``(d_qkvzba, d_weight, d_A_log, d_dt_bias)``. + + Five Triton kernels + one C++ ``causal_conv1d_bwd_function`` call, + fanned out on five side streams so memory-bound work overlaps while + the conv backward runs on the default stream. See module docstring + for the overall design. + """ + + seq_len, batch, _ = qkvzba.shape + qk_channels = num_key_heads * key_head_dim + v_channels = num_value_heads * value_head_dim + conv_dim = 2 * qk_channels + v_channels + z_offset = 2 * qk_channels + v_channels + k_w = conv1d_weight.shape[-1] + device = qkvzba.device + + # Rebuild the conv input as a NON-contiguous (b, c, s) view of qkvzba. + # ``causal_conv1d_fn`` / ``_bwd_function`` accept inputs where either + # ``stride(1) == 1`` or ``stride(2) == 1``; the permuted view of qkvzba + # satisfies the former (channel stride is 1 in the original (s, b, c) + # layout), so we can skip a 256 MB ``.contiguous()`` copy. + qkvzba_conv = qkvzba[:, :, :conv_dim].permute(1, 2, 0) + weight_2d = conv1d_weight.view(conv1d_weight.shape[0], k_w) + + # ``silu_qk_save`` is the (b, 2*qk_channels, s) bf16 buffer the + # forward wrote ``silu(conv(x))`` into for QK. Reuse it directly as + # the silu input to the l2norm backward. + silu_conv = silu_qk_save + + # Allocate d_silu_conv channel-last (stride(1)==1) — that's what + # ``causal_conv1d_channellast_bwd_kernel`` consumes natively. + d_silu_conv = torch.empty( + (batch, seq_len, conv_dim), dtype=qkvzba.dtype, device=device + ).permute(0, 2, 1) + + # Use the same stream slots as the forward for the matching scopes. + qk_stream = _get_side_stream(device, slot=_QK_STREAM_SLOT) + v_stream = _get_side_stream(device, slot=_V_STREAM_SLOT) + g_beta_stream = _get_side_stream(device, slot=_G_BETA_STREAM_SLOT) + z_stream = _get_side_stream(device, slot=_Z_STREAM_SLOT) + + # Q + K: l2norm + REPEAT backward writes into d_silu_conv's Q/K slices. + _triton_qk_l2norm_repeat_backward( + dq, + dk, + silu_conv, + d_silu_conv, + num_key_heads=num_key_heads, + num_value_heads=num_value_heads, + key_head_dim=key_head_dim, + stream=qk_stream, + ) + + # V: no l2norm and no REPEAT in forward, so d_silu_conv's V slice is + # just dv re-laid-out from (b, s, num_v_heads, value_head_dim) to + # (b, v_channels, s). + _triton_v_layout_to_conv( + dv, + d_silu_conv, + v_channel_offset=2 * qk_channels, + num_value_heads=num_value_heads, + value_head_dim=value_head_dim, + stream=v_stream, + ) + + # g + beta backward fully stores d_qkvzba's alpha + beta slices, plus + # per-head d_A_log / d_dt_bias. Conv and z slices are filled by the + # causal-conv and z kernels, so d_qkvzba does not need a pre-zero. + d_qkvzba = torch.empty_like(qkvzba) + _, d_A_log_fp32, d_dt_bias_fp32 = _triton_g_beta_backward( + qkvzba, + A_log, + dt_bias, + dg, + dbeta, + num_value_heads=num_value_heads, + key_head_dim=key_head_dim, + value_head_dim=value_head_dim, + num_key_heads=num_key_heads, + d_qkvzba_out=d_qkvzba, + stream=g_beta_stream, + ) + + # Z slice gradient: stream dgate into d_qkvzba's z slice. + _triton_z_layout_to_qkvzba( + dgate, + d_qkvzba, + z_channel_offset=z_offset, + num_value_heads=num_value_heads, + value_head_dim=value_head_dim, + stream=z_stream, + ) + + # Join only streams that wrote into d_silu_conv before causal_conv1d_bwd_function. + # g/beta and z write disjoint outputs and can continue overlapping with conv bwd. + default_stream = torch.cuda.current_stream(device) + _wait_for_streams(default_stream, qk_stream, v_stream) + + # Pre-allocate d_x_conv as a strided view INTO d_qkvzba's conv slice. + # d_qkvzba memory layout is (s, b, total_channels) contiguous, so + # element [s, b, c] sits at offset s*b_stride + b*c_stride + c. + # Re-interpreting that storage as (b, conv_dim, s) lets + # causal_conv1d_bwd_function write d_x directly into the right cells. + seq_stride = qkvzba.stride(0) + batch_stride = qkvzba.stride(1) + d_x_conv_view = d_qkvzba.as_strided( + (batch, conv_dim, seq_len), + (batch_stride, 1, seq_stride), + ) + + # Hand-tuned C++ conv backward. Internally folds the silu' factor and + # computes both d_x and d_w in fp32; writes d_x directly into the + # view above. + if _causal_conv1d_bwd_function is None: + raise RuntimeError( + "Fused pre-gated-delta-rule backward requires the 'causal_conv1d' package. " + "Install it, or use pre_gated_delta_rule_impl='unfused'." + ) + _, d_weight_fp32, _, _ = _causal_conv1d_bwd_function( + qkvzba_conv, + weight_2d, + None, # no bias + d_silu_conv, + seq_idx, + None, # initial_states + None, # dfinal_states + d_x_conv_view, # dx pre-allocated into d_qkvzba's conv slice + False, # return_dinitial_states + True, # activation (silu) + ) + + d_weight = d_weight_fp32.view(*conv1d_weight.shape).to(conv1d_weight.dtype) + default_stream.wait_stream(g_beta_stream) + d_A_log = d_A_log_fp32.to(A_log.dtype) + d_dt_bias = d_dt_bias_fp32.to(dt_bias.dtype) + default_stream.wait_stream(z_stream) + + return d_qkvzba, d_weight, d_A_log, d_dt_bias + + +class _FusedPreGatedDeltaRuleFunction(torch.autograd.Function): + """Thin :class:`torch.autograd.Function` wrapper around the fused path. + + Stashes the forward inputs + the saved ``silu_qk_save`` intermediate + in ``ctx`` and dispatches to :func:`_triton_pre_gated_delta_rule_forward` + / :func:`_triton_pre_gated_delta_rule_backward`. The actual kernel + logic lives in those two free functions so it's easy to read and + reuse outside the autograd machinery. + """ + + @staticmethod + def forward( + ctx, + qkvzba, + conv1d_weight, + A_log, + dt_bias, + cu_seqlens, + seq_idx, + num_key_heads, + num_value_heads, + key_head_dim, + value_head_dim, + ): + ctx.num_key_heads = num_key_heads + ctx.num_value_heads = num_value_heads + ctx.key_head_dim = key_head_dim + ctx.value_head_dim = value_head_dim + query, key, value, gate, beta, g, silu_qk_save = ( + _triton_pre_gated_delta_rule_forward( + qkvzba, + conv1d_weight, + A_log, + dt_bias, + num_key_heads=num_key_heads, + num_value_heads=num_value_heads, + key_head_dim=key_head_dim, + value_head_dim=value_head_dim, + cu_seqlens=cu_seqlens, + ) + ) + ctx.has_seq_idx = seq_idx is not None + if ctx.has_seq_idx: + ctx.save_for_backward(qkvzba, conv1d_weight, A_log, dt_bias, silu_qk_save, seq_idx) + else: + ctx.save_for_backward(qkvzba, conv1d_weight, A_log, dt_bias, silu_qk_save) + return query, key, value, gate, beta, g + + @staticmethod + def backward(ctx, dq, dk, dv, dgate, dbeta, dg): + if ctx.has_seq_idx: + qkvzba, conv1d_weight, A_log, dt_bias, silu_qk_save, seq_idx = ctx.saved_tensors + else: + qkvzba, conv1d_weight, A_log, dt_bias, silu_qk_save = ctx.saved_tensors + seq_idx = None + d_qkvzba, d_weight, d_A_log, d_dt_bias = _triton_pre_gated_delta_rule_backward( + qkvzba, + conv1d_weight, + silu_qk_save, + dq, + dk, + dv, + dgate, + dbeta, + dg, + A_log, + dt_bias, + num_key_heads=ctx.num_key_heads, + num_value_heads=ctx.num_value_heads, + key_head_dim=ctx.key_head_dim, + value_head_dim=ctx.value_head_dim, + seq_idx=seq_idx, + ) + # Match forward inputs: (qkvzba, conv1d_weight, A_log, dt_bias, + # cu_seqlens, seq_idx, num_key_heads, num_value_heads, + # key_head_dim, value_head_dim). + # Non-tensor args get None. + return ( + d_qkvzba, + d_weight, + d_A_log, + d_dt_bias, + None, + None, + None, + None, + None, + None, + ) + + +def fused_streamed_pre_gated_delta_rule( + qkvzba: Tensor, + conv1d_weight: Tensor, + conv1d_bias: Optional[Tensor], + A_log: Tensor, + dt_bias: Tensor, + *, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, + use_qk_l2norm: bool = True, + cu_seqlens: Optional[Tensor] = None, + seq_idx: Optional[Tensor] = None, +) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + """Streamed fused pre-gated-delta-rule entry point. + + Args: + qkvzba: ``[seq_len, batch, in_proj_dim]`` projection output. Must be + on CUDA. + conv1d_weight: ``[conv_dim, 1, k_w]`` depthwise conv weight. + conv1d_bias: Must be ``None`` (conv bias is not supported). + A_log: ``[num_value_heads]`` raw decay parameter. + dt_bias: ``[num_value_heads]`` time-step bias. + num_key_heads / num_value_heads / key_head_dim / value_head_dim: GDN + architecture parameters. ``num_value_heads`` must be an integer + multiple of ``num_key_heads``. + use_qk_l2norm: Must be ``True``; the fused backward closes over the + l2norm path. + cu_seqlens: Optional packed THD cumulative sequence lengths. When set, + ``qkvzba`` must have ``batch == 1`` and ``cu_seqlens[-1] == seq_len``. + seq_idx: Optional precomputed token-to-sequence map with shape + ``[1, seq_len]``. Used by causal-conv backward in packed THD mode. + + Returns: + ``(query, key, value, gate, beta, g)`` matching the unfused + :meth:`GatedDeltaNet.pre_gated_delta_rule` API. + """ + + assert qkvzba.is_cuda, ( + "fused_pre_gated_delta_rule requires CUDA inputs; " + f"got qkvzba.device={qkvzba.device}." + ) + assert conv1d_bias is None, ( + "Conv bias is not supported by fused_pre_gated_delta_rule " + "(production GDN config has none)." + ) + assert use_qk_l2norm, ( + "use_qk_l2norm=False is not supported by fused_pre_gated_delta_rule " + "(the backward closes over the l2norm path)." + ) + assert num_value_heads % num_key_heads == 0, ( + f"{num_value_heads=} must be a multiple of {num_key_heads=}." + ) + if cu_seqlens is not None: + assert cu_seqlens.is_cuda, ( + "Packed fused_pre_gated_delta_rule requires CUDA cu_seqlens; " + f"got cu_seqlens.device={cu_seqlens.device}." + ) + assert cu_seqlens.dtype == torch.int32, ( + "Packed fused_pre_gated_delta_rule requires int32 cu_seqlens; " + f"got {cu_seqlens.dtype=}." + ) + assert cu_seqlens.dim() == 1, ( + "Packed fused_pre_gated_delta_rule expects 1-D cu_seqlens; " + f"got {cu_seqlens.shape=}." + ) + assert qkvzba.shape[1] == 1, ( + "Packed THD fused_pre_gated_delta_rule expects batch dimension 1; " + f"got qkvzba.shape={qkvzba.shape}." + ) + assert cu_seqlens.shape[0] >= 2, ( + "Packed fused_pre_gated_delta_rule requires at least one packed sequence; " + f"got {cu_seqlens.shape=}." + ) + assert cu_seqlens[0].item() == 0, ( + "Packed fused_pre_gated_delta_rule requires cu_seqlens[0] == 0, " + f"got {cu_seqlens[0].item()}." + ) + assert torch.all(cu_seqlens[1:] >= cu_seqlens[:-1]).item(), ( + "Packed fused_pre_gated_delta_rule requires monotonically non-decreasing " + f"cu_seqlens, got {cu_seqlens}." + ) + assert cu_seqlens[-1].item() == qkvzba.shape[0], ( + "Packed fused_pre_gated_delta_rule requires cu_seqlens[-1] to match " + f"seq_len, got {cu_seqlens[-1].item()} vs {qkvzba.shape[0]}." + ) + cu_seqlens = cu_seqlens.contiguous() + seq_idx = _resolve_packed_seq_idx(cu_seqlens, seq_idx, qkvzba.shape[0]) + else: + assert seq_idx is None, "seq_idx requires cu_seqlens for packed THD mode." + + return _FusedPreGatedDeltaRuleFunction.apply( + qkvzba, + conv1d_weight, + A_log, + dt_bias, + cu_seqlens, + seq_idx, + num_key_heads, + num_value_heads, + key_head_dim, + value_head_dim, + ) + + +fused_pre_gated_delta_rule = fused_streamed_pre_gated_delta_rule diff --git a/megatron/core/models/common/embeddings/rope_utils.py b/megatron/core/models/common/embeddings/rope_utils.py index c97f738771b..0c89cdbf084 100644 --- a/megatron/core/models/common/embeddings/rope_utils.py +++ b/megatron/core/models/common/embeddings/rope_utils.py @@ -16,6 +16,7 @@ from megatron.core import parallel_state logger = logging.getLogger(__name__) +_ROPE_FUSION_FALLBACK_WARNINGS: set[str] = set() try: from megatron.core.extensions.transformer_engine import fused_apply_rotary_pos_emb @@ -29,6 +30,26 @@ fused_apply_rotary_pos_emb_thd = None +try: + from megatron.core.fusions.fused_mrope import ( + can_launch_fused_mrope_thd, + fused_apply_mrope, + fused_apply_mrope_thd, + get_fused_mrope_thd_unavailable_reason, + get_fused_mrope_unavailable_reason, + is_fused_mrope_available, + mrope_freqs_to_rotary_emb, + ) +except ImportError: + can_launch_fused_mrope_thd = None + fused_apply_mrope = None + fused_apply_mrope_thd = None + get_fused_mrope_thd_unavailable_reason = None + get_fused_mrope_unavailable_reason = None + is_fused_mrope_available = None + mrope_freqs_to_rotary_emb = None + + try: from flash_attn.layers.rotary import apply_rotary_emb as apply_rotary_emb_flash except ImportError: @@ -41,10 +62,101 @@ 'apply_rotary_pos_emb_with_cos_sin', 'fused_apply_rotary_pos_emb', 'fused_apply_rotary_pos_emb_thd', + 'can_launch_fused_mrope_thd', + 'fused_apply_mrope', + 'fused_apply_mrope_thd', + 'get_fused_mrope_thd_unavailable_reason', + 'get_fused_mrope_unavailable_reason', + 'is_fused_mrope_available', + 'mrope_freqs_to_rotary_emb', 'get_pos_emb_on_this_cp_rank', ] +def _is_raw_mrope_freqs(t: Tensor, freqs: Tensor, config: TransformerConfig) -> bool: + """Return whether freqs is the raw 3-axis mRoPE tensor for fused apply.""" + if config.mrope_section is None or freqs.dim() != 4 or freqs.shape[0] != 3: + return False + if sum(config.mrope_section) != freqs.shape[-1] or freqs.shape[-1] * 2 > t.shape[-1]: + return False + if t.dim() == 4: + return freqs.shape[1] == t.shape[1] and freqs.shape[2] == t.shape[0] + if t.dim() == 3: + return freqs.shape[1] == 1 + return False + + +def _is_raw_mrope_freqs_thd( + t: Tensor, freqs: Tensor, cu_seqlens: Tensor, config: TransformerConfig, cp_size: int +) -> bool: + """Return whether freqs is raw mRoPE for THD layout, or fail on raw-like bad shapes.""" + if config.mrope_section is None or freqs.dim() != 4 or freqs.shape[0] != 3: + return False + if t.dim() != 3: + raise ValueError( + f"raw mRoPE THD expects t with shape [tokens, heads, head_dim], got {tuple(t.shape)}" + ) + if sum(config.mrope_section) != freqs.shape[-1] or freqs.shape[-1] * 2 > t.shape[-1]: + return False + + if freqs.shape[1] != 1: + raise ValueError( + "raw mRoPE THD freqs must have singleton batch dimension with shape " + f"[3, 1, total_seqlen, rotary_dim / 2], got {tuple(freqs.shape)}" + ) + if cp_size > 1 and freqs.shape[2] % cp_size != 0: + raise ValueError( + "raw mRoPE THD freqs sequence length must be divisible by context parallel size, " + f"got freqs.shape[2]={freqs.shape[2]}, cp_size={cp_size}" + ) + expected_total_seqlen = t.shape[0] * cp_size + if freqs.shape[2] != expected_total_seqlen: + raise ValueError( + "raw mRoPE THD freqs sequence length must match local tokens times cp_size, " + f"got freqs.shape[2]={freqs.shape[2]}, tokens={t.shape[0]}, cp_size={cp_size}" + ) + if cu_seqlens.dim() != 1: + raise ValueError(f"raw mRoPE THD cu_seqlens must be 1D, got {tuple(cu_seqlens.shape)}") + return True + + +def _raw_mrope_freqs_to_emb(freqs: Tensor, config: TransformerConfig) -> Tensor: + assert mrope_freqs_to_rotary_emb is not None, "mRoPE frequency conversion is unavailable." + return mrope_freqs_to_rotary_emb( + freqs, + config.mrope_section, + interleaved_mrope=config.mrope_interleaved, + rotary_interleaved=config.rotary_interleaved, + ) + + +def _warn_rope_fusion_fallback_once(key: str, message: str) -> None: + if key in _ROPE_FUSION_FALLBACK_WARNINGS: + return + _ROPE_FUSION_FALLBACK_WARNINGS.add(key) + warnings.warn(message, stacklevel=2) + + +def _fused_mrope_unavailable_warning_key(reason: str, thd: bool = False) -> str: + prefix = "triton-mrope-thd-unavailable" if thd else "triton-mrope-unavailable" + reason_lower = reason.lower() + if "triton is not available" in reason_lower: + category = "import" + elif "cuda tensors" in reason_lower or "same device" in reason_lower: + category = "device" + elif "dtype" in reason_lower or "float32" in reason_lower: + category = "dtype" + elif "stride" in reason_lower or "contiguous" in reason_lower: + category = "stride" + elif "capability" in reason_lower: + category = "capability" + elif "rotary_interleaved" in reason_lower: + category = "rotary-interleaved" + else: + category = "other" + return f"{prefix}-{category}" + + def get_pos_emb_on_this_cp_rank( pos_emb: Tensor, seq_dim: int, cp_group: torch.distributed.ProcessGroup ) -> Tensor: @@ -181,20 +293,21 @@ def _get_thd_freqs_on_this_cp_rank( compatibility. """ if cp_size > 1: - cp_seg = x.size(0) // 2 + first_cp_seg = (x.size(0) + 1) // 2 + second_cp_seg = x.size(0) // 2 full_seqlen = cp_size * x.size(0) # Apply offset to both forward and backward segments for context parallelism - # offset=0: traditional behavior, freqs[0:cp_seg] and freqs[...] - # offset>0: exact mapping, freqs[offset+0:offset+cp_seg] and freqs[offset+...] + # offset=0: traditional behavior, freqs[0:first_cp_seg] and freqs[...] + # offset>0: exact mapping, freqs[offset+0:offset+first_cp_seg] and freqs[offset+...] return torch.cat( [ - freqs[offset + cp_rank * cp_seg : offset + (cp_rank + 1) * cp_seg], + freqs[offset + cp_rank * first_cp_seg : offset + (cp_rank + 1) * first_cp_seg], freqs[ offset + full_seqlen - - (cp_rank + 1) * cp_seg : offset + - (cp_rank + 1) * second_cp_seg : offset + full_seqlen - - cp_rank * cp_seg + - cp_rank * second_cp_seg ], ] ) @@ -205,6 +318,84 @@ def _get_thd_freqs_on_this_cp_rank( return freqs[offset : offset + x.size(0)] +def _get_thd_raw_mrope_freqs_on_this_cp_rank( + cp_rank: int, cp_size: int, x: Tensor, freqs: Tensor, offset: int = 0 +) -> Tensor: + """Get raw mRoPE frequency slices for this CP rank in THD layout.""" + if cp_size > 1: + first_cp_seg = (x.size(0) + 1) // 2 + second_cp_seg = x.size(0) // 2 + full_seqlen = cp_size * x.size(0) + return torch.cat( + [ + freqs[ + :, :, offset + cp_rank * first_cp_seg : offset + (cp_rank + 1) * first_cp_seg + ], + freqs[ + :, + :, + offset + + full_seqlen + - (cp_rank + 1) * second_cp_seg : offset + + full_seqlen + - cp_rank * second_cp_seg, + ], + ], + dim=2, + ) + else: + return freqs[:, :, offset : offset + x.size(0)] + + +def _get_thd_cp_splits(cu_seqlens: Tensor, cp_size: int) -> tuple[list[int], list[int]]: + """Return global sequence offsets and per-rank sequence lengths for THD CP fallback.""" + cu_seqlens_list = cu_seqlens.tolist() + local_seqlens = [] + for seq_start, seq_end in zip(cu_seqlens_list[:-1], cu_seqlens_list[1:]): + seq_len = seq_end - seq_start + if cp_size > 1 and seq_len % cp_size != 0: + raise ValueError( + "THD sequence lengths must be divisible by context parallel size, " + f"got sequence length {seq_len}, cp_size={cp_size}" + ) + local_seqlens.append(seq_len // cp_size) + return cu_seqlens_list, local_seqlens + + +def _pack_thd_raw_mrope_freqs( + t: Tensor, + cu_seqlens: Tensor, + freqs: Tensor, + cp_group: torch.distributed.ProcessGroup, + total_seqlen: Optional[int] = None, +) -> Tensor: + """Pack raw mRoPE freqs into the same local token order as THD tensor ``t``.""" + cp_size = cp_group.size() + cp_rank = cp_group.rank() + cu_seqlens_list, seqlens = _get_thd_cp_splits(cu_seqlens, cp_size) + sequence_splits = torch.split(t, seqlens) + if total_seqlen is None: + total_seqlen = cu_seqlens_list[-1] + assert freqs.size(2) == total_seqlen, ( + f"raw mRoPE THD freqs sequence length {freqs.size(2)} must match " + f"cu_seqlens[-1] = {total_seqlen}" + ) + + freq_slices = [] + for i, x in enumerate(sequence_splits): + seq_start_offset = cu_seqlens_list[i] + freq_slices.append( + _get_thd_raw_mrope_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs, seq_start_offset) + ) + + packed_freqs = torch.cat(freq_slices, dim=2) + assert packed_freqs.shape[2] == t.shape[0], ( + f"packed raw mRoPE freqs sequence length {packed_freqs.shape[2]} " + f"does not match THD tensor length {t.shape[0]}" + ) + return packed_freqs.contiguous() + + def _apply_rotary_pos_emb_thd( t: Tensor, cu_seqlens: Tensor, @@ -240,21 +431,21 @@ def _apply_rotary_pos_emb_thd( raise ValueError("cp_group must be provided for THD format RoPE") cp_size = cp_group.size() cp_rank = cp_group.rank() - seqlens = ((cu_seqlens[1:] - cu_seqlens[:-1]) // cp_size).tolist() + cu_seqlens_list, seqlens = _get_thd_cp_splits(cu_seqlens, cp_size) # Handle two different frequency tensor formats: - # 1. If freqs.size(0) == cu_seqlens[-1]: freqs contains all positions across all sequences + # 1. If freqs.size(0) == cu_seqlens_list[-1]: freqs contains all positions across all sequences # -> Use offset-based mapping for exact positional correspondence # 2. Otherwise: freqs contains only max sequence length positions # -> Use traditional mapping without offsets (map first :seqlen part) - if freqs.dim() >= 1 and freqs.size(0) == cu_seqlens[-1]: + if freqs.dim() >= 1 and freqs.size(0) == cu_seqlens_list[-1]: # CASE 1: Exact mapping with offsets # Build packed freqs in one pass, then apply once to the whole packed tensor sequence_splits = torch.split(t, seqlens) freq_slices = [] for i, x in enumerate(sequence_splits): # cu_seqlens[i] is the starting offset of this sequence in the original batch - seq_start_offset = cu_seqlens[i].item() + seq_start_offset = cu_seqlens_list[i] freq_slices.append( _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs, seq_start_offset) ) @@ -311,40 +502,254 @@ def apply_rotary_pos_emb( if cp_group is None: cp_group = parallel_state.get_context_parallel_group() + is_raw_mrope_freqs = ( + _is_raw_mrope_freqs(t, freqs, config) + if cu_seqlens is None + else _is_raw_mrope_freqs_thd(t, freqs, cu_seqlens, config, cp_group.size()) + ) + if config.apply_rope_fusion: if cu_seqlens is None: + force_unfused_mrope = False + if is_raw_mrope_freqs: + unavailable_reason = None + can_try_fused_mrope = ( + fused_apply_mrope is not None + and get_fused_mrope_unavailable_reason is not None + and not mla_rotary_interleaved + and not inverse + and mscale == 1.0 + ) + if can_try_fused_mrope: + unavailable_reason = get_fused_mrope_unavailable_reason( + t, freqs, config.rotary_interleaved + ) + use_fused_mrope = can_try_fused_mrope and unavailable_reason is None + if use_fused_mrope: + return fused_apply_mrope( + t, + freqs, + config.mrope_section, + interleaved_mrope=config.mrope_interleaved, + rotary_interleaved=config.rotary_interleaved, + ) + + if unavailable_reason is not None: + _warn_rope_fusion_fallback_once( + _fused_mrope_unavailable_warning_key(unavailable_reason), + f"Triton fused mRoPE is unavailable: {unavailable_reason}. " + "Using unfused implementation.", + ) + force_unfused_mrope = True + unavailable_is_rotary_interleaved = ( + unavailable_reason is not None + and "rotary_interleaved" in unavailable_reason.lower() + ) + if mscale != 1.0: + _warn_rope_fusion_fallback_once( + "triton-mrope-mscale", + f"mscale={mscale} is not supported by Triton fused mRoPE. " + "Using unfused implementation.", + ) + force_unfused_mrope = True + if mla_rotary_interleaved: + _warn_rope_fusion_fallback_once( + "triton-mrope-mla-rotary-interleaved", + "Triton fused mRoPE does not support MLA-style interleaving in RoPE. " + "Using unfused implementation.", + ) + force_unfused_mrope = True + if inverse: + _warn_rope_fusion_fallback_once( + "triton-mrope-inverse", + "inverse RoPE is not supported by Triton fused mRoPE. " + "Using unfused implementation.", + ) + force_unfused_mrope = True + if config.rotary_interleaved and not unavailable_is_rotary_interleaved: + _warn_rope_fusion_fallback_once( + "triton-mrope-rotary-interleaved", + "Triton fused mRoPE currently supports rotary_interleaved=False. " + "Using unfused implementation.", + ) + force_unfused_mrope = True + freqs = _raw_mrope_freqs_to_emb(freqs, config) + is_raw_mrope_freqs = False + if force_unfused_mrope: + return _apply_rotary_pos_emb_bshd( + t, + freqs, + rotary_interleaved=config.rotary_interleaved, + mla_rotary_interleaved=mla_rotary_interleaved, + mscale=mscale, + inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, + ) + # NOTE: TE backends do not support mRoPE in bshd format when bs > 1. use_unfused = False if config.mrope_section is not None and freqs.shape[1] > 1: # TODO: Add a check in TransformerConfig and remove this unfused implementation. - warnings.warn( - "apply_rope_fusion does not support mRoPE in bshd format when bs > 1. " - "Please set apply_rope_fusion to false. This will become an error in v0.16." + _warn_rope_fusion_fallback_once( + "te-mrope-bshd-batch", + "Transformer Engine fused RoPE does not support mRoPE in bshd format when " + "bs > 1 without raw mRoPE freqs. Using unfused implementation.", ) use_unfused = True if mscale != 1.0: - warnings.warn( + _warn_rope_fusion_fallback_once( + "te-rope-mscale", f"mscale={mscale} is not supported by TE's fused RoPE. " - "Using unfused implementation." + "Using unfused implementation.", ) use_unfused = True if mla_rotary_interleaved: - warnings.warn( - "apply_rope_fusion does not support MLA-style interleaving in RoPE." - "Using unfused implementation." + _warn_rope_fusion_fallback_once( + "te-rope-mla-rotary-interleaved", + "apply_rope_fusion does not support MLA-style interleaving in RoPE. " + "Using unfused implementation.", ) use_unfused = True if inverse: - warnings.warn( + _warn_rope_fusion_fallback_once( + "te-rope-inverse", "inverse RoPE is not supported by TE's fused RoPE. " - "Using unfused implementation." + "Using unfused implementation.", + ) + use_unfused = True + if fused_apply_rotary_pos_emb is None: + _warn_rope_fusion_fallback_once( + "te-rope-unavailable", + "Transformer Engine fused RoPE is unavailable. Using unfused implementation.", ) use_unfused = True if not use_unfused: - assert fused_apply_rotary_pos_emb is not None, "apply_rope_fusion is not available." return fused_apply_rotary_pos_emb(t, freqs, interleaved=config.rotary_interleaved) else: - assert fused_apply_rotary_pos_emb_thd is not None, "apply_rope_fusion is not available." + if is_raw_mrope_freqs: + use_fused_mrope_thd = ( + fused_apply_mrope_thd is not None + and can_launch_fused_mrope_thd is not None + and get_fused_mrope_thd_unavailable_reason is not None + and mscale == 1.0 + and not mla_rotary_interleaved + and not inverse + and not config.rotary_interleaved + ) + if use_fused_mrope_thd: + unavailable_reason = get_fused_mrope_thd_unavailable_reason( + t, + cu_seqlens, + freqs, + rotary_interleaved=config.rotary_interleaved, + cp_size=cp_group.size(), + cp_rank=cp_group.rank(), + ) + if unavailable_reason is None: + return fused_apply_mrope_thd( + t, + cu_seqlens, + freqs, + config.mrope_section, + interleaved_mrope=config.mrope_interleaved, + rotary_interleaved=config.rotary_interleaved, + cp_size=cp_group.size(), + cp_rank=cp_group.rank(), + ) + _warn_rope_fusion_fallback_once( + _fused_mrope_unavailable_warning_key(unavailable_reason, thd=True), + f"Triton fused mRoPE for THD layout is unavailable: " + f"{unavailable_reason}. Using unfused implementation.", + ) + else: + has_unsupported_option = False + if mscale != 1.0: + _warn_rope_fusion_fallback_once( + "triton-mrope-thd-mscale", + f"mscale={mscale} is not supported by Triton fused mRoPE for THD " + "layout. Using unfused implementation.", + ) + has_unsupported_option = True + if mla_rotary_interleaved: + _warn_rope_fusion_fallback_once( + "triton-mrope-thd-mla-rotary-interleaved", + "Triton fused mRoPE for THD layout does not support MLA-style " + "interleaving in RoPE. Using unfused implementation.", + ) + has_unsupported_option = True + if inverse: + _warn_rope_fusion_fallback_once( + "triton-mrope-thd-inverse", + "inverse RoPE is not supported by Triton fused mRoPE for THD layout. " + "Using unfused implementation.", + ) + has_unsupported_option = True + if config.rotary_interleaved: + _warn_rope_fusion_fallback_once( + "triton-mrope-thd-rotary-interleaved", + "Triton fused mRoPE for THD layout currently supports " + "rotary_interleaved=False. Using unfused implementation.", + ) + has_unsupported_option = True + if not has_unsupported_option: + _warn_rope_fusion_fallback_once( + "triton-mrope-thd-unavailable", + "Triton fused mRoPE for THD layout is unavailable. " + "Using unfused implementation.", + ) + freqs = _raw_mrope_freqs_to_emb(freqs, config) + return _apply_rotary_pos_emb_thd( + t, + cu_seqlens, + freqs, + rotary_interleaved=config.rotary_interleaved, + mla_rotary_interleaved=mla_rotary_interleaved, + mscale=mscale, + cp_group=cp_group, + inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, + ) + use_unfused_thd = False + if mscale != 1.0: + _warn_rope_fusion_fallback_once( + "te-rope-thd-mscale", + f"mscale={mscale} is not supported by TE's fused RoPE for THD layout. " + "Using unfused implementation.", + ) + use_unfused_thd = True + if mla_rotary_interleaved: + _warn_rope_fusion_fallback_once( + "te-rope-thd-mla-rotary-interleaved", + "TE fused RoPE for THD layout does not support MLA-style interleaving " + "in RoPE. Using unfused implementation.", + ) + use_unfused_thd = True + if inverse: + _warn_rope_fusion_fallback_once( + "te-rope-thd-inverse", + "inverse RoPE is not supported by TE's fused RoPE for THD layout. " + "Using unfused implementation.", + ) + use_unfused_thd = True + if fused_apply_rotary_pos_emb_thd is None: + _warn_rope_fusion_fallback_once( + "te-rope-thd-unavailable", + "Transformer Engine fused RoPE for THD layout is unavailable. " + "Using unfused implementation.", + ) + use_unfused_thd = True + if use_unfused_thd: + return _apply_rotary_pos_emb_thd( + t, + cu_seqlens, + freqs, + rotary_interleaved=config.rotary_interleaved, + mla_rotary_interleaved=mla_rotary_interleaved, + mscale=mscale, + cp_group=cp_group, + inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, + ) return fused_apply_rotary_pos_emb_thd( t, cu_seqlens, @@ -354,6 +759,9 @@ def apply_rotary_pos_emb( interleaved=config.rotary_interleaved, ) # use unfused implementation + if is_raw_mrope_freqs: + freqs = _raw_mrope_freqs_to_emb(freqs, config) + if cu_seqlens is None: return _apply_rotary_pos_emb_bshd( t, diff --git a/megatron/core/models/common/embeddings/rotary_pos_embedding.py b/megatron/core/models/common/embeddings/rotary_pos_embedding.py index 77eb94a34bf..0476c62b144 100644 --- a/megatron/core/models/common/embeddings/rotary_pos_embedding.py +++ b/megatron/core/models/common/embeddings/rotary_pos_embedding.py @@ -382,6 +382,8 @@ def forward( position_ids: torch.Tensor, mrope_section: List[int], cp_group: Optional[torch.distributed.ProcessGroup] = None, + return_raw_freqs: bool = False, + packed_seq: bool = False, ) -> Tensor: """Forward pass of multimodal RoPE embedding. @@ -391,9 +393,14 @@ def forward( height and width in rope calculation. cp_group (torch.distributed.ProcessGroup, optional): Context parallel group. Defaults to None. + return_raw_freqs (bool, optional): If True, return the raw per-axis frequencies with + shape [3, batchsize, seqlens, dim / 2] for fused mRoPE application. + packed_seq (bool, optional): Whether the sequence uses THD packing. Packed sequences + keep full position frequencies because THD RoPE applies CP partitioning later. Returns: - Tensor: Embeddings after applying RoPE. + Tensor: Embeddings after applying RoPE, or raw per-axis frequencies when + return_raw_freqs is True. """ seq = position_ids.to(device=self.inv_freq.device, dtype=self.inv_freq.dtype) @@ -407,6 +414,13 @@ def forward( # shape (3, bs, seq_length, dim) freqs = (inv_freq_expanded @ seq_expanded).transpose(2, 3) + if cp_group is None: + cp_group = self.cp_group + if return_raw_freqs: + if cp_group is not None and cp_group.size() > 1 and not packed_seq: + freqs = get_pos_emb_on_this_cp_rank(freqs, 2, cp_group) + return freqs.contiguous() + # first part even vector components, second part odd vector components, # 2 * dim in dimension size if self.interleaved_mrope: @@ -417,9 +431,9 @@ def forward( emb = torch.cat((freqs, freqs), dim=-1) # shape (bs, seq_length, 2 * dim) else: bs = freqs.shape[0] - emb = torch.stack((freqs.view(bs, -1, 1), freqs.view(bs, -1, 1)), dim=-1).view( - bs, freqs.shape[1], -1 - ) + emb = torch.stack( + (freqs.reshape(bs, -1, 1), freqs.reshape(bs, -1, 1)), dim=-1 + ).view(bs, freqs.shape[1], -1) else: # Original section-based layout (Qwen2-VL style). if not self.rotary_interleaved: @@ -427,8 +441,8 @@ def forward( else: bs = freqs.shape[1] emb = torch.stack( - (freqs.view(3, bs, -1, 1), freqs.view(3, bs, -1, 1)), dim=-1 - ).view(3, bs, freqs.shape[0], -1) + (freqs.reshape(3, bs, -1, 1), freqs.reshape(3, bs, -1, 1)), dim=-1 + ).view(3, bs, freqs.shape[2], -1) # generate freqs with mrope_section: cycle T/H/W per section chunk mrope_section_doubled = list(mrope_section) * 2 emb = torch.cat( @@ -437,9 +451,7 @@ def forward( # shape (seq_length, bs, 1, 2 * dim) emb = emb[..., None, :].transpose(0, 1).contiguous() - if cp_group is None: - cp_group = self.cp_group - if cp_group is not None and cp_group.size() > 1: + if cp_group is not None and cp_group.size() > 1 and not packed_seq: # slice rotary_pos_emb along sequence dimension and select the parition of the current # CP rank emb = get_pos_emb_on_this_cp_rank(emb, 0, cp_group) diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index 0038851ac5a..c57c60cea28 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/megatron/core/models/gpt/fine_grained_callables.py @@ -635,6 +635,14 @@ def submodule_moe_forward(node: ScheduleNode, dispatched_tokens: torch.Tensor): # as a gradient hook of expert_output layer.pre_mlp_norm_checkpoint.discard_output_and_register_recompute(expert_output) + # Trigger the shared-expert recompute from expert_output too (output freed in + # submodule_combine_forward). Registering on the same tensor AFTER the pre_mlp_norm + # recompute orders it after pre_mlp_layernorm_output is restored and before the attn node's + # shared-expert backward. + shared_experts_checkpoint = getattr(layer.mlp, "shared_experts_checkpoint", None) + if shared_experts_checkpoint is not None: + shared_experts_checkpoint.register_recompute_hook(expert_output) + return expert_output def submodule_combine_forward(node: ScheduleNode, output: torch.Tensor): @@ -684,6 +692,13 @@ def submodule_combine_forward(node: ScheduleNode, output: torch.Tensor): if not node.is_mtp and final_layernorm and node.is_last_layer: output = final_layernorm(output) output = make_viewless_tensor(inp=output, requires_grad=True, keep_graph=True) + + # postprocess() has consumed the shared-expert output; free its storage now (the recompute + # hook was registered on expert_output in submodule_moe_forward). + shared_experts_checkpoint = getattr(layer.mlp, "shared_experts_checkpoint", None) + if shared_experts_checkpoint is not None: + shared_experts_checkpoint.discard_output() + layer.mlp.shared_experts_checkpoint = None return output @copy_signature(layer._forward_mlp, handle_first_dst_param='preserve') diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index 01df346c05c..fb9d0c83136 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -153,6 +153,7 @@ def __init__( ignore_virtual=False, vp_stage=vp_stage, ) + self._fused_mrope_available = False self.fuse_linear_cross_entropy = ( self.config.cross_entropy_loss_fusion @@ -215,6 +216,13 @@ def __init__( assert ( self.mrope_section is not None ), "mrope require mrope_section setting, but we got None from TransformerConfig" + if self.config.apply_rope_fusion and not self.config.rotary_interleaved: + try: + from megatron.core.fusions.fused_mrope import is_fused_mrope_available + + self._fused_mrope_available = is_fused_mrope_available() + except ImportError: + self._fused_mrope_available = False # Cache for RoPE tensors which do not change between iterations. self.rotary_pos_emb_cache = {} @@ -417,10 +425,25 @@ def _preprocess( ) elif self.position_embedding_type == 'mrope' and not self.config.multi_latent_attention: if not InferenceMode.is_active() or not self.config.flash_decode: + packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + use_fused_mrope = False + use_raw_mrope_freqs = ( + self.config.apply_rope_fusion and not self.config.rotary_interleaved + ) + if self.config.fused_single_qkv_rope: + use_raw_mrope_freqs = False + # Inference indexes rotary_pos_emb as seq-major materialized embeddings. + # Raw mRoPE freqs are axis-major and are only safe for the normal decoder path. + if in_inference_mode: + use_raw_mrope_freqs = False + if use_raw_mrope_freqs: + use_fused_mrope = self._fused_mrope_available rotary_pos_emb = self.rotary_pos_emb( position_ids, self.mrope_section, cp_group=packed_seq_params.cp_group if packed_seq_params is not None else None, + return_raw_freqs=use_fused_mrope, + packed_seq=packed_seq, ) else: # Flash decoding uses precomputed cos and sin for RoPE diff --git a/megatron/core/parallel_state.py b/megatron/core/parallel_state.py index 863b5d55d9d..f2c9b0b5181 100644 --- a/megatron/core/parallel_state.py +++ b/megatron/core/parallel_state.py @@ -154,6 +154,12 @@ def get_nccl_options(pg_name, nccl_comm_cfgs): nccl_comm_cfgs (dict): nccl communicator configurations When an option (e.g., max_ctas) is not found in the config, use the NCCL default setting. """ + # The fake distributed backend (--fake-process-group) cannot accept + # ProcessGroupNCCL.Options; PyTorch's FakeProcessGroup._create_internal + # rejects them with a TypeError. Return None so callers create the + # fake sub-groups without NCCL-specific options. + if torch.distributed.is_initialized() and torch.distributed.get_backend() == "fake": + return None if pg_name in nccl_comm_cfgs: # When fields in nccl_options.config are not specified, NCCL applies default settings. # The default values for Hopper GPUs are as follows: diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index 4dd963eb928..127d3e390de 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -6,6 +6,7 @@ # LICENSE file in the root directory of this source tree. import logging +import os from dataclasses import dataclass, replace from typing import List, Optional, Tuple, Union @@ -14,9 +15,16 @@ import torch.nn.functional as F from torch import Tensor +from megatron.core import tensor_parallel from megatron.core.dist_checkpointing import ShardedTensor from megatron.core.dist_checkpointing.mapping import ReplicaId, ShardedTensorFactory from megatron.core.fp8_utils import get_fp8_align_size +from megatron.core.fusions.fused_mega_pre_gated_delta_rule import ( + fused_mega_pre_gated_delta_rule, +) +from megatron.core.fusions.fused_pre_gated_delta_rule import ( + fused_streamed_pre_gated_delta_rule, +) from megatron.core.inference.contexts import BaseInferenceContext from megatron.core.jit import jit_fuser from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group @@ -28,6 +36,7 @@ _undo_attention_load_balancing, ) from megatron.core.tensor_parallel import get_cuda_rng_tracker +from megatron.core.tensor_parallel.random import CheckpointManager from megatron.core.transformer import TransformerConfig from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.module import MegatronModule @@ -43,7 +52,14 @@ try: from fla.modules.convolution import causal_conv1d from fla.modules.l2norm import l2norm - from fla.ops.gated_delta_rule import chunk_gated_delta_rule + + if os.environ.get("MCORE_GDN_USE_OPT_WRAPPER", "0") == "1": + try: + from mcore_gdn_opt.gated_delta_rule import chunk_gated_delta_rule + except ImportError: + from fla.ops.gated_delta_rule import chunk_gated_delta_rule + else: + from fla.ops.gated_delta_rule import chunk_gated_delta_rule HAVE_FLA = True except ImportError: @@ -130,6 +146,11 @@ def __init__( self.cp_size = self.pg_collection.cp.size() self.tp_size = self.pg_collection.tp.size() self.sp_size = self.tp_size if config.sequence_parallel else 1 + self.pre_gated_delta_rule_impl = config.pre_gated_delta_rule_impl + if self.pre_gated_delta_rule_impl != "unfused": + assert ( + self.cp_size == 1 + ), "Fused pre_gated_delta_rule does not support context parallelism yet." # Attributes from config self.config = config @@ -244,6 +265,15 @@ def __init__( hidden_size=self.value_head_dim, eps=self.config.layernorm_epsilon, ) + self.recompute_norm_out = False + self.recompute_qkv = False + if self.config.recompute_granularity == "selective": + self.recompute_norm_out = "gdn_norm_out" in self.config.recompute_modules + # gdn_qkv: recompute the whole QKV proj+prep block as a discard-output checkpoint. + self.recompute_qkv = "gdn_qkv" in self.config.recompute_modules + + # Per-forward CheckpointManager for the GDN discard-output recompute (gdn_qkv/gdn_norm_out). + self.gdn_recompute_manager = None self.out_proj = build_module( submodules.out_proj, @@ -363,6 +393,112 @@ def forward( cu_seqlens_q = None cu_seqlens_kv = None + # gdn_qkv (QKV proj+prep) and gdn_norm_out (gated norm) are discard-output checkpoints; the + # QKV output `gate` feeds the gated-norm block, so when both are on the CheckpointManager + # replays them in forward order (qkv -> norm_out) from one grad hook on `out`. + recompute_qkv = self.recompute_qkv and self.training + recompute_norm_out = self.recompute_norm_out and self.training + self.gdn_recompute_manager = ( + CheckpointManager() if (recompute_qkv or recompute_norm_out) else None + ) + + # QKV projection + prep block (in_proj -> CP a2a -> conv1d -> _prepare_qkv -> g/beta). + def _qkv_proj_and_prepare(hidden_states): + return self._compute_qkv_for_gated_delta_rule( + hidden_states, batch, seq_len, cu_seqlens_q, packed_seq_params + ) + + if recompute_qkv: + # Discard the QKV outputs now; regenerate them in backward. Synchronous recompute + # (no async reload), so it is safe with the fla/compiled gated_delta_rule backward. + query, key, value, g, beta, gate = tensor_parallel.CheckpointWithoutOutput( + fp8=(self.config.fp8 or self.config.fp4), + ckpt_manager=self.gdn_recompute_manager, + ).checkpoint(_qkv_proj_and_prepare, hidden_states) + else: + query, key, value, g, beta, gate = _qkv_proj_and_prepare(hidden_states) + + # seq_len was reassigned to the post-CP-a2a sequence length inside the block; recover it + # from a produced tensor so the downstream gated-norm reshape uses the correct value. + seq_len = value.shape[1] + + nvtx_range_push(suffix="gated_delta_rule") + core_attn_out, last_recurrent_state = self.gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + cu_seqlens=cu_seqlens_q, + ) + nvtx_range_pop(suffix="gated_delta_rule") + + def _gated_norm_and_a2a(core_attn_out: torch.Tensor, gate: torch.Tensor): + # RMSNorm + nvtx_range_push(suffix="gated_norm") + norm_out_hp = self._apply_gated_norm(core_attn_out, gate) + nvtx_range_pop(suffix="gated_norm") + + # Transpose: b s x --> s b x + # From bshd back to sbhd format + norm_out_hp = norm_out_hp.reshape(batch, seq_len, -1) + norm_out_hp = norm_out_hp.transpose(0, 1).contiguous() + + # CP all to all: HP to CP + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + unpacked_norm_out = _unpack_sequence(norm_out_hp, cu_seqlens_q, dim=0) + outputs = [] + for norm_out_i in unpacked_norm_out: + norm_out_i = tensor_a2a_hp2cp( + norm_out_i, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp + ) + outputs.append(norm_out_i) + norm_out = torch.cat(outputs, dim=0) + else: + norm_out = tensor_a2a_hp2cp( + norm_out_hp, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp + ) + + return norm_out + + if recompute_norm_out: + norm_out = tensor_parallel.CheckpointWithoutOutput( + ckpt_manager=self.gdn_recompute_manager + ).checkpoint(_gated_norm_and_a2a, core_attn_out, gate) + else: + norm_out = _gated_norm_and_a2a(core_attn_out, gate) + + # Output projection + nvtx_range_push(suffix="out_proj") + out, out_bias = self.out_proj(norm_out) + nvtx_range_pop(suffix="out_proj") + + # Discard the checkpointed outputs (now consumed) and register the unified recompute hook on + # `out` — its grad is computed first in backward, before the backwards that need them. + if self.gdn_recompute_manager is not None: + self.gdn_recompute_manager.discard_all_outputs_and_register_unified_recompute(out) + self.gdn_recompute_manager = None + + return out, out_bias + + def _compute_qkv_for_gated_delta_rule( + self, hidden_states, batch, seq_len, cu_seqlens_q, packed_seq_params + ): + """QKV projection + preparation block for the gated delta rule. + + Runs in_proj, CP all-to-all, conv1d, _prepare_qkv and g/beta, producing the tensors consumed + by ``self.gated_delta_rule`` plus the ``gate`` for the gated norm. Extracted so it can be + checkpointed when ``recompute_modules`` contains ``"gdn_qkv"``. + + Returns: + Tuple of (query, key, value, g, beta, gate). + """ + cp_group = resolve_cp_group(self.pg_collection.cp, packed_seq_params) + cp_size = cp_group.size() + # Input projection nvtx_range_push(suffix="in_proj") qkvzba, _ = self.in_proj(hidden_states) @@ -405,6 +541,49 @@ def forward( ], ) + # Fused pre-gated-delta-rule path: a single fused kernel replaces the conv1d -> + # _prepare_qkv -> g/beta block below. The fused wrappers consume the post-CP-a2a + # qkvzba (s b x) directly and return (query, key, value, gate, beta, g); reorder to the + # (query, key, value, g, beta, gate) layout this method returns. + if self.pre_gated_delta_rule_impl != "unfused": + seq_idx = ( + packed_seq_params.seq_idx + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + else None + ) + if self.pre_gated_delta_rule_impl == "fused_streamed": + nvtx_range_push(suffix="fused_streamed_pre_gated_delta_rule") + query, key, value, gate, beta, g = self._fused_streamed_pre_gated_delta_rule( + qkvzba, cu_seqlens_q=cu_seqlens_q, seq_idx=seq_idx + ) + nvtx_range_pop(suffix="fused_streamed_pre_gated_delta_rule") + else: + assert self.pre_gated_delta_rule_impl == "fused_mega" + nvtx_range_push(suffix="fused_mega_pre_gated_delta_rule") + query, key, value, gate, beta, g = self._fused_mega_pre_gated_delta_rule( + qkvzba, cu_seqlens_q=cu_seqlens_q, seq_idx=seq_idx + ) + nvtx_range_pop(suffix="fused_mega_pre_gated_delta_rule") + return query, key, value, g, beta, gate + + # Unfused path: conv1d -> _prepare_qkv -> g/beta on the post-CP-a2a qkvzba. + query, key, value, gate, beta, g = self.pre_gated_delta_rule( + qkvzba, batch, seq_len, cu_seqlens_q=cu_seqlens_q, cp_group=cp_group, cp_size=cp_size + ) + return query, key, value, g, beta, gate + + def pre_gated_delta_rule( + self, qkvzba, batch, seq_len, cu_seqlens_q=None, *, cp_group, cp_size + ): + """Unfused pre-gated-delta-rule on the post-CP-a2a qkvzba. + + Runs the split -> conv1d -> _prepare_qkv -> g/beta block (the reference path that the + fused wrappers replace). Returns (query, key, value, gate, beta, g). + + ``cp_group``/``cp_size`` are the dynamic context-parallel group/size resolved by the + caller (see ``resolve_cp_group``), threaded in so this method stays free of global + process-group reads. + """ # Transpose: s b x --> b s x # From sbhd to bshd format qkvzba = qkvzba.transpose(0, 1) @@ -498,47 +677,45 @@ def forward( g, beta = self._compute_g_and_beta(A_log_local_cp, dt_bias_local_cp, alpha, beta) nvtx_range_pop(suffix="g_and_beta") - nvtx_range_push(suffix="gated_delta_rule") - core_attn_out, last_recurrent_state = self.gated_delta_rule( - query, - key, - value, - g=g, - beta=beta, - initial_state=None, - output_final_state=False, - use_qk_l2norm_in_kernel=False, - cu_seqlens=cu_seqlens_q, - ) - nvtx_range_pop(suffix="gated_delta_rule") + return query, key, value, gate, beta, g - # RMSNorm - nvtx_range_push(suffix="gated_norm") - norm_out = self._apply_gated_norm(core_attn_out, gate) - nvtx_range_pop(suffix="gated_norm") + def _fused_streamed_pre_gated_delta_rule(self, qkvzba, cu_seqlens_q=None, seq_idx=None): + """Call the streamed fused pre-GDR wrapper. Returns (query, key, value, gate, beta, g).""" - # Transpose: b s x --> s b x - # From bshd back to sbhd format - norm_out = norm_out.reshape(batch, seq_len, -1) - norm_out = norm_out.transpose(0, 1).contiguous() - - # CP all to all: HP to CP - if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': - unpacked_norm_out = _unpack_sequence(norm_out, cu_seqlens_q, dim=0) - outputs = [] - for norm_out_i in unpacked_norm_out: - norm_out_i = tensor_a2a_hp2cp(norm_out_i, seq_dim=0, head_dim=-1, cp_group=cp_group) - outputs.append(norm_out_i) - norm_out = torch.cat(outputs, dim=0) - else: - norm_out = tensor_a2a_hp2cp(norm_out, seq_dim=0, head_dim=-1, cp_group=cp_group) + assert self.cp_size == 1, "Fused pre_gated_delta_rule does not support CP yet." + return fused_streamed_pre_gated_delta_rule( + qkvzba, + self.conv1d.weight, + self.conv1d.bias if self.conv_bias else None, + self.A_log, + self.dt_bias, + num_key_heads=self.qk_dim_local_tp // self.key_head_dim, + num_value_heads=self.v_dim_local_tp // self.value_head_dim, + key_head_dim=self.key_head_dim, + value_head_dim=self.value_head_dim, + use_qk_l2norm=self.use_qk_l2norm, + cu_seqlens=cu_seqlens_q, + seq_idx=seq_idx, + ) - # Output projection - nvtx_range_push(suffix="out_proj") - out, out_bias = self.out_proj(norm_out) - nvtx_range_pop(suffix="out_proj") + def _fused_mega_pre_gated_delta_rule(self, qkvzba, cu_seqlens_q=None, seq_idx=None): + """Call the mega fused pre-GDR wrapper. Returns (query, key, value, gate, beta, g).""" - return out, out_bias + assert self.cp_size == 1, "Fused pre_gated_delta_rule does not support CP yet." + return fused_mega_pre_gated_delta_rule( + qkvzba, + self.conv1d.weight, + self.conv1d.bias if self.conv_bias else None, + self.A_log, + self.dt_bias, + num_key_heads=self.qk_dim_local_tp // self.key_head_dim, + num_value_heads=self.v_dim_local_tp // self.value_head_dim, + key_head_dim=self.key_head_dim, + value_head_dim=self.value_head_dim, + use_qk_l2norm=self.use_qk_l2norm, + cu_seqlens=cu_seqlens_q, + seq_idx=seq_idx, + ) @jit_fuser def _apply_gated_norm(self, x, gate): diff --git a/megatron/core/tensor_parallel/random.py b/megatron/core/tensor_parallel/random.py index 4ad95ce2253..27066ecd310 100644 --- a/megatron/core/tensor_parallel/random.py +++ b/megatron/core/tensor_parallel/random.py @@ -921,30 +921,43 @@ def _recompute(self, _): self.outputs = None self.ctx = None - def discard_output_and_register_recompute(self, hook_tensor): - """ - Release the output tensor storages and register the recompute function as a grad hook of - the hook_tensor. + def discard_output(self): + """Free the output storages (metadata kept for backward). - Note: the caller should make sure that the output tensors are no longer used - in the forward pass and the gradient of the hook_tensor is computed before the recomputed - tensors are used. + Pair with :meth:`register_recompute_hook` when the output is freed in a different place + than where the recompute hook is registered; otherwise use + :meth:`discard_output_and_register_recompute`. """ - # When ckpt_manager is set, this is a no-op. - # Manager handles all discarding and hook registration uniformly. from megatron.core.transformer.cuda_graphs import is_graph_warmup if self.ckpt_manager is not None or is_graph_warmup(): return - - # use resize to release the output tensor memory and still keep the metadata in the tensors. - # the metadata is still needed for backward + # resize keeps tensor metadata (needed for backward) while releasing the memory. for output in self.outputs: output.untyped_storage().resize_(0) - # register the recomputation as a backward hook, when the the gradient of the hook_tensor - # is computed, the recomputation will be triggered. The hook_tensor should be selected - # carefully to ensure that the tensors are recomputed before it is used by other backward - # computations. + def register_recompute_hook(self, hook_tensor): + """Trigger the recompute from ``hook_tensor``'s grad hook. + + ``hook_tensor`` must have its grad computed before the discarded outputs are needed in + backward (and, if the recompute reads other discarded activations, after those are + restored). + """ + from megatron.core.transformer.cuda_graphs import is_graph_warmup + + if self.ckpt_manager is not None or is_graph_warmup(): + return if hook_tensor.requires_grad: hook_tensor.register_hook(self._recompute) + + def discard_output_and_register_recompute(self, hook_tensor): + """ + Release the output tensor storages and register the recompute function as a grad hook of + the hook_tensor. + + Note: the caller should make sure that the output tensors are no longer used + in the forward pass and the gradient of the hook_tensor is computed before the recomputed + tensors are used. + """ + self.discard_output() + self.register_recompute_hook(hook_tensor) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 3e61eb12a5f..3434bfd7507 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -486,6 +486,7 @@ def _build_per_layer_rotary_pos_emb(self, rotary_base: float) -> None: rotary_interleaved=self.config.rotary_interleaved, seq_len_interpolation_factor=seq_len_interpolation_factor, rotary_base=rotary_base, + interleaved_mrope=self.config.mrope_interleaved, ) self.mrope_section = self.config.mrope_section assert ( diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index 359b7c4a4bd..a46480960dd 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -255,6 +255,18 @@ def __init__( config.recompute_granularity == 'selective' and "shared_experts" in config.recompute_modules ) + # Discard the shared-expert OUTPUT (CheckpointWithoutOutput) instead of a standard + # checkpoint that keeps it. The recompute is wired up where ordering is correct (after any + # pre_mlp_layernorm recompute) by the caller: the fine-grained callables for A2A overlap, + # or TransformerLayer._forward_post_mlp otherwise. Disabled under MoE cudagraph partial + # capture, where the output is a graph output and must keep its storage. + self.shared_experts_recompute_discard_output = ( + self.shared_experts_recompute + and not bool(getattr(config, "cuda_graph_modules", None)) + ) + # The active CheckpointWithoutOutput, handed to the caller that frees the output and + # registers the recompute hook (None when not using discard-output recompute). + self.shared_experts_checkpoint = None self.tp_group = pg_collection.tp self.tp_ep_group = pg_collection.tp_ep @@ -534,7 +546,19 @@ def shared_experts_compute(self, hidden_states: torch.Tensor): shared_expert_output = None if self.use_shared_expert and not self.shared_expert_overlap: # Compute the shared expert separately when not overlapped with communication. - if self.shared_experts_recompute: + if self.shared_experts_recompute_discard_output and self.training: + # Recompute-with-discarded-output: run the shared expert under no_grad, then free + # its output in postprocess() and regenerate it (with its backward graph) from a + # grad hook. CheckpointWithoutOutput handles fp8/fp4 internally via its fp8 flag. + self.shared_experts_checkpoint = tensor_parallel.CheckpointWithoutOutput( + fp8=(self.config.fp8 or self.config.fp4) + ) + shared_expert_output = self.shared_experts_checkpoint.checkpoint( + apply_module(self.shared_experts), hidden_states + ) + elif self.shared_experts_recompute: + # Standard checkpoint fallback (e.g. MoE cudagraph partial capture or eval): keep + # the output, recompute only the intermediates. if self.config.fp8 or self.config.fp4: shared_expert_output = te_checkpoint( apply_module(self.shared_experts), diff --git a/megatron/core/transformer/moe/shared_experts.py b/megatron/core/transformer/moe/shared_experts.py index 3ffc572ae5e..d5e25dbd7b6 100644 --- a/megatron/core/transformer/moe/shared_experts.py +++ b/megatron/core/transformer/moe/shared_experts.py @@ -28,6 +28,8 @@ is_te_min_version, is_torch_min_version, make_sharded_tensor_for_checkpoint, + nvtx_range_pop, + nvtx_range_push, ) if HAVE_TE: @@ -188,11 +190,13 @@ def __init__( def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: """Forward function""" + nvtx_range_push("SharedExpert.forward") output, _ = super().forward(hidden_states) if self.use_shared_expert_gate: logits = torch.nn.functional.linear(hidden_states, self.gate_weight) gate_score = torch.nn.functional.sigmoid(logits) output = output * gate_score + nvtx_range_pop("SharedExpert.forward") return output def _reset_parameters(self): @@ -238,6 +242,7 @@ def pre_forward_comm(self, input, wait_current_stream=True): if wait_current_stream: self.wait_current_stream() with torch.cuda.stream(self.stream): + nvtx_range_push("SharedExpert.pre_forward_comm") if self.use_shared_expert_gate: logits = torch.nn.functional.linear(input, self.gate_weight) self.gate_score = torch.nn.functional.sigmoid(logits) @@ -250,6 +255,7 @@ def pre_forward_comm(self, input, wait_current_stream=True): input, group=self.tp_group ) set_tensor_grad_fn_sequence_sr(self.cached_fc1_input, torch.iinfo(torch.int).max) + nvtx_range_pop("SharedExpert.pre_forward_comm") @overlap_state_check( SharedExpertState.PRE_FORWARD_COMM_DONE, SharedExpertState.FC1_FORWARD_DONE @@ -261,6 +267,7 @@ def linear_fc1_forward_and_act(self, overlapped_comm_output=None): It is only useful when --moe-shared-expert-overlap is set and may be changed. """ with torch.cuda.stream(self.stream): + nvtx_range_push("SharedExpert.linear_fc1_forward_and_act") # [s, b, 4 * h/p] intermediate_parallel, bias_parallel = apply_module(self.linear_fc1)( self.cached_fc1_input @@ -308,6 +315,7 @@ def glu(x): intermediate_parallel = self.activation_func(intermediate_parallel) self.cached_fc2_input = intermediate_parallel + nvtx_range_pop("SharedExpert.linear_fc1_forward_and_act") # Tensor sequence number is used to control the backward order. # Decrease the sequence number of the expert output to make the comm launched first # in the backward order. @@ -327,9 +335,11 @@ def linear_fc2_forward(self, overlapped_comm_output=None): if overlapped_comm_output is not None: set_tensor_grad_fn_sequence_sr(overlapped_comm_output, torch.iinfo(torch.int).max) with torch.cuda.stream(self.stream): + nvtx_range_push("SharedExpert.linear_fc2_forward") # [s, b, h] self.cached_fc2_output, _ = apply_module(self.linear_fc2)(self.cached_fc2_input) self.cached_fc2_input = None + nvtx_range_pop("SharedExpert.linear_fc2_forward") @overlap_state_check( SharedExpertState.FC2_FORWARD_DONE, SharedExpertState.POST_FORWARD_COMM_DONE @@ -341,6 +351,7 @@ def post_forward_comm(self): It is only useful when --moe-shared-expert-overlap is set and may be changed. """ with torch.cuda.stream(self.stream): + nvtx_range_push("SharedExpert.post_forward_comm") if self.config.sequence_parallel: self.cached_output = reduce_scatter_to_sequence_parallel_region( self.cached_fc2_output, group=self.tp_group @@ -351,6 +362,7 @@ def post_forward_comm(self): ) self.cached_fc2_output = None set_tensor_grad_fn_sequence_sr(self.cached_output, torch.iinfo(torch.int).max) + nvtx_range_pop("SharedExpert.post_forward_comm") @overlap_state_check(SharedExpertState.POST_FORWARD_COMM_DONE, SharedExpertState.IDLE) def get_output(self): @@ -360,6 +372,7 @@ def get_output(self): It is only useful when --moe-shared-expert-overlap is set and may be changed. """ with torch.cuda.stream(self.stream): + nvtx_range_push("SharedExpert.get_output") if self.use_shared_expert_gate: assert self.gate_score is not None output = self.cached_output * self.gate_score @@ -367,6 +380,7 @@ def get_output(self): else: output = self.cached_output self.cached_output = None + nvtx_range_pop("SharedExpert.get_output") torch.cuda.current_stream().wait_stream(self.stream) return output diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index caf9e8d26b6..88171323a73 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -379,6 +379,9 @@ class TransformerConfig(ModelParallelConfig): linear_num_value_heads: Optional[int] = 32 """Number of value and gate heads for the gated delta net.""" + pre_gated_delta_rule_impl: Literal["unfused", "fused_streamed", "fused_mega"] = "unfused" + """Pre-gated-delta-rule implementation for GatedDeltaNet.""" + #################### # initialization #################### @@ -557,7 +560,7 @@ class TransformerConfig(ModelParallelConfig): recompute_modules: Optional[List[str]] = None """The submodules to recompute. choices: "core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", - "shared_experts", "mhc". + "shared_experts", "mhc", "gdn_norm_out". default: ["core_attn"]. "core_attn": recompute the core attention part of the transformer layer. "moe_act": recompute the MoE MLP activation function. @@ -569,8 +572,9 @@ class TransformerConfig(ModelParallelConfig): "mhc": recompute HyperConnection intermediate activations via CheckpointWithoutOutput + CheckpointManager. Requires enable_hyper_connections=True. Cannot be used with "mlp". - "moe_act", "layernorm", "mla_up_proj", and "mhc" use output-discarding checkpointing, - "core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing. + "gdn_norm_out": recompute the GatedDeltaNet output norm and HP-to-CP all-to-all. + "moe_act", "layernorm", "mla_up_proj", "mhc", and "gdn_norm_out" use output-discarding + checkpointing, "core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing. """ #################### @@ -1425,6 +1429,21 @@ def __post_init__(self): self.experimental_attention_variant = self.linear_attention_type self.linear_attention_type = None + valid_pre_gdr_impls = ("unfused", "fused_streamed", "fused_mega") + if self.pre_gated_delta_rule_impl not in valid_pre_gdr_impls: + raise ValueError( + "pre_gated_delta_rule_impl must be one of " + f"{valid_pre_gdr_impls}, got {self.pre_gated_delta_rule_impl!r}." + ) + if ( + self.pre_gated_delta_rule_impl != "unfused" + and self.experimental_attention_variant != "gated_delta_net" + ): + raise ValueError( + "pre_gated_delta_rule_impl can select a fused path only when " + "experimental_attention_variant='gated_delta_net'." + ) + if self.experimental_attention_variant in ["gated_delta_net"]: assert ( self.linear_attention_freq is not None @@ -1859,6 +1878,8 @@ def __post_init__(self): "moe", "shared_experts", "mhc", + "gdn_norm_out", + "gdn_qkv", } invalid_modules = set(self.recompute_modules) - allowed_modules assert not invalid_modules, ( @@ -1877,6 +1898,24 @@ def __post_init__(self): "multi_latent_attention." ) + if ( + "gdn_norm_out" in self.recompute_modules + and self.experimental_attention_variant != "gated_delta_net" + ): + raise ValueError( + "gdn_norm_out in recompute_modules is only supported with " + "experimental_attention_variant='gated_delta_net'." + ) + + if ( + "gdn_qkv" in self.recompute_modules + and self.experimental_attention_variant != "gated_delta_net" + ): + raise ValueError( + "gdn_qkv in recompute_modules is only supported with " + "experimental_attention_variant='gated_delta_net'." + ) + if "core_attn" in self.recompute_modules: warnings.warn( "If you are using transformer_engine as the transformer implementation, " @@ -2314,7 +2353,18 @@ def __post_init__(self): "It is experimental and may change in future versions." ) else: - if self.rotary_interleaved: + fused_mrope_available = False + # Triton fused mRoPE supports split-half RoPE only. Keep rotary_interleaved + # configs on the TE validation path so the TE >= 2.3 check still applies. + if self.mrope_section is not None and not self.rotary_interleaved: + try: + from megatron.core.fusions.fused_mrope import is_fused_mrope_available + + fused_mrope_available = is_fused_mrope_available() + except ImportError: + fused_mrope_available = False + + if self.rotary_interleaved and not fused_mrope_available: if not is_te_min_version("2.3.0"): raise ValueError( "rotary_interleaved does not work with apply_rope_fusion for " @@ -2326,9 +2376,14 @@ def __post_init__(self): fused_apply_rotary_pos_emb_thd, ) - if fused_apply_rotary_pos_emb is None and fused_apply_rotary_pos_emb_thd is None: + if ( + fused_apply_rotary_pos_emb is None + and fused_apply_rotary_pos_emb_thd is None + and not fused_mrope_available + ): raise ValueError( - "apply_rope_fusion is not available. Please install TE >= 1.4." + "apply_rope_fusion is not available. Please install TE >= 1.4 " + "or Triton for fused mRoPE." ) if self.fused_single_qkv_rope: diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 67bb04837ac..fe1f17c6573 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -957,6 +957,19 @@ def _forward_post_mlp( mlp_output_with_bias[0] ) + # Shared-expert discard-output recompute (non-overlap path; the A2A-overlap path uses the + # fine-grained callables and never reaches here). The output was consumed by the MoE + # postprocess add, so free it and register the recompute on mlp_output_with_bias[0] AFTER + # the pre_mlp_norm recompute above (same hook tensor) — this orders it after the shared + # expert's input pre_mlp_layernorm_output is restored and before its backward. + if self.is_moe_layer: + shared_experts_checkpoint = getattr(self.mlp, "shared_experts_checkpoint", None) + if shared_experts_checkpoint is not None: + shared_experts_checkpoint.discard_output_and_register_recompute( + mlp_output_with_bias[0] + ) + self.mlp.shared_experts_checkpoint = None + # TODO: could we move `bias_dropout_add_exec_handler` itself # inside the module provided in the `bias_dropout_add_spec` module? nvtx_range_push(suffix="mlp_bda") diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 29e21544599..aeaeff206f3 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1638,7 +1638,7 @@ def validate_args(args, defaults={}): # Legacy RoPE arguments if args.use_rotary_position_embeddings: args.position_embedding_type = 'rope' - if args.position_embedding_type != 'rope': + if args.position_embedding_type not in ('rope', 'mrope'): args.apply_rope_fusion = False # Would just need to add 'NoPE' as a position_embedding_type to support this, but for now @@ -4680,6 +4680,10 @@ def _add_mla_args(parser): def _add_experimental_attention_variant_args(parser): group = parser.add_argument_group(title="experimental_attention_variant") + # NOTE: --pre-gated-delta-rule-impl is auto-generated from the + # TransformerConfig.pre_gated_delta_rule_impl field by ArgumentGroupFactory + # (see _add_transformer_engine_args / build_group), so it must NOT be + # registered manually here — doing so raises an argparse conflict. # Linear attention group.add_argument( '--linear-attention-freq', diff --git a/megatron/training/training.py b/megatron/training/training.py index 8c9bba1d16c..c72b7f708dd 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -1318,6 +1318,25 @@ def pretrain( # Model, optimizer, and learning rate. timers('model-and-optimizer-setup', log_level=0).start(barrier=True) + + # Enable CUDA memory event recording BEFORE model/optimizer init, so the + # snapshot at dump time contains full allocation event traces (frames + + # device_traces), not just segment state. Gated on the existing + # `--record-memory-history` flag. + if args.record_memory_history and ( + is_last_rank() or torch.distributed.get_backend() == 'fake' + ): + try: + torch.cuda.memory._record_memory_history( + enabled='all', + context='all', + stacks='python', + max_entries=100000, + ) + print_rank_0("[memory_snapshot] enabled torch.cuda.memory event recording (mode=all)") + except Exception as _e: # noqa: BLE001 + print_rank_0(f"[memory_snapshot] _record_memory_history failed: {_e}") + model, optimizer, opt_param_scheduler = setup_model_and_optimizer( model_provider, model_type, checkpointing_context=checkpointing_context ) diff --git a/tests/unit_tests/a2a_overlap/test_fsdp_1f1b_overlap.py b/tests/unit_tests/a2a_overlap/test_fsdp_1f1b_overlap.py index 4e5ddc7eb02..c7005c72884 100644 --- a/tests/unit_tests/a2a_overlap/test_fsdp_1f1b_overlap.py +++ b/tests/unit_tests/a2a_overlap/test_fsdp_1f1b_overlap.py @@ -74,9 +74,12 @@ def test_fsdp_1f1b_training_step( [[], ["attn_norm", "core_attn", "attn_proj", "mlp_norm", "expert_fc1", "moe_act"]], ) def test_fsdp_1f1b_memory_opt(self, recompute_modules, offload_modules): + # Configure shared experts so recompute_modules=[..., "shared_experts"] actually + # exercises the shared-experts discard-output recompute (a no-op without them). self._run_test_helper( dispatcher_type="alltoall", sharding_strategy="optim_grads_params", + shared_expert_intermediate_size=512, recompute_modules=recompute_modules, offload_modules=offload_modules, ) diff --git a/tests/unit_tests/fusions/test_fused_mrope.py b/tests/unit_tests/fusions/test_fused_mrope.py new file mode 100644 index 00000000000..b033f6b9bed --- /dev/null +++ b/tests/unit_tests/fusions/test_fused_mrope.py @@ -0,0 +1,1311 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import warnings +from types import SimpleNamespace + +import pytest +import torch + +import megatron.core.models.common.embeddings.rope_utils as rope_utils +from megatron.core import parallel_state +from megatron.core.fusions.fused_mrope import ( + fused_apply_mrope, + fused_apply_mrope_thd, + get_fused_mrope_thd_unavailable_reason, + get_fused_mrope_unavailable_reason, + is_fused_mrope_available, + mrope_freqs_to_rotary_emb, +) +from megatron.core.models.common.embeddings import apply_rotary_pos_emb +from megatron.core.models.common.embeddings.rope_utils import ( + _ROPE_FUSION_FALLBACK_WARNINGS, + _apply_rotary_pos_emb_bshd, + _apply_rotary_pos_emb_thd, +) +from megatron.core.models.common.embeddings.rotary_pos_embedding import MultimodalRotaryEmbedding +from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.core.transformer.transformer_config import TransformerConfig +from tests.unit_tests.test_utilities import Utils + + +class FakeCPGroup: + def __init__(self, size=1, rank=0): + self._size = size + self._rank = rank + + def size(self): + return self._size + + def rank(self): + return self._rank + + +class FakeDynamicInferenceContext: + def is_dynamic_batching(self): + return True + + def is_static_batching(self): + return False + + +class FakeStaticInferenceContext: + def is_dynamic_batching(self): + return False + + def is_static_batching(self): + return True + + +@pytest.fixture(autouse=True) +def clear_rope_fusion_fallback_warnings(): + _ROPE_FUSION_FALLBACK_WARNINGS.clear() + yield + _ROPE_FUSION_FALLBACK_WARNINGS.clear() + + +def _dtype_tols(dtype): + if dtype == torch.bfloat16: + return dict(rtol=2.0e-2, atol=5.0e-2) + if dtype == torch.float16: + return dict(rtol=3.0e-3, atol=1.0e-2) + return dict(rtol=1.0e-6, atol=1.0e-6) + + +def _make_inputs( + dtype=torch.bfloat16, + requires_grad=False, + head_dim=20, + rotary_dim=16, + mrope_section=None, + interleaved_mrope=False, + batch=2, +): + seq = 32 + heads = 3 + if mrope_section is None: + mrope_section = [3, 3, 2] if interleaved_mrope else [2, 3, 3] + + generator = torch.Generator(device="cuda").manual_seed(1234) + t = torch.randn( + seq, + batch, + heads, + head_dim, + dtype=dtype, + device="cuda", + generator=generator, + requires_grad=requires_grad, + ) + freqs = torch.randn( + 3, batch, seq, rotary_dim // 2, dtype=torch.float32, device="cuda", generator=generator + ) + return t, freqs, mrope_section + + +def _make_position_ids(seq, batch): + base = torch.arange(seq, device="cuda", dtype=torch.long) + batch_offsets = torch.arange(batch, device="cuda", dtype=torch.long) + return ( + torch.stack((base, base * 2 + 3, base * 3 + 5), dim=0)[:, None, :] + + batch_offsets[None, :, None] + ).contiguous() + + +def _make_thd_inputs( + dtype=torch.bfloat16, + requires_grad=False, + interleaved_mrope=False, + cp_size=1, + padded_seq_lens=(12, 16), + head_dim=20, + rotary_dim=16, + mrope_section=None, +): + total_seq = sum(padded_seq_lens) + local_seq = total_seq // cp_size + heads = 3 + if mrope_section is None: + mrope_section = [3, 3, 2] if interleaved_mrope else [2, 3, 3] + + generator = torch.Generator(device="cuda").manual_seed(5678) + t = torch.randn( + local_seq, + heads, + head_dim, + dtype=dtype, + device="cuda", + generator=generator, + requires_grad=requires_grad, + ) + freqs = torch.randn( + 3, 1, total_seq, rotary_dim // 2, dtype=torch.float32, device="cuda", generator=generator + ) + cu_seqlens = torch.tensor([0, padded_seq_lens[0], total_seq], dtype=torch.int32, device="cuda") + return t, freqs, cu_seqlens, mrope_section + + +def _make_mrope_config( + num_attention_heads, mrope_section, interleaved_mrope=False, rotary_interleaved=False +): + return TransformerConfig( + num_attention_heads=num_attention_heads, + num_layers=1, + apply_rope_fusion=True, + mrope_section=mrope_section, + mrope_interleaved=interleaved_mrope, + rotary_interleaved=rotary_interleaved, + ) + + +def _fallback_warnings(recorded_warnings): + return [ + warning + for warning in recorded_warnings + if issubclass(warning.category, UserWarning) + and "Using unfused implementation" in str(warning.message) + ] + + +def _thd_cp_freq_indices(cu_seqlens_cpu, cp_size, cp_rank): + indices = [] + for global_start, global_end in zip(cu_seqlens_cpu[:-1], cu_seqlens_cpu[1:]): + local_seq_len = (global_end - global_start) // cp_size + first_cp_seg = (local_seq_len + 1) // 2 + second_cp_seg = local_seq_len // 2 + indices.extend( + range( + global_start + cp_rank * first_cp_seg, global_start + (cp_rank + 1) * first_cp_seg + ) + ) + indices.extend( + range(global_end - (cp_rank + 1) * second_cp_seg, global_end - cp_rank * second_cp_seg) + ) + return indices + + +def _assert_thd_cp_freq_index_coverage(cu_seqlens_cpu, cp_size): + expected = [] + actual = [] + for global_start, global_end in zip(cu_seqlens_cpu[:-1], cu_seqlens_cpu[1:]): + expected.extend(range(global_start, global_end)) + for cp_rank in range(cp_size): + actual.extend(_thd_cp_freq_indices(cu_seqlens_cpu, cp_size, cp_rank)) + assert sorted(actual) == expected + assert len(set(actual)) == len(actual) + + +@pytest.mark.parametrize("use_packed_seq", [False, True]) +def test_gpt_mrope_eval_requests_raw_freqs_when_fusion_available(use_packed_seq): + captured_kwargs = {} + + def fake_rotary_pos_emb(*args, **kwargs): + captured_kwargs.update(kwargs) + return "raw-mrope-freqs" + + model = SimpleNamespace( + training=False, + pre_process=False, + mtp_process=False, + position_embedding_type="mrope", + config=SimpleNamespace( + multi_latent_attention=False, + flash_decode=False, + apply_rope_fusion=True, + rotary_interleaved=False, + cuda_graph_impl=None, + fused_single_qkv_rope=False, + ), + rotary_pos_emb=fake_rotary_pos_emb, + mrope_section=[2, 3, 3], + _fused_mrope_available=True, + ) + packed_seq_params = ( + SimpleNamespace(qkv_format="thd", cp_group=FakeCPGroup()) if use_packed_seq else None + ) + + output = GPTModel._preprocess( + model, + input_ids=torch.zeros(1, 4, dtype=torch.long), + position_ids=torch.zeros(3, 1, 4, dtype=torch.long), + decoder_input=torch.zeros(4, 1, 12), + packed_seq_params=packed_seq_params, + ) + + assert output[1] == "raw-mrope-freqs" + assert captured_kwargs["return_raw_freqs"] is True + assert captured_kwargs["packed_seq"] is use_packed_seq + + +def test_gpt_mrope_eval_keeps_materialized_freqs_with_fused_single_qkv_rope(): + captured_kwargs = {} + + def fake_rotary_pos_emb(*args, **kwargs): + captured_kwargs.update(kwargs) + return "materialized-mrope-freqs" + + model = SimpleNamespace( + training=False, + pre_process=False, + mtp_process=False, + position_embedding_type="mrope", + config=SimpleNamespace( + multi_latent_attention=False, + flash_decode=False, + apply_rope_fusion=True, + rotary_interleaved=False, + cuda_graph_impl=None, + fused_single_qkv_rope=True, + ), + rotary_pos_emb=fake_rotary_pos_emb, + mrope_section=[2, 3, 3], + _fused_mrope_available=True, + ) + + output = GPTModel._preprocess( + model, + input_ids=torch.zeros(1, 4, dtype=torch.long), + position_ids=torch.zeros(3, 1, 4, dtype=torch.long), + decoder_input=torch.zeros(4, 1, 12), + ) + + assert output[1] == "materialized-mrope-freqs" + assert captured_kwargs["return_raw_freqs"] is False + + +def test_gpt_mrope_dynamic_inference_keeps_materialized_freqs(): + captured_kwargs = {} + + def fake_rotary_pos_emb(*args, **kwargs): + captured_kwargs.update(kwargs) + return "materialized-mrope-freqs" + + model = SimpleNamespace( + training=False, + pre_process=False, + mtp_process=False, + position_embedding_type="mrope", + config=SimpleNamespace( + multi_latent_attention=False, + flash_decode=False, + apply_rope_fusion=True, + rotary_interleaved=False, + cuda_graph_impl=None, + fused_single_qkv_rope=False, + ), + rotary_pos_emb=fake_rotary_pos_emb, + mrope_section=[2, 3, 3], + _fused_mrope_available=True, + ) + + output = GPTModel._preprocess( + model, + input_ids=torch.zeros(1, 4, dtype=torch.long), + position_ids=torch.zeros(3, 1, 4, dtype=torch.long), + decoder_input=torch.zeros(4, 1, 12), + inference_context=FakeDynamicInferenceContext(), + ) + + assert output[1] == "materialized-mrope-freqs" + assert captured_kwargs["return_raw_freqs"] is False + + +def test_gpt_mrope_static_inference_keeps_materialized_freqs(): + captured_kwargs = {} + + def fake_rotary_pos_emb(*args, **kwargs): + captured_kwargs.update(kwargs) + return "materialized-mrope-freqs" + + model = SimpleNamespace( + training=False, + pre_process=False, + mtp_process=False, + position_embedding_type="mrope", + config=SimpleNamespace( + multi_latent_attention=False, + flash_decode=False, + apply_rope_fusion=True, + rotary_interleaved=False, + cuda_graph_impl=None, + fused_single_qkv_rope=False, + ), + rotary_pos_emb=fake_rotary_pos_emb, + mrope_section=[2, 3, 3], + _fused_mrope_available=True, + ) + + output = GPTModel._preprocess( + model, + input_ids=torch.zeros(1, 4, dtype=torch.long), + position_ids=torch.zeros(3, 1, 4, dtype=torch.long), + decoder_input=torch.zeros(4, 1, 12), + inference_context=FakeStaticInferenceContext(), + ) + + assert output[1] == "materialized-mrope-freqs" + assert captured_kwargs["return_raw_freqs"] is False + + +def test_is_fused_mrope_available_requires_cuda(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + assert not is_fused_mrope_available() + + +def test_transformer_config_rejects_fused_mrope_without_cuda_or_te(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(rope_utils, "fused_apply_rotary_pos_emb", None) + monkeypatch.setattr(rope_utils, "fused_apply_rotary_pos_emb_thd", None) + + with pytest.raises(ValueError, match="apply_rope_fusion is not available"): + TransformerConfig( + num_attention_heads=1, num_layers=1, apply_rope_fusion=True, mrope_section=[1, 1, 1] + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("interleaved_mrope", [False, True]) +@pytest.mark.parametrize("head_dim", [16, 20]) +def test_fused_mrope_matches_unfused_forward_backward(interleaved_mrope, head_dim): + t_ref, freqs, mrope_section = _make_inputs( + requires_grad=True, head_dim=head_dim, interleaved_mrope=interleaved_mrope + ) + t_fused = t_ref.detach().clone().requires_grad_(True) + + emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + ref = _apply_rotary_pos_emb_bshd(t_ref, emb, rotary_interleaved=False) + out = fused_apply_mrope( + t_fused, freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + + tols = _dtype_tols(t_ref.dtype) + torch.testing.assert_close(ref.float(), out.float(), **tols) + + grad = torch.randn_like(ref) + ref.backward(grad) + out.backward(grad) + torch.testing.assert_close(t_ref.grad.float(), t_fused.grad.float(), **tols) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("interleaved_mrope", [False, True]) +def test_apply_rotary_pos_emb_bshd_eval_uses_triton_without_te(interleaved_mrope, monkeypatch): + t, freqs, mrope_section = _make_inputs(interleaved_mrope=interleaved_mrope, batch=1) + config = _make_mrope_config(t.shape[2], mrope_section, interleaved_mrope) + + emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + ref = _apply_rotary_pos_emb_bshd(t, emb, rotary_interleaved=False) + + fused_calls = 0 + orig_fused_apply_mrope = rope_utils.fused_apply_mrope + + def wrapped_fused_apply_mrope(*args, **kwargs): + nonlocal fused_calls + fused_calls += 1 + return orig_fused_apply_mrope(*args, **kwargs) + + monkeypatch.setattr(rope_utils, "fused_apply_rotary_pos_emb", None) + monkeypatch.setattr(rope_utils, "fused_apply_mrope", wrapped_fused_apply_mrope) + with torch.no_grad(), warnings.catch_warnings(record=True) as recorded_warnings: + warnings.simplefilter("always") + out = apply_rotary_pos_emb(t, freqs, config, cp_group=FakeCPGroup()) + + assert fused_calls == 1 + assert not _fallback_warnings(recorded_warnings) + assert not out.requires_grad + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize( + "fallback_kwargs, warning_match", + [ + ({"mscale": 1.25}, "mscale=1.25 is not supported by Triton fused mRoPE"), + ({"inverse": True}, "inverse RoPE is not supported by Triton fused mRoPE"), + ], +) +def test_apply_rotary_pos_emb_raw_mrope_fallbacks_match_unfused(fallback_kwargs, warning_match): + t, freqs, mrope_section = _make_inputs() + config = TransformerConfig( + num_attention_heads=t.shape[2], + num_layers=1, + apply_rope_fusion=True, + mrope_section=mrope_section, + ) + + with pytest.warns(UserWarning, match=warning_match): + out = apply_rotary_pos_emb(t, freqs, config, cp_group=FakeCPGroup(), **fallback_kwargs) + + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + ref = _apply_rotary_pos_emb_bshd(t, emb, rotary_interleaved=False, **fallback_kwargs) + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + with warnings.catch_warnings(record=True) as repeated_warnings: + warnings.simplefilter("always") + out_again = apply_rotary_pos_emb( + t, freqs, config, cp_group=FakeCPGroup(), **fallback_kwargs + ) + assert not repeated_warnings + torch.testing.assert_close(ref.float(), out_again.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize( + "fallback_kwargs, config_kwargs, expected_warning_key, warning_text", + [ + ( + {"mscale": 1.25}, + {}, + "triton-mrope-mscale", + "mscale=1.25 is not supported by Triton fused mRoPE", + ), + ( + {"inverse": True}, + {}, + "triton-mrope-inverse", + "inverse RoPE is not supported by Triton fused mRoPE", + ), + ( + {"mla_rotary_interleaved": True}, + {}, + "triton-mrope-mla-rotary-interleaved", + "does not support MLA-style interleaving", + ), + ( + {}, + {"rotary_interleaved": True}, + "triton-mrope-unavailable-rotary-interleaved", + "rotary_interleaved=True is not supported", + ), + ], +) +def test_apply_rotary_pos_emb_raw_mrope_fallback_emits_single_warning( + fallback_kwargs, config_kwargs, expected_warning_key, warning_text +): + t, freqs, mrope_section = _make_inputs() + config = _make_mrope_config(t.shape[2], mrope_section, **config_kwargs) + + with warnings.catch_warnings(record=True) as recorded_warnings: + warnings.simplefilter("always") + out = apply_rotary_pos_emb(t, freqs, config, cp_group=FakeCPGroup(), **fallback_kwargs) + + fallback_warnings = _fallback_warnings(recorded_warnings) + assert len(fallback_warnings) == 1 + assert warning_text in str(fallback_warnings[0].message) + assert _ROPE_FUSION_FALLBACK_WARNINGS == {expected_warning_key} + + emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, rotary_interleaved=config.rotary_interleaved + ) + ref = _apply_rotary_pos_emb_bshd( + t, emb, rotary_interleaved=config.rotary_interleaved, **fallback_kwargs + ) + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +def test_interleaved_mrope_rejects_inconsistent_sections(): + freqs = torch.randn(3, 2, 8, 8, dtype=torch.float32) + + with pytest.raises(AssertionError, match="interleaved mRoPE"): + mrope_freqs_to_rotary_emb(freqs, [2, 3, 3], interleaved_mrope=True) + + +def test_raw_mrope_cpu_falls_back_to_unfused(): + t = torch.randn(8, 1, 3, 20, dtype=torch.float32) + freqs = torch.randn(3, 1, 8, 8, dtype=torch.float32) + mrope_section = [2, 3, 3] + config = SimpleNamespace( + apply_rope_fusion=True, + mrope_section=mrope_section, + mrope_interleaved=False, + rotary_interleaved=False, + ) + + unavailable_reason = get_fused_mrope_unavailable_reason(t, freqs) + assert unavailable_reason is not None + with pytest.warns( + UserWarning, match="(CUDA tensors|Triton is not available).*Using unfused implementation" + ): + out = apply_rotary_pos_emb(t, freqs, config, cp_group=FakeCPGroup()) + assert _ROPE_FUSION_FALLBACK_WARNINGS in ( + {"triton-mrope-unavailable-device"}, + {"triton-mrope-unavailable-import"}, + ) + + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + ref = _apply_rotary_pos_emb_bshd(t, emb, rotary_interleaved=False) + torch.testing.assert_close(ref, out) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_raw_mrope_unsupported_dtype_falls_back_to_unfused(): + t, freqs, mrope_section = _make_inputs(dtype=torch.float64) + config = TransformerConfig( + num_attention_heads=t.shape[2], + num_layers=1, + apply_rope_fusion=True, + mrope_section=mrope_section, + ) + + assert "dtype" in get_fused_mrope_unavailable_reason(t, freqs) + with pytest.warns(UserWarning, match="dtype.*Using unfused implementation"): + out = apply_rotary_pos_emb(t, freqs, config, cp_group=FakeCPGroup()) + assert _ROPE_FUSION_FALLBACK_WARNINGS == {"triton-mrope-unavailable-dtype"} + + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + ref = _apply_rotary_pos_emb_bshd(t, emb, rotary_interleaved=False) + torch.testing.assert_close(ref, out) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("interleaved_mrope", [False, True]) +def test_apply_rotary_pos_emb_dispatches_raw_mrope(interleaved_mrope): + t, freqs, mrope_section = _make_inputs(interleaved_mrope=interleaved_mrope) + config = TransformerConfig( + num_attention_heads=t.shape[2], + num_layers=1, + apply_rope_fusion=True, + mrope_section=mrope_section, + mrope_interleaved=interleaved_mrope, + ) + + out = apply_rotary_pos_emb(t, freqs, config, cp_group=FakeCPGroup()) + + emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + ref = _apply_rotary_pos_emb_bshd(t, emb, rotary_interleaved=False) + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_raw_mrope_unsupported_freq_dtype_warning_key_is_dtype(): + t, freqs, mrope_section = _make_inputs() + freqs = freqs.to(torch.float16) + config = TransformerConfig( + num_attention_heads=t.shape[2], + num_layers=1, + apply_rope_fusion=True, + mrope_section=mrope_section, + ) + + assert "float32" in get_fused_mrope_unavailable_reason(t, freqs) + with pytest.warns(UserWarning, match="float32.*Using unfused implementation"): + out = apply_rotary_pos_emb(t, freqs, config, cp_group=FakeCPGroup()) + assert _ROPE_FUSION_FALLBACK_WARNINGS == {"triton-mrope-unavailable-dtype"} + + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + ref = _apply_rotary_pos_emb_bshd(t, emb, rotary_interleaved=False) + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +def test_apply_rotary_pos_emb_raw_mrope_checks_triton_availability_once(monkeypatch): + t = torch.randn(4, 1, 2, 8, dtype=torch.float32) + freqs = torch.randn(3, 1, 4, 4, dtype=torch.float32) + config = SimpleNamespace( + apply_rope_fusion=True, + mrope_section=[1, 1, 2], + mrope_interleaved=False, + rotary_interleaved=False, + ) + + calls = 0 + + def fake_unavailable_reason(*args, **kwargs): + nonlocal calls + calls += 1 + return None + + monkeypatch.setattr(rope_utils, "get_fused_mrope_unavailable_reason", fake_unavailable_reason) + monkeypatch.setattr(rope_utils, "fused_apply_mrope", lambda *args, **kwargs: t + 1) + + out = apply_rotary_pos_emb(t, freqs, config, cp_group=FakeCPGroup()) + + assert calls == 1 + torch.testing.assert_close(out, t + 1) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("layout", ["bshd", "thd"]) +def test_materialized_mrope_falls_back_without_te_fused_rope(monkeypatch, layout): + if layout == "bshd": + t, freqs, mrope_section = _make_inputs(batch=1) + cu_seqlens = None + monkeypatch.setattr(rope_utils, "fused_apply_rotary_pos_emb", None) + else: + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs() + monkeypatch.setattr(rope_utils, "fused_apply_rotary_pos_emb_thd", None) + + config = _make_mrope_config(t.shape[-2], mrope_section) + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + + with pytest.warns(UserWarning, match="Transformer Engine fused RoPE.*unavailable"): + out = apply_rotary_pos_emb(t, emb, config, cu_seqlens, cp_group=FakeCPGroup()) + + if layout == "bshd": + ref = _apply_rotary_pos_emb_bshd(t, emb, rotary_interleaved=False) + else: + ref = _apply_rotary_pos_emb_thd(t, cu_seqlens, emb, cp_group=FakeCPGroup()) + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize( + "fallback_kwargs, expected_warning_key, warning_text", + [ + ( + {"mscale": 1.25}, + "te-rope-thd-mscale", + "mscale=1.25 is not supported by TE's fused RoPE for THD layout", + ), + ( + {"inverse": True}, + "te-rope-thd-inverse", + "inverse RoPE is not supported by TE's fused RoPE for THD layout", + ), + ( + {"mla_rotary_interleaved": True}, + "te-rope-thd-mla-rotary-interleaved", + "does not support MLA-style interleaving", + ), + ], +) +def test_materialized_thd_mrope_option_fallbacks_do_not_call_te( + monkeypatch, fallback_kwargs, expected_warning_key, warning_text +): + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs() + config = _make_mrope_config(t.shape[1], mrope_section) + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + + def unexpected_te_thd_call(*args, **kwargs): + raise AssertionError("TE THD fused RoPE should not be called") + + monkeypatch.setattr(rope_utils, "fused_apply_rotary_pos_emb_thd", unexpected_te_thd_call) + with warnings.catch_warnings(record=True) as recorded_warnings: + warnings.simplefilter("always") + out = apply_rotary_pos_emb( + t, emb, config, cu_seqlens, cp_group=FakeCPGroup(), **fallback_kwargs + ) + + fallback_warnings = _fallback_warnings(recorded_warnings) + assert len(fallback_warnings) == 1 + assert warning_text in str(fallback_warnings[0].message) + assert _ROPE_FUSION_FALLBACK_WARNINGS == {expected_warning_key} + + ref = _apply_rotary_pos_emb_thd(t, cu_seqlens, emb, cp_group=FakeCPGroup(), **fallback_kwargs) + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("interleaved_mrope", [False, True]) +@pytest.mark.parametrize("cp_size, cp_rank", [(1, 0), (2, 0), (2, 1)]) +def test_fused_mrope_thd_matches_unfused_forward_backward( + interleaved_mrope, cp_size, cp_rank, monkeypatch +): + t_ref, freqs, cu_seqlens, mrope_section = _make_thd_inputs( + requires_grad=True, interleaved_mrope=interleaved_mrope, cp_size=cp_size + ) + t_fused = t_ref.detach().clone().requires_grad_(True) + cp_group = FakeCPGroup(size=cp_size, rank=cp_rank) + config = TransformerConfig( + num_attention_heads=t_ref.shape[1], + num_layers=1, + context_parallel_size=cp_size, + apply_rope_fusion=True, + mrope_section=mrope_section, + mrope_interleaved=interleaved_mrope, + ) + + emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + ref = _apply_rotary_pos_emb_thd(t_ref, cu_seqlens, emb, cp_group=cp_group) + + fused_calls = 0 + orig_fused_apply_mrope_thd = rope_utils.fused_apply_mrope_thd + + def wrapped_fused_apply_mrope_thd(*args, **kwargs): + nonlocal fused_calls + fused_calls += 1 + return orig_fused_apply_mrope_thd(*args, **kwargs) + + def unexpected_pack(*args, **kwargs): + raise AssertionError("raw THD mRoPE fusion should not materialize packed freqs") + + monkeypatch.setattr(rope_utils, "fused_apply_mrope_thd", wrapped_fused_apply_mrope_thd) + monkeypatch.setattr(rope_utils, "_pack_thd_raw_mrope_freqs", unexpected_pack) + out = apply_rotary_pos_emb(t_fused, freqs, config, cu_seqlens, cp_group=cp_group) + assert fused_calls == 1 + + tols = _dtype_tols(t_ref.dtype) + torch.testing.assert_close(ref.float(), out.float(), **tols) + + grad = torch.randn_like(ref) + ref.backward(grad) + out.backward(grad) + torch.testing.assert_close(t_ref.grad.float(), t_fused.grad.float(), **tols) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("interleaved_mrope", [False, True]) +def test_apply_rotary_pos_emb_thd_eval_uses_triton_without_te(interleaved_mrope, monkeypatch): + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs(interleaved_mrope=interleaved_mrope) + config = _make_mrope_config(t.shape[1], mrope_section, interleaved_mrope) + + emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + ref = _apply_rotary_pos_emb_thd(t, cu_seqlens, emb, cp_group=FakeCPGroup()) + + fused_calls = 0 + orig_fused_apply_mrope_thd = rope_utils.fused_apply_mrope_thd + + def wrapped_fused_apply_mrope_thd(*args, **kwargs): + nonlocal fused_calls + fused_calls += 1 + return orig_fused_apply_mrope_thd(*args, **kwargs) + + def unexpected_pack(*args, **kwargs): + raise AssertionError("raw THD mRoPE fusion should not materialize packed freqs") + + monkeypatch.setattr(rope_utils, "fused_apply_rotary_pos_emb_thd", None) + monkeypatch.setattr(rope_utils, "fused_apply_mrope_thd", wrapped_fused_apply_mrope_thd) + monkeypatch.setattr(rope_utils, "_pack_thd_raw_mrope_freqs", unexpected_pack) + with torch.no_grad(), warnings.catch_warnings(record=True) as recorded_warnings: + warnings.simplefilter("always") + out = apply_rotary_pos_emb(t, freqs, config, cu_seqlens, cp_group=FakeCPGroup()) + + assert fused_calls == 1 + assert not _fallback_warnings(recorded_warnings) + assert not out.requires_grad + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_apply_rotary_pos_emb_thd_fused_dispatch_does_not_read_cuda_scalars(monkeypatch): + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs() + config = _make_mrope_config(t.shape[1], mrope_section) + + def unexpected_item(_tensor): + raise AssertionError("fused raw THD mRoPE dispatch should not call Tensor.item()") + + monkeypatch.setattr(torch.Tensor, "item", unexpected_item) + out = apply_rotary_pos_emb(t, freqs, config, cu_seqlens, cp_group=FakeCPGroup()) + + assert out.shape == t.shape + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize( + "fallback_kwargs, config_kwargs, expected_warning_key, warning_text", + [ + ( + {"mscale": 1.25}, + {}, + "triton-mrope-thd-mscale", + "mscale=1.25 is not supported by Triton fused mRoPE for THD layout", + ), + ( + {"inverse": True}, + {}, + "triton-mrope-thd-inverse", + "inverse RoPE is not supported by Triton fused mRoPE for THD layout", + ), + ( + {"mla_rotary_interleaved": True}, + {}, + "triton-mrope-thd-mla-rotary-interleaved", + "does not support MLA-style interleaving", + ), + ( + {}, + {"rotary_interleaved": True}, + "triton-mrope-thd-rotary-interleaved", + "currently supports rotary_interleaved=False", + ), + ], +) +def test_apply_rotary_pos_emb_thd_raw_mrope_fallback_emits_option_warning( + fallback_kwargs, config_kwargs, expected_warning_key, warning_text +): + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs() + config = _make_mrope_config(t.shape[1], mrope_section, **config_kwargs) + + with warnings.catch_warnings(record=True) as recorded_warnings: + warnings.simplefilter("always") + out = apply_rotary_pos_emb( + t, freqs, config, cu_seqlens, cp_group=FakeCPGroup(), **fallback_kwargs + ) + + fallback_warnings = _fallback_warnings(recorded_warnings) + assert len(fallback_warnings) == 1 + assert warning_text in str(fallback_warnings[0].message) + assert _ROPE_FUSION_FALLBACK_WARNINGS == {expected_warning_key} + + emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, rotary_interleaved=config.rotary_interleaved + ) + ref = _apply_rotary_pos_emb_thd( + t, + cu_seqlens, + emb, + rotary_interleaved=config.rotary_interleaved, + cp_group=FakeCPGroup(), + **fallback_kwargs, + ) + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_thd_raw_mrope_rejects_sequence_length_mismatch(): + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs() + config = _make_mrope_config(t.shape[1], mrope_section) + bad_freqs = freqs[:, :, :-1, :].contiguous() + + with pytest.raises(ValueError, match="sequence length must match local tokens"): + apply_rotary_pos_emb(t, bad_freqs, config, cu_seqlens, cp_group=FakeCPGroup()) + + +def test_thd_raw_mrope_rejects_global_sequence_length_not_divisible_by_cp(): + t = torch.randn(2, 3, 20, dtype=torch.float32) + freqs = torch.randn(3, 1, 5, 8, dtype=torch.float32) + cu_seqlens = torch.tensor([0, 5], dtype=torch.int32) + config = SimpleNamespace( + apply_rope_fusion=True, + mrope_section=[2, 3, 3], + mrope_interleaved=False, + rotary_interleaved=False, + ) + + with pytest.raises(ValueError, match="divisible by context parallel size"): + apply_rotary_pos_emb(t, freqs, config, cu_seqlens, cp_group=FakeCPGroup(size=2)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_thd_raw_mrope_unavailable_reason_rejects_global_sequence_length_not_divisible_by_cp(): + t = torch.randn(3, 3, 20, dtype=torch.bfloat16, device="cuda") + freqs = torch.randn(3, 1, 5, 8, dtype=torch.float32, device="cuda") + cu_seqlens = torch.tensor([0, 5], dtype=torch.int32, device="cuda") + + assert "divisible by context parallel size" in get_fused_mrope_thd_unavailable_reason( + t, cu_seqlens, freqs, cp_size=2, cp_rank=0 + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("cp_rank", [0, 1]) +def test_thd_raw_mrope_cp_odd_local_sequence_lengths_match_manual_reference(cp_rank): + cp_size = 2 + t_ref, freqs, cu_seqlens, mrope_section = _make_thd_inputs( + requires_grad=True, cp_size=cp_size, padded_seq_lens=(10, 14) + ) + t_fused = t_ref.detach().clone().requires_grad_(True) + config = _make_mrope_config(t_ref.shape[1], mrope_section) + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + + cu_seqlens_cpu = cu_seqlens.cpu().tolist() + _assert_thd_cp_freq_index_coverage(cu_seqlens_cpu, cp_size) + packed_freqs = emb[_thd_cp_freq_indices(cu_seqlens_cpu, cp_size, cp_rank)] + + ref = _apply_rotary_pos_emb_bshd(t_ref.unsqueeze(1), packed_freqs).squeeze(1) + out = apply_rotary_pos_emb( + t_fused, freqs, config, cu_seqlens, cp_group=FakeCPGroup(size=cp_size, rank=cp_rank) + ) + + tols = _dtype_tols(t_ref.dtype) + torch.testing.assert_close(ref.float(), out.float(), **tols) + + grad = torch.randn_like(ref) + ref.backward(grad) + out.backward(grad) + torch.testing.assert_close(t_ref.grad.float(), t_fused.grad.float(), **tols) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("cp_rank", [0, 1]) +def test_thd_raw_mrope_fallback_supports_odd_local_sequence_lengths(cp_rank): + cp_size = 2 + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs( + cp_size=cp_size, padded_seq_lens=(10, 14) + ) + config = _make_mrope_config(t.shape[1], mrope_section) + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + cu_seqlens_cpu = cu_seqlens.cpu().tolist() + packed_freqs = emb[_thd_cp_freq_indices(cu_seqlens_cpu, cp_size, cp_rank)] + + with pytest.warns(UserWarning, match="mscale=1.25.*Using unfused implementation"): + out = apply_rotary_pos_emb( + t, + freqs, + config, + cu_seqlens, + mscale=1.25, + cp_group=FakeCPGroup(size=cp_size, rank=cp_rank), + ) + + ref = _apply_rotary_pos_emb_bshd(t.unsqueeze(1), packed_freqs, mscale=1.25).squeeze(1) + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_thd_raw_mrope_rejects_batch_dimension_greater_than_one(): + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs() + config = _make_mrope_config(t.shape[1], mrope_section) + bad_freqs = freqs.expand(-1, 2, -1, -1).contiguous() + + with pytest.raises(ValueError, match="singleton batch dimension"): + apply_rotary_pos_emb(t, bad_freqs, config, cu_seqlens, cp_group=FakeCPGroup()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_thd_raw_mrope_rejects_non_thd_tensor_shape(): + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs() + config = _make_mrope_config(t.shape[1], mrope_section) + + with pytest.raises(ValueError, match="raw mRoPE THD expects t"): + apply_rotary_pos_emb( + t[..., :8].unsqueeze(1), freqs, config, cu_seqlens, cp_group=FakeCPGroup() + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_fused_mrope_thd_public_api_matches_unfused(): + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs() + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + + ref = _apply_rotary_pos_emb_thd(t, cu_seqlens, emb, cp_group=FakeCPGroup()) + assert ( + get_fused_mrope_thd_unavailable_reason(t, cu_seqlens, freqs, cp_size=1, cp_rank=0) is None + ) + out = fused_apply_mrope_thd(t, cu_seqlens, freqs, mrope_section) + + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("padded_seq_lens", [(28,), (8, 10, 10)]) +def test_fused_mrope_thd_matches_unfused_for_different_sequence_counts(padded_seq_lens): + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs(padded_seq_lens=padded_seq_lens) + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + + ref = _apply_rotary_pos_emb_thd(t, cu_seqlens, emb, cp_group=FakeCPGroup()) + out = fused_apply_mrope_thd(t, cu_seqlens, freqs, mrope_section) + + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_fused_mrope_thd_fp32_compute_matches_explicit_cast_forward_backward(): + t_ref, freqs, cu_seqlens, mrope_section = _make_thd_inputs(requires_grad=True) + t_fused = t_ref.detach().clone().requires_grad_(True) + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + + ref = _apply_rotary_pos_emb_thd(t_ref.float(), cu_seqlens, emb, cp_group=FakeCPGroup()).to( + t_ref.dtype + ) + out = fused_apply_mrope_thd(t_fused, cu_seqlens, freqs, mrope_section, fp32_compute=True) + + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t_ref.dtype)) + + grad = torch.randn_like(ref) + ref.backward(grad) + out.backward(grad) + torch.testing.assert_close(t_ref.grad.float(), t_fused.grad.float(), **_dtype_tols(t_ref.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("return_raw_freqs", [False, True]) +def test_mrope_packed_seq_keeps_global_freqs_with_context_parallel(return_raw_freqs): + class FakeCPGroup2: + def size(self): + return 2 + + def rank(self): + return 0 + + seq = 16 + batch = 1 + head_dim = 20 + rotary_dim = 16 + mrope_section = [2, 3, 3] + cp_group = FakeCPGroup2() + position_ids = _make_position_ids(seq, batch) + rope = MultimodalRotaryEmbedding( + head_dim, rotary_percent=rotary_dim / head_dim, cp_group=cp_group + ) + + unpacked_freqs = rope( + position_ids, + mrope_section, + cp_group=cp_group, + return_raw_freqs=return_raw_freqs, + packed_seq=False, + ) + packed_freqs = rope( + position_ids, + mrope_section, + cp_group=cp_group, + return_raw_freqs=return_raw_freqs, + packed_seq=True, + ) + + seq_dim = 2 if return_raw_freqs else 0 + assert unpacked_freqs.shape[seq_dim] == seq // cp_group.size() + assert packed_freqs.shape[seq_dim] == seq + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.skipif(Utils.world_size < 2, reason="CP test requires at least 2 distributed ranks") +@pytest.mark.parametrize("interleaved_mrope", [False, True]) +def test_raw_mrope_fusion_matches_unfused_with_context_parallel(interleaved_mrope): + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=2) + try: + cp_group = parallel_state.get_context_parallel_group() + seq = 32 + batch = 2 + heads = 3 + head_dim = 20 + rotary_dim = 16 + mrope_section = [3, 3, 2] if interleaved_mrope else [2, 3, 3] + position_ids = _make_position_ids(seq, batch) + + rope = MultimodalRotaryEmbedding( + head_dim, + rotary_percent=rotary_dim / head_dim, + cp_group=cp_group, + interleaved_mrope=interleaved_mrope, + ) + raw_freqs = rope(position_ids, mrope_section, cp_group=cp_group, return_raw_freqs=True) + materialized_emb = rope(position_ids, mrope_section, cp_group=cp_group) + raw_freqs_emb = mrope_freqs_to_rotary_emb( + raw_freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + torch.testing.assert_close(raw_freqs_emb, materialized_emb) + + local_seq = seq // cp_group.size() + assert raw_freqs.shape == (3, batch, local_seq, rotary_dim // 2) + assert materialized_emb.shape == (local_seq, batch, 1, rotary_dim) + + generator = torch.Generator(device="cuda").manual_seed(4321) + t_ref = torch.randn( + local_seq, + batch, + heads, + head_dim, + dtype=torch.bfloat16, + device="cuda", + generator=generator, + requires_grad=True, + ) + t_fused = t_ref.detach().clone().requires_grad_(True) + + config = TransformerConfig( + num_attention_heads=heads, + num_layers=1, + context_parallel_size=cp_group.size(), + apply_rope_fusion=True, + mrope_section=mrope_section, + mrope_interleaved=interleaved_mrope, + ) + + ref = _apply_rotary_pos_emb_bshd(t_ref, materialized_emb, rotary_interleaved=False) + out = apply_rotary_pos_emb(t_fused, raw_freqs, config, cp_group=cp_group) + tols = _dtype_tols(t_ref.dtype) + torch.testing.assert_close(ref.float(), out.float(), **tols) + + grad = torch.randn_like(ref) + ref.backward(grad) + out.backward(grad) + torch.testing.assert_close(t_ref.grad.float(), t_fused.grad.float(), **tols) + finally: + Utils.destroy_model_parallel() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.skipif( + Utils.world_size < 2, reason="THD CP test requires at least 2 distributed ranks" +) +@pytest.mark.parametrize("interleaved_mrope", [False, True]) +def test_raw_mrope_thd_fusion_matches_unfused_with_context_parallel(interleaved_mrope): + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=2) + try: + cp_group = parallel_state.get_context_parallel_group() + t_ref, _, cu_seqlens, mrope_section = _make_thd_inputs( + requires_grad=True, interleaved_mrope=interleaved_mrope, cp_size=cp_group.size() + ) + t_fused = t_ref.detach().clone().requires_grad_(True) + config = TransformerConfig( + num_attention_heads=t_ref.shape[1], + num_layers=1, + context_parallel_size=cp_group.size(), + apply_rope_fusion=True, + mrope_section=mrope_section, + mrope_interleaved=interleaved_mrope, + ) + total_seq = int(cu_seqlens[-1].item()) + position_ids = _make_position_ids(total_seq, 1) + rope = MultimodalRotaryEmbedding( + t_ref.shape[-1], + rotary_percent=16 / t_ref.shape[-1], + cp_group=cp_group, + interleaved_mrope=interleaved_mrope, + ) + freqs = rope( + position_ids, mrope_section, cp_group=cp_group, return_raw_freqs=True, packed_seq=True + ) + emb = rope(position_ids, mrope_section, cp_group=cp_group, packed_seq=True) + + raw_freqs_emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + assert freqs.shape == (3, 1, total_seq, 8) + assert emb.shape == (total_seq, 1, 1, 16) + torch.testing.assert_close(raw_freqs_emb, emb) + + ref = _apply_rotary_pos_emb_thd(t_ref, cu_seqlens, emb, cp_group=cp_group) + out = apply_rotary_pos_emb(t_fused, freqs, config, cu_seqlens, cp_group=cp_group) + tols = _dtype_tols(t_ref.dtype) + torch.testing.assert_close(ref.float(), out.float(), **tols) + + grad = torch.randn_like(ref) + ref.backward(grad) + out.backward(grad) + torch.testing.assert_close(t_ref.grad.float(), t_fused.grad.float(), **tols) + finally: + Utils.destroy_model_parallel() + + +# --------------------------------------------------------------------------- +# Real Qwen3.5-VL deployment shapes. +# +# The parametrized tests above use head_dim=16/20 with rotary_dim=16 (~80% of +# channels rotated). The real Qwen3.5-VL config is head_dim=256 with +# rotary_percent=0.25 -> rotary_dim=64 (only 25% rotated, 75% pass-through) and +# mrope_section=[11,11,10] (interleaved). Exercise those exact shapes so a kernel +# regression in the large-pass-through / large-section regime is caught. +# --------------------------------------------------------------------------- + +# (head_dim, rotary_dim, mrope_section, interleaved_mrope) +_REAL_BSHD_SHAPES = [ + (256, 64, [11, 11, 10], True), # Qwen3.5-VL LLM decoder (75% pass-through) + (256, 64, [10, 11, 11], False), # same, section (non-interleaved) layout + (256, 256, [43, 43, 42], True), # full rotary (no pass-through) +] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("head_dim,rotary_dim,mrope_section,interleaved_mrope", _REAL_BSHD_SHAPES) +def test_fused_mrope_matches_unfused_real_shapes( + head_dim, rotary_dim, mrope_section, interleaved_mrope +): + t_ref, freqs, mrope_section = _make_inputs( + requires_grad=True, + head_dim=head_dim, + rotary_dim=rotary_dim, + mrope_section=mrope_section, + interleaved_mrope=interleaved_mrope, + ) + t_fused = t_ref.detach().clone().requires_grad_(True) + + emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + ref = _apply_rotary_pos_emb_bshd(t_ref, emb, rotary_interleaved=False) + out = fused_apply_mrope( + t_fused, freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + + tols = _dtype_tols(t_ref.dtype) + torch.testing.assert_close(ref.float(), out.float(), **tols) + + grad = torch.randn_like(ref) + ref.backward(grad) + out.backward(grad) + torch.testing.assert_close(t_ref.grad.float(), t_fused.grad.float(), **tols) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("interleaved_mrope", [False, True]) +def test_fused_mrope_thd_matches_unfused_real_shapes(interleaved_mrope): + # Real Qwen3.5-VL head_dim=256, rotary_dim=64 in THD packed layout. + section = [11, 11, 10] if interleaved_mrope else [10, 11, 11] + t_ref, freqs, cu_seqlens, mrope_section = _make_thd_inputs( + requires_grad=True, + interleaved_mrope=interleaved_mrope, + head_dim=256, + rotary_dim=64, + mrope_section=section, + ) + t_fused = t_ref.detach().clone().requires_grad_(True) + cp_group = FakeCPGroup(size=1, rank=0) + + emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + ref = _apply_rotary_pos_emb_thd(t_ref, cu_seqlens, emb, cp_group=cp_group) + out = fused_apply_mrope_thd( + t_fused, cu_seqlens, freqs, mrope_section, + interleaved_mrope=interleaved_mrope, rotary_interleaved=False, cp_size=1, cp_rank=0, + ) + + tols = _dtype_tols(t_ref.dtype) + torch.testing.assert_close(ref.float(), out.float(), **tols) + + grad = torch.randn_like(ref) + ref.backward(grad) + out.backward(grad) + torch.testing.assert_close(t_ref.grad.float(), t_fused.grad.float(), **tols) + + +def test_thd_unavailable_reason_rejects_non_cp_divisible_subsequence(): + # Per-sequence CP divisibility: total length is divisible by cp_size but an + # individual packed sub-sequence is not. The fused THD launch path + # (apply_rotary_pos_emb -> fused_apply_mrope_thd) must reject this so it falls + # back to the unfused path (which splits per-sequence correctly), instead of + # silently computing wrong CP token indices. + cp_size = 2 + # sub-sequence lengths 10 and 14 -> both even (OK); 9 and 15 -> total 24 even + # but each odd (must be rejected). + cu_seqlens = torch.tensor([0, 9, 24], dtype=torch.int32, device="cuda") + local_tokens = 24 // cp_size + t = torch.randn(local_tokens, 3, 20, dtype=torch.bfloat16, device="cuda") + freqs = torch.randn(3, 1, 24, 8, dtype=torch.float32, device="cuda") + reason = get_fused_mrope_thd_unavailable_reason( + t, cu_seqlens, freqs, rotary_interleaved=False, cp_size=cp_size, cp_rank=0 + ) + assert reason is not None and "sub-sequence" in reason, reason + + # Control: all sub-sequences divisible by cp_size -> launchable (reason None). + cu_ok = torch.tensor([0, 10, 24], dtype=torch.int32, device="cuda") + reason_ok = get_fused_mrope_thd_unavailable_reason( + t, cu_ok, freqs, rotary_interleaved=False, cp_size=cp_size, cp_rank=0 + ) + assert reason_ok is None, reason_ok diff --git a/tests/unit_tests/ssm/bench_gdn_cuda_opt.py b/tests/unit_tests/ssm/bench_gdn_cuda_opt.py new file mode 100644 index 00000000000..c1f82ac1443 --- /dev/null +++ b/tests/unit_tests/ssm/bench_gdn_cuda_opt.py @@ -0,0 +1,457 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Direct GatedDeltaNet CUDA optimization correctness and performance runner. + +This runner intentionally uses installed packages and normal project imports. +Install `mcore_gdn_opt` and FLA in editable mode before running it. +""" + +import argparse +import importlib.util +import os +import statistics +from contextlib import nullcontext +from dataclasses import dataclass + +import torch +import torch.nn.functional as F + +from megatron.core import parallel_state +from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_experimental_attention_variant_module_spec, +) +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer import TransformerConfig +from tests.unit_tests.test_utilities import Utils + + +FLAGS = ( + "MCORE_GDN_USE_OPT_WRAPPER", + "MCORE_GDN_OPT_BACKEND", + "MCORE_GDN_OPT_WARN_FALLBACK", + "MCORE_GDN_OPT_ENABLE_FWD_H", + "MCORE_GDN_OPT_ENABLE_WY_BWD", + "MCORE_GDN_OPT_ENABLE_DV_DHU", + "MCORE_GDN_OPT_ENABLE_DHU", + "MCORE_GDN_OPT_ENABLE_DQKWG", + "MCORE_GDN_OPT_ENABLE_DHU_DQKWG", + "FLA_CUTE_FWD_H", + "CHUNK_DELTA_FWD_USE_BWD_PORT", + "FLA_CUTE_WY_BWD", + "FLA_CUTE_BWD_DV_DHU", + "FLA_CUTE_BWD_DHU", + "FLA_CUTE_BWD_DQKWG", + "FLA_CUTE_BWD_DHU_DQKWG", + "FLA_CUTE_BWD_DHU_DQKWG_KERNEL", + "FLA_CUTE_BWD_DHU_DQKWG_DIRECT", +) + + +SCENARIOS = { + "baseline": ("Triton baseline", {}), + "wrapper_fla": ( + "MCore wrapper forced FLA", + {"MCORE_GDN_USE_OPT_WRAPPER": "1", "MCORE_GDN_OPT_BACKEND": "fla"}, + ), + "wrapper_auto": ( + "MCore wrapper auto", + {"MCORE_GDN_USE_OPT_WRAPPER": "1", "MCORE_GDN_OPT_BACKEND": "auto"}, + ), + "wrapper_cuda": ( + "MCore wrapper forced CUDA", + {"MCORE_GDN_USE_OPT_WRAPPER": "1", "MCORE_GDN_OPT_BACKEND": "cuda"}, + ), + "wy": ( + "CUDA wy_bwd", + { + "MCORE_GDN_USE_OPT_WRAPPER": "1", + "MCORE_GDN_OPT_BACKEND": "cuda", + "MCORE_GDN_OPT_ENABLE_FWD_H": "0", + "MCORE_GDN_OPT_ENABLE_DV_DHU": "0", + "MCORE_GDN_OPT_ENABLE_DHU": "0", + "MCORE_GDN_OPT_ENABLE_DQKWG": "0", + "MCORE_GDN_OPT_ENABLE_DHU_DQKWG": "0", + }, + ), + "dv_dhu": ( + "CUDA dv_local+delta_h fused", + { + "MCORE_GDN_USE_OPT_WRAPPER": "1", + "MCORE_GDN_OPT_BACKEND": "cuda", + "MCORE_GDN_OPT_ENABLE_FWD_H": "0", + "MCORE_GDN_OPT_ENABLE_WY_BWD": "0", + "MCORE_GDN_OPT_ENABLE_DHU": "0", + "MCORE_GDN_OPT_ENABLE_DQKWG": "0", + "MCORE_GDN_OPT_ENABLE_DHU_DQKWG": "0", + }, + ), + "dhu": ( + "CUDA delta_h", + { + "MCORE_GDN_USE_OPT_WRAPPER": "1", + "MCORE_GDN_OPT_BACKEND": "cuda", + "MCORE_GDN_OPT_ENABLE_FWD_H": "0", + "MCORE_GDN_OPT_ENABLE_WY_BWD": "0", + "MCORE_GDN_OPT_ENABLE_DV_DHU": "0", + "MCORE_GDN_OPT_ENABLE_DQKWG": "0", + "MCORE_GDN_OPT_ENABLE_DHU_DQKWG": "0", + }, + ), + "dqkwg": ( + "CUDA dqkwg", + { + "MCORE_GDN_USE_OPT_WRAPPER": "1", + "MCORE_GDN_OPT_BACKEND": "cuda", + "MCORE_GDN_OPT_ENABLE_FWD_H": "0", + "MCORE_GDN_OPT_ENABLE_WY_BWD": "0", + "MCORE_GDN_OPT_ENABLE_DV_DHU": "0", + "MCORE_GDN_OPT_ENABLE_DHU": "0", + "MCORE_GDN_OPT_ENABLE_DHU_DQKWG": "0", + }, + ), + "fused": ( + "CUDA wy+dhu+dqkwg fused", + { + "MCORE_GDN_USE_OPT_WRAPPER": "1", + "MCORE_GDN_OPT_BACKEND": "cuda", + "MCORE_GDN_OPT_ENABLE_FWD_H": "0", + "MCORE_GDN_OPT_ENABLE_DV_DHU": "0", + "MCORE_GDN_OPT_ENABLE_DHU": "0", + "MCORE_GDN_OPT_ENABLE_DQKWG": "0", + }, + ), + "separate": ( + "CUDA all three separate", + { + "MCORE_GDN_USE_OPT_WRAPPER": "1", + "MCORE_GDN_OPT_BACKEND": "cuda", + "MCORE_GDN_OPT_ENABLE_FWD_H": "0", + "MCORE_GDN_OPT_ENABLE_DV_DHU": "0", + "MCORE_GDN_OPT_ENABLE_DHU_DQKWG": "0", + }, + ), + "dv_dhu_dqkwg": ( + "CUDA dv_local+delta_h fused + dqkwg", + { + "MCORE_GDN_USE_OPT_WRAPPER": "1", + "MCORE_GDN_OPT_BACKEND": "cuda", + "MCORE_GDN_OPT_ENABLE_FWD_H": "0", + "MCORE_GDN_OPT_ENABLE_WY_BWD": "0", + "MCORE_GDN_OPT_ENABLE_DHU": "0", + "MCORE_GDN_OPT_ENABLE_DHU_DQKWG": "0", + }, + ), + "all_four": ( + "CUDA all four", + { + "MCORE_GDN_USE_OPT_WRAPPER": "1", + "MCORE_GDN_OPT_BACKEND": "cuda", + "MCORE_GDN_OPT_ENABLE_DV_DHU": "0", + "MCORE_GDN_OPT_ENABLE_DHU_DQKWG": "0", + }, + ), + "all_four_dv_dhu": ( + "CUDA fwd_h+wy+dv_dhu+dqkwg", + { + "MCORE_GDN_USE_OPT_WRAPPER": "1", + "MCORE_GDN_OPT_BACKEND": "cuda", + "MCORE_GDN_OPT_ENABLE_DHU": "0", + "MCORE_GDN_OPT_ENABLE_DHU_DQKWG": "0", + }, + ), +} + + +@dataclass +class AccuracyRow: + name: str + status: str + output_max_abs: float + input_grad_max_abs: float + worst_param: str + worst_param_max_abs: float + + +@dataclass +class PerfRow: + name: str + mean_us: float + median_us: float + min_us: float + max_us: float + speedup: float + + +def set_env(overrides): + for flag in FLAGS: + os.environ.pop(flag, None) + if "MCORE_GDN_USE_OPT_WRAPPER" not in overrides: + os.environ["MCORE_GDN_USE_OPT_WRAPPER"] = "0" + if "MCORE_GDN_OPT_BACKEND" not in overrides: + os.environ["MCORE_GDN_OPT_BACKEND"] = "fla" + os.environ.update(overrides) + + +def set_model_dispatch(model): + if os.environ.get("MCORE_GDN_USE_OPT_WRAPPER", "0") == "1": + from mcore_gdn_opt.gated_delta_rule import chunk_gated_delta_rule + else: + from fla.ops.gated_delta_rule import chunk_gated_delta_rule + + model.gated_delta_rule = chunk_gated_delta_rule + + +def validate_dispatch_sources(scenario_items): + if any("MCORE_GDN_OPT_BACKEND" in env for _, (_, env) in scenario_items): + for module_name in ( + "mcore_gdn_opt.gated_delta_rule.chunk", + "mcore_gdn_opt.gated_delta_rule.backward", + ): + spec = importlib.util.find_spec(module_name) + if spec is None or spec.origin is None: + raise RuntimeError(f"cannot locate required mcore_gdn_opt module {module_name!r}") + print(f"MCORE_GDN_OPT_DISPATCH_SOURCE module={module_name} path={spec.origin}", flush=True) + + +def nvtx_range(label, enabled=True): + if enabled and torch.cuda.is_available(): + return torch.cuda.nvtx.range(label) + return nullcontext() + + +def scenario_label(index, name): + safe_name = name.replace(" ", "_").replace("+", "plus").replace("/", "_") + return f"gdn_only/{index:02d}_{safe_name}" + + +def make_model(dtype): + from megatron.core.ssm.gated_delta_net import GatedDeltaNet + + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, context_parallel_size=1 + ) + model_parallel_cuda_manual_seed(123) + pg_collection = ProcessGroupCollection( + tp=parallel_state.get_tensor_model_parallel_group(), + cp=parallel_state.get_context_parallel_group(), + ) + cfg = TransformerConfig( + hidden_size=128, + linear_conv_kernel_dim=2, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_num_key_heads=64, + linear_num_value_heads=64, + num_layers=1, + normalization="RMSNorm", + use_cpu_initialization=True, + layernorm_zero_centered_gamma=True, + num_attention_heads=64, + activation_func=F.silu, + bf16=(dtype == torch.bfloat16), + fp16=(dtype == torch.float16), + experimental_attention_variant="gated_delta_net", + linear_attention_freq=[1], + transformer_impl="transformer_engine", + ) + submodules = get_experimental_attention_variant_module_spec(config=cfg).submodules + return GatedDeltaNet( + cfg, + submodules=submodules, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=pg_collection, + ).cuda().to(dtype) + + +def zero_grads(model): + model.zero_grad(set_to_none=True) + + +def compute_loss(output, loss): + if loss == "sum": + return output.float().sum() + if loss == "square_mean": + return output.float().square().mean() + raise ValueError(f"unknown loss: {loss}") + + +def run_once(model, x, env, loss, nvtx_label=None, use_nvtx=True): + set_env(env) + set_model_dispatch(model) + print( + "RUN_ONCE " + f"label={nvtx_label or 'none'} " + f"use_wrapper={os.environ.get('MCORE_GDN_USE_OPT_WRAPPER', '')} " + f"backend={os.environ.get('MCORE_GDN_OPT_BACKEND', '')}", + flush=True, + ) + zero_grads(model) + inp = x.detach().clone().requires_grad_(True) + with nvtx_range(nvtx_label, enabled=use_nvtx and nvtx_label is not None): + out, _ = model(inp, attention_mask=None) + compute_loss(out, loss).backward() + torch.cuda.synchronize() + grads = { + name: param.grad.detach().float().clone().cpu() + for name, param in model.named_parameters() + if param.grad is not None + } + return out.detach().float().clone().cpu(), inp.grad.detach().float().clone().cpu(), grads + + +def diff_max_abs(actual, expected): + return float((actual - expected).abs().max().item()) + + +def allclose(actual, expected, atol, rtol): + return bool(torch.isfinite(actual).all().item()) and bool( + torch.allclose(actual, expected, atol=atol, rtol=rtol) + ) + + +def check_accuracy(model, x, scenario_items, loss, atol, rtol, use_nvtx=True): + base_name, base_env = SCENARIOS["baseline"] + base_out, base_grad, base_params = run_once( + model, x, base_env, loss, "gdn_only/00_accuracy_reference/Triton_baseline", use_nvtx + ) + rows = [] + for scenario_idx, (_, (name, env)) in enumerate(scenario_items, start=1): + label = f"{scenario_label(scenario_idx, name)}/accuracy" + out, grad, params = run_once(model, x, env, loss, label, use_nvtx) + output_ok = allclose(out, base_out, atol, rtol) + grad_ok = allclose(grad, base_grad, atol, rtol) + worst_param = "" + worst_param_abs = 0.0 + params_ok = True + for param_name, expected in base_params.items(): + actual = params[param_name] + params_ok = params_ok and allclose(actual, expected, atol, rtol) + param_abs = diff_max_abs(actual, expected) + if param_abs > worst_param_abs: + worst_param = param_name + worst_param_abs = param_abs + rows.append( + AccuracyRow( + name=name, + status="PASS" if output_ok and grad_ok and params_ok else "FAIL", + output_max_abs=diff_max_abs(out, base_out), + input_grad_max_abs=diff_max_abs(grad, base_grad), + worst_param=worst_param, + worst_param_max_abs=worst_param_abs, + ) + ) + return rows + + +def fwd_bwd(model, x, env, loss, nvtx_label=None, use_nvtx=True): + set_env(env) + set_model_dispatch(model) + zero_grads(model) + inp = x.detach().requires_grad_(True) + with nvtx_range(nvtx_label, enabled=use_nvtx and nvtx_label is not None): + out, _ = model(inp, attention_mask=None) + compute_loss(out, loss).backward() + + +def benchmark(model, x, scenario_items, loss, warmup, repeats, rounds, use_nvtx=True): + rows = [] + baseline_us = None + for scenario_idx, (_, (name, env)) in enumerate(scenario_items, start=1): + base_label = scenario_label(scenario_idx, name) + for warmup_idx in range(warmup): + fwd_bwd(model, x, env, loss, f"{base_label}/warmup_{warmup_idx:02d}", use_nvtx) + torch.cuda.synchronize() + samples = [] + for round_idx in range(rounds): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + with nvtx_range(f"{base_label}/round_{round_idx:02d}/measured_{repeats}iters", enabled=use_nvtx): + start.record() + for iter_idx in range(repeats): + fwd_bwd(model, x, env, loss, f"{base_label}/round_{round_idx:02d}/iter_{iter_idx:02d}", use_nvtx) + end.record() + torch.cuda.synchronize() + samples.append(start.elapsed_time(end) * 1000.0 / repeats) + mean_us = statistics.mean(samples) + if baseline_us is None: + baseline_us = mean_us + rows.append( + PerfRow( + name=name, + mean_us=mean_us, + median_us=statistics.median(samples), + min_us=min(samples), + max_us=max(samples), + speedup=baseline_us / mean_us, + ) + ) + return rows + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dtype", choices=("bf16", "fp16"), default="bf16") + parser.add_argument("--loss", choices=("sum", "square_mean"), default="square_mean") + parser.add_argument("--scenarios", default="baseline,fused,separate,all_four") + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--repeats", type=int, default=20) + parser.add_argument("--rounds", type=int, default=3) + parser.add_argument("--atol", type=float, default=5e-3) + parser.add_argument("--rtol", type=float, default=5e-3) + parser.add_argument("--fail-on-accuracy", action="store_true") + parser.add_argument("--no-nvtx", dest="use_nvtx", action="store_false", default=True) + return parser.parse_args() + + +def main(): + args = parse_args() + keys = [key.strip() for key in args.scenarios.split(",") if key.strip()] + if "baseline" not in keys: + keys.insert(0, "baseline") + unknown = [key for key in keys if key not in SCENARIOS] + if unknown: + raise ValueError(f"unknown scenarios: {unknown}; choices={sorted(SCENARIOS)}") + scenario_items = [(key, SCENARIOS[key]) for key in keys] + validate_dispatch_sources(scenario_items) + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 + + torch.manual_seed(123) + set_env({}) + print( + f"DEVICE {torch.cuda.get_device_name(0)} SHAPE B=2 T=8192 H=64 D=128 " + f"dtype={args.dtype} loss={args.loss}" + ) + try: + model = make_model(dtype).eval() + x = torch.randn(8192, 2, 128, device="cuda", dtype=dtype) + accuracy_rows = check_accuracy(model, x, scenario_items, args.loss, args.atol, args.rtol, args.use_nvtx) + for row in accuracy_rows: + print( + f"ACCURACY name={row.name!r} status={row.status} " + f"output_max_abs={row.output_max_abs:.9f} " + f"input_grad_max_abs={row.input_grad_max_abs:.9f} " + f"worst_param={row.worst_param} " + f"worst_param_max_abs={row.worst_param_max_abs:.9f}" + ) + perf_rows = benchmark(model, x, scenario_items, args.loss, args.warmup, args.repeats, args.rounds, args.use_nvtx) + for row in perf_rows: + print( + f"PERF name={row.name!r} mean_us={row.mean_us:.3f} " + f"median_us={row.median_us:.3f} min_us={row.min_us:.3f} " + f"max_us={row.max_us:.3f} speedup_vs_baseline={row.speedup:.3f}" + ) + if args.fail_on_accuracy and any(row.status != "PASS" for row in accuracy_rows): + raise SystemExit(1) + finally: + set_env({}) + Utils.destroy_model_parallel() + + +if __name__ == "__main__": + main() diff --git a/tests/unit_tests/ssm/test_bench_gdn_cuda_opt_scenarios.py b/tests/unit_tests/ssm/test_bench_gdn_cuda_opt_scenarios.py new file mode 100644 index 00000000000..3d583985ef2 --- /dev/null +++ b/tests/unit_tests/ssm/test_bench_gdn_cuda_opt_scenarios.py @@ -0,0 +1,43 @@ +import ast +from pathlib import Path + + +BENCH = Path(__file__).with_name("bench_gdn_cuda_opt.py") + + +def _literal_assignment(name): + tree = ast.parse(BENCH.read_text()) + for node in tree.body: + if isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Name) and target.id == name: + return ast.literal_eval(node.value) + raise AssertionError(f"{name} assignment not found") + + +def test_optimized_scenarios_route_through_mcore_wrapper(): + scenarios = _literal_assignment("SCENARIOS") + optimized = [ + "wy", + "dv_dhu", + "dhu", + "dqkwg", + "fused", + "separate", + "dv_dhu_dqkwg", + "all_four", + "all_four_dv_dhu", + ] + + for key in optimized: + env = scenarios[key][1] + assert env["MCORE_GDN_USE_OPT_WRAPPER"] == "1", key + assert env["MCORE_GDN_OPT_BACKEND"] == "cuda", key + assert not any(flag.startswith("FLA_CUTE_") for flag in env), key + + +def test_benchmark_does_not_require_patched_fla_sources(): + text = BENCH.read_text() + + assert "patched flash-linear-attention" not in text + assert "FLA_DISPATCH_SOURCE" not in text diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index f490e7cfdb8..44cd551c19a 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -1,5 +1,6 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import copy from functools import partial from unittest import mock @@ -16,6 +17,7 @@ get_transformer_block_with_experimental_attention_variant_spec, ) from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.gated_delta_net import GatedDeltaNet from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed @@ -45,6 +47,49 @@ HAVE_FLA = False +def _make_gdn_config(**overrides): + config_kwargs = { + "hidden_size": 128, + "linear_conv_kernel_dim": 2, + "linear_key_head_dim": 32, + "linear_value_head_dim": 32, + "linear_num_key_heads": 4, + "linear_num_value_heads": 8, + "num_layers": 1, + "normalization": "RMSNorm", + "use_cpu_initialization": True, + "layernorm_zero_centered_gamma": True, + "num_attention_heads": 8, + "activation_func": F.silu, + "bf16": True, + "experimental_attention_variant": "gated_delta_net", + "linear_attention_freq": [1], + "transformer_impl": "transformer_engine", + } + config_kwargs.update(overrides) + return TransformerConfig(**config_kwargs) + + +@pytest.mark.parametrize("pre_gated_delta_rule_impl", ["unfused", "fused_streamed", "fused_mega"]) +def test_pre_gated_delta_rule_impl_accepts_gdn_modes(pre_gated_delta_rule_impl): + config = _make_gdn_config(pre_gated_delta_rule_impl=pre_gated_delta_rule_impl) + assert config.pre_gated_delta_rule_impl == pre_gated_delta_rule_impl + + +def test_pre_gated_delta_rule_impl_rejects_invalid_value(): + with pytest.raises(ValueError, match="pre_gated_delta_rule_impl must be one of"): + _make_gdn_config(pre_gated_delta_rule_impl="fused") + + +def test_pre_gated_delta_rule_impl_requires_gdn_variant(): + with pytest.raises(ValueError, match="experimental_attention_variant='gated_delta_net'"): + _make_gdn_config( + experimental_attention_variant=None, + linear_attention_freq=None, + pre_gated_delta_rule_impl="fused_streamed", + ) + + @pytest.mark.parametrize( ("tp_size", "sp", "cp_size"), [(1, False, 1), (2, False, 1), (2, True, 1), (1, False, 2), (2, False, 2), (2, True, 2)], @@ -142,6 +187,171 @@ def test_gpu_forward(self): output.dtype == hidden_states.dtype ), f"Output dtype {output.dtype=} mismatch with {hidden_states.dtype=}" + def test_selective_recompute_norm_out(self): + tp_group = parallel_state.get_tensor_model_parallel_group() + cp_group = parallel_state.get_context_parallel_group() + pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group) + + def build_gdn(config): + gdn_submodules = get_experimental_attention_variant_module_spec( + config=config + ).submodules + gdn = GatedDeltaNet( + config, + submodules=gdn_submodules, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=pg_collection, + ) + return gdn.cuda().bfloat16() + + def run(gdn, hidden_states): + output, _ = gdn(hidden_states, None) + output.float().sum().backward() + grads = { + name: param.grad.detach() + for name, param in gdn.named_parameters() + if param.grad is not None + } + input_grad = hidden_states.grad.detach().clone() + return output.detach(), grads, input_grad + + micro_batch_size = 2 + seq_length = 64 + base_config = copy.deepcopy(self.transformer_config) + rec_config = copy.deepcopy(self.transformer_config) + rec_config.recompute_granularity = "selective" + rec_config.recompute_modules = ["gdn_norm_out"] + + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + hidden_states = torch.randn( + ( + seq_length // self.sp_size // self.cp_size, + micro_batch_size, + self.gdn.config.hidden_size, + ), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + requires_grad=True, + ) + + # --- Baseline (no recompute) --- + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + base_gdn = build_gdn(base_config) + assert base_gdn.recompute_norm_out is False + base_output, base_grads, base_input_grad = run(base_gdn, hidden_states) + hidden_states.grad = None + del base_gdn + torch.cuda.empty_cache() + + # --- Recompute --- + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + rec_gdn = build_gdn(rec_config) + assert rec_gdn.recompute_norm_out is True + rec_output, rec_grads, rec_input_grad = run(rec_gdn, hidden_states) + + rank = torch.distributed.get_rank() + assert torch.equal(rec_output, base_output), f"Output not identical ({rank=})" + assert torch.equal(rec_input_grad, base_input_grad), f"Input grad not identical ({rank=})" + assert set(rec_grads.keys()) == set(base_grads.keys()) + for name in base_grads: + assert torch.equal( + rec_grads[name], base_grads[name] + ), f"Grad not identical for {name} ({rank=})" + + def test_selective_recompute_gdn_qkv(self): + """gdn_qkv discard-output recompute must be numerically exact. + + recompute_modules=["gdn_qkv"] recomputes the whole QKV projection + + preparation block (in_proj -> CP a2a -> conv1d -> _prepare_qkv -> g/beta) + as a discard-output checkpoint. Output, parameter grads and input grad + must match the no-recompute baseline bit-for-bit. + """ + tp_group = parallel_state.get_tensor_model_parallel_group() + cp_group = parallel_state.get_context_parallel_group() + pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group) + + def build_gdn(config): + gdn_submodules = get_experimental_attention_variant_module_spec( + config=config + ).submodules + gdn = GatedDeltaNet( + config, + submodules=gdn_submodules, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=pg_collection, + ) + return gdn.cuda().bfloat16() + + def run(gdn, hidden_states): + output, _ = gdn(hidden_states, None) + output.float().sum().backward() + grads = { + name: param.grad.detach() + for name, param in gdn.named_parameters() + if param.grad is not None + } + input_grad = hidden_states.grad.detach().clone() + return output.detach(), grads, input_grad + + micro_batch_size = 2 + seq_length = 64 + base_config = copy.deepcopy(self.transformer_config) + rec_config = copy.deepcopy(self.transformer_config) + rec_config.recompute_granularity = "selective" + rec_config.recompute_modules = ["gdn_qkv"] + + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + hidden_states = torch.randn( + ( + seq_length // self.sp_size // self.cp_size, + micro_batch_size, + self.gdn.config.hidden_size, + ), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + requires_grad=True, + ) + + # --- Baseline (no recompute) --- + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + base_gdn = build_gdn(base_config) + assert base_gdn.recompute_qkv is False + base_output, base_grads, base_input_grad = run(base_gdn, hidden_states) + hidden_states.grad = None + del base_gdn + torch.cuda.empty_cache() + + # --- Recompute --- + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + rec_gdn = build_gdn(rec_config) + assert rec_gdn.recompute_qkv is True + rec_output, rec_grads, rec_input_grad = run(rec_gdn, hidden_states) + + rank = torch.distributed.get_rank() + assert torch.equal(rec_output, base_output), f"Output not identical ({rank=})" + assert torch.equal(rec_input_grad, base_input_grad), f"Input grad not identical ({rank=})" + assert set(rec_grads.keys()) == set(base_grads.keys()) + for name in base_grads: + assert torch.equal( + rec_grads[name], base_grads[name] + ), f"Grad not identical for {name} ({rank=})" + def test_jit_compiled_helpers(self): import torch._dynamo @@ -309,6 +519,462 @@ def test_gpu_forward_thd_padding_correctness(self): self.gdn(hidden_states_thd, None, packed_seq_params=actual_mismatch_params) +@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") +@pytest.mark.internal +class TestFusedPreGatedDeltaRule: + + @pytest.fixture(scope='function', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=1, + ) + model_parallel_cuda_manual_seed(123) + + tp_group = parallel_state.get_tensor_model_parallel_group() + cp_group = parallel_state.get_context_parallel_group() + self.pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group) + + self.unfused_gdn = self._build_gdn(pre_gated_delta_rule_impl="unfused") + self.fused_gdn = self._build_gdn(pre_gated_delta_rule_impl="fused_streamed") + self.fused_gdn.load_state_dict(self.unfused_gdn.state_dict()) + + def teardown_method(self): + Utils.destroy_model_parallel() + + def _build_gdn( + self, + pre_gated_delta_rule_impl: str, + *, + deterministic_mode: bool = True, + conv_kernel_dim: int = 2, + ): + transformer_config = TransformerConfig( + hidden_size=256, + linear_conv_kernel_dim=conv_kernel_dim, + linear_key_head_dim=64, + linear_value_head_dim=64, + linear_num_key_heads=4, + linear_num_value_heads=8, + num_layers=1, + normalization="RMSNorm", + use_cpu_initialization=True, + layernorm_zero_centered_gamma=True, + num_attention_heads=8, + activation_func=F.silu, + bf16=True, + tensor_model_parallel_size=1, + context_parallel_size=1, + experimental_attention_variant="gated_delta_net", + linear_attention_freq=[1], + transformer_impl="transformer_engine", + deterministic_mode=deterministic_mode, + pre_gated_delta_rule_impl=pre_gated_delta_rule_impl, + ) + gdn_submodules = get_experimental_attention_variant_module_spec( + config=transformer_config + ).submodules + gdn = GatedDeltaNet( + transformer_config, + submodules=gdn_submodules, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=self.pg_collection, + ) + return gdn.cuda().bfloat16() + + def _packed_pre_gated_delta_rule_reference(self, gdn, qkvzba, cu_seqlens): + """Run the dense torch reference independently on each packed sequence.""" + + segment_outputs = [[] for _ in range(6)] + for start, end in zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist()): + outputs = gdn.pre_gated_delta_rule(qkvzba[start:end], batch=1, seq_len=end - start) + for output_list, output in zip(segment_outputs, outputs): + output_list.append(output) + return tuple(torch.cat(outputs, dim=1) for outputs in segment_outputs) + + def _assert_pre_gated_delta_rule_outputs_close( + self, + fused_outputs, + unfused_outputs, + *, + atol: float, + rtol: float, + output_tolerances=None, + ): + """Compare named pre-GDR outputs with optional per-output tolerances.""" + + output_names = ("query", "key", "value", "gate", "beta", "g") + output_tolerances = output_tolerances or {} + for name, fused, unfused in zip(output_names, fused_outputs, unfused_outputs): + output_atol, output_rtol = output_tolerances.get(name, (atol, rtol)) + torch.testing.assert_close( + fused, + unfused, + atol=output_atol, + rtol=output_rtol, + msg=lambda msg, output_name=name: f"{output_name} mismatch: {msg}", + ) + + def test_fused_and_unfused_forward_match(self): + hidden_states = torch.randn( + (32, 2, self.unfused_gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + with torch.no_grad(): + unfused_output, unfused_bias = self.unfused_gdn(hidden_states, None) + fused_output, fused_bias = self.fused_gdn(hidden_states, None) + + torch.testing.assert_close(fused_output, unfused_output, atol=1e-3, rtol=1e-3) + assert fused_bias == unfused_bias + + @pytest.mark.parametrize("pre_gated_delta_rule_impl", ["fused_streamed", "fused_mega"]) + def test_fused_and_unfused_forward_thd_match(self, pre_gated_delta_rule_impl): + unfused_gdn = self._build_gdn( + pre_gated_delta_rule_impl="unfused", + deterministic_mode=False, + conv_kernel_dim=4, + ) + fused_gdn = self._build_gdn( + pre_gated_delta_rule_impl=pre_gated_delta_rule_impl, + deterministic_mode=False, + conv_kernel_dim=4, + ) + fused_gdn.load_state_dict(unfused_gdn.state_dict()) + + hidden_states = torch.randn( + (32, 1, unfused_gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + cu_seqlens = torch.tensor([0, 1, 4, 11, 32], device=torch.cuda.current_device(), dtype=torch.int32) + packed_seq_params = PackedSeqParams( + qkv_format='thd', + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=21, + max_seqlen_kv=21, + total_tokens=hidden_states.shape[0], + ) + assert packed_seq_params.seq_idx is not None + + with torch.no_grad(): + unfused_output, unfused_bias = unfused_gdn( + hidden_states, None, packed_seq_params=packed_seq_params + ) + fused_output, fused_bias = fused_gdn( + hidden_states, None, packed_seq_params=packed_seq_params + ) + + torch.testing.assert_close(fused_output, unfused_output, atol=2e-3, rtol=2e-3) + assert fused_bias == unfused_bias + + def test_fused_and_unfused_forward_thd_padding_match(self): + unfused_gdn = self._build_gdn( + pre_gated_delta_rule_impl="unfused", + deterministic_mode=False, + conv_kernel_dim=4, + ) + fused_gdn = self._build_gdn( + pre_gated_delta_rule_impl="fused_streamed", + deterministic_mode=False, + conv_kernel_dim=4, + ) + fused_gdn.load_state_dict(unfused_gdn.state_dict()) + + hidden_states = torch.randn( + (12, 1, unfused_gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + cu_seqlens = torch.tensor([0, 1, 4, 9], device=torch.cuda.current_device(), dtype=torch.int32) + cu_seqlens_padded = torch.tensor( + [0, 2, 6, 12], device=torch.cuda.current_device(), dtype=torch.int32 + ) + packed_seq_params = PackedSeqParams( + qkv_format='thd', + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=6, + max_seqlen_kv=6, + total_tokens=hidden_states.shape[0], + ) + assert packed_seq_params.seq_idx is not None + + with torch.no_grad(): + unfused_output, unfused_bias = unfused_gdn( + hidden_states, None, packed_seq_params=packed_seq_params + ) + fused_output, fused_bias = fused_gdn( + hidden_states, None, packed_seq_params=packed_seq_params + ) + + torch.testing.assert_close(fused_output, unfused_output, atol=2e-3, rtol=2e-3) + assert fused_bias == unfused_bias + + def test_fused_and_unfused_pre_gated_delta_rule_match(self): + batch = 2 + seq_len = 32 + hidden_states = torch.randn( + (seq_len, batch, self.unfused_gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + with torch.no_grad(): + qkvzba, _ = self.unfused_gdn.in_proj(hidden_states) + unfused_outputs = self.unfused_gdn.pre_gated_delta_rule(qkvzba, batch, seq_len) + fused_outputs = self.fused_gdn._fused_streamed_pre_gated_delta_rule(qkvzba) + + self._assert_pre_gated_delta_rule_outputs_close( + fused_outputs, + unfused_outputs, + atol=1e-3, + rtol=1e-3, + # g uses Triton exp/log softplus in the fused path and torch softplus + # in the reference path, so its direct intermediate parity needs a + # slightly looser relative tolerance than the layout/conv outputs. + output_tolerances={"g": (1e-3, 3e-3)}, + ) + + def test_fused_and_unfused_packed_pre_gated_delta_rule_forward_match(self): + reference_gdn = self._build_gdn( + pre_gated_delta_rule_impl="unfused", + deterministic_mode=True, + conv_kernel_dim=4, + ) + fused_gdn = self._build_gdn( + pre_gated_delta_rule_impl="fused_streamed", + deterministic_mode=False, + conv_kernel_dim=4, + ) + fused_gdn.load_state_dict(reference_gdn.state_dict()) + + batch = 1 + cu_seqlens = torch.tensor([0, 1, 4, 6, 11], device=torch.cuda.current_device(), dtype=torch.int32) + seq_len = cu_seqlens[-1].item() + qkvzba = torch.randn( + (seq_len, batch, reference_gdn.in_proj_dim), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + with torch.no_grad(): + unfused_outputs = self._packed_pre_gated_delta_rule_reference( + reference_gdn, qkvzba, cu_seqlens + ) + fused_outputs = fused_gdn._fused_streamed_pre_gated_delta_rule( + qkvzba, cu_seqlens_q=cu_seqlens + ) + + self._assert_pre_gated_delta_rule_outputs_close( + fused_outputs, unfused_outputs, atol=2e-3, rtol=2e-3 + ) + + def test_fused_and_unfused_packed_pre_gated_delta_rule_backward_match(self): + reference_gdn = self._build_gdn( + pre_gated_delta_rule_impl="unfused", + deterministic_mode=True, + conv_kernel_dim=4, + ) + fused_gdn = self._build_gdn( + pre_gated_delta_rule_impl="fused_streamed", + deterministic_mode=False, + conv_kernel_dim=4, + ) + fused_gdn.load_state_dict(reference_gdn.state_dict()) + + batch = 1 + cu_seqlens = torch.tensor([0, 1, 4, 6, 11], device=torch.cuda.current_device(), dtype=torch.int32) + seq_len = cu_seqlens[-1].item() + qkvzba = torch.randn( + (seq_len, batch, reference_gdn.in_proj_dim), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + qkvzba_unfused = qkvzba.detach().clone().requires_grad_(True) + qkvzba_fused = qkvzba.detach().clone().requires_grad_(True) + + reference_gdn.zero_grad(set_to_none=True) + fused_gdn.zero_grad(set_to_none=True) + + unfused_outputs = self._packed_pre_gated_delta_rule_reference( + reference_gdn, qkvzba_unfused, cu_seqlens + ) + fused_outputs = fused_gdn._fused_streamed_pre_gated_delta_rule( + qkvzba_fused, cu_seqlens_q=cu_seqlens + ) + grad_outputs = [torch.randn_like(output.float()) for output in unfused_outputs] + + unfused_loss = sum( + (output.float() * grad).sum() for output, grad in zip(unfused_outputs, grad_outputs) + ) + fused_loss = sum( + (output.float() * grad).sum() for output, grad in zip(fused_outputs, grad_outputs) + ) + unfused_loss.backward() + fused_loss.backward() + + torch.testing.assert_close(qkvzba_fused.grad, qkvzba_unfused.grad, atol=3e-2, rtol=3e-2) + torch.testing.assert_close( + fused_gdn.conv1d.weight.grad, + reference_gdn.conv1d.weight.grad, + atol=3e-2, + rtol=3e-2, + ) + torch.testing.assert_close(fused_gdn.A_log.grad, reference_gdn.A_log.grad, atol=3e-2, rtol=3e-2) + torch.testing.assert_close( + fused_gdn.dt_bias.grad, reference_gdn.dt_bias.grad, atol=3e-2, rtol=3e-2 + ) + + def test_fused_packed_conv_forward_boundary_isolation(self): + from megatron.core.fusions.fused_pre_gated_delta_rule import ( + fused_streamed_pre_gated_delta_rule, + ) + + seq_len = 5 + boundary = 3 + num_key_heads = 1 + # Keep qkvzba.stride(0) aligned for causal_conv1d's channel-last + # backward guard; the boundary condition under test is independent + # of the value-head repeat factor. + num_value_heads = 4 + key_head_dim = 32 + value_head_dim = 32 + conv_width = 4 + qk_channels = num_key_heads * key_head_dim + v_channels = num_value_heads * value_head_dim + v_offset = 2 * qk_channels + k_offset = qk_channels + total_channels = 2 * qk_channels + 2 * v_channels + 2 * num_value_heads + device = torch.cuda.current_device() + + qkvzba = torch.zeros((seq_len, 1, total_channels), device=device, dtype=torch.bfloat16) + qkvzba[boundary - 1, 0, :qk_channels] = 10.0 + qkvzba[boundary - 1, 0, k_offset : k_offset + qk_channels] = 10.0 + qkvzba[boundary - 1, 0, v_offset : v_offset + v_channels] = 10.0 + conv_weight = torch.zeros((2 * qk_channels + v_channels, 1, conv_width), device=device) + conv_weight[:qk_channels, 0, conv_width - 2] = 1.0 + conv_weight[k_offset : k_offset + qk_channels, 0, conv_width - 2] = 1.0 + conv_weight[v_offset : v_offset + v_channels, 0, conv_width - 2] = 1.0 + A_log = torch.zeros((num_value_heads,), device=device, dtype=torch.bfloat16) + dt_bias = torch.zeros((num_value_heads,), device=device, dtype=torch.bfloat16) + cu_seqlens = torch.tensor([0, boundary, seq_len], device=device, dtype=torch.int32) + + query, key, value, _, _, _ = fused_streamed_pre_gated_delta_rule( + qkvzba, + conv_weight.to(torch.bfloat16), + None, + A_log, + dt_bias, + num_key_heads=num_key_heads, + num_value_heads=num_value_heads, + key_head_dim=key_head_dim, + value_head_dim=value_head_dim, + cu_seqlens=cu_seqlens, + ) + + torch.testing.assert_close( + query[0, boundary], + torch.zeros_like(query[0, boundary]), + atol=0.0, + rtol=0.0, + ) + torch.testing.assert_close( + key[0, boundary], + torch.zeros_like(key[0, boundary]), + atol=0.0, + rtol=0.0, + ) + torch.testing.assert_close( + value[0, boundary], + torch.zeros_like(value[0, boundary]), + atol=0.0, + rtol=0.0, + ) + + def test_fused_packed_conv_backward_boundary_isolation(self): + from megatron.core.fusions.fused_pre_gated_delta_rule import ( + fused_streamed_pre_gated_delta_rule, + ) + + seq_len = 5 + boundary = 3 + num_key_heads = 1 + # Keep qkvzba.stride(0) aligned for causal_conv1d's channel-last + # backward guard; the boundary condition under test is independent + # of the value-head repeat factor. + num_value_heads = 4 + key_head_dim = 32 + value_head_dim = 32 + conv_width = 4 + qk_channels = num_key_heads * key_head_dim + v_channels = num_value_heads * value_head_dim + v_offset = 2 * qk_channels + k_offset = qk_channels + total_channels = 2 * qk_channels + 2 * v_channels + 2 * num_value_heads + device = torch.cuda.current_device() + + qkvzba = torch.zeros( + (seq_len, 1, total_channels), device=device, dtype=torch.bfloat16, requires_grad=True + ) + conv_weight = torch.zeros( + (2 * qk_channels + v_channels, 1, conv_width), + device=device, + dtype=torch.bfloat16, + requires_grad=True, + ) + with torch.no_grad(): + conv_weight[:qk_channels, 0, conv_width - 2] = 1.0 + conv_weight[k_offset : k_offset + qk_channels, 0, conv_width - 2] = 1.0 + conv_weight[v_offset : v_offset + v_channels, 0, conv_width - 2] = 1.0 + A_log = torch.zeros((num_value_heads,), device=device, dtype=torch.bfloat16, requires_grad=True) + dt_bias = torch.zeros( + (num_value_heads,), device=device, dtype=torch.bfloat16, requires_grad=True + ) + cu_seqlens = torch.tensor([0, boundary, seq_len], device=device, dtype=torch.int32) + + query, key, value, gate, beta, g = fused_streamed_pre_gated_delta_rule( + qkvzba, + conv_weight, + None, + A_log, + dt_bias, + num_key_heads=num_key_heads, + num_value_heads=num_value_heads, + key_head_dim=key_head_dim, + value_head_dim=value_head_dim, + cu_seqlens=cu_seqlens, + ) + + loss = ( + query[0, boundary].float().sum() + + key[0, boundary].float().sum() + + value[0, boundary].float().sum() + ) + loss = loss + 0.0 * ( + gate.float().sum() + + beta.float().sum() + + g.float().sum() + ) + loss.backward() + leaked_q_grad = qkvzba.grad[boundary - 1, 0, :qk_channels] + leaked_k_grad = qkvzba.grad[boundary - 1, 0, k_offset : k_offset + qk_channels] + leaked_grad = qkvzba.grad[boundary - 1, 0, v_offset : v_offset + v_channels] + torch.testing.assert_close(leaked_q_grad, torch.zeros_like(leaked_q_grad), atol=0.0, rtol=0.0) + torch.testing.assert_close(leaked_k_grad, torch.zeros_like(leaked_k_grad), atol=0.0, rtol=0.0) + torch.testing.assert_close(leaked_grad, torch.zeros_like(leaked_grad), atol=0.0, rtol=0.0) + + @pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") @pytest.mark.internal class TestGDNCuSeqlensResolve: diff --git a/tests/unit_tests/ssm/test_gated_delta_net_cuda_opt.py b/tests/unit_tests/ssm/test_gated_delta_net_cuda_opt.py new file mode 100644 index 00000000000..46309db6559 --- /dev/null +++ b/tests/unit_tests/ssm/test_gated_delta_net_cuda_opt.py @@ -0,0 +1,88 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Focused GatedDeltaNet CUDA optimization coverage. + +This keeps the optimized-kernel correctness and optional perf check separate +from the generic GatedDeltaNet unit tests. +""" + +import os + +import pytest +import torch + +from tests.unit_tests.ssm import bench_gdn_cuda_opt as runner + +try: + import fla # noqa: F401 + + HAVE_FLA = True +except ImportError: + HAVE_FLA = False + + +def _scenario_items(): + keys = [ + key.strip() + for key in os.environ.get( + "MCORE_GDN_UNIT_TEST_SCENARIOS", "baseline,all_four_dv_dhu" + ).split(",") + if key.strip() + ] + if "baseline" not in keys: + keys.insert(0, "baseline") + unknown = [key for key in keys if key not in runner.SCENARIOS] + if unknown: + raise ValueError(f"unknown GDN CUDA opt scenarios: {unknown}") + return [(key, runner.SCENARIOS[key]) for key in keys] + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) +@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available.") +@pytest.mark.internal +def test_gated_delta_net_cuda_opt_correctness_and_optional_perf(dtype): + scenario_items = _scenario_items() + runner.validate_dispatch_sources(scenario_items) + + torch.manual_seed(123) + runner.set_env({}) + try: + model = runner.make_model(dtype).eval() + seq_len = int(os.environ.get("MCORE_GDN_UNIT_TEST_T", "8192")) + batch = int(os.environ.get("MCORE_GDN_UNIT_TEST_B", "2")) + x = torch.randn(seq_len, batch, 128, device="cuda", dtype=dtype) + + atol = float(os.environ.get("MCORE_GDN_UNIT_TEST_ATOL", "5e-3")) + rtol = float(os.environ.get("MCORE_GDN_UNIT_TEST_RTOL", "5e-3")) + loss = os.environ.get("MCORE_GDN_UNIT_TEST_LOSS", "sum") + accuracy_rows = runner.check_accuracy( + model, x, scenario_items, loss=loss, atol=atol, rtol=rtol, use_nvtx=False + ) + failed = [row for row in accuracy_rows if row.status != "PASS"] + assert not failed, "\n".join( + f"{row.name}: output={row.output_max_abs:.9f} " + f"input_grad={row.input_grad_max_abs:.9f} " + f"{row.worst_param}={row.worst_param_max_abs:.9f}" + for row in failed + ) + + if os.environ.get("MCORE_GDN_UNIT_TEST_PERF", "0") == "1": + perf_rows = runner.benchmark( + model, + x, + scenario_items, + loss=loss, + warmup=int(os.environ.get("MCORE_GDN_UNIT_TEST_WARMUP", "5")), + repeats=int(os.environ.get("MCORE_GDN_UNIT_TEST_REPEATS", "20")), + rounds=int(os.environ.get("MCORE_GDN_UNIT_TEST_ROUNDS", "3")), + use_nvtx=True, + ) + for row in perf_rows: + print( + f"PERF {row.name}: mean_us={row.mean_us:.3f} " + f"speedup_vs_baseline={row.speedup:.3f}" + ) + finally: + runner.set_env({}) + runner.Utils.destroy_model_parallel() diff --git a/third_party/mcore_gdn_opt b/third_party/mcore_gdn_opt new file mode 160000 index 00000000000..12605c515bd --- /dev/null +++ b/third_party/mcore_gdn_opt @@ -0,0 +1 @@ +Subproject commit 12605c515bdbc1f239df991a9cea570f7e79d234