| """Fixed production kernels for Diagonal Kalman Delta Networks. |
| |
| Vendored from https://github.com/ngocbh/kalman-delta-networks |
| (``lit_gpt/kdn_ops/diag_kdn_chunk.py`` at commit |
| 53e437be8591521a508bb6a04cd5b26594a3d5ec) for arXiv:2609.07816. |
| |
| Upstream notes: generated mechanically from the reachable qualified |
| covariance-state S3 implementation at frozen reference commit |
| 95de7b2fc8b2ebec478761b1f32b32b7afa32474. The |
| FLA-derived kernels retain the upstream MIT license and Copyright 2023-2026 |
| FLA contributors. |
| |
| The pinned flash-linear-attention build has no ``USE_CUDA_GRAPH`` switch, so |
| autotuning benchmarks run without CUDA graphs. |
| """ |
| from __future__ import annotations |
|
|
| import math |
|
|
| import torch |
| import triton |
| import triton.language as tl |
| from torch.autograd.function import once_differentiable |
|
|
| from fla.ops.utils.cache import fla_cache_autotune |
| from fla.ops.utils.constant import RCP_LN2 |
| from fla.ops.utils.op import exp, exp2, gather |
| from fla.ops.utils.softplus import softplus |
| from fla.utils import ( |
| IS_GATHER_SUPPORTED, |
| IS_NVIDIA_HOPPER, |
| autotune_cache_kwargs, |
| check_shared_mem, |
| ) |
|
|
| USE_CUDA_GRAPH = False |
|
|
|
|
| _DIAG_KDN_PRECISION_ARM = "P1CBWU32_INTRADQFP32_WYDKDM_TF32_BWDOUTBF16" |
| _DIAG_KDN_MEMORY_CHUNK_SIZE = 32 |
| _DIAG_KDN_GAIN_CHUNK_SIZE = 256 |
| _DIAG_KDN_FRONTEND_LOCAL_CHUNK_SIZE = 64 |
| _DIAG_KDN_READOUT_FORWARD = "FBT" |
| _DIAG_KDN_READOUT_BACKWARD = "BQ" |
| _DIAG_KDN_VECTOR_DTYPE = torch.bfloat16 |
| _DIAG_KDN_STATE_DTYPE = torch.float32 |
| _DIAG_KDN_SCORE_DOT_PRECISION = "tf32x3" |
| _DIAG_KDN_SOLVE_DOT_PRECISION = "tf32x3" |
| _DIAG_KDN_STATE_DOT_PRECISION = "tf32x3" |
| _DIAG_KDN_STATE_RESIDUAL_DOT_PRECISION = "tf32" |
| _DIAG_KDN_BACKWARD_DOT_PRECISION = "tf32x3" |
| _DIAG_KDN_INTRA_QUERY_DTYPE = torch.float32 |
| _DIAG_KDN_FINAL_GRADIENT_DTYPE = torch.bfloat16 |
|
|
| |
| _frozen_intra_SOLVE_TRIL_DOT_PRECISION = tl.constexpr("tf32") |
|
|
| |
|
|
| _ROW_BLOCK = 4 |
|
|
| _frozen_norm_NUM_WARPS = 4 |
|
|
| _frozen_norm_NUM_STAGES = 1 |
|
|
| @triton.jit |
| def _frozen_norm_div_rn_fp32(numerator, denominator): |
| """IEEE round-to-nearest FP32 division, never reciprocal multiplication.""" |
| return tl.inline_asm_elementwise('div.rn.f32 $0, $1, $2;', '=f,f,f', [numerator, denominator], dtype=tl.float32, is_pure=True, pack=1) |
|
|
| @triton.jit |
| def _frozen_norm_sqrt_rn_fp32(value): |
| """Non-FTZ IEEE round-to-nearest FP32 square root.""" |
| return tl.inline_asm_elementwise('sqrt.rn.f32 $0, $1;', '=f,f', [value], dtype=tl.float32, is_pure=True, pack=1) |
|
|
| @triton.jit |
| def _frozen_norm_mul_rn_fp32(left, right): |
| """FP32 multiply barrier used by each measured autograd edge.""" |
| return tl.inline_asm_elementwise('mul.rn.f32 $0, $1, $2;', '=f,f,f', [left, right], dtype=tl.float32, is_pure=True, pack=1) |
|
|
| @triton.jit |
| def _frozen_norm_add_rn_fp32(left, right): |
| """FP32 add barrier preventing edge multiplication/addition contraction.""" |
| return tl.inline_asm_elementwise('add.rn.f32 $0, $1, $2;', '=f,f,f', [left, right], dtype=tl.float32, is_pure=True, pack=1) |
|
|
| @triton.jit |
| def _frozen_norm_fma_rn_fp32(left, right, accumulator): |
| """Non-FTZ FP32 FMA used by PyTorch's NormTwo reduction.""" |
| return tl.inline_asm_elementwise('fma.rn.f32 $0, $1, $2, $3;', '=f,f,f,f', [left, right, accumulator], dtype=tl.float32, is_pure=True, pack=1) |
|
|
| @triton.jit |
| def _frozen_norm_shfl_down_fp32(value, OFFSET: tl.constexpr): |
| value_bits = tl.cast(value, tl.int32, bitcast=True) |
| offset = tl.full(value.shape, OFFSET, tl.int32) |
| shuffled_bits = tl.inline_asm_elementwise('shfl.sync.down.b32 $0, $1, $2, 0x1f, 0xffffffff;', '=r,r,r', [value_bits, offset], dtype=tl.int32, is_pure=True, pack=1) |
| return tl.cast(shuffled_bits, tl.float32, bitcast=True) |
|
|
| @triton.jit |
| def _frozen_norm_broadcast_lane0_fp32(value): |
| value_bits = tl.cast(value, tl.int32, bitcast=True) |
| source_lane = tl.zeros(value.shape, tl.int32) |
| broadcast_bits = tl.inline_asm_elementwise('shfl.sync.idx.b32 $0, $1, $2, 0x1f, 0xffffffff;', '=r,r,r', [value_bits, source_lane], dtype=tl.int32, is_pure=True, pack=1) |
| return tl.cast(broadcast_bits, tl.float32, bitcast=True) |
|
|
| @triton.jit |
| def _frozen_norm_combine_vt0(left0, left1, left2, left3): |
| value = _frozen_norm_add_rn_fp32(left0, left1) |
| value = _frozen_norm_add_rn_fp32(value, left2) |
| return _frozen_norm_add_rn_fp32(value, left3) |
|
|
| @triton.jit |
| def _frozen_norm_torch_reduce_cuh_accumulators(accumulator0, accumulator1, accumulator2, accumulator3, K: tl.constexpr, SMALL_ROWS: tl.constexpr, REDUCTION_WIDTH: tl.constexpr): |
| """PyTorch 2.8 Reduce.cuh topology, returned as a lane-0 broadcast.""" |
| positive_zero = tl.zeros(accumulator0.shape, tl.float32) |
| if REDUCTION_WIDTH == 128: |
| tl.static_assert(K == 128 and SMALL_ROWS) |
| p0 = _frozen_norm_combine_vt0(accumulator0, positive_zero, positive_zero, positive_zero) |
| p64 = _frozen_norm_combine_vt0(accumulator2, positive_zero, positive_zero, positive_zero) |
| p32 = _frozen_norm_combine_vt0(accumulator1, positive_zero, positive_zero, positive_zero) |
| p96 = _frozen_norm_combine_vt0(accumulator3, positive_zero, positive_zero, positive_zero) |
| left = _frozen_norm_add_rn_fp32(p0, p64) |
| right = _frozen_norm_add_rn_fp32(p32, p96) |
| value = _frozen_norm_add_rn_fp32(left, right) |
| elif REDUCTION_WIDTH == 64: |
| tl.static_assert(SMALL_ROWS) |
| p0 = _frozen_norm_combine_vt0(accumulator0, accumulator2, positive_zero, positive_zero) |
| p32 = _frozen_norm_combine_vt0(accumulator1, accumulator3, positive_zero, positive_zero) |
| value = _frozen_norm_add_rn_fp32(p0, p32) |
| else: |
| tl.static_assert(REDUCTION_WIDTH == 32 and (not SMALL_ROWS)) |
| value = _frozen_norm_combine_vt0(accumulator0, accumulator1, accumulator2, accumulator3) |
| value = _frozen_norm_add_rn_fp32(value, _frozen_norm_shfl_down_fp32(value, OFFSET=1)) |
| value = _frozen_norm_add_rn_fp32(value, _frozen_norm_shfl_down_fp32(value, OFFSET=2)) |
| value = _frozen_norm_add_rn_fp32(value, _frozen_norm_shfl_down_fp32(value, OFFSET=4)) |
| value = _frozen_norm_add_rn_fp32(value, _frozen_norm_shfl_down_fp32(value, OFFSET=8)) |
| value = _frozen_norm_add_rn_fp32(value, _frozen_norm_shfl_down_fp32(value, OFFSET=16)) |
| return _frozen_norm_broadcast_lane0_fp32(value) |
|
|
| @triton.jit |
| def _frozen_norm_normalize_row_forward(x0, x1, x2, x3, K: tl.constexpr, SMALL_ROWS: tl.constexpr, REDUCTION_WIDTH: tl.constexpr): |
| """Exact PyTorch NormTwo row reduction on a [ROW_BLOCK,32] tile.""" |
| positive_zero = tl.zeros(x0.shape, tl.float32) |
| accumulator0 = _frozen_norm_fma_rn_fp32(x0, x0, positive_zero) |
| accumulator1 = _frozen_norm_fma_rn_fp32(x1, x1, positive_zero) |
| if K == 128: |
| accumulator2 = _frozen_norm_fma_rn_fp32(x2, x2, positive_zero) |
| accumulator3 = _frozen_norm_fma_rn_fp32(x3, x3, positive_zero) |
| else: |
| accumulator2 = positive_zero |
| accumulator3 = positive_zero |
| square_sum = _frozen_norm_torch_reduce_cuh_accumulators(accumulator0, accumulator1, accumulator2, accumulator3, K=K, SMALL_ROWS=SMALL_ROWS, REDUCTION_WIDTH=REDUCTION_WIDTH) |
| return _frozen_norm_sqrt_rn_fp32(square_sum) |
|
|
| @triton.jit |
| def _frozen_norm_normalize_row_backward(denominator_lane0, denominator_lane1, denominator_lane2, denominator_lane3, K: tl.constexpr, SMALL_ROWS: tl.constexpr, REDUCTION_WIDTH: tl.constexpr): |
| """Exact PyTorch sum row reduction for the denominator VJP.""" |
| positive_zero = tl.zeros(denominator_lane0.shape, tl.float32) |
| accumulator0 = _frozen_norm_add_rn_fp32(positive_zero, denominator_lane0) |
| accumulator1 = _frozen_norm_add_rn_fp32(positive_zero, denominator_lane1) |
| if K == 128: |
| accumulator2 = _frozen_norm_add_rn_fp32(positive_zero, denominator_lane2) |
| accumulator3 = _frozen_norm_add_rn_fp32(positive_zero, denominator_lane3) |
| else: |
| accumulator2 = positive_zero |
| accumulator3 = positive_zero |
| return _frozen_norm_torch_reduce_cuh_accumulators(accumulator0, accumulator1, accumulator2, accumulator3, K=K, SMALL_ROWS=SMALL_ROWS, REDUCTION_WIDTH=REDUCTION_WIDTH) |
|
|
| @triton.jit |
| def _frozen_norm_store_forward_component(x, denominator, output_ptr, offsets, mask): |
| normalized = _frozen_norm_div_rn_fp32(x, denominator) |
| memory = tl.cast(normalized, tl.bfloat16, fp_downcast_rounding='rtne') |
| tl.store(output_ptr + offsets, memory, mask=mask) |
|
|
| @triton.jit |
| def _frozen_norm_store_key_forward_component(x, denominator, gain_output_ptr, memory_output_ptr, offsets, mask): |
| normalized = _frozen_norm_div_rn_fp32(x, denominator) |
| memory = tl.cast(normalized, tl.bfloat16, fp_downcast_rounding='rtne') |
| tl.store(gain_output_ptr + offsets, normalized, mask=mask) |
| tl.store(memory_output_ptr + offsets, memory, mask=mask) |
|
|
| @triton.jit |
| def _frozen_norm_store_backward_component(x, cotangent, denominator, norm, denominator_grad, output_ptr, offsets, mask, INPUT_IS_BF16: tl.constexpr, FINAL_ROUND: tl.constexpr): |
| direct_fp32 = _frozen_norm_div_rn_fp32(cotangent, denominator) |
| safe_norm = tl.where(norm == 0.0, 1.0, norm) |
| unit = _frozen_norm_div_rn_fp32(x, safe_norm) |
| unit = tl.where(norm == 0.0, 0.0, unit) |
| norm_fp32 = tl.where(norm >= 1e-12, _frozen_norm_mul_rn_fp32(denominator_grad, unit), 0.0) |
| if INPUT_IS_BF16: |
| if FINAL_ROUND: |
| combined_fp32 = _frozen_norm_add_rn_fp32(direct_fp32, norm_fp32) |
| combined_raw = tl.cast(combined_fp32, tl.bfloat16, fp_downcast_rounding='rtne') |
| tl.store(output_ptr + offsets, combined_raw, mask=mask) |
| else: |
| direct_raw = tl.cast(direct_fp32, tl.bfloat16, fp_downcast_rounding='rtne') |
| norm_raw = tl.cast(norm_fp32, tl.bfloat16, fp_downcast_rounding='rtne') |
| tl.store(output_ptr + offsets, direct_raw + norm_raw, mask=mask) |
| else: |
| tl.store(output_ptr + offsets, _frozen_norm_add_rn_fp32(direct_fp32, norm_fp32), mask=mask) |
|
|
| @triton.jit |
| def _frozen_norm_diag_exact_norm_fwd_kernel(q_ptr, k_ptr, q_memory_ptr, k_gain_ptr, k_memory_ptr, q_norm_ptr, k_norm_ptr, n_rows, K: tl.constexpr, ROW_BLOCK: tl.constexpr, PROCESS_K: tl.constexpr, SMALL_ROWS: tl.constexpr, REDUCTION_WIDTH: tl.constexpr): |
| """One launch for q rows; the reviewed helper is reused by the future k arm.""" |
| row = tl.program_id(0) * ROW_BLOCK + tl.arange(0, ROW_BLOCK) |
| lane = tl.arange(0, 32) |
| offsets0 = row[:, None] * K + lane[None, :] |
| offsets1 = offsets0 + 32 |
| offsets2 = offsets0 + 64 |
| offsets3 = offsets0 + 96 |
| row_mask = row < n_rows |
| mask0 = row_mask[:, None] |
| mask1 = row_mask[:, None] & (lane[None, :] + 32 < K) |
| mask2 = row_mask[:, None] & (lane[None, :] + 64 < K) |
| mask3 = row_mask[:, None] & (lane[None, :] + 96 < K) |
| q0 = tl.load(q_ptr + offsets0, mask=mask0, other=0.0).to(tl.float32) |
| q1 = tl.load(q_ptr + offsets1, mask=mask1, other=0.0).to(tl.float32) |
| q2 = tl.load(q_ptr + offsets2, mask=mask2, other=0.0).to(tl.float32) |
| q3 = tl.load(q_ptr + offsets3, mask=mask3, other=0.0).to(tl.float32) |
| q_norm = _frozen_norm_normalize_row_forward(q0, q1, q2, q3, K=K, SMALL_ROWS=SMALL_ROWS, REDUCTION_WIDTH=REDUCTION_WIDTH) |
| denominator = tl.where(q_norm != q_norm, q_norm, tl.maximum(q_norm, 1e-12)) |
| _frozen_norm_store_forward_component(q0, denominator, q_memory_ptr, offsets0, mask0) |
| _frozen_norm_store_forward_component(q1, denominator, q_memory_ptr, offsets1, mask1) |
| if K == 128: |
| _frozen_norm_store_forward_component(q2, denominator, q_memory_ptr, offsets2, mask2) |
| _frozen_norm_store_forward_component(q3, denominator, q_memory_ptr, offsets3, mask3) |
| lane0_mask = row_mask[:, None] & (lane[None, :] == 0) |
| tl.store(q_norm_ptr + row[:, None] + lane[None, :], q_norm, mask=lane0_mask) |
| if PROCESS_K: |
| k0 = tl.load(k_ptr + offsets0, mask=mask0, other=0.0).to(tl.float32) |
| k1 = tl.load(k_ptr + offsets1, mask=mask1, other=0.0).to(tl.float32) |
| k2 = tl.load(k_ptr + offsets2, mask=mask2, other=0.0).to(tl.float32) |
| k3 = tl.load(k_ptr + offsets3, mask=mask3, other=0.0).to(tl.float32) |
| k_norm = _frozen_norm_normalize_row_forward(k0, k1, k2, k3, K=K, SMALL_ROWS=SMALL_ROWS, REDUCTION_WIDTH=REDUCTION_WIDTH) |
| k_denominator = tl.where(k_norm != k_norm, k_norm, tl.maximum(k_norm, 1e-12)) |
| _frozen_norm_store_key_forward_component(k0, k_denominator, k_gain_ptr, k_memory_ptr, offsets0, mask0) |
| _frozen_norm_store_key_forward_component(k1, k_denominator, k_gain_ptr, k_memory_ptr, offsets1, mask1) |
| if K == 128: |
| _frozen_norm_store_key_forward_component(k2, k_denominator, k_gain_ptr, k_memory_ptr, offsets2, mask2) |
| _frozen_norm_store_key_forward_component(k3, k_denominator, k_gain_ptr, k_memory_ptr, offsets3, mask3) |
| tl.store(k_norm_ptr + row[:, None] + lane[None, :], k_norm, mask=lane0_mask) |
|
|
| @triton.jit |
| def _frozen_norm_diag_exact_norm_bwd_kernel(q_ptr, k_ptr, dq_memory_ptr, dk_gain_ptr, dk_memory_ptr, q_norm_ptr, k_norm_ptr, dq_ptr, dk_ptr, B, T, H, n_rows, dq_stride_b, dq_stride_t, dq_stride_h, dq_stride_k, dkg_stride_b, dkg_stride_t, dkg_stride_h, dkg_stride_k, dkm_stride_b, dkm_stride_t, dkm_stride_h, dkm_stride_k, K: tl.constexpr, ROW_BLOCK: tl.constexpr, PROCESS_Q: tl.constexpr, PROCESS_K: tl.constexpr, HAS_GAIN: tl.constexpr, HAS_MEMORY: tl.constexpr, INPUT_IS_BF16: tl.constexpr, FINAL_ROUND: tl.constexpr, SMALL_ROWS: tl.constexpr, REDUCTION_WIDTH: tl.constexpr): |
| """One launch with direct strided cotangent loads for the shared row VJP.""" |
| row = tl.program_id(0) * ROW_BLOCK + tl.arange(0, ROW_BLOCK) |
| lane = tl.arange(0, 32) |
| offsets0 = row[:, None] * K + lane[None, :] |
| offsets1 = offsets0 + 32 |
| offsets2 = offsets0 + 64 |
| offsets3 = offsets0 + 96 |
| row_mask = row < n_rows |
| mask0 = row_mask[:, None] |
| mask1 = row_mask[:, None] & (lane[None, :] + 32 < K) |
| mask2 = row_mask[:, None] & (lane[None, :] + 64 < K) |
| mask3 = row_mask[:, None] & (lane[None, :] + 96 < K) |
| if PROCESS_Q: |
| h_index = row % H |
| bt_index = row // H |
| t_index = bt_index % T |
| b_index = bt_index // T |
| dq_base = b_index.to(tl.int64)[:, None] * dq_stride_b + t_index.to(tl.int64)[:, None] * dq_stride_t + h_index.to(tl.int64)[:, None] * dq_stride_h |
| k0 = lane.to(tl.int64)[None, :] |
| dq_offsets0 = dq_base + k0 * dq_stride_k |
| dq_offsets1 = dq_base + (k0 + 32) * dq_stride_k |
| dq_offsets2 = dq_base + (k0 + 64) * dq_stride_k |
| dq_offsets3 = dq_base + (k0 + 96) * dq_stride_k |
| q0 = tl.load(q_ptr + offsets0, mask=mask0, other=0.0).to(tl.float32) |
| q1 = tl.load(q_ptr + offsets1, mask=mask1, other=0.0).to(tl.float32) |
| q2 = tl.load(q_ptr + offsets2, mask=mask2, other=0.0).to(tl.float32) |
| q3 = tl.load(q_ptr + offsets3, mask=mask3, other=0.0).to(tl.float32) |
| u0 = tl.load(dq_memory_ptr + dq_offsets0, mask=mask0, other=0.0).to(tl.float32) |
| u1 = tl.load(dq_memory_ptr + dq_offsets1, mask=mask1, other=0.0).to(tl.float32) |
| u2 = tl.load(dq_memory_ptr + dq_offsets2, mask=mask2, other=0.0).to(tl.float32) |
| u3 = tl.load(dq_memory_ptr + dq_offsets3, mask=mask3, other=0.0).to(tl.float32) |
| lane0_mask = row_mask[:, None] & (lane[None, :] == 0) |
| q_norm_seed = tl.load(q_norm_ptr + row[:, None] + lane[None, :], mask=lane0_mask, other=0.0).to(tl.float32) |
| q_norm = _frozen_norm_broadcast_lane0_fp32(q_norm_seed) |
| denominator = tl.where(q_norm != q_norm, q_norm, tl.maximum(q_norm, 1e-12)) |
| q0_over_d = _frozen_norm_div_rn_fp32(q0, denominator) |
| q1_over_d = _frozen_norm_div_rn_fp32(q1, denominator) |
| z0 = _frozen_norm_mul_rn_fp32(-u0, _frozen_norm_div_rn_fp32(q0_over_d, denominator)) |
| z1 = _frozen_norm_mul_rn_fp32(-u1, _frozen_norm_div_rn_fp32(q1_over_d, denominator)) |
| if K == 128: |
| q2_over_d = _frozen_norm_div_rn_fp32(q2, denominator) |
| q3_over_d = _frozen_norm_div_rn_fp32(q3, denominator) |
| z2 = _frozen_norm_mul_rn_fp32(-u2, _frozen_norm_div_rn_fp32(q2_over_d, denominator)) |
| z3 = _frozen_norm_mul_rn_fp32(-u3, _frozen_norm_div_rn_fp32(q3_over_d, denominator)) |
| else: |
| z2 = tl.zeros(z0.shape, tl.float32) |
| z3 = tl.zeros(z0.shape, tl.float32) |
| denominator_grad = _frozen_norm_normalize_row_backward(z0, z1, z2, z3, K=K, SMALL_ROWS=SMALL_ROWS, REDUCTION_WIDTH=REDUCTION_WIDTH) |
| _frozen_norm_store_backward_component(q0, u0, denominator, q_norm, denominator_grad, dq_ptr, offsets0, mask0, INPUT_IS_BF16=INPUT_IS_BF16, FINAL_ROUND=FINAL_ROUND) |
| _frozen_norm_store_backward_component(q1, u1, denominator, q_norm, denominator_grad, dq_ptr, offsets1, mask1, INPUT_IS_BF16=INPUT_IS_BF16, FINAL_ROUND=FINAL_ROUND) |
| if K == 128: |
| _frozen_norm_store_backward_component(q2, u2, denominator, q_norm, denominator_grad, dq_ptr, offsets2, mask2, INPUT_IS_BF16=INPUT_IS_BF16, FINAL_ROUND=FINAL_ROUND) |
| _frozen_norm_store_backward_component(q3, u3, denominator, q_norm, denominator_grad, dq_ptr, offsets3, mask3, INPUT_IS_BF16=INPUT_IS_BF16, FINAL_ROUND=FINAL_ROUND) |
| if PROCESS_K: |
| h_index = row % H |
| bt_index = row // H |
| t_index = bt_index % T |
| b_index = bt_index // T |
| dkg_base = b_index.to(tl.int64)[:, None] * dkg_stride_b + t_index.to(tl.int64)[:, None] * dkg_stride_t + h_index.to(tl.int64)[:, None] * dkg_stride_h |
| dkm_base = b_index.to(tl.int64)[:, None] * dkm_stride_b + t_index.to(tl.int64)[:, None] * dkm_stride_t + h_index.to(tl.int64)[:, None] * dkm_stride_h |
| key_index = lane.to(tl.int64)[None, :] |
| dkg_offsets0 = dkg_base + key_index * dkg_stride_k |
| dkg_offsets1 = dkg_base + (key_index + 32) * dkg_stride_k |
| dkg_offsets2 = dkg_base + (key_index + 64) * dkg_stride_k |
| dkg_offsets3 = dkg_base + (key_index + 96) * dkg_stride_k |
| dkm_offsets0 = dkm_base + key_index * dkm_stride_k |
| dkm_offsets1 = dkm_base + (key_index + 32) * dkm_stride_k |
| dkm_offsets2 = dkm_base + (key_index + 64) * dkm_stride_k |
| dkm_offsets3 = dkm_base + (key_index + 96) * dkm_stride_k |
| k0 = tl.load(k_ptr + offsets0, mask=mask0, other=0.0).to(tl.float32) |
| k1 = tl.load(k_ptr + offsets1, mask=mask1, other=0.0).to(tl.float32) |
| k2 = tl.load(k_ptr + offsets2, mask=mask2, other=0.0).to(tl.float32) |
| k3 = tl.load(k_ptr + offsets3, mask=mask3, other=0.0).to(tl.float32) |
| lane0_mask = row_mask[:, None] & (lane[None, :] == 0) |
| k_norm_seed = tl.load(k_norm_ptr + row[:, None] + lane[None, :], mask=lane0_mask, other=0.0).to(tl.float32) |
| k_norm = _frozen_norm_broadcast_lane0_fp32(k_norm_seed) |
| k_denominator = tl.where(k_norm != k_norm, k_norm, tl.maximum(k_norm, 1e-12)) |
| if HAS_GAIN: |
| kg0 = tl.load(dk_gain_ptr + dkg_offsets0, mask=mask0, other=0.0).to(tl.float32) |
| kg1 = tl.load(dk_gain_ptr + dkg_offsets1, mask=mask1, other=0.0).to(tl.float32) |
| kg2 = tl.load(dk_gain_ptr + dkg_offsets2, mask=mask2, other=0.0).to(tl.float32) |
| kg3 = tl.load(dk_gain_ptr + dkg_offsets3, mask=mask3, other=0.0).to(tl.float32) |
| ku0 = kg0 |
| ku1 = kg1 |
| ku2 = kg2 |
| ku3 = kg3 |
| if HAS_MEMORY: |
| km0 = tl.load(dk_memory_ptr + dkm_offsets0, mask=mask0, other=0.0).to(tl.float32) |
| km1 = tl.load(dk_memory_ptr + dkm_offsets1, mask=mask1, other=0.0).to(tl.float32) |
| km2 = tl.load(dk_memory_ptr + dkm_offsets2, mask=mask2, other=0.0).to(tl.float32) |
| km3 = tl.load(dk_memory_ptr + dkm_offsets3, mask=mask3, other=0.0).to(tl.float32) |
| ku0 = _frozen_norm_add_rn_fp32(kg0, km0) |
| ku1 = _frozen_norm_add_rn_fp32(kg1, km1) |
| ku2 = _frozen_norm_add_rn_fp32(kg2, km2) |
| ku3 = _frozen_norm_add_rn_fp32(kg3, km3) |
| else: |
| km0 = tl.load(dk_memory_ptr + dkm_offsets0, mask=mask0, other=0.0).to(tl.float32) |
| km1 = tl.load(dk_memory_ptr + dkm_offsets1, mask=mask1, other=0.0).to(tl.float32) |
| km2 = tl.load(dk_memory_ptr + dkm_offsets2, mask=mask2, other=0.0).to(tl.float32) |
| km3 = tl.load(dk_memory_ptr + dkm_offsets3, mask=mask3, other=0.0).to(tl.float32) |
| ku0 = km0 |
| ku1 = km1 |
| ku2 = km2 |
| ku3 = km3 |
| kz0 = _frozen_norm_mul_rn_fp32(-ku0, _frozen_norm_div_rn_fp32(_frozen_norm_div_rn_fp32(k0, k_denominator), k_denominator)) |
| kz1 = _frozen_norm_mul_rn_fp32(-ku1, _frozen_norm_div_rn_fp32(_frozen_norm_div_rn_fp32(k1, k_denominator), k_denominator)) |
| if K == 128: |
| kz2 = _frozen_norm_mul_rn_fp32(-ku2, _frozen_norm_div_rn_fp32(_frozen_norm_div_rn_fp32(k2, k_denominator), k_denominator)) |
| kz3 = _frozen_norm_mul_rn_fp32(-ku3, _frozen_norm_div_rn_fp32(_frozen_norm_div_rn_fp32(k3, k_denominator), k_denominator)) |
| else: |
| kz2 = tl.zeros(kz0.shape, tl.float32) |
| kz3 = tl.zeros(kz0.shape, tl.float32) |
| k_denominator_grad = _frozen_norm_normalize_row_backward(kz0, kz1, kz2, kz3, K=K, SMALL_ROWS=SMALL_ROWS, REDUCTION_WIDTH=REDUCTION_WIDTH) |
| _frozen_norm_store_backward_component(k0, ku0, k_denominator, k_norm, k_denominator_grad, dk_ptr, offsets0, mask0, INPUT_IS_BF16=INPUT_IS_BF16, FINAL_ROUND=FINAL_ROUND) |
| _frozen_norm_store_backward_component(k1, ku1, k_denominator, k_norm, k_denominator_grad, dk_ptr, offsets1, mask1, INPUT_IS_BF16=INPUT_IS_BF16, FINAL_ROUND=FINAL_ROUND) |
| if K == 128: |
| _frozen_norm_store_backward_component(k2, ku2, k_denominator, k_norm, k_denominator_grad, dk_ptr, offsets2, mask2, INPUT_IS_BF16=INPUT_IS_BF16, FINAL_ROUND=FINAL_ROUND) |
| _frozen_norm_store_backward_component(k3, ku3, k_denominator, k_norm, k_denominator_grad, dk_ptr, offsets3, mask3, INPUT_IS_BF16=INPUT_IS_BF16, FINAL_ROUND=FINAL_ROUND) |
|
|
| def _frozen_norm_reduction_width(n_rows: int, k: int) -> int: |
| return 32 if n_rows >= 16 else 64 if n_rows >= 8 else k |
|
|
| def _frozen_norm_launch_diag_norm_boundary_forward(q_raw: torch.Tensor, k_raw: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| q_memory = torch.empty_like(q_raw, dtype=torch.bfloat16, memory_format=torch.contiguous_format) |
| k_gain = torch.empty_like(k_raw, dtype=torch.float32, memory_format=torch.contiguous_format) |
| k_memory = torch.empty_like(k_raw, dtype=torch.bfloat16, memory_format=torch.contiguous_format) |
| q_norm = torch.empty(q_raw.shape[:-1], device=q_raw.device, dtype=torch.float32) |
| k_norm = torch.empty(k_raw.shape[:-1], device=k_raw.device, dtype=torch.float32) |
| n_rows = q_raw.numel() // q_raw.shape[-1] |
| if n_rows != 0: |
| grid = (triton.cdiv(n_rows, _ROW_BLOCK),) |
| _frozen_norm_diag_exact_norm_fwd_kernel[grid](q_raw, k_raw, q_memory, k_gain, k_memory, q_norm, k_norm, n_rows, K=q_raw.shape[-1], ROW_BLOCK=_ROW_BLOCK, PROCESS_K=True, SMALL_ROWS=n_rows < 16, REDUCTION_WIDTH=_frozen_norm_reduction_width(n_rows, q_raw.shape[-1]), num_warps=_frozen_norm_NUM_WARPS, num_stages=_frozen_norm_NUM_STAGES) |
| return (q_memory, k_gain, k_memory, q_norm, k_norm) |
|
|
|
|
| |
|
|
| @triton.jit(do_not_specialize=['T', 'info_scale', 'info_scale_base']) |
| def _frozen_gain_kla_kappa_passA_kernel(k_ptr, a_ptr, omega_ptr, r_ptr, cm_ptr, T, info_scale, log_info_scale_ptr, info_scale_base, NT: tl.constexpr, H: tl.constexpr, K: tl.constexpr, BK: tl.constexpr, BT: tl.constexpr, USE_LOG_SCALE: tl.constexpr): |
| i_nt = tl.program_id(0) |
| pid = tl.program_id(1) |
| i_b = pid // H |
| i_h = pid % H |
| if USE_LOG_SCALE: |
| log_info_scale = tl.load(log_info_scale_ptr + i_h).to(tl.float32) |
| head_info_scale = info_scale_base * tl.exp(log_info_scale) |
| tl.device_assert((head_info_scale > 0.0) & (head_info_scale <= 3.402823466e+38), 'effective info scale must be positive and finite') |
| i_k = tl.arange(0, BK) |
| mask = i_k < K |
| base = i_b * T * H * K + i_h * K |
| base_r = i_b * T * H + i_h |
| t0 = i_nt * BT |
| t1 = tl.minimum(t0 + BT, T) |
| mA = tl.full((BK,), 1.0, tl.float32) |
| mB = tl.full((BK,), 0.0, tl.float32) |
| mC = tl.full((BK,), 0.0, tl.float32) |
| mD = tl.full((BK,), 1.0, tl.float32) |
| for t in range(t0, t1): |
| off = base + t * H * K + i_k |
| k = tl.load(k_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| a = tl.load(a_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| omega = tl.load(omega_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| r = tl.load(r_ptr + base_r + t * H).to(tl.float32) |
| a2 = a * a |
| if USE_LOG_SCALE: |
| u = head_info_scale * k * k / r |
| else: |
| u = info_scale * k * k / r |
| tA = a2 |
| tB = omega |
| tC = u * a2 |
| tD = 1.0 + u * omega |
| nA = tB * mC + tA * mA |
| nB = tB * mD + tA * mB |
| nC = tD * mC + tC * mA |
| nD = tD * mD + tC * mB |
| m = tl.maximum(tl.maximum(tl.abs(nD), tl.abs(nC)), tl.maximum(tl.abs(nB), tl.abs(nA))) |
| inv = 1.0 / tl.maximum(m, 1e-30) |
| mA = nA * inv |
| mB = nB * inv |
| mC = nC * inv |
| mD = nD * inv |
| cm_base = ((i_b * NT + i_nt) * H + i_h) * K * 4 |
| tl.store(cm_ptr + cm_base + i_k * 4 + 0, mA, mask=mask) |
| tl.store(cm_ptr + cm_base + i_k * 4 + 1, mB, mask=mask) |
| tl.store(cm_ptr + cm_base + i_k * 4 + 2, mC, mask=mask) |
| tl.store(cm_ptr + cm_base + i_k * 4 + 3, mD, mask=mask) |
|
|
| @triton.jit |
| def _frozen_gain_kla_kappa_passB_kernel(cm_ptr, mu_ptr, carry_ptr, NT: tl.constexpr, H: tl.constexpr, K: tl.constexpr, BK: tl.constexpr): |
| pid = tl.program_id(0) |
| i_b = pid // H |
| i_h = pid % H |
| i_k = tl.arange(0, BK) |
| mask = i_k < K |
| mu = tl.load(mu_ptr + i_h).to(tl.float32) |
| p_num = tl.full((BK,), 1.0, tl.float32) |
| p_den = tl.full((BK,), 0.0, tl.float32) + mu |
| for nt in range(0, NT): |
| carry_base = ((i_b * NT + nt) * H + i_h) * K * 2 |
| tl.store(carry_ptr + carry_base + i_k * 2 + 0, p_num, mask=mask) |
| tl.store(carry_ptr + carry_base + i_k * 2 + 1, p_den, mask=mask) |
| cm_base = ((i_b * NT + nt) * H + i_h) * K * 4 |
| mA = tl.load(cm_ptr + cm_base + i_k * 4 + 0, mask=mask, other=0.0) |
| mB = tl.load(cm_ptr + cm_base + i_k * 4 + 1, mask=mask, other=0.0) |
| mC = tl.load(cm_ptr + cm_base + i_k * 4 + 2, mask=mask, other=0.0) |
| mD = tl.load(cm_ptr + cm_base + i_k * 4 + 3, mask=mask, other=0.0) |
| p_num_new = mB * p_den + mA * p_num |
| p_den_new = mD * p_den + mC * p_num |
| m = tl.maximum(tl.abs(p_den_new), tl.abs(p_num_new)) |
| inv = 1.0 / tl.maximum(m, 1e-30) |
| p_den = p_den_new * inv |
| p_num = p_num_new * inv |
|
|
| @triton.jit(do_not_specialize=['T', 'info_scale', 'info_scale_base']) |
| def _frozen_gain_kla_kappa_passC_kappa_only_kernel(k_ptr, a_ptr, omega_ptr, r_ptr, carry_ptr, kappa_ptr, n_ptr, d_ptr, T, info_scale, log_info_scale_ptr, info_scale_base, NT: tl.constexpr, H: tl.constexpr, K: tl.constexpr, BK: tl.constexpr, BT: tl.constexpr, SAVE_STATE: tl.constexpr, USE_LOG_SCALE: tl.constexpr): |
| """Pass C specialization which neither accepts nor writes beta_ch.""" |
| i_nt = tl.program_id(0) |
| pid = tl.program_id(1) |
| i_b = pid // H |
| i_h = pid % H |
| if USE_LOG_SCALE: |
| log_info_scale = tl.load(log_info_scale_ptr + i_h).to(tl.float32) |
| head_info_scale = info_scale_base * tl.exp(log_info_scale) |
| tl.device_assert((head_info_scale > 0.0) & (head_info_scale <= 3.402823466e+38), 'effective info scale must be positive and finite') |
| i_k = tl.arange(0, BK) |
| mask = i_k < K |
| base = i_b * T * H * K + i_h * K |
| base_r = i_b * T * H + i_h |
| t0 = i_nt * BT |
| t1 = tl.minimum(t0 + BT, T) |
| carry_base = ((i_b * NT + i_nt) * H + i_h) * K * 2 |
| p_num = tl.load(carry_ptr + carry_base + i_k * 2, mask=mask, other=1.0).to(tl.float32) |
| p_den = tl.load(carry_ptr + carry_base + i_k * 2 + 1, mask=mask, other=0.0).to(tl.float32) |
| for t in range(t0, t1): |
| off = base + t * H * K + i_k |
| k = tl.load(k_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| a = tl.load(a_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| omega = tl.load(omega_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| r = tl.load(r_ptr + base_r + t * H).to(tl.float32) |
| if SAVE_STATE: |
| tl.store(n_ptr + off, p_den, mask=mask) |
| tl.store(d_ptr + off, p_num, mask=mask) |
| a2 = a * a |
| p_hat = a2 * p_num / p_den + omega |
| ksq = k * k |
| denom = r + tl.sum(tl.where(mask, p_hat * ksq, 0.0), axis=0) |
| kappa = p_hat * k / denom |
| tl.store(kappa_ptr + off, kappa.to(kappa_ptr.dtype.element_ty), mask=mask) |
| if USE_LOG_SCALE: |
| u = head_info_scale * ksq / r |
| else: |
| u = info_scale * ksq / r |
| p_num_new = omega * p_den + a2 * p_num |
| p_den_new = (1.0 + u * omega) * p_den + u * a2 * p_num |
| m = tl.maximum(tl.abs(p_den_new), tl.abs(p_num_new)) |
| inv = 1.0 / tl.maximum(m, 1e-30) |
| p_den = p_den_new * inv |
| p_num = p_num_new * inv |
|
|
| @triton.jit(do_not_specialize=['T', 'info_scale', 'info_scale_base']) |
| def _frozen_gain_kla_kappa_bwd_fill_kernel(k_ptr, a_ptr, omega_ptr, r_ptr, carry_ptr, p_excl_ptr, T, info_scale, log_info_scale_ptr, info_scale_base, NT: tl.constexpr, H: tl.constexpr, K: tl.constexpr, BK: tl.constexpr, BT: tl.constexpr, USE_LOG_SCALE: tl.constexpr): |
| i_nt = tl.program_id(0) |
| pid = tl.program_id(1) |
| i_b = pid // H |
| i_h = pid % H |
| if USE_LOG_SCALE: |
| log_info_scale = tl.load(log_info_scale_ptr + i_h).to(tl.float32) |
| head_info_scale = info_scale_base * tl.exp(log_info_scale) |
| tl.device_assert((head_info_scale > 0.0) & (head_info_scale <= 3.402823466e+38), 'effective info scale must be positive and finite') |
| i_k = tl.arange(0, BK) |
| mask = i_k < K |
| base = i_b * T * H * K + i_h * K |
| base_r = i_b * T * H + i_h |
| t0 = i_nt * BT |
| t1 = tl.minimum(t0 + BT, T) |
| carry_base = ((i_b * NT + i_nt) * H + i_h) * K * 2 |
| if USE_LOG_SCALE: |
| if log_info_scale != 0.0: |
| p_num64 = tl.load(carry_ptr + carry_base + i_k * 2 + 0, mask=mask, other=1.0).to(tl.float64) |
| p_den64 = tl.load(carry_ptr + carry_base + i_k * 2 + 1, mask=mask, other=0.0).to(tl.float64) |
| Kf64 = head_info_scale.to(tl.float64) |
| for t in range(t0, t1): |
| off = base + t * H * K + i_k |
| tl.store(p_excl_ptr + off, p_num64 / p_den64, mask=mask) |
| k64 = tl.load(k_ptr + off, mask=mask, other=0.0).to(tl.float64) |
| a64 = tl.load(a_ptr + off, mask=mask, other=0.0).to(tl.float64) |
| omega64 = tl.load(omega_ptr + off, mask=mask, other=0.0).to(tl.float64) |
| r64 = tl.load(r_ptr + base_r + t * H).to(tl.float64) |
| a264 = a64 * a64 |
| u64 = Kf64 * k64 * k64 / r64 |
| p_num_new64 = omega64 * p_den64 + a264 * p_num64 |
| p_den_new64 = (1.0 + u64 * omega64) * p_den64 + u64 * a264 * p_num64 |
| magnitude64 = tl.maximum(tl.abs(p_den_new64), tl.abs(p_num_new64)) |
| inv64 = 1.0 / tl.maximum(magnitude64, 1e-30) |
| p_den64 = tl.where(mask, p_den_new64 * inv64, p_den64) |
| p_num64 = tl.where(mask, p_num_new64 * inv64, p_num64) |
| else: |
| p_num = tl.load(carry_ptr + carry_base + i_k * 2 + 0, mask=mask, other=1.0).to(tl.float32) |
| p_den = tl.load(carry_ptr + carry_base + i_k * 2 + 1, mask=mask, other=0.0).to(tl.float32) |
| for t in range(t0, t1): |
| off = base + t * H * K + i_k |
| tl.store(p_excl_ptr + off, p_num / p_den, mask=mask) |
| k = tl.load(k_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| a = tl.load(a_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| omega = tl.load(omega_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| r = tl.load(r_ptr + base_r + t * H).to(tl.float32) |
| a2 = a * a |
| u = head_info_scale * k * k / r |
| p_num_new = omega * p_den + a2 * p_num |
| p_den_new = (1.0 + u * omega) * p_den + u * a2 * p_num |
| magnitude = tl.maximum(tl.abs(p_den_new), tl.abs(p_num_new)) |
| inv = 1.0 / tl.maximum(magnitude, 1e-30) |
| p_den = p_den_new * inv |
| p_num = p_num_new * inv |
| else: |
| p_num = tl.load(carry_ptr + carry_base + i_k * 2 + 0, mask=mask, other=1.0).to(tl.float32) |
| p_den = tl.load(carry_ptr + carry_base + i_k * 2 + 1, mask=mask, other=0.0).to(tl.float32) |
| for t in range(t0, t1): |
| off = base + t * H * K + i_k |
| tl.store(p_excl_ptr + off, p_num / p_den, mask=mask) |
| k = tl.load(k_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| a = tl.load(a_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| omega = tl.load(omega_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| r = tl.load(r_ptr + base_r + t * H).to(tl.float32) |
| a2 = a * a |
| u = info_scale * k * k / r |
| p_num_new = omega * p_den + a2 * p_num |
| p_den_new = (1.0 + u * omega) * p_den + u * a2 * p_num |
| magnitude = tl.maximum(tl.abs(p_den_new), tl.abs(p_num_new)) |
| inv = 1.0 / tl.maximum(magnitude, 1e-30) |
| p_den = p_den_new * inv |
| p_num = p_num_new * inv |
|
|
| @triton.jit(do_not_specialize=['T', 'info_scale', 'info_scale_base']) |
| def _frozen_gain_kla_kappa_bwd_passA_kappa_only_kernel(k_ptr, a_ptr, omega_ptr, r_ptr, dkappa_ptr, p_excl_ptr, ab_ptr, T, info_scale, log_info_scale_ptr, info_scale_base, NT: tl.constexpr, H: tl.constexpr, K: tl.constexpr, BK: tl.constexpr, BT: tl.constexpr, USE_LOG_SCALE: tl.constexpr): |
| """Parallel reverse-map construction without a dbeta pointer.""" |
| i_nt = tl.program_id(0) |
| pid = tl.program_id(1) |
| i_b = pid // H |
| i_h = pid % H |
| if USE_LOG_SCALE: |
| log_info_scale = tl.load(log_info_scale_ptr + i_h).to(tl.float32) |
| Kf = info_scale_base * tl.exp(log_info_scale) |
| tl.device_assert((Kf > 0.0) & (Kf <= 3.402823466e+38), 'effective info scale must be positive and finite') |
| else: |
| Kf = info_scale |
| i_k = tl.arange(0, BK) |
| mask = i_k < K |
| base = i_b * T * H * K + i_h * K |
| base_r = i_b * T * H + i_h |
| t0 = i_nt * BT |
| t1 = tl.minimum(t0 + BT, T) |
| A_net = tl.full((BK,), 1.0, tl.float32) |
| B_net = tl.full((BK,), 0.0, tl.float32) |
| for t in range(t1 - 1, t0 - 1, -1): |
| off = base + t * H * K + i_k |
| pprev = tl.load(p_excl_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| k = tl.load(k_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| a = tl.load(a_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| omega = tl.load(omega_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| r = tl.load(r_ptr + base_r + t * H).to(tl.float32) |
| dki = tl.load(dkappa_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| a2 = a * a |
| ksq = k * k |
| p_hat = a2 * pprev + omega |
| u = Kf * ksq / r |
| h = 1.0 + u * p_hat |
| D = r + tl.sum(tl.where(mask, p_hat * ksq, 0.0), axis=0) |
| kap = p_hat * k / D |
| G = tl.sum(tl.where(mask, dki * kap, 0.0), axis=0) |
| Pbar_gain = dki * k / D - ksq * G / D |
| A_t = a2 / (h * h) |
| B_t = a2 * Pbar_gain |
| A_net_new = A_t * A_net |
| B_net = A_t * B_net + B_t |
| A_net = A_net_new |
| ab_base = ((i_b * NT + i_nt) * H + i_h) * K * 2 |
| tl.store(ab_ptr + ab_base + i_k * 2, A_net, mask=mask) |
| tl.store(ab_ptr + ab_base + i_k * 2 + 1, B_net, mask=mask) |
|
|
| @triton.jit |
| def _frozen_gain_kla_kappa_bwd_passB_kernel(ab_ptr, carry_ptr, NT: tl.constexpr, H: tl.constexpr, K: tl.constexpr, BK: tl.constexpr): |
| pid = tl.program_id(0) |
| i_b = pid // H |
| i_h = pid % H |
| i_k = tl.arange(0, BK) |
| mask = i_k < K |
| lam = tl.zeros((BK,), tl.float32) |
| for nt in range(NT - 1, -1, -1): |
| carry_base = ((i_b * NT + nt) * H + i_h) * K |
| tl.store(carry_ptr + carry_base + i_k, lam, mask=mask) |
| ab_base = ((i_b * NT + nt) * H + i_h) * K * 2 |
| A_chunk = tl.load(ab_ptr + ab_base + i_k * 2 + 0, mask=mask, other=0.0) |
| B_chunk = tl.load(ab_ptr + ab_base + i_k * 2 + 1, mask=mask, other=0.0) |
| lam = A_chunk * lam + B_chunk |
|
|
| @triton.jit(do_not_specialize=['T', 'info_scale', 'info_scale_base']) |
| def _frozen_gain_kla_kappa_bwd_passC_kappa_only_kernel(k_ptr, a_ptr, omega_ptr, r_ptr, dkappa_ptr, p_excl_ptr, carry_ptr, omega_raw_ptr, r_raw_ptr, dk_ptr, dalpha_ptr, domega_raw_ptr, dr_raw_ptr, dmu_ptr, dlog_ptr, T, info_scale, log_info_scale_ptr, info_scale_base, NT: tl.constexpr, H: tl.constexpr, K: tl.constexpr, BK: tl.constexpr, BT: tl.constexpr, RAW_ACTIVATION_VJP: tl.constexpr, USE_LOG_SCALE: tl.constexpr, NEED_DMU: tl.constexpr, NEED_DLOG: tl.constexpr): |
| """Parallel gradient emission without a dbeta pointer.""" |
| i_nt = tl.program_id(0) |
| pid = tl.program_id(1) |
| i_b = pid // H |
| i_h = pid % H |
| if USE_LOG_SCALE: |
| log_info_scale = tl.load(log_info_scale_ptr + i_h).to(tl.float32) |
| Kf = info_scale_base * tl.exp(log_info_scale) |
| tl.device_assert((Kf > 0.0) & (Kf <= 3.402823466e+38), 'effective info scale must be positive and finite') |
| else: |
| Kf = info_scale |
| i_k = tl.arange(0, BK) |
| mask = i_k < K |
| base = i_b * T * H * K + i_h * K |
| base_r = i_b * T * H + i_h |
| t0 = i_nt * BT |
| t1 = tl.minimum(t0 + BT, T) |
| carry_base = ((i_b * NT + i_nt) * H + i_h) * K |
| lam = tl.load(carry_ptr + carry_base + i_k, mask=mask, other=0.0).to(tl.float32) |
| dlog = 0.0 |
| for t in range(t1 - 1, t0 - 1, -1): |
| off = base + t * H * K + i_k |
| pprev = tl.load(p_excl_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| k = tl.load(k_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| a = tl.load(a_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| omega = tl.load(omega_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| r = tl.load(r_ptr + base_r + t * H).to(tl.float32) |
| dki = tl.load(dkappa_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| a2 = a * a |
| ksq = k * k |
| p_hat = a2 * pprev + omega |
| u = Kf * ksq / r |
| h = 1.0 + u * p_hat |
| pnext = p_hat / h |
| D = r + tl.sum(tl.where(mask, p_hat * ksq, 0.0), axis=0) |
| kap = p_hat * k / D |
| beta = p_hat / D |
| G = tl.sum(tl.where(mask, dki * kap, 0.0), axis=0) |
| Pbar_gain = dki * k / D - ksq * G / D |
| kbar_gain = dki * beta - 2.0 * p_hat * k * G / D |
| Pbar_tot = Pbar_gain + lam / (h * h) |
| ubar = -lam * pnext * pnext |
| dk = kbar_gain + ubar * (2.0 * Kf * k / r) |
| dalpha = 2.0 * a * pprev * Pbar_tot |
| dr = -G / D + tl.sum(tl.where(mask, ubar * (-Kf * ksq / (r * r)), 0.0), axis=0) |
| tl.store(dk_ptr + off, dk.to(dk_ptr.dtype.element_ty), mask=mask) |
| tl.store(dalpha_ptr + off, dalpha.to(dalpha_ptr.dtype.element_ty), mask=mask) |
| if RAW_ACTIVATION_VJP: |
| omega_raw = tl.load(omega_raw_ptr + off, mask=mask, other=0.0).to(tl.float32) |
| r_raw = tl.load(r_raw_ptr + base_r + t * H).to(tl.float32) |
| domega_raw = Pbar_tot * tl.sigmoid(omega_raw) |
| dr_raw = dr * tl.sigmoid(r_raw) |
| tl.store(domega_raw_ptr + off, domega_raw.to(domega_raw_ptr.dtype.element_ty), mask=mask) |
| tl.store(dr_raw_ptr + base_r + t * H, dr_raw.to(dr_raw_ptr.dtype.element_ty)) |
| else: |
| tl.store(domega_raw_ptr + off, Pbar_tot.to(domega_raw_ptr.dtype.element_ty), mask=mask) |
| tl.store(dr_raw_ptr + base_r + t * H, dr.to(dr_raw_ptr.dtype.element_ty)) |
| if NEED_DLOG: |
| dlog += tl.sum(tl.where(mask, ubar * u, 0.0), axis=0) |
| lam = tl.where(mask, a2 * Pbar_tot, 0.0) |
| if NEED_DMU: |
| if i_nt == 0: |
| p_initial = tl.load(p_excl_ptr + base + i_k, mask=mask, other=0.0).to(tl.float32) |
| tl.atomic_add(dmu_ptr + i_h, tl.sum(tl.where(mask, -lam * p_initial * p_initial, 0.0), axis=0)) |
| if NEED_DLOG: |
| tl.atomic_add(dlog_ptr + i_h, dlog) |
|
|
|
|
| |
|
|
| @triton.jit |
| def _frozen_dot_whitelisted_dot(lhs, rhs, BF16_DOT_OPERANDS: tl.constexpr, DOT_PRECISION: tl.constexpr): |
| if BF16_DOT_OPERANDS: |
| lhs = lhs.to(tl.bfloat16) |
| rhs = rhs.to(tl.bfloat16) |
| return tl.dot(lhs, rhs, input_precision=DOT_PRECISION, out_dtype=tl.float32) |
|
|
|
|
| |
|
|
| |
|
|
|
|
| |
|
|
| _frozen_kda_NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] |
|
|
| @triton.heuristics({'USE_G': lambda args: args['g'] is not None, 'USE_GK': lambda args: args['gk'] is not None, 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, 'SAVE_NEW_VALUE': lambda args: args['v_new'] is not None, 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) |
| @triton.autotune(configs=[triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) for num_warps in [2, 4] for num_stages in ([2, 3, 4] if check_shared_mem('ampere') else [2, 1]) for BV in ([32, 64] if check_shared_mem('ada') else [32])], key=['H', 'K', 'V', 'BT', 'USE_EXP2', 'TRANSPOSE_STATE', 'DOT_PRECISION', 'BF16_DOT_OPERANDS', 'RESIDUAL_DOT_PRECISION', 'RESIDUAL_BF16_DOT_OPERANDS'], use_cuda_graph=USE_CUDA_GRAPH, **autotune_cache_kwargs) |
| @triton.jit(do_not_specialize=['T']) |
| def _frozen_kda_chunk_gated_delta_rule_fwd_kernel_h_blockdim64_diag(k, v, w, v_new, g, gk, h, h0, ht, cu_seqlens, chunk_offsets, T, H: tl.constexpr, Hq: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BV: tl.constexpr, USE_G: tl.constexpr, USE_GK: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, STORE_FINAL_STATE: tl.constexpr, SAVE_NEW_VALUE: tl.constexpr, USE_EXP2: tl.constexpr, TRANSPOSE_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, BF16_DOT_OPERANDS: tl.constexpr, RESIDUAL_DOT_PRECISION: tl.constexpr, RESIDUAL_BF16_DOT_OPERANDS: tl.constexpr): |
| i_v, i_nh = (tl.program_id(0), tl.program_id(1)) |
| i_n, i_h = (i_nh // H, i_nh % H) |
| if IS_VARLEN: |
| bos, eos = (tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)) |
| T = eos - bos |
| NT = tl.cdiv(T, BT) |
| boh = tl.load(chunk_offsets + i_n).to(tl.int32) |
| else: |
| bos, eos = (i_n * T, i_n * T + T) |
| NT = tl.cdiv(T, BT) |
| boh = i_n * NT |
| if TRANSPOSE_STATE: |
| b_h1 = tl.zeros([BV, 64], dtype=tl.float32) |
| if K > 64: |
| b_h2 = tl.zeros([BV, 64], dtype=tl.float32) |
| if K > 128: |
| b_h3 = tl.zeros([BV, 64], dtype=tl.float32) |
| if K > 192: |
| b_h4 = tl.zeros([BV, 64], dtype=tl.float32) |
| else: |
| b_h1 = tl.zeros([64, BV], dtype=tl.float32) |
| if K > 64: |
| b_h2 = tl.zeros([64, BV], dtype=tl.float32) |
| if K > 128: |
| b_h3 = tl.zeros([64, BV], dtype=tl.float32) |
| if K > 192: |
| b_h4 = tl.zeros([64, BV], dtype=tl.float32) |
| h += (boh * H + i_h).to(tl.int64) * K * V |
| v += (bos * H + i_h).to(tl.int64) * V |
| k += (bos * Hq + i_h // (H // Hq)).to(tl.int64) * K |
| w += (bos * H + i_h).to(tl.int64) * K |
| if SAVE_NEW_VALUE: |
| v_new += (bos * H + i_h).to(tl.int64) * V |
| if USE_INITIAL_STATE: |
| h0 = h0 + i_nh * K * V |
| if STORE_FINAL_STATE: |
| ht = ht + i_nh * K * V |
| if USE_INITIAL_STATE: |
| if TRANSPOSE_STATE: |
| p_h0_1 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) |
| else: |
| p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) |
| b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) |
| if K > 64: |
| if TRANSPOSE_STATE: |
| p_h0_2 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) |
| else: |
| p_h0_2 = tl.make_block_ptr(h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) |
| b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) |
| if K > 128: |
| if TRANSPOSE_STATE: |
| p_h0_3 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) |
| else: |
| p_h0_3 = tl.make_block_ptr(h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) |
| b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) |
| if K > 192: |
| if TRANSPOSE_STATE: |
| p_h0_4 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) |
| else: |
| p_h0_4 = tl.make_block_ptr(h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) |
| b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) |
| for i_t in range(NT): |
| i_t_int64 = i_t.to(tl.int64) |
| if TRANSPOSE_STATE: |
| p_h1 = tl.make_block_ptr(h + i_t_int64 * H * K * V, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) |
| else: |
| p_h1 = tl.make_block_ptr(h + i_t_int64 * H * K * V, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) |
| tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) |
| if K > 64: |
| if TRANSPOSE_STATE: |
| p_h2 = tl.make_block_ptr(h + i_t_int64 * H * K * V, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) |
| else: |
| p_h2 = tl.make_block_ptr(h + i_t_int64 * H * K * V, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) |
| tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) |
| if K > 128: |
| if TRANSPOSE_STATE: |
| p_h3 = tl.make_block_ptr(h + i_t_int64 * H * K * V, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) |
| else: |
| p_h3 = tl.make_block_ptr(h + i_t_int64 * H * K * V, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) |
| tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) |
| if K > 192: |
| if TRANSPOSE_STATE: |
| p_h4 = tl.make_block_ptr(h + i_t_int64 * H * K * V, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) |
| else: |
| p_h4 = tl.make_block_ptr(h + i_t_int64 * H * K * V, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) |
| tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) |
| p_w = tl.make_block_ptr(w, (T, K), (H * K, 1), (i_t * BT, 0), (BT, 64), (1, 0)) |
| b_w = tl.load(p_w, boundary_check=(0, 1)).to(tl.float32) |
| if TRANSPOSE_STATE: |
| b_v = _frozen_dot_whitelisted_dot(b_w, tl.trans(b_h1).to(b_w.dtype), RESIDUAL_BF16_DOT_OPERANDS, RESIDUAL_DOT_PRECISION) |
| else: |
| b_v = _frozen_dot_whitelisted_dot(b_w, b_h1.to(b_w.dtype), RESIDUAL_BF16_DOT_OPERANDS, RESIDUAL_DOT_PRECISION) |
| if K > 64: |
| p_w = tl.make_block_ptr(w, (T, K), (H * K, 1), (i_t * BT, 64), (BT, 64), (1, 0)) |
| b_w = tl.load(p_w, boundary_check=(0, 1)).to(tl.float32) |
| if TRANSPOSE_STATE: |
| b_v += _frozen_dot_whitelisted_dot(b_w, tl.trans(b_h2).to(b_w.dtype), RESIDUAL_BF16_DOT_OPERANDS, RESIDUAL_DOT_PRECISION) |
| else: |
| b_v += _frozen_dot_whitelisted_dot(b_w, b_h2.to(b_w.dtype), RESIDUAL_BF16_DOT_OPERANDS, RESIDUAL_DOT_PRECISION) |
| if K > 128: |
| p_w = tl.make_block_ptr(w, (T, K), (H * K, 1), (i_t * BT, 128), (BT, 64), (1, 0)) |
| b_w = tl.load(p_w, boundary_check=(0, 1)).to(tl.float32) |
| if TRANSPOSE_STATE: |
| b_v += _frozen_dot_whitelisted_dot(b_w, tl.trans(b_h3).to(b_w.dtype), RESIDUAL_BF16_DOT_OPERANDS, RESIDUAL_DOT_PRECISION) |
| else: |
| b_v += _frozen_dot_whitelisted_dot(b_w, b_h3.to(b_w.dtype), RESIDUAL_BF16_DOT_OPERANDS, RESIDUAL_DOT_PRECISION) |
| if K > 192: |
| p_w = tl.make_block_ptr(w, (T, K), (H * K, 1), (i_t * BT, 192), (BT, 64), (1, 0)) |
| b_w = tl.load(p_w, boundary_check=(0, 1)).to(tl.float32) |
| if TRANSPOSE_STATE: |
| b_v += _frozen_dot_whitelisted_dot(b_w, tl.trans(b_h4).to(b_w.dtype), RESIDUAL_BF16_DOT_OPERANDS, RESIDUAL_DOT_PRECISION) |
| else: |
| b_v += _frozen_dot_whitelisted_dot(b_w, b_h4.to(b_w.dtype), RESIDUAL_BF16_DOT_OPERANDS, RESIDUAL_DOT_PRECISION) |
| p_v = tl.make_block_ptr(v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) |
| b_v = tl.load(p_v, boundary_check=(0, 1)).to(tl.float32) - b_v |
| if SAVE_NEW_VALUE: |
| p_v = tl.make_block_ptr(v_new, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) |
| tl.store(p_v, b_v.to(p_v.dtype.element_ty), boundary_check=(0, 1)) |
| last_idx = min((i_t + 1) * BT, T) - 1 |
| if USE_G: |
| m_t = i_t * BT + tl.arange(0, BT) < T |
| b_g_last = tl.load(g + (bos * H + last_idx * H + i_h).to(tl.int64)).to(tl.float32) |
| p_g = tl.make_block_ptr(g + (bos * H + i_h).to(tl.int64), (T,), (H,), (i_t * BT,), (BT,), (0,)) |
| b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) |
| if USE_EXP2: |
| b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] |
| b_g_last = exp2(b_g_last) |
| else: |
| b_v = b_v * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] |
| b_g_last = exp(b_g_last) |
| b_h1 *= b_g_last |
| if K > 64: |
| b_h2 *= b_g_last |
| if K > 128: |
| b_h3 *= b_g_last |
| if K > 192: |
| b_h4 *= b_g_last |
| if USE_GK: |
| o_k1 = tl.arange(0, 64) |
| b_gk_last1 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k1, mask=o_k1 < K, other=0.0).to(tl.float32) |
| if TRANSPOSE_STATE: |
| if USE_EXP2: |
| b_h1 *= exp2(b_gk_last1)[None, :] |
| else: |
| b_h1 *= exp(b_gk_last1)[None, :] |
| elif USE_EXP2: |
| b_h1 *= exp2(b_gk_last1)[:, None] |
| else: |
| b_h1 *= exp(b_gk_last1)[:, None] |
| if K > 64: |
| o_k2 = 64 + o_k1 |
| b_gk_last2 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k2, mask=o_k2 < K, other=0.0).to(tl.float32) |
| if TRANSPOSE_STATE: |
| if USE_EXP2: |
| b_h2 *= exp2(b_gk_last2)[None, :] |
| else: |
| b_h2 *= exp(b_gk_last2)[None, :] |
| elif USE_EXP2: |
| b_h2 *= exp2(b_gk_last2)[:, None] |
| else: |
| b_h2 *= exp(b_gk_last2)[:, None] |
| if K > 128: |
| o_k3 = 128 + o_k1 |
| b_gk_last3 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k3, mask=o_k3 < K, other=0.0).to(tl.float32) |
| if TRANSPOSE_STATE: |
| if USE_EXP2: |
| b_h3 *= exp2(b_gk_last3)[None, :] |
| else: |
| b_h3 *= exp(b_gk_last3)[None, :] |
| elif USE_EXP2: |
| b_h3 *= exp2(b_gk_last3)[:, None] |
| else: |
| b_h3 *= exp(b_gk_last3)[:, None] |
| if K > 192: |
| o_k4 = 192 + o_k1 |
| b_gk_last4 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k4, mask=o_k4 < K, other=0.0).to(tl.float32) |
| if TRANSPOSE_STATE: |
| if USE_EXP2: |
| b_h4 *= exp2(b_gk_last4)[None, :] |
| else: |
| b_h4 *= exp(b_gk_last4)[None, :] |
| elif USE_EXP2: |
| b_h4 *= exp2(b_gk_last4)[:, None] |
| else: |
| b_h4 *= exp(b_gk_last4)[:, None] |
| b_v = b_v.to(tl.float32) |
| p_k = tl.make_block_ptr(k, (K, T), (1, Hq * K), (0, i_t * BT), (64, BT), (0, 1)) |
| b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) |
| if TRANSPOSE_STATE: |
| b_h1 += tl.trans(tl.dot(b_k, b_v, input_precision=DOT_PRECISION, out_dtype=tl.float32)) |
| else: |
| b_h1 += tl.dot(b_k, b_v, input_precision=DOT_PRECISION, out_dtype=tl.float32) |
| if K > 64: |
| p_k = tl.make_block_ptr(k, (K, T), (1, Hq * K), (64, i_t * BT), (64, BT), (0, 1)) |
| b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) |
| if TRANSPOSE_STATE: |
| b_h2 += tl.trans(tl.dot(b_k, b_v, input_precision=DOT_PRECISION, out_dtype=tl.float32)) |
| else: |
| b_h2 += tl.dot(b_k, b_v, input_precision=DOT_PRECISION, out_dtype=tl.float32) |
| if K > 128: |
| p_k = tl.make_block_ptr(k, (K, T), (1, Hq * K), (128, i_t * BT), (64, BT), (0, 1)) |
| b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) |
| if TRANSPOSE_STATE: |
| b_h3 += tl.trans(tl.dot(b_k, b_v, input_precision=DOT_PRECISION, out_dtype=tl.float32)) |
| else: |
| b_h3 += tl.dot(b_k, b_v, input_precision=DOT_PRECISION, out_dtype=tl.float32) |
| if K > 192: |
| p_k = tl.make_block_ptr(k, (K, T), (1, Hq * K), (192, i_t * BT), (64, BT), (0, 1)) |
| b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) |
| if TRANSPOSE_STATE: |
| b_h4 += tl.trans(tl.dot(b_k, b_v, input_precision=DOT_PRECISION, out_dtype=tl.float32)) |
| else: |
| b_h4 += tl.dot(b_k, b_v, input_precision=DOT_PRECISION, out_dtype=tl.float32) |
| if STORE_FINAL_STATE: |
| if TRANSPOSE_STATE: |
| p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) |
| else: |
| p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) |
| tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) |
| if K > 64: |
| if TRANSPOSE_STATE: |
| p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) |
| else: |
| p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) |
| tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) |
| if K > 128: |
| if TRANSPOSE_STATE: |
| p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) |
| else: |
| p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) |
| tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) |
| if K > 192: |
| if TRANSPOSE_STATE: |
| p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) |
| else: |
| p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) |
| tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) |
|
|
| @triton.heuristics({'USE_G': lambda args: args['g'] is not None, 'USE_GK': lambda args: args['gk'] is not None, 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) |
| @triton.autotune(configs=[triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) for num_warps in [2, 4] for num_stages in ([2, 3, 4] if check_shared_mem('ampere') else [1]) for BV in ([32, 64] if check_shared_mem('ada') else [32])], key=['H', 'K', 'V', 'BT', 'BV', 'USE_G', 'USE_EXP2', 'TRANSPOSE_STATE', 'DOT_PRECISION', 'BF16_QG_DO_OPERANDS', 'BF16_W_DV_OPERANDS'], use_cuda_graph=USE_CUDA_GRAPH, **autotune_cache_kwargs) |
| @triton.jit(do_not_specialize=['T']) |
| def _frozen_kda_chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64_diag(q, k, w, g, gk, dht, dh0, do, dh, dv, dv2, cu_seqlens, chunk_offsets, scale, T, H: tl.constexpr, Hq: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BV: tl.constexpr, USE_G: tl.constexpr, USE_GK: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, USE_FINAL_STATE_GRADIENT: tl.constexpr, USE_EXP2: tl.constexpr, TRANSPOSE_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, BF16_QG_DO_OPERANDS: tl.constexpr, BF16_W_DV_OPERANDS: tl.constexpr): |
| i_v, i_nh = (tl.program_id(0), tl.program_id(1)) |
| i_n, i_h = (i_nh // H, i_nh % H) |
| if IS_VARLEN: |
| bos, eos = (tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)) |
| T = eos - bos |
| NT = tl.cdiv(T, BT) |
| boh = tl.load(chunk_offsets + i_n).to(tl.int32) |
| else: |
| bos, eos = (i_n * T, i_n * T + T) |
| NT = tl.cdiv(T, BT) |
| boh = i_n * NT |
| if TRANSPOSE_STATE: |
| b_dh1 = tl.zeros([BV, 64], dtype=tl.float32) |
| if K > 64: |
| b_dh2 = tl.zeros([BV, 64], dtype=tl.float32) |
| if K > 128: |
| b_dh3 = tl.zeros([BV, 64], dtype=tl.float32) |
| if K > 192: |
| b_dh4 = tl.zeros([BV, 64], dtype=tl.float32) |
| else: |
| b_dh1 = tl.zeros([64, BV], dtype=tl.float32) |
| if K > 64: |
| b_dh2 = tl.zeros([64, BV], dtype=tl.float32) |
| if K > 128: |
| b_dh3 = tl.zeros([64, BV], dtype=tl.float32) |
| if K > 192: |
| b_dh4 = tl.zeros([64, BV], dtype=tl.float32) |
| q += (bos * Hq + i_h // (H // Hq)).to(tl.int64) * K |
| k += (bos * Hq + i_h // (H // Hq)).to(tl.int64) * K |
| w += (bos * H + i_h).to(tl.int64) * K |
| do += (bos * H + i_h).to(tl.int64) * V |
| dv += (bos * H + i_h).to(tl.int64) * V |
| dv2 += (bos * H + i_h).to(tl.int64) * V |
| dh += (boh * H + i_h).to(tl.int64) * K * V |
| if USE_GK: |
| gk += (bos * H + i_h).to(tl.int64) * K |
| if USE_INITIAL_STATE: |
| dh0 += i_nh * K * V |
| if USE_FINAL_STATE_GRADIENT: |
| dht += i_nh * K * V |
| if USE_FINAL_STATE_GRADIENT: |
| if TRANSPOSE_STATE: |
| p_dht1 = tl.make_block_ptr(dht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) |
| else: |
| p_dht1 = tl.make_block_ptr(dht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) |
| b_dh1 += tl.load(p_dht1, boundary_check=(0, 1)) |
| if K > 64: |
| if TRANSPOSE_STATE: |
| p_dht2 = tl.make_block_ptr(dht, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) |
| else: |
| p_dht2 = tl.make_block_ptr(dht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) |
| b_dh2 += tl.load(p_dht2, boundary_check=(0, 1)) |
| if K > 128: |
| if TRANSPOSE_STATE: |
| p_dht3 = tl.make_block_ptr(dht, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) |
| else: |
| p_dht3 = tl.make_block_ptr(dht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) |
| b_dh3 += tl.load(p_dht3, boundary_check=(0, 1)) |
| if K > 192: |
| if TRANSPOSE_STATE: |
| p_dht4 = tl.make_block_ptr(dht, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) |
| else: |
| p_dht4 = tl.make_block_ptr(dht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) |
| b_dh4 += tl.load(p_dht4, boundary_check=(0, 1)) |
| for i_t in range(NT - 1, -1, -1): |
| i_t_int64 = i_t.to(tl.int64) |
| if TRANSPOSE_STATE: |
| p_dh1 = tl.make_block_ptr(dh + i_t_int64 * H * K * V, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) |
| else: |
| p_dh1 = tl.make_block_ptr(dh + i_t_int64 * H * K * V, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) |
| tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) |
| if K > 64: |
| if TRANSPOSE_STATE: |
| p_dh2 = tl.make_block_ptr(dh + i_t_int64 * H * K * V, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) |
| else: |
| p_dh2 = tl.make_block_ptr(dh + i_t_int64 * H * K * V, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) |
| tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) |
| if K > 128: |
| if TRANSPOSE_STATE: |
| p_dh3 = tl.make_block_ptr(dh + i_t_int64 * H * K * V, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) |
| else: |
| p_dh3 = tl.make_block_ptr(dh + i_t_int64 * H * K * V, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) |
| tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) |
| if K > 192: |
| if TRANSPOSE_STATE: |
| p_dh4 = tl.make_block_ptr(dh + i_t_int64 * H * K * V, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) |
| else: |
| p_dh4 = tl.make_block_ptr(dh + i_t_int64 * H * K * V, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) |
| tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) |
| last_idx = min((i_t + 1) * BT, T) - 1 |
| if USE_G: |
| bg_last = tl.load(g + (bos + last_idx) * H + i_h).to(tl.float32) |
| p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) |
| b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) |
| if USE_EXP2: |
| bg_last_exp = exp2(bg_last) |
| b_g_exp = exp2(b_g) |
| else: |
| bg_last_exp = exp(bg_last) |
| b_g_exp = exp(b_g) |
| p_dv = tl.make_block_ptr(dv, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) |
| p_dv2 = tl.make_block_ptr(dv2, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) |
| p_do = tl.make_block_ptr(do, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) |
| b_do = tl.load(p_do, boundary_check=(0, 1)).to(tl.float32) |
| p_k = tl.make_block_ptr(k, (T, K), (Hq * K, 1), (i_t * BT, 0), (BT, 64), (1, 0)) |
| b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) |
| if USE_GK: |
| o_k1 = tl.arange(0, 64) |
| b_gk_last1 = tl.load(gk + last_idx * H * K + o_k1, mask=o_k1 < K, other=0.0).to(tl.float32) |
| if TRANSPOSE_STATE: |
| b_dv = tl.dot(b_k, tl.trans(b_dh1).to(b_k.dtype), input_precision=DOT_PRECISION, out_dtype=tl.float32) |
| else: |
| b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype), input_precision=DOT_PRECISION, out_dtype=tl.float32) |
| if K > 64: |
| p_k = tl.make_block_ptr(k, (T, K), (Hq * K, 1), (i_t * BT, 64), (BT, 64), (1, 0)) |
| b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) |
| if USE_GK: |
| o_k2 = 64 + o_k1 |
| b_gk_last2 = tl.load(gk + last_idx * H * K + o_k2, mask=o_k2 < K, other=0.0).to(tl.float32) |
| if TRANSPOSE_STATE: |
| b_dv += tl.dot(b_k, tl.trans(b_dh2).to(b_k.dtype), input_precision=DOT_PRECISION, out_dtype=tl.float32) |
| else: |
| b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype), input_precision=DOT_PRECISION, out_dtype=tl.float32) |
| if K > 128: |
| p_k = tl.make_block_ptr(k, (T, K), (Hq * K, 1), (i_t * BT, 128), (BT, 64), (1, 0)) |
| b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) |
| if USE_GK: |
| o_k3 = 128 + o_k1 |
| b_gk_last3 = tl.load(gk + last_idx * H * K + o_k3, mask=o_k3 < K, other=0.0).to(tl.float32) |
| if TRANSPOSE_STATE: |
| b_dv += tl.dot(b_k, tl.trans(b_dh3).to(b_k.dtype), input_precision=DOT_PRECISION, out_dtype=tl.float32) |
| else: |
| b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype), input_precision=DOT_PRECISION, out_dtype=tl.float32) |
| if K > 192: |
| p_k = tl.make_block_ptr(k, (T, K), (Hq * K, 1), (i_t * BT, 192), (BT, 64), (1, 0)) |
| b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) |
| if USE_GK: |
| o_k4 = 192 + o_k1 |
| b_gk_last4 = tl.load(gk + last_idx * H * K + o_k4, mask=o_k4 < K, other=0.0).to(tl.float32) |
| if TRANSPOSE_STATE: |
| b_dv += tl.dot(b_k, tl.trans(b_dh4).to(b_k.dtype), input_precision=DOT_PRECISION, out_dtype=tl.float32) |
| else: |
| b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype), input_precision=DOT_PRECISION, out_dtype=tl.float32) |
| if USE_G: |
| m_t = i_t * BT + tl.arange(0, BT) < T |
| if USE_EXP2: |
| b_dv *= tl.where(m_t, exp2(bg_last - b_g), 0)[:, None] |
| else: |
| b_dv *= tl.where(m_t, exp(bg_last - b_g), 0)[:, None] |
| b_dv += tl.load(p_dv, boundary_check=(0, 1)) |
| tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) |
| p_w = tl.make_block_ptr(w, (K, T), (1, H * K), (0, i_t * BT), (64, BT), (0, 1)) |
| p_q = tl.make_block_ptr(q, (K, T), (1, Hq * K), (0, i_t * BT), (64, BT), (0, 1)) |
| b_w = tl.load(p_w, boundary_check=(0, 1)).to(tl.float32) |
| b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) |
| if USE_G: |
| b_dh1 *= bg_last_exp |
| b_q = b_q * b_g_exp[None, :] |
| if USE_GK: |
| if TRANSPOSE_STATE: |
| if USE_EXP2: |
| b_dh1 *= exp2(b_gk_last1)[None, :] |
| else: |
| b_dh1 *= exp(b_gk_last1)[None, :] |
| elif USE_EXP2: |
| b_dh1 *= exp2(b_gk_last1[:, None]) |
| else: |
| b_dh1 *= exp(b_gk_last1[:, None]) |
| if TRANSPOSE_STATE: |
| b_dh1 += tl.trans(_frozen_dot_whitelisted_dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype), BF16_QG_DO_OPERANDS, DOT_PRECISION) * scale - _frozen_dot_whitelisted_dot(b_w, b_dv.to(b_w.dtype), BF16_W_DV_OPERANDS, DOT_PRECISION)) |
| else: |
| b_dh1 += _frozen_dot_whitelisted_dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype), BF16_QG_DO_OPERANDS, DOT_PRECISION) * scale - _frozen_dot_whitelisted_dot(b_w, b_dv.to(b_w.dtype), BF16_W_DV_OPERANDS, DOT_PRECISION) |
| if K > 64: |
| p_q = tl.make_block_ptr(q, (K, T), (1, Hq * K), (64, i_t * BT), (64, BT), (0, 1)) |
| p_w = tl.make_block_ptr(w, (K, T), (1, H * K), (64, i_t * BT), (64, BT), (0, 1)) |
| b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) |
| b_w = tl.load(p_w, boundary_check=(0, 1)).to(tl.float32) |
| if USE_G: |
| b_dh2 *= bg_last_exp |
| b_q = b_q * b_g_exp[None, :] |
| if USE_GK: |
| if TRANSPOSE_STATE: |
| if USE_EXP2: |
| b_dh2 *= exp2(b_gk_last2)[None, :] |
| else: |
| b_dh2 *= exp(b_gk_last2)[None, :] |
| elif USE_EXP2: |
| b_dh2 *= exp2(b_gk_last2[:, None]) |
| else: |
| b_dh2 *= exp(b_gk_last2[:, None]) |
| if TRANSPOSE_STATE: |
| b_dh2 += tl.trans(_frozen_dot_whitelisted_dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype), BF16_QG_DO_OPERANDS, DOT_PRECISION) * scale - _frozen_dot_whitelisted_dot(b_w, b_dv.to(b_w.dtype), BF16_W_DV_OPERANDS, DOT_PRECISION)) |
| else: |
| b_dh2 += _frozen_dot_whitelisted_dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype), BF16_QG_DO_OPERANDS, DOT_PRECISION) * scale - _frozen_dot_whitelisted_dot(b_w, b_dv.to(b_w.dtype), BF16_W_DV_OPERANDS, DOT_PRECISION) |
| if K > 128: |
| p_q = tl.make_block_ptr(q, (K, T), (1, Hq * K), (128, i_t * BT), (64, BT), (0, 1)) |
| p_w = tl.make_block_ptr(w, (K, T), (1, H * K), (128, i_t * BT), (64, BT), (0, 1)) |
| b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) |
| b_w = tl.load(p_w, boundary_check=(0, 1)).to(tl.float32) |
| if USE_G: |
| b_dh3 *= bg_last_exp |
| b_q = b_q * b_g_exp[None, :] |
| if USE_GK: |
| if TRANSPOSE_STATE: |
| if USE_EXP2: |
| b_dh3 *= exp2(b_gk_last3)[None, :] |
| else: |
| b_dh3 *= exp(b_gk_last3)[None, :] |
| elif USE_EXP2: |
| b_dh3 *= exp2(b_gk_last3[:, None]) |
| else: |
| b_dh3 *= exp(b_gk_last3[:, None]) |
| if TRANSPOSE_STATE: |
| b_dh3 += tl.trans(_frozen_dot_whitelisted_dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype), BF16_QG_DO_OPERANDS, DOT_PRECISION) * scale - _frozen_dot_whitelisted_dot(b_w, b_dv.to(b_w.dtype), BF16_W_DV_OPERANDS, DOT_PRECISION)) |
| else: |
| b_dh3 += _frozen_dot_whitelisted_dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype), BF16_QG_DO_OPERANDS, DOT_PRECISION) * scale - _frozen_dot_whitelisted_dot(b_w, b_dv.to(b_w.dtype), BF16_W_DV_OPERANDS, DOT_PRECISION) |
| if K > 192: |
| p_q = tl.make_block_ptr(q, (K, T), (1, Hq * K), (192, i_t * BT), (64, BT), (0, 1)) |
| p_w = tl.make_block_ptr(w, (K, T), (1, H * K), (192, i_t * BT), (64, BT), (0, 1)) |
| b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) |
| b_w = tl.load(p_w, boundary_check=(0, 1)).to(tl.float32) |
| if USE_G: |
| b_dh4 *= bg_last_exp |
| b_q = b_q * b_g_exp[None, :] |
| if USE_GK: |
| if TRANSPOSE_STATE: |
| if USE_EXP2: |
| b_dh4 *= exp2(b_gk_last4)[None, :] |
| else: |
| b_dh4 *= exp(b_gk_last4)[None, :] |
| elif USE_EXP2: |
| b_dh4 *= exp2(b_gk_last4[:, None]) |
| else: |
| b_dh4 *= exp(b_gk_last4[:, None]) |
| if TRANSPOSE_STATE: |
| b_dh4 += tl.trans(_frozen_dot_whitelisted_dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype), BF16_QG_DO_OPERANDS, DOT_PRECISION) * scale - _frozen_dot_whitelisted_dot(b_w, b_dv.to(b_w.dtype), BF16_W_DV_OPERANDS, DOT_PRECISION)) |
| else: |
| b_dh4 += _frozen_dot_whitelisted_dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype), BF16_QG_DO_OPERANDS, DOT_PRECISION) * scale - _frozen_dot_whitelisted_dot(b_w, b_dv.to(b_w.dtype), BF16_W_DV_OPERANDS, DOT_PRECISION) |
| if USE_INITIAL_STATE: |
| if TRANSPOSE_STATE: |
| p_dh0 = tl.make_block_ptr(dh0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) |
| else: |
| p_dh0 = tl.make_block_ptr(dh0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) |
| tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) |
| if K > 64: |
| if TRANSPOSE_STATE: |
| p_dh1 = tl.make_block_ptr(dh0, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) |
| else: |
| p_dh1 = tl.make_block_ptr(dh0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) |
| tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) |
| if K > 128: |
| if TRANSPOSE_STATE: |
| p_dh2 = tl.make_block_ptr(dh0, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) |
| else: |
| p_dh2 = tl.make_block_ptr(dh0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) |
| tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) |
| if K > 192: |
| if TRANSPOSE_STATE: |
| p_dh3 = tl.make_block_ptr(dh0, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) |
| else: |
| p_dh3 = tl.make_block_ptr(dh0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) |
| tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) |
|
|
| @triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) |
| @triton.autotune(configs=[triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in _frozen_kda_NUM_WARPS for num_stages in [2, 3, 4]], key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'DOT_PRECISION', 'BF16_DOT_OPERANDS', 'BF16_DAQK_DO_VNEW_OPERANDS', 'DVNEW_AQK_DO_DOT_PRECISION', 'BF16_DVNEW_AQK_DO_OPERANDS'], **autotune_cache_kwargs) |
| @triton.jit(do_not_specialize=['T']) |
| def _frozen_kda_chunk_kda_bwd_kernel_dAv_diag(q, k, v, A, do, dv, dA, cu_seqlens, chunk_indices, scale, T, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, BF16_DOT_OPERANDS: tl.constexpr, BF16_DAQK_DO_VNEW_OPERANDS: tl.constexpr, DVNEW_AQK_DO_DOT_PRECISION: tl.constexpr, BF16_DVNEW_AQK_DO_OPERANDS: tl.constexpr): |
| i_t, i_bh = (tl.program_id(0), tl.program_id(1)) |
| i_b, i_h = (i_bh // H, i_bh % H) |
| if IS_VARLEN: |
| i_n, i_t = (tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)) |
| bos, eos = (tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)) |
| T = eos - bos |
| else: |
| bos, eos = (i_b * T, i_b * T + T) |
| q += (bos * H + i_h) * K |
| k += (bos * H + i_h) * K |
| v += (bos * H + i_h) * V |
| do += (bos * H + i_h) * V |
| dv += (bos * H + i_h) * V |
| dA += (bos * H + i_h) * BT |
| p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1)) |
| b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) |
| o_t = i_t * BT + tl.arange(0, BT) |
| m_t = o_t < T |
| m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t) |
| b_A = tl.where(m_A, b_A, 0).to(tl.float32) |
| b_dA = tl.zeros([BT, BT], dtype=tl.float32) |
| for i_v in range(tl.cdiv(V, BV)): |
| p_v = tl.make_block_ptr(v, (V, T), (1, H * V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) |
| p_do = tl.make_block_ptr(do, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) |
| p_dv = tl.make_block_ptr(dv, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) |
| b_v = tl.load(p_v, boundary_check=(0, 1)).to(tl.float32) |
| b_do = tl.load(p_do, boundary_check=(0, 1)).to(tl.float32) |
| b_dA += _frozen_dot_whitelisted_dot(b_do, b_v, BF16_DAQK_DO_VNEW_OPERANDS, DOT_PRECISION) |
| b_dv = _frozen_dot_whitelisted_dot(b_A, b_do, BF16_DVNEW_AQK_DO_OPERANDS, DVNEW_AQK_DO_DOT_PRECISION) |
| tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) |
| p_dA = tl.make_block_ptr(dA, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) |
| b_dA = tl.where(o_t[:, None] >= o_t, b_dA * scale, 0.0) |
| tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) |
|
|
|
|
| |
|
|
| @triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) |
| @fla_cache_autotune(configs=[triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) for BK in [32, 64] for BV in [64, 128] for num_warps in [2, 4, 8] for num_stages in [2, 3, 4]], key=['BT', 'HV', 'TRANSPOSE_STATE', 'DOT_PRECISION', 'BF16_INTER_DOT_OPERANDS', 'BF16_INTRA_DOT_OPERANDS'], **autotune_cache_kwargs) |
| @triton.jit(do_not_specialize=['T']) |
| def _frozen_readout_chunk_gla_fwd_kernel_o_diag(q, v, g, h, o, A, cu_seqlens, chunk_indices, scale, T, H: tl.constexpr, HV: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, USE_EXP2: tl.constexpr, TRANSPOSE_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, BF16_INTER_DOT_OPERANDS: tl.constexpr, BF16_INTRA_DOT_OPERANDS: tl.constexpr): |
| i_v, i_t, i_bh = (tl.program_id(0), tl.program_id(1), tl.program_id(2)) |
| i_b, i_hv = (i_bh // HV, i_bh % HV) |
| i_h = i_hv // (HV // H) |
| if IS_VARLEN: |
| i_tg = i_t.to(tl.int64) |
| i_n, i_t = (tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)) |
| bos, eos = (tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)) |
| T = eos - bos |
| NT = tl.cdiv(T, BT) |
| else: |
| NT = tl.cdiv(T, BT) |
| i_tg = (i_b * NT + i_t).to(tl.int64) |
| bos, eos = ((i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64)) |
| m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] |
| q += (bos * H + i_h) * K |
| g += (bos * HV + i_hv) * K |
| v += (bos * HV + i_hv) * V |
| o += (bos * HV + i_hv) * V |
| h += (i_tg * HV + i_hv).to(tl.int64) * K * V |
| A += (bos * HV + i_hv) * BT |
| b_o = tl.zeros([BT, BV], dtype=tl.float32) |
| for i_k in range(tl.cdiv(K, BK)): |
| p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) |
| p_g = tl.make_block_ptr(g, (T, K), (HV * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) |
| if TRANSPOSE_STATE: |
| p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) |
| else: |
| p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) |
| b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) |
| b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) |
| if USE_EXP2: |
| b_qg = (b_q * exp2(b_g)).to(tl.float32) |
| else: |
| b_qg = (b_q * exp(b_g)).to(tl.float32) |
| b_h = tl.load(p_h, boundary_check=(0, 1)).to(tl.float32) |
| if i_k >= 0: |
| if TRANSPOSE_STATE: |
| b_o += _frozen_dot_whitelisted_dot(b_qg, tl.trans(b_h).to(b_qg.dtype), BF16_INTER_DOT_OPERANDS, DOT_PRECISION) |
| else: |
| b_o += _frozen_dot_whitelisted_dot(b_qg, b_h.to(b_qg.dtype), BF16_INTER_DOT_OPERANDS, DOT_PRECISION) |
| b_o *= scale |
| p_v = tl.make_block_ptr(v, (T, V), (HV * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) |
| p_o = tl.make_block_ptr(o, (T, V), (HV * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) |
| p_A = tl.make_block_ptr(A, (T, BT), (HV * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) |
| b_v = tl.load(p_v, boundary_check=(0, 1)).to(tl.float32) |
| b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) |
| b_A = tl.where(m_s, b_A, 0.0).to(tl.float32) |
| b_o += _frozen_dot_whitelisted_dot(b_A, b_v, BF16_INTRA_DOT_OPERANDS, DOT_PRECISION) |
| tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) |
|
|
|
|
| |
|
|
| NUM_WARPS_INTRA = [1, 2, 4] if IS_NVIDIA_HOPPER else [1, 2, 4, 8] |
|
|
| @triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) |
| @triton.autotune(configs=[triton.Config({'BH': BH}, num_warps=num_warps) for BH in [1, 2, 4, 8] for num_warps in [1, 2, 4, 8]], key=['K', 'H'], **autotune_cache_kwargs) |
| @triton.jit(do_not_specialize=['T', 'N']) |
| def _frozen_intra_chunk_kalman_fwd_kernel_intra_token_parallel(q, k, kappa, g, Aqk, Akk, scale, cu_seqlens, N, T, H: tl.constexpr, K: tl.constexpr, BT: tl.constexpr, BC: tl.constexpr, BH: tl.constexpr, IS_VARLEN: tl.constexpr): |
| i_tg, i_hg = (tl.program_id(0), tl.program_id(1)) |
| if IS_VARLEN: |
| i_n = 0 |
| left, right = (0, N) |
| for _ in range(20): |
| if left < right: |
| mid = (left + right) // 2 |
| if i_tg < tl.load(cu_seqlens + mid + 1).to(tl.int32): |
| right = mid |
| else: |
| left = mid + 1 |
| i_n = left |
| bos, eos = (tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)) |
| T = eos - bos |
| i_t = i_tg - bos |
| else: |
| bos = i_tg // T * T |
| i_t = i_tg % T |
| if i_t >= T: |
| return |
| i_c = i_t // BT |
| i_s = i_t % BT // BC |
| i_tc = i_c * BT |
| i_ts = i_tc + i_s * BC |
| q += bos * H * K |
| k += bos * H * K |
| kappa += bos * H * K |
| g += bos * H * K |
| Aqk += bos * H * BT |
| Akk += bos * H * BC |
| BK: tl.constexpr = triton.next_power_of_2(K) |
| o_h = tl.arange(0, BH) |
| o_k = tl.arange(0, BK) |
| m_h = i_hg * BH + o_h < H |
| m_k = o_k < K |
| p_q = tl.make_block_ptr(q + i_t * H * K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) |
| p_k = tl.make_block_ptr(k + i_t * H * K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) |
| p_g = tl.make_block_ptr(g + i_t * H * K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) |
| b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) |
| b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) |
| b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) |
| for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))): |
| p_kappaj = tl.make_block_ptr(kappa + j * H * K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) |
| p_gj = tl.make_block_ptr(g + j * H * K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) |
| b_kappaj = tl.load(p_kappaj, boundary_check=(0, 1)).to(tl.float32) |
| b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32) |
| b_kappagj = b_kappaj * exp2(b_g - b_gj) |
| b_kappagj = tl.where(m_k[None, :], b_kappagj, 0.0) |
| b_Aqk = tl.sum(b_q * b_kappagj, axis=1) * scale |
| b_Akk = tl.sum(b_k * b_kappagj, axis=1) * tl.where(j < i_t, 1.0, 0.0) |
| tl.store(Aqk + i_t * H * BT + (i_hg * BH + o_h) * BT + j % BT, b_Aqk.to(Aqk.dtype.element_ty), mask=m_h) |
| tl.store(Akk + i_t * H * BC + (i_hg * BH + o_h) * BC + j - i_ts, b_Akk.to(Akk.dtype.element_ty), mask=m_h) |
|
|
| @triton.heuristics({'STORE_M': lambda args: args['Mraw'] is not None}) |
| @triton.autotune(configs=[triton.Config({'BK': BK}, num_warps=num_warps) for BK in [32, 64] for num_warps in [1, 2, 4]], key=['H', 'K', 'BC', 'SCORE_DOT_PRECISION', 'SOLVE_DOT_PRECISION'], **autotune_cache_kwargs) |
| @triton.jit(do_not_specialize=['T']) |
| def _frozen_intra_chunk_kalman_fwd_kernel_inter_solve_fused_bt32_explicit(q, k, kappa, g, Aqk, Akkd, Akk, Mraw, scale, T, H: tl.constexpr, K: tl.constexpr, BT: tl.constexpr, BC: tl.constexpr, NC: tl.constexpr, BK: tl.constexpr, STORE_M: tl.constexpr, SCORE_DOT_PRECISION: tl.constexpr, SOLVE_DOT_PRECISION: tl.constexpr): |
| """Two-subchunk BT32 solve with the same FP32/TF32x3 contract as BT64.""" |
| tl.static_assert(BT == 32) |
| tl.static_assert(BC == 16) |
| tl.static_assert(NC == 2) |
| tl.static_assert(BT == BC * NC) |
| i_t, i_bh = (tl.program_id(0), tl.program_id(1)) |
| i_b, i_h = (i_bh // H, i_bh % H) |
| bos = i_b * T |
| if i_t * BT >= T: |
| return |
| i_tc0 = i_t * BT |
| i_tc1 = i_t * BT + BC |
| q += (bos * H + i_h) * K |
| k += (bos * H + i_h) * K |
| kappa += (bos * H + i_h) * K |
| g += (bos * H + i_h) * K |
| Aqk += (bos * H + i_h) * BT |
| Akk += (bos * H + i_h) * BT |
| Akkd += (bos * H + i_h) * BC |
| if STORE_M: |
| Mraw += (bos * H + i_h) * BT |
| o_i = tl.arange(0, BC) |
| m_tc1 = i_tc1 + o_i < T |
| b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32) |
| b_M10 = tl.zeros([BC, BC], dtype=tl.float32) |
| for i_k in range(tl.cdiv(K, BK)): |
| o_k = i_k * BK + tl.arange(0, BK) |
| m_k = o_k < K |
| p_kappa0 = tl.make_block_ptr(kappa, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) |
| p_g0 = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) |
| b_kappa0 = tl.load(p_kappa0, boundary_check=(0, 1), padding_option='zero').to(tl.float32) |
| b_g0 = tl.load(p_g0, boundary_check=(0, 1), padding_option='zero').to(tl.float32) |
| if i_tc1 < T: |
| p_q1 = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) |
| p_k1 = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) |
| p_g1 = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) |
| b_q1 = tl.load(p_q1, boundary_check=(0, 1), padding_option='zero').to(tl.float32) |
| b_k1 = tl.load(p_k1, boundary_check=(0, 1), padding_option='zero').to(tl.float32) |
| b_g1 = tl.load(p_g1, boundary_check=(0, 1), padding_option='zero').to(tl.float32) |
| b_gn1 = tl.load(g + i_tc1 * H * K + o_k, mask=m_k, other=0).to(tl.float32) |
| b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0) |
| b_kgt = tl.trans(b_kappa0 * exp2(b_gn1[None, :] - b_g0)) |
| b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt, input_precision=SCORE_DOT_PRECISION, out_dtype=tl.float32) |
| b_M10 += tl.dot(b_k1 * b_gqn, b_kgt, input_precision=SCORE_DOT_PRECISION, out_dtype=tl.float32) |
| b_Aqk10 = tl.where(m_tc1[:, None], b_Aqk10, 0.0) |
| b_M10 = tl.where(m_tc1[:, None], b_M10, 0.0) |
| if i_tc1 < T: |
| p_Aqk10 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) |
| tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) |
| p_Akkd00 = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc0, 0), (BC, BC), (1, 0)) |
| p_Akkd11 = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc1, 0), (BC, BC), (1, 0)) |
| b_Ai00 = tl.load(p_Akkd00, boundary_check=(0, 1), padding_option='zero').to(tl.float32) |
| b_Ai11 = tl.load(p_Akkd11, boundary_check=(0, 1), padding_option='zero').to(tl.float32) |
| m_A = o_i[:, None] > o_i[None, :] |
| m_I = o_i[:, None] == o_i[None, :] |
| b_Ai00 = -tl.where(m_A, b_Ai00, 0) |
| b_Ai11 = -tl.where(m_A, b_Ai11, 0) |
| if STORE_M: |
| p_M00 = tl.make_block_ptr(Mraw, (T, BT), (H * BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) |
| p_M11 = tl.make_block_ptr(Mraw, (T, BT), (H * BT, 1), (i_tc1, BC), (BC, BC), (1, 0)) |
| p_M10 = tl.make_block_ptr(Mraw, (T, BT), (H * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) |
| tl.store(p_M00, (-b_Ai00).to(Mraw.dtype.element_ty), boundary_check=(0, 1)) |
| tl.store(p_M11, (-b_Ai11).to(Mraw.dtype.element_ty), boundary_check=(0, 1)) |
| tl.store(p_M10, b_M10.to(Mraw.dtype.element_ty), boundary_check=(0, 1)) |
| for i in range(2, min(BC, T - i_tc0)): |
| b_a00 = -tl.load(Akkd + (i_tc0 + i) * H * BC + o_i) |
| b_a00 = tl.where(o_i < i, b_a00, 0.0) |
| b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0) |
| b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) |
| for i in range(BC + 2, min(2 * BC, T - i_tc0)): |
| b_a11 = -tl.load(Akkd + (i_tc0 + i) * H * BC + o_i) |
| b_a11 = tl.where(o_i < i - BC, b_a11, 0.0) |
| b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0) |
| b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11) |
| b_Ai00 += m_I |
| b_Ai11 += m_I |
| b_Ai10 = -tl.dot(tl.dot(b_Ai11, b_M10, input_precision=SOLVE_DOT_PRECISION, out_dtype=tl.float32), b_Ai00, input_precision=SOLVE_DOT_PRECISION, out_dtype=tl.float32) |
| b_Ai10 = tl.where(m_tc1[:, None], b_Ai10, 0.0) |
| p_Akk00 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) |
| p_Akk10 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) |
| p_Akk11 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc1, BC), (BC, BC), (1, 0)) |
| tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) |
| tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) |
| tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) |
|
|
| @triton.heuristics({'STORE_QG': lambda args: args['qg'] is not None, 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) |
| @triton.autotune(configs=[triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in [2, 4, 8] for num_stages in [2, 3, 4]], key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN', 'DOT_PRECISION', 'BF16_DOT_OPERANDS'], **autotune_cache_kwargs) |
| @triton.jit(do_not_specialize=['T']) |
| def _frozen_intra_recompute_w_u_fwd_kalman_kernel_explicit(q, k, kappa, qg, kg, v, w, u, A, gk, cu_seqlens, chunk_indices, T, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, STORE_QG: tl.constexpr, IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, BF16_DOT_OPERANDS: tl.constexpr): |
| i_t, i_bh = (tl.program_id(0), tl.program_id(1)) |
| i_b, i_h = (i_bh // H, i_bh % H) |
| if IS_VARLEN: |
| i_n, i_t = (tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)) |
| bos, eos = (tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)) |
| T = eos - bos |
| else: |
| bos, eos = (i_b * T, i_b * T + T) |
| p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) |
| b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) |
| for i_v in range(tl.cdiv(V, BV)): |
| p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) |
| p_u = tl.make_block_ptr(u + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) |
| b_v = tl.load(p_v, boundary_check=(0, 1)).to(tl.float32) |
| b_u = _frozen_dot_whitelisted_dot(b_A, b_v, BF16_DOT_OPERANDS, DOT_PRECISION) |
| tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) |
| for i_k in range(tl.cdiv(K, BK)): |
| p_w = tl.make_block_ptr(w + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) |
| p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) |
| b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) |
| p_gk = tl.make_block_ptr(gk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) |
| b_gk = tl.load(p_gk, boundary_check=(0, 1)).to(tl.float32) |
| b_kb = b_k * exp2(b_gk) |
| if STORE_QG: |
| p_q = tl.make_block_ptr(q + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) |
| p_qg = tl.make_block_ptr(qg + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) |
| b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) |
| b_qg = b_q * exp2(b_gk) |
| tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) |
| last_idx = min(i_t * BT + BT, T) - 1 |
| o_k = i_k * BK + tl.arange(0, BK) |
| m_k = o_k < K |
| b_gn = tl.load(gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0).to(tl.float32) |
| p_kappa = tl.make_block_ptr(kappa + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) |
| b_kappa = tl.load(p_kappa, boundary_check=(0, 1)).to(tl.float32) |
| b_kg = b_kappa * tl.where((i_t * BT + tl.arange(0, BT) < T)[:, None], exp2(b_gn[None, :] - b_gk), 0) |
| p_kg = tl.make_block_ptr(kg + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) |
| tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) |
| b_w = _frozen_dot_whitelisted_dot(b_A, b_kb.to(tl.float32), BF16_DOT_OPERANDS, DOT_PRECISION) |
| tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) |
|
|
| @triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) |
| @triton.autotune(configs=[triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in NUM_WARPS_INTRA for num_stages in [2, 3, 4]], key=['BK', 'NC', 'BT', 'DOT_PRECISION', 'BF16_DOT_OPERANDS', 'BF16_DK_ROW_OPERANDS', 'BF16_DQ_ROW_OPERANDS', 'DQ_ROW_DOT_PRECISION', 'K_TILES_PER_CTA'], **autotune_cache_kwargs) |
| @triton.jit(do_not_specialize=['B', 'T']) |
| def _frozen_intra_chunk_kalman_bwd_kernel_intra_explicit(q, k, kappa, g, dAqk, dAkk, dq, dq2, dk, dk2, dkappa, dkappa2, dg, dg2, cu_seqlens, chunk_indices, B, T, H: tl.constexpr, K: tl.constexpr, BT: tl.constexpr, BC: tl.constexpr, BK: tl.constexpr, NC: tl.constexpr, IS_VARLEN: tl.constexpr, SAFE_GATE: tl.constexpr, USE_GATHER: tl.constexpr, DOT_PRECISION: tl.constexpr, BF16_DOT_OPERANDS: tl.constexpr, BF16_DK_ROW_OPERANDS: tl.constexpr, BF16_DQ_ROW_OPERANDS: tl.constexpr, DQ_ROW_DOT_PRECISION: tl.constexpr, K_TILES_PER_CTA: tl.constexpr): |
| i_kc, i_t, i_bh = (tl.program_id(0), tl.program_id(1), tl.program_id(2)) |
| i_b, i_h = (i_bh // H, i_bh % H) |
| i_k_group, i_i = (i_kc // NC, i_kc % NC) |
| all = B * T |
| if IS_VARLEN: |
| i_n, i_t = (tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)) |
| bos, eos = (tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)) |
| else: |
| bos, eos = (i_b * T, i_b * T + T) |
| T = eos - bos |
| i_ti = i_t * BT + i_i * BC |
| if i_ti >= T: |
| return |
| q += (bos * H + i_h) * K |
| k += (bos * H + i_h) * K |
| kappa += (bos * H + i_h) * K |
| g += (bos * H + i_h) * K |
| dAqk += (bos * H + i_h) * BT |
| dAkk += (bos * H + i_h) * BT |
| dq += (bos * H + i_h) * K |
| dq2 += (bos * H + i_h) * K |
| dk += (bos * H + i_h) * K |
| dk2 += (bos * H + i_h) * K |
| dkappa += (bos * H + i_h) * K |
| dkappa2 += (bos * H + i_h) * K |
| dg += (bos * H + i_h) * K |
| dg2 += (bos * H + i_h) * K |
| for i_k_offset in range(0, K_TILES_PER_CTA): |
| i_k = i_k_group * K_TILES_PER_CTA + i_k_offset |
| o_k = i_k * BK + tl.arange(0, BK) |
| m_k = o_k < K |
| p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) |
| b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) |
| b_dq2 = tl.zeros([BC, BK], dtype=tl.float32) |
| b_dk2 = tl.zeros([BC, BK], dtype=tl.float32) |
| if i_i > 0: |
| p_gn = g + i_ti * H * K + o_k |
| b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] |
| for i_j in range(0, i_i): |
| p_kappa = tl.make_block_ptr(kappa, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) |
| p_gk = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) |
| p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (H * BT, 1), (i_ti, i_j * BC), (BC, BC), (1, 0)) |
| p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (H * BT, 1), (i_ti, i_j * BC), (BC, BC), (1, 0)) |
| b_kappa = tl.load(p_kappa, boundary_check=(0, 1)).to(tl.float32) |
| b_gk = tl.load(p_gk, boundary_check=(0, 1)).to(tl.float32) |
| b_kappag = b_kappa * exp2(b_gn - b_gk) |
| b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) |
| b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) |
| b_dq2 += _frozen_dot_whitelisted_dot(b_dAqk, b_kappag, BF16_DQ_ROW_OPERANDS, DQ_ROW_DOT_PRECISION) |
| b_dk2 += _frozen_dot_whitelisted_dot(b_dAkk, b_kappag, BF16_DK_ROW_OPERANDS, DOT_PRECISION) |
| b_gqn = exp2(b_g - b_gn) |
| b_dq2 *= b_gqn |
| b_dk2 *= b_gqn |
| o_i = tl.arange(0, BC) |
| m_dA = i_ti + o_i < T |
| o_dA = (i_ti + o_i) * H * BT + i_i * BC |
| p_kappaj = kappa + i_ti * H * K + o_k |
| p_gkj = g + i_ti * H * K + o_k |
| p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) |
| p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) |
| b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) |
| b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) |
| if SAFE_GATE: |
| p_kappa_diag = tl.make_block_ptr(kappa, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) |
| b_kappa_diag = tl.load(p_kappa_diag, boundary_check=(0, 1)).to(tl.float32) |
| if USE_GATHER: |
| b_gn = gather(b_g, tl.full([1, BK], min(BC // 2, T - i_ti - 1), dtype=tl.int16), axis=0) |
| else: |
| p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H * K + o_k |
| b_gn = tl.load(p_gn, mask=m_k, other=0)[None, :] |
| p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (H * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) |
| p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (H * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) |
| b_dAqk_diag_qk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32) |
| b_dAkk_diag_qk = tl.load(p_dAkk, boundary_check=(0, 1)).to(tl.float32) |
| m_i_diag_qk = (o_i[:, None] >= o_i[None, :]) & (i_ti + o_i[:, None] < T) & (i_ti + o_i[None, :] < T) |
| m_j_diag_qk = i_ti + o_i[:, None] < T |
| b_dAqk_diag_qk = tl.where(m_i_diag_qk, b_dAqk_diag_qk, 0.0) |
| b_dAkk_diag_qk = tl.where(m_i_diag_qk, b_dAkk_diag_qk, 0.0) |
| b_g_diag_qk = tl.where(m_j_diag_qk, b_g - b_gn, 0.0) |
| exp_b_g_diag_qk = tl.where(m_j_diag_qk, exp2(b_g_diag_qk), 0.0) |
| exp_neg_b_g_diag_qk = tl.where(m_j_diag_qk, exp2(-b_g_diag_qk), 0.0) |
| b_kappa_exp_diag_qk = b_kappa_diag * exp_neg_b_g_diag_qk |
| b_dq2 += _frozen_dot_whitelisted_dot(b_dAqk_diag_qk, b_kappa_exp_diag_qk, BF16_DQ_ROW_OPERANDS, DQ_ROW_DOT_PRECISION) * exp_b_g_diag_qk |
| b_dk2 += _frozen_dot_whitelisted_dot(b_dAkk_diag_qk, b_kappa_exp_diag_qk, BF16_DK_ROW_OPERANDS, DOT_PRECISION) * exp_b_g_diag_qk |
| else: |
| for j in range(0, min(BC, T - i_t * BT - i_i * BC)): |
| b_dAqk = tl.load(dAqk + o_dA + j, mask=m_dA, other=0) |
| b_dAkk = tl.load(dAkk + o_dA + j, mask=m_dA, other=0) |
| b_kappaj = tl.load(p_kappaj, mask=m_k, other=0).to(tl.float32) |
| b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) |
| m_i = o_i[:, None] >= j |
| b_gqk = exp2(b_g - b_gkj[None, :]) |
| b_dq2 += tl.where(m_i, b_dAqk[:, None] * b_kappaj[None, :] * b_gqk, 0.0) |
| b_dk2 += tl.where(m_i, b_dAkk[:, None] * b_kappaj[None, :] * b_gqk, 0.0) |
| p_kappaj += H * K |
| p_gkj += H * K |
| p_dq = tl.make_block_ptr(dq, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) |
| p_dq2 = tl.make_block_ptr(dq2, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) |
| b_dg2 = b_q * b_dq2 |
| b_dq2 = b_dq2 + tl.load(p_dq, boundary_check=(0, 1)) |
| tl.store(p_dq2, b_dq2.to(p_dq2.dtype.element_ty), boundary_check=(0, 1)) |
| tl.debug_barrier() |
| b_dkt = tl.zeros([BC, BK], dtype=tl.float32) |
| active_nc = min(NC, tl.cdiv(T - i_t * BT, BC)) |
| if i_i < active_nc - 1: |
| p_gn = g + (min(i_ti + BC, T) - 1) * H * K + o_k |
| b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] |
| for i_j in range(i_i + 1, active_nc): |
| p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) |
| p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) |
| p_gk = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) |
| p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, H * BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) |
| p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, H * BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) |
| b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) |
| b_k_row = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) |
| b_gk = tl.load(p_gk, boundary_check=(0, 1)).to(tl.float32) |
| b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) |
| b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) |
| o_j = i_t * BT + i_j * BC + o_i |
| m_j = o_j < T |
| b_gkn = exp2(b_gk - b_gn) |
| b_qg = b_q * tl.where(m_j[:, None], b_gkn, 0) |
| b_kg_row = b_k_row * tl.where(m_j[:, None], b_gkn, 0) |
| b_dkt += tl.dot(b_dAqk, b_qg, input_precision=DOT_PRECISION, out_dtype=tl.float32) |
| b_dkt += tl.dot(b_dAkk, b_kg_row, input_precision=DOT_PRECISION, out_dtype=tl.float32) |
| b_dkt *= exp2(b_gn - b_g) |
| o_dA = i_ti * H * BT + i_i * BC + o_i |
| p_qj = q + i_ti * H * K + o_k |
| p_kj = k + i_ti * H * K + o_k |
| p_gkj = g + i_ti * H * K + o_k |
| if SAFE_GATE: |
| if USE_GATHER: |
| b_gn = gather(b_g, tl.full([1, BK], min(BC // 2, T - i_ti - 1), dtype=tl.int16), axis=0) |
| else: |
| p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H * K + o_k |
| b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] |
| p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) |
| b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) |
| p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, H * BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) |
| p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, H * BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) |
| b_dAqk_diag_kk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32) |
| b_dAkk_diag_kk = tl.load(p_dAkk, boundary_check=(0, 1)).to(tl.float32) |
| m_i_diag_kk = (o_i[:, None] <= o_i[None, :]) & (i_ti + o_i[:, None] < T) & (i_ti + o_i[None, :] < T) |
| m_j_diag_kk = i_ti + o_i[:, None] < T |
| b_dAqk_diag_kk = tl.where(m_i_diag_kk, b_dAqk_diag_kk, 0.0) |
| b_dAkk_diag_kk = tl.where(m_i_diag_kk, b_dAkk_diag_kk, 0.0) |
| b_g_diag_kk = tl.where(m_j_diag_kk, b_g - b_gn, 0.0) |
| exp_b_g_diag_kk = tl.where(m_j_diag_kk, exp2(b_g_diag_kk), 0.0) |
| exp_neg_b_g_diag_kk = tl.where(m_j_diag_kk, exp2(-b_g_diag_kk), 0.0) |
| b_q_exp = b_q * exp_b_g_diag_kk |
| b_k_exp = b_k * exp_b_g_diag_kk |
| b_dkt += tl.dot(b_dAqk_diag_kk, b_q_exp, input_precision=DOT_PRECISION, out_dtype=tl.float32) * exp_neg_b_g_diag_kk |
| b_dkt += tl.dot(b_dAkk_diag_kk, b_k_exp, input_precision=DOT_PRECISION, out_dtype=tl.float32) * exp_neg_b_g_diag_kk |
| else: |
| for j in range(0, min(BC, T - i_t * BT - i_i * BC)): |
| b_dAqk = tl.load(dAqk + o_dA + j * H * BT) |
| b_dAkk = tl.load(dAkk + o_dA + j * H * BT) |
| b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) |
| b_kj_row = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) |
| b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) |
| m_i = o_i[:, None] <= j |
| b_gkq = exp2(b_gkj[None, :] - b_g) |
| b_dkt += tl.where(m_i, b_dAqk[:, None] * b_qj[None, :] * b_gkq, 0.0) |
| b_dkt += tl.where(m_i, b_dAkk[:, None] * b_kj_row[None, :] * b_gkq, 0.0) |
| p_qj += H * K |
| p_kj += H * K |
| p_gkj += H * K |
| p_kappa_row = tl.make_block_ptr(kappa, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) |
| b_kappa_row = tl.load(p_kappa_row, boundary_check=(0, 1)) |
| p_dk = tl.make_block_ptr(dk, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) |
| p_dk2 = tl.make_block_ptr(dk2, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) |
| p_dkappa = tl.make_block_ptr(dkappa, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) |
| p_dkappa2 = tl.make_block_ptr(dkappa2, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) |
| p_dg = tl.make_block_ptr(dg, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) |
| p_dg2 = tl.make_block_ptr(dg2, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) |
| b_dg2 += b_dk2 * b_k - b_dkt * b_kappa_row + tl.load(p_dg, boundary_check=(0, 1)) |
| b_dk2 += tl.load(p_dk, boundary_check=(0, 1)) |
| b_dkt += tl.load(p_dkappa, boundary_check=(0, 1)) |
| tl.store(p_dk2, b_dk2.to(p_dk2.dtype.element_ty), boundary_check=(0, 1)) |
| tl.store(p_dkappa2, b_dkt.to(p_dkappa2.dtype.element_ty), boundary_check=(0, 1)) |
| tl.store(p_dg2, b_dg2.to(p_dg2.dtype.element_ty), boundary_check=(0, 1)) |
|
|
|
|
| |
|
|
| BT = 32 |
|
|
| BK = 64 |
|
|
| BV = 64 |
|
|
| @triton.jit(do_not_specialize=['T']) |
| def _frozen_wy_chunk_kalman_bwd_kernel_wy_buffered_dw_producer(q, k, kappa, v_new, g, A, h, do, dh, dv, dq, dk, dkappa, dv2, dg, dw_scratch, scale, T, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, DOT_PRECISION: tl.constexpr, DKAPPA_DM_DOT_PRECISION: tl.constexpr, DQ_DO_H_DOT_PRECISION: tl.constexpr, BF16_DOT_OPERANDS: tl.constexpr, BF16_DW_DU_H_OPERANDS: tl.constexpr, BF16_DK_ABAR_DW_OPERANDS: tl.constexpr, BF16_DQ_DO_H_OPERANDS: tl.constexpr): |
| """Produce the five non-dM outputs and post-negation FP32 dw.""" |
| tl.static_assert(BT == 32) |
| tl.static_assert(BK == 64) |
| tl.static_assert(BV == 64) |
| tl.static_assert(K == V) |
| tl.static_assert(K == 64 or K == 128) |
| tl.static_assert(DOT_PRECISION == 'tf32x3') |
| tl.static_assert(DKAPPA_DM_DOT_PRECISION == 'tf32') |
| tl.static_assert(DQ_DO_H_DOT_PRECISION == 'tf32x3' or DQ_DO_H_DOT_PRECISION == 'tf32') |
| tl.static_assert(BF16_DOT_OPERANDS) |
| tl.static_assert(BF16_DW_DU_H_OPERANDS) |
| tl.static_assert(BF16_DK_ABAR_DW_OPERANDS) |
| i_t, i_bh = (tl.program_id(0), tl.program_id(1)) |
| i_b, i_h = (i_bh // H, i_bh % H) |
| NT = tl.cdiv(T, BT) |
| i_tg = (i_b * NT + i_t).to(tl.int64) |
| bos = (i_b * T).to(tl.int64) |
| o_t = i_t * BT + tl.arange(0, BT) |
| m_t = o_t < T |
| m_last = o_t == min(T, i_t * BT + BT) - 1 |
| q += (bos * H + i_h) * K |
| k += (bos * H + i_h) * K |
| kappa += (bos * H + i_h) * K |
| v_new += (bos * H + i_h) * V |
| g += (bos * H + i_h) * K |
| A += (bos * H + i_h) * BT |
| h += (i_tg * H + i_h) * K * V |
| do += (bos * H + i_h) * V |
| dh += (i_tg * H + i_h) * K * V |
| dv += (bos * H + i_h) * V |
| dq += (bos * H + i_h) * K |
| dk += (bos * H + i_h) * K |
| dkappa += (bos * H + i_h) * K |
| dv2 += (bos * H + i_h) * V |
| dg += (bos * H + i_h) * K |
| dw_scratch += (bos * H + i_h) * K |
| p_A_p = tl.make_block_ptr(A, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1)) |
| b_A_p = tl.load(p_A_p, boundary_check=(0, 1)).to(tl.float32) |
| for i_k_p in range(tl.cdiv(K, BK)): |
| o_k = i_k_p * BK + tl.arange(0, BK) |
| m_k = o_k < K |
| b_dq_p = tl.zeros([BT, BK], dtype=tl.float32) |
| b_dkappa_p = tl.zeros([BT, BK], dtype=tl.float32) |
| b_dw_p = tl.zeros([BT, BK], dtype=tl.float32) |
| b_dgk_p = tl.zeros([BK], dtype=tl.float32) |
| for i_v_p in range(tl.cdiv(V, BV)): |
| p_v_new_p = tl.make_block_ptr(v_new, (T, V), (H * V, 1), (i_t * BT, i_v_p * BV), (BT, BV), (1, 0)) |
| p_do_p = tl.make_block_ptr(do, (T, V), (H * V, 1), (i_t * BT, i_v_p * BV), (BT, BV), (1, 0)) |
| p_h_p = tl.make_block_ptr(h, (V, K), (1, V), (i_v_p * BV, i_k_p * BK), (BV, BK), (0, 1)) |
| p_dh_p = tl.make_block_ptr(dh, (V, K), (1, V), (i_v_p * BV, i_k_p * BK), (BV, BK), (0, 1)) |
| p_dv_p = tl.make_block_ptr(dv, (T, V), (H * V, 1), (i_t * BT, i_v_p * BV), (BT, BV), (1, 0)) |
| b_v_new_p = tl.load(p_v_new_p, boundary_check=(0, 1)).to(tl.float32) |
| b_do_p = tl.load(p_do_p, boundary_check=(0, 1)).to(tl.float32) |
| b_h_p = tl.load(p_h_p, boundary_check=(0, 1)).to(tl.float32) |
| b_dh_p = tl.load(p_dh_p, boundary_check=(0, 1)).to(tl.float32) |
| b_dv_p = tl.load(p_dv_p, boundary_check=(0, 1)).to(tl.float32) |
| b_dgk_p += tl.sum(b_h_p * b_dh_p, axis=0) |
| b_dq_p += _frozen_dot_whitelisted_dot(b_do_p, b_h_p, BF16_DQ_DO_H_OPERANDS, DQ_DO_H_DOT_PRECISION) |
| b_dkappa_p += tl.dot(b_v_new_p, b_dh_p, input_precision=DKAPPA_DM_DOT_PRECISION, out_dtype=tl.float32) |
| b_dw_p += _frozen_dot_whitelisted_dot(b_dv_p, b_h_p, BF16_DW_DU_H_OPERANDS, DOT_PRECISION) |
| tl.debug_barrier() |
| if i_k_p == 0: |
| p_dv2_p = tl.make_block_ptr(dv2, (T, V), (H * V, 1), (i_t * BT, i_v_p * BV), (BT, BV), (1, 0)) |
| b_dv2_p = _frozen_dot_whitelisted_dot(b_A_p, b_dv_p, BF16_DOT_OPERANDS, DOT_PRECISION) |
| tl.store(p_dv2_p, b_dv2_p.to(p_dv2_p.dtype.element_ty), boundary_check=(0, 1)) |
| p_k_p = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT, i_k_p * BK), (BT, BK), (1, 0)) |
| p_g_p = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT, i_k_p * BK), (BT, BK), (1, 0)) |
| b_k_p = tl.load(p_k_p, boundary_check=(0, 1)).to(tl.float32) |
| b_g_p = tl.load(p_g_p, boundary_check=(0, 1)).to(tl.float32) |
| p_gn_p = g + (min(T, i_t * BT + BT) - 1).to(tl.int64) * H * K + o_k |
| b_gn_p = tl.load(p_gn_p, mask=m_k, other=0.0).to(tl.float32) |
| b_g_exp_p = exp2(b_g_p) |
| b_dgk_p *= exp2(b_gn_p) |
| b_dq_p = b_dq_p * b_g_exp_p * scale |
| b_dkappa_p *= tl.where(m_t[:, None], exp2(b_gn_p[None, :] - b_g_p), 0.0) |
| b_kg_p = b_k_p * b_g_exp_p |
| b_dw_p = -b_dw_p.to(tl.float32) |
| b_dkgb_p = _frozen_dot_whitelisted_dot(b_A_p, b_dw_p, BF16_DK_ABAR_DW_OPERANDS, DOT_PRECISION) |
| p_dw_scratch_p = tl.make_block_ptr(dw_scratch, (T, K), (H * K, 1), (i_t * BT, i_k_p * BK), (BT, BK), (1, 0)) |
| tl.store(p_dw_scratch_p, b_dw_p.to(p_dw_scratch_p.dtype.element_ty), boundary_check=(0, 1)) |
| p_q_p = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT, i_k_p * BK), (BT, BK), (1, 0)) |
| p_kappa_p = tl.make_block_ptr(kappa, (T, K), (H * K, 1), (i_t * BT, i_k_p * BK), (BT, BK), (1, 0)) |
| b_q_p = tl.load(p_q_p, boundary_check=(0, 1)).to(tl.float32) |
| b_kappa_p = tl.load(p_kappa_p, boundary_check=(0, 1)).to(tl.float32) |
| b_kdk_p = b_kappa_p * b_dkappa_p |
| b_dgk_p += tl.sum(b_kdk_p, axis=0) |
| b_dg_p = b_q_p * b_dq_p - b_kdk_p + m_last[:, None] * b_dgk_p + b_kg_p * b_dkgb_p |
| b_dk_p = b_dkgb_p * b_g_exp_p |
| p_dq_p = tl.make_block_ptr(dq, (T, K), (H * K, 1), (i_t * BT, i_k_p * BK), (BT, BK), (1, 0)) |
| p_dk_p = tl.make_block_ptr(dk, (T, K), (H * K, 1), (i_t * BT, i_k_p * BK), (BT, BK), (1, 0)) |
| p_dkappa_p = tl.make_block_ptr(dkappa, (T, K), (H * K, 1), (i_t * BT, i_k_p * BK), (BT, BK), (1, 0)) |
| p_dg_p = tl.make_block_ptr(dg, (T, K), (H * K, 1), (i_t * BT, i_k_p * BK), (BT, BK), (1, 0)) |
| tl.store(p_dq_p, b_dq_p.to(p_dq_p.dtype.element_ty), boundary_check=(0, 1)) |
| tl.store(p_dk_p, b_dk_p.to(p_dk_p.dtype.element_ty), boundary_check=(0, 1)) |
| tl.store(p_dkappa_p, b_dkappa_p.to(p_dkappa_p.dtype.element_ty), boundary_check=(0, 1)) |
| tl.store(p_dg_p, b_dg_p.to(p_dg_p.dtype.element_ty), boundary_check=(0, 1)) |
|
|
| @triton.jit(do_not_specialize=['T']) |
| def _frozen_wy_chunk_kalman_bwd_kernel_wy_buffered_dw_consumer(k, v, g, A, dv, dw_scratch, dM, T, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, DKAPPA_DM_DOT_PRECISION: tl.constexpr): |
| """Consume buffered post-negation dw to form the exact Phase-D dM.""" |
| tl.static_assert(BT == 32) |
| tl.static_assert(BK == 64) |
| tl.static_assert(BV == 64) |
| tl.static_assert(K == V) |
| tl.static_assert(K == 64 or K == 128) |
| tl.static_assert(DKAPPA_DM_DOT_PRECISION == 'tf32') |
| i_t, i_bh = (tl.program_id(0), tl.program_id(1)) |
| i_b, i_h = (i_bh // H, i_bh % H) |
| bos = (i_b * T).to(tl.int64) |
| o_t = i_t * BT + tl.arange(0, BT) |
| m_t = o_t < T |
| k += (bos * H + i_h) * K |
| v += (bos * H + i_h) * V |
| g += (bos * H + i_h) * K |
| A += (bos * H + i_h) * BT |
| dv += (bos * H + i_h) * V |
| dw_scratch += (bos * H + i_h) * K |
| dM += (bos * H + i_h) * BT |
| b_dM_c = tl.zeros([BT, BT], dtype=tl.float32) |
| for i_v_c in range(tl.cdiv(V, BV)): |
| p_v_c = tl.make_block_ptr(v, (T, V), (H * V, 1), (i_t * BT, i_v_c * BV), (BT, BV), (1, 0)) |
| p_dv_c = tl.make_block_ptr(dv, (T, V), (H * V, 1), (i_t * BT, i_v_c * BV), (BT, BV), (1, 0)) |
| b_v_c = tl.load(p_v_c, boundary_check=(0, 1)).to(tl.float32) |
| b_dv_c = tl.load(p_dv_c, boundary_check=(0, 1)).to(tl.float32) |
| b_dM_c += tl.dot(b_dv_c, tl.trans(b_v_c), input_precision=DKAPPA_DM_DOT_PRECISION, out_dtype=tl.float32) |
| for i_k_c in range(tl.cdiv(K, BK)): |
| p_dw_scratch_c = tl.make_block_ptr(dw_scratch, (T, K), (H * K, 1), (i_t * BT, i_k_c * BK), (BT, BK), (1, 0)) |
| b_dw_c = tl.load(p_dw_scratch_c, boundary_check=(0, 1)).to(tl.float32) |
| p_k_c = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT, i_k_c * BK), (BT, BK), (1, 0)) |
| p_g_c = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT, i_k_c * BK), (BT, BK), (1, 0)) |
| b_k_c = tl.load(p_k_c, boundary_check=(0, 1)).to(tl.float32) |
| b_g_c = tl.load(p_g_c, boundary_check=(0, 1)).to(tl.float32) |
| b_kg_c = b_k_c * exp2(b_g_c) |
| b_dM_c += tl.dot(b_dw_c, tl.trans(b_kg_c.to(tl.float32)), input_precision=DKAPPA_DM_DOT_PRECISION, out_dtype=tl.float32) |
| m_M = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) |
| b_dM_c = tl.where(m_M, b_dM_c, 0.0) |
| p_A_c = tl.make_block_ptr(A, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1)) |
| b_A_c = tl.load(p_A_c, boundary_check=(0, 1)).to(tl.float32) |
| b_dM_c = tl.dot(b_dM_c.to(tl.float32), b_A_c, input_precision=DKAPPA_DM_DOT_PRECISION, out_dtype=tl.float32) |
| b_dM_c = tl.dot(b_A_c, b_dM_c.to(tl.float32), input_precision=DKAPPA_DM_DOT_PRECISION, out_dtype=tl.float32) |
| b_dM_c = tl.where(m_M, -b_dM_c, 0.0) |
| p_dM_c = tl.make_block_ptr(dM, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) |
| tl.store(p_dM_c, b_dM_c.to(p_dM_c.dtype.element_ty), boundary_check=(0, 1)) |
|
|
|
|
| |
|
|
| @triton.jit |
| def _frozen_frontend_diag_frontend_fwd_kernel(retention_raw_ptr, omega_raw_ptr, r_raw_ptr, A_log_ptr, dt_bias_ptr, initial_precision_param_ptr, first_ptr, alpha_ptr, omega_ptr, r_ptr, initial_precision_ptr, omega_min, r_min, T, NT, H: tl.constexpr, K: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr, OMEGA_COUPLING: tl.constexpr, EMIT_RAW_G: tl.constexpr): |
| i_bc = tl.program_id(0) |
| i_h = tl.program_id(1) |
| i_b = i_bc // NT |
| i_c = i_bc % NT |
| offs_t_1d = i_c * BT + tl.arange(0, BT) |
| offs_k_1d = tl.arange(0, BK) |
| offs_t = offs_t_1d[:, None] |
| offs_k = offs_k_1d[None, :] |
| mask_bthk = (offs_t_1d[:, None] < T) & (offs_k_1d[None, :] < K) |
| mask_bth = offs_t_1d < T |
| offsets_bthk = ((i_b * T + offs_t) * H + i_h) * K + offs_k |
| offsets_bth = (i_b * T + offs_t_1d) * H + i_h |
| f_value = tl.load(retention_raw_ptr + offsets_bthk, mask=mask_bthk, other=0.0).to(tl.float32) |
| z_value = tl.load(omega_raw_ptr + offsets_bthk, mask=mask_bthk, other=0.0).to(tl.float32) |
| y_value = tl.load(r_raw_ptr + offsets_bth, mask=mask_bth, other=0.0).to(tl.float32) |
| A_log_value = tl.load(A_log_ptr + i_h).to(tl.float32) |
| dt_bias_value = tl.load(dt_bias_ptr + i_h * K + offs_k_1d, mask=offs_k_1d < K, other=0.0).to(tl.float32) |
| rate_value = tl.exp(A_log_value) |
| x_value = f_value + dt_bias_value[None, :] |
| softplus_value = softplus(x_value) |
| ell_value = -rate_value * softplus_value |
| alpha_value = tl.exp(ell_value) |
| clamped_alpha_value = tl.maximum(alpha_value, 1e-06) |
| memory_alpha_value = tl.where(alpha_value != alpha_value, alpha_value, clamped_alpha_value) |
| ell_mem_value = tl.log(memory_alpha_value) |
| masked_ell_mem_value = tl.where(mask_bthk, ell_mem_value, 0.0) |
| prefix_value = tl.cumsum(masked_ell_mem_value, axis=0) * 1.4426950408889634 |
| omega_base_value = softplus(z_value) |
| omega_value = omega_min + omega_base_value |
| r_value = r_min + softplus(y_value) |
| tl.store(first_ptr + offsets_bthk, ell_mem_value if EMIT_RAW_G else prefix_value, mask=mask_bthk) |
| tl.store(alpha_ptr + offsets_bthk, alpha_value, mask=mask_bthk) |
| tl.store(omega_ptr + offsets_bthk, omega_value, mask=mask_bthk) |
| tl.store(r_ptr + offsets_bth, r_value, mask=mask_bth) |
| if i_bc == 0: |
| mu_value = tl.load(initial_precision_param_ptr + i_h).to(tl.float32) |
| mu_output_value = 0.1 + softplus(mu_value) |
| tl.store(initial_precision_ptr + i_h, mu_output_value) |
|
|
| @triton.jit |
| def _frozen_frontend_diag_frontend_uncoupled_omega_local_vjp_kernel(omega_raw_ptr, d_omega_ptr, domega_local_ptr, T, NT, H: tl.constexpr, K: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr): |
| i_bc = tl.program_id(0) |
| i_h = tl.program_id(1) |
| i_b = i_bc // NT |
| i_c = i_bc % NT |
| offs_t_1d = i_c * BT + tl.arange(0, BT) |
| offs_k_1d = tl.arange(0, BK) |
| offs_t = offs_t_1d[:, None] |
| offs_k = offs_k_1d[None, :] |
| offsets = ((i_b * T + offs_t) * H + i_h) * K + offs_k |
| mask = (offs_t_1d[:, None] < T) & (offs_k_1d[None, :] < K) |
| z_value = tl.load(omega_raw_ptr + offsets, mask=mask, other=0.0).to(tl.float32) |
| d_omega_value = tl.load(d_omega_ptr + offsets, mask=mask, other=0.0).to(tl.float32) |
| sigmoid_value = tl.sigmoid(z_value) |
| dz_local_value = d_omega_value * sigmoid_value |
| tl.store(domega_local_ptr + offsets, dz_local_value, mask=mask) |
|
|
| @triton.jit |
| def _frozen_frontend_diag_frontend_r_local_vjp_kernel(r_raw_ptr, d_r_ptr, dr_local_ptr, T, NT, H: tl.constexpr, BT: tl.constexpr): |
| i_bc = tl.program_id(0) |
| i_h = tl.program_id(1) |
| i_b = i_bc // NT |
| i_c = i_bc % NT |
| offs_t = i_c * BT + tl.arange(0, BT) |
| offsets = (i_b * T + offs_t) * H + i_h |
| mask = offs_t < T |
| y_value = tl.load(r_raw_ptr + offsets, mask=mask, other=0.0).to(tl.float32) |
| d_r_value = tl.load(d_r_ptr + offsets, mask=mask, other=0.0).to(tl.float32) |
| sigmoid_value = tl.sigmoid(y_value) |
| dy_local_value = d_r_value * sigmoid_value |
| tl.store(dr_local_ptr + offsets, dy_local_value, mask=mask) |
|
|
| @triton.jit |
| def _frozen_frontend_diag_frontend_mu_local_vjp_kernel(initial_precision_param_ptr, d_initial_precision_ptr, dinitial_precision_local_ptr, H, BH: tl.constexpr): |
| offs_h = tl.program_id(0) * BH + tl.arange(0, BH) |
| mask = offs_h < H |
| mu_value = tl.load(initial_precision_param_ptr + offs_h, mask=mask, other=0.0).to(tl.float32) |
| d_mu_value = tl.load(d_initial_precision_ptr + offs_h, mask=mask, other=0.0).to(tl.float32) |
| sigmoid_value = tl.sigmoid(mu_value) |
| dmu_local_value = d_mu_value * sigmoid_value |
| tl.store(dinitial_precision_local_ptr + offs_h, dmu_local_value, mask=mask) |
|
|
| @triton.jit |
| def _frozen_frontend_diag_frontend_dA_reduction_vjp_kernel(dA_partial_ptr, dA_log_ptr, NP, H: tl.constexpr, BNP: tl.constexpr): |
| i_h = tl.program_id(0) |
| offs_n = tl.arange(0, BNP) |
| dA_partial_value = tl.load(dA_partial_ptr + offs_n * H + i_h, mask=offs_n < NP, other=0.0).to(tl.float32) |
| dA_log_value = tl.sum(dA_partial_value, axis=0) |
| tl.store(dA_log_ptr + i_h, dA_log_value) |
|
|
| @triton.jit |
| def _frozen_frontend_diag_frontend_ddt_reduction_vjp_kernel(ddt_partial_ptr, ddt_bias_ptr, NP, H: tl.constexpr, K: tl.constexpr, BNP: tl.constexpr): |
| i_hk = tl.program_id(0) |
| i_h = i_hk // K |
| i_k = i_hk % K |
| offs_n = tl.arange(0, BNP) |
| ddt_partial_value = tl.load(ddt_partial_ptr + (offs_n * H + i_h) * K + i_k, mask=offs_n < NP, other=0.0).to(tl.float32) |
| ddt_bias_value = tl.sum(ddt_partial_value, axis=0) |
| tl.store(ddt_bias_ptr + i_hk, ddt_bias_value) |
|
|
| @triton.jit |
| def _frozen_frontend_diag_frontend_alpha_presence_vjp_kernel(retention_raw_ptr, A_log_ptr, dt_bias_ptr, alpha_ptr, d_first_ptr, d_alpha_gain_ptr, dretention_raw_ptr, dA_partial_ptr, ddt_partial_ptr, stride_d_first_b, stride_d_first_t, stride_d_first_h, stride_d_first_k, T, NT, H: tl.constexpr, K: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr, ALPHA_FLOOR: tl.constexpr, HAS_DFIRST: tl.constexpr, HAS_DALPHA: tl.constexpr, NEED_DF_RAW: tl.constexpr, NEED_DA_PARTIAL: tl.constexpr, NEED_DDT_PARTIAL: tl.constexpr, ROUND_MEMORY_COTANGENT: tl.constexpr): |
| i_bc = tl.program_id(0) |
| i_h = tl.program_id(1) |
| i_b = i_bc // NT |
| i_c = i_bc % NT |
| offs_t_1d = i_c * BT + tl.arange(0, BT) |
| offs_k_1d = tl.arange(0, BK) |
| offs_t = offs_t_1d[:, None] |
| offs_k = offs_k_1d[None, :] |
| offsets = ((i_b * T + offs_t) * H + i_h) * K + offs_k |
| mask_t = offs_t_1d < T |
| mask_k = offs_k_1d < K |
| mask = mask_t[:, None] & mask_k[None, :] |
| f_value = tl.load(retention_raw_ptr + offsets, mask=mask, other=0.0).to(tl.float32) |
| A_log_value = tl.load(A_log_ptr + i_h).to(tl.float32) |
| dt_bias_value = tl.load(dt_bias_ptr + i_h * K + offs_k_1d, mask=mask_k, other=0.0).to(tl.float32) |
| alpha_value = tl.load(alpha_ptr + offsets, mask=mask, other=1.0).to(tl.float32) |
| dalpha_value = tl.zeros((BT, BK), tl.float32) |
| if HAS_DFIRST: |
| d_first_offsets = i_b * stride_d_first_b + offs_t * stride_d_first_t + i_h * stride_d_first_h + offs_k * stride_d_first_k |
| loaded_d_first_value = tl.load(d_first_ptr + d_first_offsets, mask=mask, other=0.0).to(tl.float32) |
| d_first_value = tl.where(mask, loaded_d_first_value, 0.0) |
| reverse_value = tl.cumsum(d_first_value, axis=0, reverse=True) |
| if ROUND_MEMORY_COTANGENT: |
| memory_cotangent = tl.cast(reverse_value, tl.bfloat16, fp_downcast_rounding='rtne').to(tl.float32) |
| else: |
| memory_cotangent = reverse_value |
| dalpha_value += tl.where(alpha_value >= ALPHA_FLOOR, memory_cotangent / alpha_value, 0.0) |
| if HAS_DALPHA: |
| dalpha_value += tl.load(d_alpha_gain_ptr + offsets, mask=mask, other=0.0).to(tl.float32) |
| rate_value = tl.exp(A_log_value) |
| x_value = f_value + dt_bias_value[None, :] |
| ell_value = -rate_value * softplus(x_value) |
| dell_value = dalpha_value * alpha_value |
| sigmoid_value = tl.sigmoid(x_value) |
| df_value = dell_value * -rate_value * sigmoid_value |
| dA_value = dell_value * ell_value |
| if NEED_DA_PARTIAL: |
| valid_dA_value = tl.where(mask, dA_value, 0.0) |
| dA_partial_value = tl.sum(tl.sum(valid_dA_value, axis=1), axis=0) |
| tl.store(dA_partial_ptr + i_bc * H + i_h, dA_partial_value) |
| if NEED_DDT_PARTIAL: |
| valid_df_value = tl.where(mask, df_value, 0.0) |
| ddt_partial_value = tl.sum(valid_df_value, axis=0) |
| partial_k_offsets = (i_bc * H + i_h) * K + offs_k_1d |
| tl.store(ddt_partial_ptr + partial_k_offsets, ddt_partial_value, mask=mask_k) |
| if NEED_DF_RAW: |
| tl.store(dretention_raw_ptr + offsets, df_value, mask=mask) |
|
|
|
|
| |
|
|
| |
|
|
| |
|
|
| def _validate_diag_kdn_norm_inputs(q_raw, k_raw): |
| for name, value in (("q_raw", q_raw), ("k_raw", k_raw)): |
| if not isinstance(value, torch.Tensor): |
| raise TypeError(f"{name} must be a torch.Tensor") |
| if value.ndim != 4: |
| raise ValueError(f"{name} must have rank 4 [B,T,H,K]") |
| if not value.is_cuda: |
| raise ValueError(f"{name} must be a CUDA tensor") |
| if value.dtype not in (torch.bfloat16, torch.float32): |
| raise TypeError(f"{name} must be BF16 or FP32") |
| if not value.is_contiguous(): |
| raise ValueError(f"{name} must be contiguous") |
| if q_raw.shape != k_raw.shape or q_raw.device != k_raw.device: |
| raise ValueError("q_raw and k_raw must have matching shape and device") |
| if q_raw.dtype != k_raw.dtype or q_raw.requires_grad != k_raw.requires_grad: |
| raise ValueError("q_raw and k_raw must have matching dtype and grad mode") |
| if q_raw.shape[-1] not in (64, 128): |
| raise ValueError("q_raw and k_raw require K in {64,128}") |
|
|
|
|
| class _DiagKdnNormFunction(torch.autograd.Function): |
| @staticmethod |
| @torch.amp.custom_fwd(device_type="cuda", cast_inputs=None) |
| def forward(ctx, q_raw, k_raw): |
| q_memory, k_gain, k_memory, q_norm, k_norm = ( |
| _frozen_norm_launch_diag_norm_boundary_forward(q_raw, k_raw) |
| ) |
| ctx.save_for_backward(q_raw, k_raw, q_norm, k_norm) |
| ctx.set_materialize_grads(False) |
| return q_memory, k_gain, k_memory |
|
|
| @staticmethod |
| @torch.amp.custom_bwd(device_type="cuda") |
| @once_differentiable |
| def backward(ctx, d_q_memory, d_k_gain, d_k_memory): |
| if not ctx.needs_input_grad[0]: |
| d_q_memory = None |
| if not ctx.needs_input_grad[1]: |
| d_k_gain = None |
| d_k_memory = None |
| if d_q_memory is None and d_k_gain is None and d_k_memory is None: |
| return None, None |
| q_raw, k_raw, q_norm, k_norm = ctx.saved_tensors |
| return _diag_kdn_norm_backward( |
| q_raw, |
| k_raw, |
| d_q_memory, |
| d_k_gain, |
| d_k_memory, |
| q_norm, |
| k_norm, |
| ) |
|
|
|
|
| def _diag_kdn_norm_backward( |
| q_raw, |
| k_raw, |
| d_q_memory, |
| d_k_gain, |
| d_k_memory, |
| q_norm, |
| k_norm, |
| ): |
| process_q = d_q_memory is not None |
| has_gain = d_k_gain is not None |
| has_memory = d_k_memory is not None |
| process_k = has_gain or has_memory |
| d_q_raw = ( |
| torch.empty_like(q_raw, memory_format=torch.contiguous_format) |
| if process_q |
| else None |
| ) |
| d_k_raw = ( |
| torch.empty_like(k_raw, memory_format=torch.contiguous_format) |
| if process_k |
| else None |
| ) |
| use_single_row = ( |
| (d_q_memory is not None and not d_q_memory.is_contiguous()) |
| or (d_k_gain is not None and not d_k_gain.is_contiguous()) |
| or (d_k_memory is not None and not d_k_memory.is_contiguous()) |
| ) |
| row_block = 1 if use_single_row else _ROW_BLOCK |
| num_warps = 1 if use_single_row else _frozen_norm_NUM_WARPS |
| n_rows = q_raw.numel() // q_raw.shape[-1] |
| if n_rows and (process_q or process_k): |
| q_source = d_q_memory if d_q_memory is not None else q_raw |
| k_gain_source = d_k_gain if d_k_gain is not None else q_raw |
| k_memory_source = d_k_memory if d_k_memory is not None else q_raw |
| q_output = d_q_raw if d_q_raw is not None else q_raw |
| k_output = d_k_raw if d_k_raw is not None else k_raw |
| _frozen_norm_diag_exact_norm_bwd_kernel[ |
| (triton.cdiv(n_rows, row_block),) |
| ]( |
| q_raw, |
| k_raw, |
| q_source, |
| k_gain_source, |
| k_memory_source, |
| q_norm, |
| k_norm, |
| q_output, |
| k_output, |
| q_raw.shape[0], |
| q_raw.shape[1], |
| q_raw.shape[2], |
| n_rows, |
| q_source.stride(0), |
| q_source.stride(1), |
| q_source.stride(2), |
| q_source.stride(3), |
| k_gain_source.stride(0), |
| k_gain_source.stride(1), |
| k_gain_source.stride(2), |
| k_gain_source.stride(3), |
| k_memory_source.stride(0), |
| k_memory_source.stride(1), |
| k_memory_source.stride(2), |
| k_memory_source.stride(3), |
| K=q_raw.shape[-1], |
| ROW_BLOCK=row_block, |
| PROCESS_Q=process_q, |
| PROCESS_K=process_k, |
| HAS_GAIN=has_gain, |
| HAS_MEMORY=has_memory, |
| INPUT_IS_BF16=q_raw.dtype == torch.bfloat16, |
| FINAL_ROUND=True, |
| SMALL_ROWS=n_rows < 16, |
| REDUCTION_WIDTH=_frozen_norm_reduction_width( |
| n_rows, q_raw.shape[-1] |
| ), |
| num_warps=num_warps, |
| num_stages=_frozen_norm_NUM_STAGES, |
| ) |
| return d_q_raw, d_k_raw |
|
|
|
|
| def _diag_kdn_normalize( |
| q_raw: torch.Tensor, |
| k_raw: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| """Normalize q/k at the certified BF16 public boundary.""" |
| _validate_diag_kdn_norm_inputs(q_raw, k_raw) |
| return _DiagKdnNormFunction.apply(q_raw, k_raw) |
|
|
|
|
| def _validate_diag_kdn_frontend_inputs( |
| retention_raw, |
| omega_raw, |
| r_raw, |
| A_log, |
| dt_bias, |
| initial_precision_param, |
| ): |
| tensors = { |
| "retention_raw": retention_raw, |
| "omega_raw": omega_raw, |
| "r_raw": r_raw, |
| "A_log": A_log, |
| "dt_bias": dt_bias, |
| "initial_precision_param": initial_precision_param, |
| } |
| for name, value in tensors.items(): |
| if not isinstance(value, torch.Tensor): |
| raise TypeError(f"{name} must be a torch.Tensor") |
| if not value.is_cuda: |
| raise ValueError(f"{name} must be a CUDA tensor") |
| if value.device != retention_raw.device: |
| raise ValueError("all DiagKDN frontend inputs must share one CUDA device") |
| if not value.is_contiguous(): |
| raise ValueError(f"{name} must be contiguous") |
| if retention_raw.ndim != 4: |
| raise ValueError("retention_raw must have rank 4 [B,T,H,K]") |
| B, T, H, K = retention_raw.shape |
| if min(B, T, H) <= 0 or K not in (64, 128): |
| raise ValueError("DiagKDN frontend requires non-empty B/T/H and K in {64,128}") |
| if omega_raw.shape != (B, T, H, K): |
| raise ValueError("omega_raw must match retention_raw") |
| if r_raw.shape != (B, T, H): |
| raise ValueError("r_raw must have shape [B,T,H]") |
| if A_log.shape != (H,) or initial_precision_param.shape != (H,): |
| raise ValueError("A_log and initial_precision_param must have shape [H]") |
| if dt_bias.shape != (H, K): |
| raise ValueError("dt_bias must have shape [H,K]") |
| for name in ("retention_raw", "omega_raw", "r_raw"): |
| if tensors[name].dtype != _DIAG_KDN_VECTOR_DTYPE: |
| raise TypeError(f"{name} must be torch.bfloat16") |
| for name in ("A_log", "dt_bias", "initial_precision_param"): |
| if tensors[name].dtype != _DIAG_KDN_STATE_DTYPE: |
| raise TypeError(f"{name} must be torch.float32") |
| return B, T, H, K |
|
|
|
|
| def _diag_kdn_frontend_forward( |
| retention_raw, |
| omega_raw, |
| r_raw, |
| A_log, |
| dt_bias, |
| initial_precision_param, |
| ): |
| B, T, H, K = retention_raw.shape |
| BT = _DIAG_KDN_MEMORY_CHUNK_SIZE |
| NT = triton.cdiv(T, BT) |
| g_cumsum = torch.empty_like(retention_raw, dtype=_DIAG_KDN_STATE_DTYPE) |
| alpha = torch.empty_like(retention_raw, dtype=_DIAG_KDN_STATE_DTYPE) |
| omega = torch.empty_like(omega_raw, dtype=_DIAG_KDN_STATE_DTYPE) |
| r = torch.empty_like(r_raw, dtype=_DIAG_KDN_STATE_DTYPE) |
| initial_precision = torch.empty_like( |
| initial_precision_param, dtype=_DIAG_KDN_STATE_DTYPE |
| ) |
| with torch.cuda.device(retention_raw.device): |
| _frozen_frontend_diag_frontend_fwd_kernel[(B * NT, H)]( |
| retention_raw, |
| omega_raw, |
| r_raw, |
| A_log, |
| dt_bias, |
| initial_precision_param, |
| g_cumsum, |
| alpha, |
| omega, |
| r, |
| initial_precision, |
| 0.0, |
| 0.01, |
| T=T, |
| NT=NT, |
| H=H, |
| K=K, |
| BT=BT, |
| BK=K, |
| OMEGA_COUPLING=False, |
| EMIT_RAW_G=False, |
| num_warps=4, |
| num_stages=2, |
| ) |
| return g_cumsum, alpha, omega, r, initial_precision |
|
|
|
|
| def _diag_kdn_frontend_backward( |
| retention_raw, |
| omega_raw, |
| r_raw, |
| A_log, |
| dt_bias, |
| initial_precision_param, |
| alpha, |
| d_g_cumsum, |
| d_alpha, |
| d_omega, |
| d_r, |
| d_initial_precision, |
| needs, |
| ): |
| B, T, H, K = retention_raw.shape |
| BT = _DIAG_KDN_MEMORY_CHUNK_SIZE |
| NT, NP = triton.cdiv(T, BT), B * triton.cdiv(T, BT) |
| local_BT = _DIAG_KDN_FRONTEND_LOCAL_CHUNK_SIZE |
| local_NT = triton.cdiv(T, local_BT) |
| need_retention, need_omega, need_r, need_A, need_dt, need_initial = needs |
| has_decay = d_g_cumsum is not None or d_alpha is not None |
| need_decay = need_retention or need_A or need_dt |
| dretention_raw = None |
| dA_partial = None |
| ddt_partial = None |
| if has_decay and need_decay: |
| dretention_raw = ( |
| torch.empty_like(retention_raw) |
| if need_retention |
| else None |
| ) |
| dA_partial = ( |
| torch.empty((NP, H), device=retention_raw.device, dtype=torch.float32) |
| if need_A |
| else None |
| ) |
| ddt_partial = ( |
| torch.empty( |
| (NP, H, K), device=retention_raw.device, dtype=torch.float32 |
| ) |
| if need_dt |
| else None |
| ) |
| with torch.cuda.device(retention_raw.device): |
| _frozen_frontend_diag_frontend_alpha_presence_vjp_kernel[(NP, H)]( |
| retention_raw, |
| A_log, |
| dt_bias, |
| alpha, |
| d_g_cumsum if d_g_cumsum is not None else alpha, |
| d_alpha if d_alpha is not None else alpha, |
| dretention_raw if dretention_raw is not None else retention_raw, |
| dA_partial if dA_partial is not None else A_log, |
| ddt_partial if ddt_partial is not None else dt_bias, |
| d_g_cumsum.stride(0) if d_g_cumsum is not None else alpha.stride(0), |
| d_g_cumsum.stride(1) if d_g_cumsum is not None else alpha.stride(1), |
| d_g_cumsum.stride(2) if d_g_cumsum is not None else alpha.stride(2), |
| d_g_cumsum.stride(3) if d_g_cumsum is not None else alpha.stride(3), |
| T, |
| NT, |
| H=H, |
| K=K, |
| BT=BT, |
| BK=K, |
| ALPHA_FLOOR=1.0e-6, |
| HAS_DFIRST=d_g_cumsum is not None, |
| HAS_DALPHA=d_alpha is not None, |
| NEED_DF_RAW=need_retention, |
| NEED_DA_PARTIAL=need_A, |
| NEED_DDT_PARTIAL=need_dt, |
| ROUND_MEMORY_COTANGENT=False, |
| num_warps=4, |
| num_stages=2, |
| ) |
| dA_log = None |
| if dA_partial is not None: |
| dA_log = torch.empty_like(A_log) |
| with torch.cuda.device(retention_raw.device): |
| _frozen_frontend_diag_frontend_dA_reduction_vjp_kernel[(H,)]( |
| dA_partial, |
| dA_log, |
| NP, |
| H=H, |
| BNP=triton.next_power_of_2(NP), |
| num_warps=4, |
| num_stages=1, |
| ) |
| ddt_bias = None |
| if ddt_partial is not None: |
| ddt_bias = torch.empty_like(dt_bias) |
| with torch.cuda.device(retention_raw.device): |
| _frozen_frontend_diag_frontend_ddt_reduction_vjp_kernel[(H * K,)]( |
| ddt_partial, |
| ddt_bias, |
| NP, |
| H=H, |
| K=K, |
| BNP=triton.next_power_of_2(NP), |
| num_warps=4, |
| num_stages=1, |
| ) |
| domega_raw = None |
| if need_omega and d_omega is not None: |
| domega_raw = torch.empty_like(omega_raw) |
| with torch.cuda.device(retention_raw.device): |
| _frozen_frontend_diag_frontend_uncoupled_omega_local_vjp_kernel[ |
| (B * local_NT, H) |
| ]( |
| omega_raw, |
| d_omega, |
| domega_raw, |
| T, |
| local_NT, |
| H=H, |
| K=K, |
| BT=local_BT, |
| BK=K, |
| num_warps=4, |
| num_stages=2, |
| ) |
| dr_raw = None |
| if need_r and d_r is not None: |
| dr_raw = torch.empty_like(r_raw) |
| with torch.cuda.device(retention_raw.device): |
| _frozen_frontend_diag_frontend_r_local_vjp_kernel[(B * local_NT, H)]( |
| r_raw, |
| d_r, |
| dr_raw, |
| T, |
| local_NT, |
| H=H, |
| BT=local_BT, |
| num_warps=4, |
| num_stages=2, |
| ) |
| dinitial_precision_param = None |
| if need_initial and d_initial_precision is not None: |
| dinitial_precision_param = torch.empty_like(initial_precision_param) |
| with torch.cuda.device(retention_raw.device): |
| _frozen_frontend_diag_frontend_mu_local_vjp_kernel[ |
| (triton.cdiv(H, local_BT),) |
| ]( |
| initial_precision_param, |
| d_initial_precision, |
| dinitial_precision_param, |
| H, |
| BH=local_BT, |
| num_warps=4, |
| num_stages=2, |
| ) |
| return ( |
| dretention_raw, |
| domega_raw, |
| dr_raw, |
| dA_log, |
| ddt_bias, |
| dinitial_precision_param, |
| ) |
|
|
|
|
| class _DiagKdnFrontendFunction(torch.autograd.Function): |
| @staticmethod |
| def forward( |
| ctx, |
| retention_raw, |
| omega_raw, |
| r_raw, |
| A_log, |
| dt_bias, |
| initial_precision_param, |
| ): |
| outputs = _diag_kdn_frontend_forward( |
| retention_raw, |
| omega_raw, |
| r_raw, |
| A_log, |
| dt_bias, |
| initial_precision_param, |
| ) |
| ctx.save_for_backward( |
| retention_raw, |
| omega_raw, |
| r_raw, |
| A_log, |
| dt_bias, |
| initial_precision_param, |
| outputs[1], |
| ) |
| ctx.set_materialize_grads(False) |
| return outputs |
|
|
| @staticmethod |
| @once_differentiable |
| def backward(ctx, d_g_cumsum, d_alpha, d_omega, d_r, d_initial_precision): |
| return _diag_kdn_frontend_backward( |
| *ctx.saved_tensors, |
| d_g_cumsum, |
| d_alpha, |
| d_omega, |
| d_r, |
| d_initial_precision, |
| ctx.needs_input_grad, |
| ) |
|
|
|
|
| def _diag_kdn_frontend_precumsum_paired( |
| retention_raw: torch.Tensor, |
| omega_raw: torch.Tensor, |
| r_raw: torch.Tensor, |
| A_log: torch.Tensor, |
| dt_bias: torch.Tensor, |
| initial_precision_param: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| """Apply the fixed uncoupled omega=0/r=0.01 BT32 frontend.""" |
| _validate_diag_kdn_frontend_inputs( |
| retention_raw, |
| omega_raw, |
| r_raw, |
| A_log, |
| dt_bias, |
| initial_precision_param, |
| ) |
| return _DiagKdnFrontendFunction.apply( |
| retention_raw, |
| omega_raw, |
| r_raw, |
| A_log, |
| dt_bias, |
| initial_precision_param, |
| ) |
|
|
|
|
| @torch.compiler.disable |
| def _diag_kdn_gain_forward( |
| k: torch.Tensor, |
| alpha: torch.Tensor, |
| omega: torch.Tensor, |
| r: torch.Tensor, |
| initial_precision: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| B, T, H, K = k.shape |
| BT = _DIAG_KDN_GAIN_CHUNK_SIZE |
| NT = (T + BT - 1) // BT |
| BK = triton.next_power_of_2(K) |
| kappa = torch.empty_like(k, dtype=_DIAG_KDN_STATE_DTYPE) |
| chunk_maps = torch.empty( |
| (B, NT, H, K, 4), dtype=_DIAG_KDN_STATE_DTYPE, device=k.device |
| ) |
| carry = torch.empty( |
| (B, NT, H, K, 2), dtype=_DIAG_KDN_STATE_DTYPE, device=k.device |
| ) |
| grid = (NT, B * H) |
| _frozen_gain_kla_kappa_passA_kernel[grid]( |
| k, |
| alpha, |
| omega, |
| r, |
| chunk_maps, |
| T, |
| float(K), |
| k, |
| 1.0, |
| NT=NT, |
| H=H, |
| K=K, |
| BK=BK, |
| BT=BT, |
| USE_LOG_SCALE=False, |
| ) |
| _frozen_gain_kla_kappa_passB_kernel[(B * H,)]( |
| chunk_maps, |
| initial_precision, |
| carry, |
| NT=NT, |
| H=H, |
| K=K, |
| BK=BK, |
| ) |
| _frozen_gain_kla_kappa_passC_kappa_only_kernel[grid]( |
| k, |
| alpha, |
| omega, |
| r, |
| carry, |
| kappa, |
| kappa, |
| kappa, |
| T, |
| float(K), |
| k, |
| 1.0, |
| NT=NT, |
| H=H, |
| K=K, |
| BK=BK, |
| BT=BT, |
| SAVE_STATE=False, |
| USE_LOG_SCALE=False, |
| ) |
| return kappa, carry |
|
|
|
|
| @torch.compiler.disable |
| def _diag_kdn_gain_backward( |
| k: torch.Tensor, |
| alpha: torch.Tensor, |
| omega: torch.Tensor, |
| r: torch.Tensor, |
| initial_precision: torch.Tensor, |
| omega_raw: torch.Tensor, |
| r_raw: torch.Tensor, |
| dkappa: torch.Tensor | None, |
| carry: torch.Tensor, |
| *, |
| need_initial_precision: bool, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: |
| B, T, H, K = k.shape |
| BT = _DIAG_KDN_GAIN_CHUNK_SIZE |
| NT = (T + BT - 1) // BT |
| BK = triton.next_power_of_2(K) |
| if carry.shape != (B, NT, H, K, 2): |
| raise RuntimeError("saved DiagKDN gain carry shape changed") |
| if carry.dtype != _DIAG_KDN_STATE_DTYPE or not carry.is_contiguous(): |
| raise TypeError("saved DiagKDN gain carry must be contiguous FP32") |
| dkappa = torch.zeros_like(k) if dkappa is None else dkappa.contiguous() |
| dk = torch.empty_like(k, dtype=_DIAG_KDN_STATE_DTYPE) |
| dalpha = torch.empty_like(alpha, dtype=_DIAG_KDN_STATE_DTYPE) |
| domega_raw = torch.empty_like(omega_raw) |
| dr_raw = torch.empty_like(r_raw) |
| dinitial_precision = ( |
| torch.zeros((H,), dtype=_DIAG_KDN_STATE_DTYPE, device=k.device) |
| if need_initial_precision |
| else None |
| ) |
| p_exclusive = torch.empty_like(k, dtype=_DIAG_KDN_STATE_DTYPE) |
| reverse_maps = torch.empty( |
| (B, NT, H, K, 2), dtype=_DIAG_KDN_STATE_DTYPE, device=k.device |
| ) |
| reverse_carry = torch.empty( |
| (B, NT, H, K), dtype=_DIAG_KDN_STATE_DTYPE, device=k.device |
| ) |
| grid = (NT, B * H) |
| _frozen_gain_kla_kappa_bwd_fill_kernel[grid]( |
| k, |
| alpha, |
| omega, |
| r, |
| carry, |
| p_exclusive, |
| T, |
| float(K), |
| k, |
| 1.0, |
| NT=NT, |
| H=H, |
| K=K, |
| BK=BK, |
| BT=BT, |
| USE_LOG_SCALE=False, |
| ) |
| _frozen_gain_kla_kappa_bwd_passA_kappa_only_kernel[grid]( |
| k, |
| alpha, |
| omega, |
| r, |
| dkappa, |
| p_exclusive, |
| reverse_maps, |
| T, |
| float(K), |
| k, |
| 1.0, |
| NT=NT, |
| H=H, |
| K=K, |
| BK=BK, |
| BT=BT, |
| USE_LOG_SCALE=False, |
| ) |
| _frozen_gain_kla_kappa_bwd_passB_kernel[(B * H,)]( |
| reverse_maps, |
| reverse_carry, |
| NT=NT, |
| H=H, |
| K=K, |
| BK=BK, |
| ) |
| _frozen_gain_kla_kappa_bwd_passC_kappa_only_kernel[grid]( |
| k, |
| alpha, |
| omega, |
| r, |
| dkappa, |
| p_exclusive, |
| reverse_carry, |
| omega_raw, |
| r_raw, |
| dk, |
| dalpha, |
| domega_raw, |
| dr_raw, |
| dinitial_precision if dinitial_precision is not None else initial_precision, |
| initial_precision, |
| T, |
| float(K), |
| k, |
| 1.0, |
| NT=NT, |
| H=H, |
| K=K, |
| BK=BK, |
| BT=BT, |
| RAW_ACTIVATION_VJP=True, |
| USE_LOG_SCALE=False, |
| NEED_DMU=need_initial_precision, |
| NEED_DLOG=False, |
| ) |
| return dk, dalpha, domega_raw, dr_raw, dinitial_precision |
|
|
|
|
| class _DiagKdnGainFunction(torch.autograd.Function): |
| @staticmethod |
| def forward(ctx, k, alpha, omega, r, initial_precision, omega_raw, r_raw): |
| kappa, carry = _diag_kdn_gain_forward( |
| k, alpha, omega, r, initial_precision |
| ) |
| ctx.save_for_backward( |
| k, alpha, omega, r, initial_precision, omega_raw, r_raw, carry |
| ) |
| ctx.set_materialize_grads(False) |
| return kappa |
|
|
| @staticmethod |
| @once_differentiable |
| def backward(ctx, dkappa): |
| ( |
| k, |
| alpha, |
| omega, |
| r, |
| initial_precision, |
| omega_raw, |
| r_raw, |
| carry, |
| ) = ctx.saved_tensors |
| needs = ctx.needs_input_grad |
| dk, dalpha, domega_raw, dr_raw, dinitial_precision = ( |
| _diag_kdn_gain_backward( |
| k, |
| alpha, |
| omega, |
| r, |
| initial_precision, |
| omega_raw, |
| r_raw, |
| dkappa, |
| carry, |
| need_initial_precision=needs[4], |
| ) |
| ) |
| return ( |
| dk if needs[0] else None, |
| dalpha if needs[1] else None, |
| None, |
| None, |
| dinitial_precision if needs[4] else None, |
| domega_raw if needs[5] else None, |
| dr_raw if needs[6] else None, |
| ) |
|
|
|
|
| def _diag_kdn_gain( |
| k_gain: torch.Tensor, |
| alpha: torch.Tensor, |
| omega: torch.Tensor, |
| r: torch.Tensor, |
| initial_precision: torch.Tensor, |
| omega_raw: torch.Tensor, |
| r_raw: torch.Tensor, |
| ) -> torch.Tensor: |
| """Return the fixed-scale FP32 diagonal gain with fused raw-logit VJPs.""" |
| tensors = { |
| "k_gain": k_gain, |
| "alpha": alpha, |
| "omega": omega, |
| "r": r, |
| "initial_precision": initial_precision, |
| "omega_raw": omega_raw, |
| "r_raw": r_raw, |
| } |
| for name, value in tensors.items(): |
| if not isinstance(value, torch.Tensor): |
| raise TypeError(f"{name} must be a torch.Tensor") |
| if not value.is_cuda: |
| raise ValueError(f"{name} must be a CUDA tensor") |
| if value.device != k_gain.device: |
| raise ValueError("all DiagKDN gain inputs must share one CUDA device") |
| if not value.is_contiguous(): |
| raise ValueError(f"{name} must be contiguous") |
| if k_gain.ndim != 4: |
| raise ValueError("k_gain must have rank 4 [B,T,H,K]") |
| B, T, H, K = k_gain.shape |
| if min(B, T, H) <= 0 or K not in (64, 128): |
| raise ValueError("DiagKDN gain requires non-empty B/T/H and K in {64,128}") |
| expected = (B, T, H, K) |
| for name in ("alpha", "omega", "omega_raw"): |
| if tensors[name].shape != expected: |
| raise ValueError(f"{name} must have shape {expected}") |
| if r.shape != (B, T, H) or r_raw.shape != (B, T, H): |
| raise ValueError("r and r_raw must have shape [B,T,H]") |
| if initial_precision.shape != (H,): |
| raise ValueError(f"initial_precision must have shape [{H}]") |
| for name in ("k_gain", "alpha", "omega", "r", "initial_precision"): |
| if tensors[name].dtype != torch.float32: |
| raise TypeError(f"{name} must be torch.float32") |
| if omega_raw.dtype != _DIAG_KDN_VECTOR_DTYPE or r_raw.dtype != _DIAG_KDN_VECTOR_DTYPE: |
| raise TypeError("omega_raw and r_raw must be torch.bfloat16") |
| return _DiagKdnGainFunction.apply( |
| k_gain, |
| alpha, |
| omega.detach(), |
| r.detach(), |
| initial_precision, |
| omega_raw, |
| r_raw, |
| ) |
|
|
|
|
| def _validate_diag_kdn_memory_inputs(q, k, kappa, v, g_cumsum, scale): |
| tensors = {"q": q, "k": k, "kappa": kappa, "v": v, "g_cumsum": g_cumsum} |
| for name, value in tensors.items(): |
| if not isinstance(value, torch.Tensor): |
| raise TypeError(f"{name} must be a torch.Tensor") |
| if value.ndim != 4: |
| raise ValueError(f"{name} must have rank 4 [B,T,H,D]") |
| if not value.is_cuda: |
| raise ValueError(f"{name} must be a CUDA tensor") |
| if value.device != q.device: |
| raise ValueError("all DiagKDN memory inputs must share one CUDA device") |
| if not value.is_contiguous(): |
| raise ValueError(f"{name} must be contiguous") |
| B, T, H, K = q.shape |
| if min(B, T, H) <= 0 or K not in (64, 128): |
| raise ValueError("DiagKDN memory requires non-empty B/T/H and K in {64,128}") |
| expected = (B, T, H, K) |
| for name in ("k", "kappa", "v", "g_cumsum"): |
| if tensors[name].shape != expected: |
| raise ValueError(f"{name} must have shape {expected}") |
| for name in ("q", "k", "v"): |
| if tensors[name].dtype != _DIAG_KDN_VECTOR_DTYPE: |
| raise TypeError(f"{name} must be torch.bfloat16") |
| if kappa.dtype != _DIAG_KDN_STATE_DTYPE or g_cumsum.dtype != _DIAG_KDN_STATE_DTYPE: |
| raise TypeError("kappa and g_cumsum must be torch.float32") |
| if type(scale) is not float or not math.isfinite(scale): |
| raise TypeError("scale must be an exact finite float") |
| return B, T, H, K |
|
|
|
|
| def _diag_kdn_recompute_forward(k, kappa, v, Abar, g_cumsum): |
| """Fixed no-varlen recompute without the optional query output.""" |
| B, T, H, K = k.shape |
| V = v.shape[-1] |
| BT, BK, BV = _DIAG_KDN_MEMORY_CHUNK_SIZE, 64, 64 |
| NT = triton.cdiv(T, BT) |
| w = torch.empty_like(k, dtype=_DIAG_KDN_STATE_DTYPE) |
| u = torch.empty_like(v, dtype=_DIAG_KDN_STATE_DTYPE) |
| kappa_fed = torch.empty_like(kappa, dtype=_DIAG_KDN_STATE_DTYPE) |
| _frozen_intra_recompute_w_u_fwd_kalman_kernel_explicit[(NT, B * H)]( |
| q=None, |
| k=k, |
| kappa=kappa, |
| qg=None, |
| kg=kappa_fed, |
| v=v, |
| w=w, |
| u=u, |
| A=Abar, |
| gk=g_cumsum, |
| cu_seqlens=None, |
| chunk_indices=None, |
| T=T, |
| H=H, |
| K=K, |
| V=V, |
| BT=BT, |
| BK=BK, |
| BV=BV, |
| DOT_PRECISION=_DIAG_KDN_SCORE_DOT_PRECISION, |
| BF16_DOT_OPERANDS=False, |
| ) |
| return w, u, kappa_fed |
|
|
|
|
| def _diag_kdn_recompute_backward(q, k, kappa, v, Abar, g_cumsum): |
| """Fixed no-varlen recompute with the FP32 gated-query history.""" |
| B, T, H, K = k.shape |
| V = v.shape[-1] |
| BT, BK, BV = _DIAG_KDN_MEMORY_CHUNK_SIZE, 64, 64 |
| NT = triton.cdiv(T, BT) |
| w = torch.empty_like(k, dtype=_DIAG_KDN_STATE_DTYPE) |
| u = torch.empty_like(v, dtype=_DIAG_KDN_STATE_DTYPE) |
| kappa_fed = torch.empty_like(kappa, dtype=_DIAG_KDN_STATE_DTYPE) |
| qg = torch.empty_like(q, dtype=_DIAG_KDN_INTRA_QUERY_DTYPE) |
| _frozen_intra_recompute_w_u_fwd_kalman_kernel_explicit[(NT, B * H)]( |
| q=q, |
| k=k, |
| kappa=kappa, |
| qg=qg, |
| kg=kappa_fed, |
| v=v, |
| w=w, |
| u=u, |
| A=Abar, |
| gk=g_cumsum, |
| cu_seqlens=None, |
| chunk_indices=None, |
| T=T, |
| H=H, |
| K=K, |
| V=V, |
| BT=BT, |
| BK=BK, |
| BV=BV, |
| DOT_PRECISION=_DIAG_KDN_SCORE_DOT_PRECISION, |
| BF16_DOT_OPERANDS=False, |
| ) |
| return w, u, kappa_fed, qg |
|
|
|
|
| def _diag_kdn_intra_forward(q, k, kappa, v, g_cumsum, scale): |
| """Build the fixed BT32 FP32 Aqk/Abar caches and WY auxiliaries.""" |
| B, T, H, K = k.shape |
| BT, BC = _DIAG_KDN_MEMORY_CHUNK_SIZE, 16 |
| NT, NC = triton.cdiv(T, BT), 2 |
| Aqk = torch.zeros(B, T, H, BT, device=k.device, dtype=_DIAG_KDN_STATE_DTYPE) |
| Abar = torch.zeros(B, T, H, BT, device=k.device, dtype=_DIAG_KDN_STATE_DTYPE) |
| diagonal = torch.empty(B, T, H, BC, device=k.device, dtype=_DIAG_KDN_STATE_DTYPE) |
|
|
| def diagonal_grid(meta): |
| return B * T, triton.cdiv(H, meta["BH"]) |
|
|
| _frozen_intra_chunk_kalman_fwd_kernel_intra_token_parallel[diagonal_grid]( |
| q=q, |
| k=k, |
| kappa=kappa, |
| g=g_cumsum, |
| Aqk=Aqk, |
| Akk=diagonal, |
| scale=scale, |
| cu_seqlens=None, |
| N=B, |
| T=T, |
| H=H, |
| K=K, |
| BT=BT, |
| BC=BC, |
| ) |
| _frozen_intra_chunk_kalman_fwd_kernel_inter_solve_fused_bt32_explicit[ |
| (NT, B * H) |
| ]( |
| q=q, |
| k=k, |
| kappa=kappa, |
| g=g_cumsum, |
| Aqk=Aqk, |
| Akkd=diagonal, |
| Akk=Abar, |
| Mraw=None, |
| scale=scale, |
| T=T, |
| H=H, |
| K=K, |
| BT=BT, |
| BC=BC, |
| NC=NC, |
| SCORE_DOT_PRECISION=_DIAG_KDN_SCORE_DOT_PRECISION, |
| SOLVE_DOT_PRECISION=_DIAG_KDN_SOLVE_DOT_PRECISION, |
| ) |
| w, u, kappa_fed = _diag_kdn_recompute_forward( |
| k, kappa, v, Abar, g_cumsum |
| ) |
| return w, u, kappa_fed, Aqk, Abar |
|
|
|
|
| def _diag_kdn_state_forward(kappa_fed, w, u, g_cumsum): |
| """Run only the fixed no-state-input/no-final-state BT32 recurrence.""" |
| B, T, Hq, K = kappa_fed.shape |
| H, V = u.shape[2], u.shape[-1] |
| BT = _DIAG_KDN_MEMORY_CHUNK_SIZE |
| NT = triton.cdiv(T, BT) |
| h = torch.empty((B, NT, H, K, V), device=u.device, dtype=_DIAG_KDN_STATE_DTYPE) |
| v_new = torch.empty_like(u, dtype=_DIAG_KDN_STATE_DTYPE) |
|
|
| def grid(meta): |
| return triton.cdiv(V, meta["BV"]), B * H |
|
|
| _frozen_kda_chunk_gated_delta_rule_fwd_kernel_h_blockdim64_diag[grid]( |
| k=kappa_fed, |
| v=u, |
| w=w, |
| v_new=v_new, |
| g=None, |
| gk=g_cumsum, |
| h=h, |
| h0=None, |
| ht=None, |
| cu_seqlens=None, |
| chunk_offsets=None, |
| T=T, |
| H=H, |
| Hq=Hq, |
| K=K, |
| V=V, |
| BT=BT, |
| USE_EXP2=True, |
| TRANSPOSE_STATE=False, |
| DOT_PRECISION=_DIAG_KDN_STATE_DOT_PRECISION, |
| BF16_DOT_OPERANDS=True, |
| RESIDUAL_DOT_PRECISION=_DIAG_KDN_STATE_RESIDUAL_DOT_PRECISION, |
| RESIDUAL_BF16_DOT_OPERANDS=False, |
| ) |
| return h, v_new |
|
|
|
|
| def _diag_kdn_readout(q, v_new, g_cumsum, Aqk, h, scale): |
| """Run only the fixed FBT no-varlen readout.""" |
| B, T, H, K = q.shape |
| V = v_new.shape[-1] |
| BT = _DIAG_KDN_MEMORY_CHUNK_SIZE |
| NT = triton.cdiv(T, BT) |
| output = torch.zeros(v_new.shape, device=v_new.device, dtype=_DIAG_KDN_STATE_DTYPE) |
|
|
| def grid(meta): |
| return triton.cdiv(V, meta["BV"]), NT, B * H |
|
|
| _frozen_readout_chunk_gla_fwd_kernel_o_diag[grid]( |
| q=q, |
| v=v_new, |
| g=g_cumsum, |
| h=h, |
| o=output, |
| A=Aqk, |
| cu_seqlens=None, |
| chunk_indices=None, |
| scale=scale, |
| T=T, |
| H=H, |
| HV=H, |
| K=K, |
| V=V, |
| BT=BT, |
| USE_EXP2=True, |
| TRANSPOSE_STATE=False, |
| DOT_PRECISION="tf32", |
| BF16_INTER_DOT_OPERANDS=False, |
| BF16_INTRA_DOT_OPERANDS=False, |
| ) |
| return output |
|
|
|
|
| def _diag_kdn_output_backward(q, kappa_fed, v_new, do, Aqk, scale): |
| """Run the fixed BQ dAqk/dv_new readout VJP.""" |
| B, T, H, K = kappa_fed.shape |
| V = do.shape[-1] |
| BT = _DIAG_KDN_MEMORY_CHUNK_SIZE |
| BK = min(max(triton.next_power_of_2(K), 16), 128) |
| BV = min(max(triton.next_power_of_2(V), 16), 128) |
| NT = triton.cdiv(T, BT) |
| dAqk = torch.empty((B, T, H, BT), device=q.device, dtype=_DIAG_KDN_STATE_DTYPE) |
| dv_new = torch.empty_like(do, dtype=_DIAG_KDN_STATE_DTYPE) |
| _frozen_kda_chunk_kda_bwd_kernel_dAv_diag[(NT, B * H)]( |
| q=q, |
| k=kappa_fed, |
| v=v_new, |
| A=Aqk, |
| do=do, |
| dv=dv_new, |
| dA=dAqk, |
| cu_seqlens=None, |
| chunk_indices=None, |
| scale=scale, |
| T=T, |
| H=H, |
| K=K, |
| V=V, |
| BT=BT, |
| BK=BK, |
| BV=BV, |
| DOT_PRECISION=_DIAG_KDN_BACKWARD_DOT_PRECISION, |
| BF16_DOT_OPERANDS=True, |
| BF16_DAQK_DO_VNEW_OPERANDS=False, |
| DVNEW_AQK_DO_DOT_PRECISION="tf32x3", |
| BF16_DVNEW_AQK_DO_OPERANDS=True, |
| ) |
| return dAqk, dv_new |
|
|
|
|
| def _diag_kdn_state_backward(qg, kappa_fed, w, g_cumsum, do, dv_new, scale): |
| """Run only the fixed no-state-input BT32 state VJP.""" |
| B, T, Hq, K = qg.shape |
| H, V = do.shape[2], do.shape[-1] |
| BT = _DIAG_KDN_MEMORY_CHUNK_SIZE |
| NT = triton.cdiv(T, BT) |
| dh = torch.empty((B, NT, H, K, V), device=qg.device, dtype=_DIAG_KDN_STATE_DTYPE) |
| du = torch.empty_like(dv_new, dtype=_DIAG_KDN_STATE_DTYPE) |
|
|
| def grid(meta): |
| return triton.cdiv(V, meta["BV"]), B * H |
|
|
| _frozen_kda_chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64_diag[grid]( |
| q=qg, |
| k=kappa_fed, |
| w=w, |
| g=None, |
| gk=g_cumsum, |
| dht=None, |
| dh0=None, |
| do=do, |
| dh=dh, |
| dv=dv_new, |
| dv2=du, |
| cu_seqlens=None, |
| chunk_offsets=None, |
| scale=scale, |
| T=T, |
| H=H, |
| Hq=Hq, |
| K=K, |
| V=V, |
| BT=BT, |
| USE_EXP2=True, |
| TRANSPOSE_STATE=False, |
| DOT_PRECISION=_DIAG_KDN_STATE_DOT_PRECISION, |
| BF16_QG_DO_OPERANDS=False, |
| BF16_W_DV_OPERANDS=True, |
| ) |
| return dh, du |
|
|
|
|
| def _diag_kdn_wy_backward(q, k, kappa, v, v_new, g, Abar, h, do, dh, du, scale): |
| """Run the sole fixed buffered-WY BQ backward.""" |
| B, T, H, K = k.shape |
| V = v.shape[-1] |
| dq = torch.empty_like(q, dtype=_DIAG_KDN_STATE_DTYPE) |
| dk = torch.empty_like(k, dtype=_DIAG_KDN_STATE_DTYPE) |
| dkappa = torch.empty_like(kappa, dtype=_DIAG_KDN_STATE_DTYPE) |
| dv = torch.empty_like(v, dtype=_DIAG_KDN_FINAL_GRADIENT_DTYPE) |
| dg = torch.empty_like(g, dtype=_DIAG_KDN_STATE_DTYPE) |
| dM = torch.empty_like(Abar, dtype=_DIAG_KDN_STATE_DTYPE) |
| dw = torch.empty_like(k, dtype=_DIAG_KDN_STATE_DTYPE) |
| grid = (triton.cdiv(T, _DIAG_KDN_MEMORY_CHUNK_SIZE), B * H) |
| _frozen_wy_chunk_kalman_bwd_kernel_wy_buffered_dw_producer[grid]( |
| q, |
| k, |
| kappa, |
| v_new, |
| g, |
| Abar, |
| h, |
| do, |
| dh, |
| du, |
| dq, |
| dk, |
| dkappa, |
| dv, |
| dg, |
| dw, |
| scale, |
| T, |
| H, |
| K, |
| V, |
| _DIAG_KDN_MEMORY_CHUNK_SIZE, |
| 64, |
| 64, |
| _DIAG_KDN_BACKWARD_DOT_PRECISION, |
| "tf32", |
| "tf32", |
| True, |
| True, |
| True, |
| False, |
| num_warps=4, |
| num_stages=4, |
| num_ctas=1, |
| ) |
| _frozen_wy_chunk_kalman_bwd_kernel_wy_buffered_dw_consumer[grid]( |
| k, |
| v, |
| g, |
| Abar, |
| du, |
| dw, |
| dM, |
| T, |
| H, |
| K, |
| V, |
| _DIAG_KDN_MEMORY_CHUNK_SIZE, |
| 64, |
| 64, |
| "tf32", |
| num_warps=4, |
| num_stages=4, |
| num_ctas=1, |
| ) |
| return dq, dk, dkappa, dv, dg, dM |
|
|
|
|
| def _diag_kdn_intra_backward(q, k, kappa, g, dAqk, dM, dq, dk, dkappa, dg): |
| """Run only the fixed no-varlen Triton intra VJP.""" |
| B, T, H, K = k.shape |
| BT, BC, BK = _DIAG_KDN_MEMORY_CHUNK_SIZE, 16, 32 |
| NT, NC, NK = triton.cdiv(T, BT), 2, triton.cdiv(K, BK) |
| dq_out = torch.empty_like(q, dtype=_DIAG_KDN_FINAL_GRADIENT_DTYPE) |
| dk_out = torch.empty_like(k, dtype=_DIAG_KDN_FINAL_GRADIENT_DTYPE) |
| dkappa_out = torch.empty_like(kappa, dtype=_DIAG_KDN_STATE_DTYPE) |
| dg_out = torch.empty_like(dg, dtype=_DIAG_KDN_STATE_DTYPE) |
| _frozen_intra_chunk_kalman_bwd_kernel_intra_explicit[ |
| (NK * NC, NT, B * H) |
| ]( |
| q=q, |
| k=k, |
| kappa=kappa, |
| g=g, |
| dAqk=dAqk, |
| dAkk=dM, |
| dq=dq, |
| dq2=dq_out, |
| dk=dk, |
| dk2=dk_out, |
| dkappa=dkappa, |
| dkappa2=dkappa_out, |
| dg=dg, |
| dg2=dg_out, |
| cu_seqlens=None, |
| chunk_indices=None, |
| B=B, |
| T=T, |
| H=H, |
| K=K, |
| BT=BT, |
| BC=BC, |
| BK=BK, |
| NC=NC, |
| SAFE_GATE=False, |
| USE_GATHER=IS_GATHER_SUPPORTED, |
| DOT_PRECISION=_DIAG_KDN_BACKWARD_DOT_PRECISION, |
| BF16_DOT_OPERANDS=True, |
| BF16_DK_ROW_OPERANDS=True, |
| BF16_DQ_ROW_OPERANDS=False, |
| DQ_ROW_DOT_PRECISION="tf32x3", |
| K_TILES_PER_CTA=1, |
| ) |
| return dq_out, dk_out, dkappa_out, dg_out |
|
|
|
|
| def _diag_kdn_memory_forward(q, k, kappa, v, g_cumsum, scale): |
| B, T, H, K = _validate_diag_kdn_memory_inputs( |
| q, k, kappa, v, g_cumsum, scale |
| ) |
| with torch.cuda.device(q.device): |
| w, u, kappa_fed, Aqk, Abar = _diag_kdn_intra_forward( |
| q, k, kappa, v, g_cumsum, scale |
| ) |
| BT = _DIAG_KDN_MEMORY_CHUNK_SIZE |
| for name, value in (("Aqk", Aqk), ("Abar", Abar)): |
| if value.shape != (B, T, H, BT): |
| raise RuntimeError(f"saved {name} shape changed") |
| if value.dtype != torch.float32 or not value.is_contiguous(): |
| raise TypeError(f"saved {name} must be contiguous FP32") |
| h, v_new = _diag_kdn_state_forward(kappa_fed, w, u, g_cumsum) |
| output = _diag_kdn_readout( |
| q, v_new, g_cumsum, Aqk, h, scale |
| ) |
| if output.shape != (B, T, H, K): |
| raise RuntimeError("fixed DiagKDN output shape changed") |
| return output, Aqk, Abar |
|
|
|
|
| class _DiagKdnMemoryFunction(torch.autograd.Function): |
| @staticmethod |
| def forward(ctx, q, k, kappa, v, g_cumsum, scale): |
| output, Aqk, Abar = _diag_kdn_memory_forward( |
| q, k, kappa, v, g_cumsum, scale |
| ) |
| ctx.save_for_backward(q, k, kappa, v, g_cumsum, Aqk, Abar) |
| ctx.scale = scale |
| ctx.set_materialize_grads(False) |
| return output |
|
|
| @staticmethod |
| @once_differentiable |
| def backward(ctx, do): |
| q, k, kappa, v, g_cumsum, Aqk, Abar = ctx.saved_tensors |
| if do is None: |
| return None, None, None, None, None, None |
| do = do.contiguous() |
| B, T, H, K = _validate_diag_kdn_memory_inputs( |
| q, k, kappa, v, g_cumsum, ctx.scale |
| ) |
| if do.shape != (B, T, H, K) or do.dtype != torch.float32: |
| raise TypeError("DiagKDN output cotangent must be FP32 BTHK") |
| with torch.cuda.device(q.device): |
| w, u, kappa_fed, qg = _diag_kdn_recompute_backward( |
| q, k, kappa, v, Abar, g_cumsum |
| ) |
| h, v_new = _diag_kdn_state_forward( |
| kappa_fed, w, u, g_cumsum |
| ) |
| do_work = do.to(torch.bfloat16) |
| dAqk, dv_new = _diag_kdn_output_backward( |
| q, kappa_fed, v_new, do_work, Aqk, ctx.scale |
| ) |
| dAqk = dAqk.to(torch.float32).contiguous() |
| dv_new = dv_new.to(torch.float32) |
| dh, du = _diag_kdn_state_backward( |
| qg, kappa_fed, w, g_cumsum, do_work, dv_new, ctx.scale |
| ) |
| dq_w, dk_read, dkappa, dv, dg_w, dM = _diag_kdn_wy_backward( |
| q, |
| k, |
| kappa, |
| v, |
| v_new, |
| g_cumsum, |
| Abar, |
| h, |
| do_work, |
| dh, |
| du, |
| ctx.scale, |
| ) |
| dq, dk_read, dkappa, dG = _diag_kdn_intra_backward( |
| q, |
| k, |
| kappa, |
| g_cumsum, |
| dAqk, |
| dM, |
| dq_w, |
| dk_read, |
| dkappa, |
| dg_w, |
| ) |
| return dq, dk_read, dkappa, dv, dG.contiguous(), None |
|
|
|
|
| def chunk_kalman( |
| q_memory: torch.Tensor, |
| k_memory: torch.Tensor, |
| kappa: torch.Tensor, |
| v: torch.Tensor, |
| g_cumsum: torch.Tensor, |
| scale: float, |
| ) -> torch.Tensor: |
| """Run the sole fixed BT32 production DiagKDN memory route.""" |
| _validate_diag_kdn_memory_inputs( |
| q_memory, k_memory, kappa, v, g_cumsum, scale |
| ) |
| return _DiagKdnMemoryFunction.apply( |
| q_memory, k_memory, kappa, v, g_cumsum, scale |
| ) |
|
|
|
|
| __all__ = ("chunk_kalman",) |
|
|
