File size: 5,969 Bytes
5ff5661 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | """Ordered native Mamba backward reductions for the bounded M4 trajectory.
Candidate policy: GPU gates and a real128 distributed restart are mandatory.
Forward equations and parameters are unchanged. Native CUDA/Triton calls remain
responsible for derivatives; only the order of shared-parameter sums changes.
"""
from __future__ import annotations
import importlib
import torch
POLICY = 'native_single_writer_ordered_reductions_v1'
_installed = None
def install():
"""Freeze native tiles and remove concurrent writes to parameter gradients."""
global _installed
if _installed is not None:
return _installed
from causal_conv1d.cpp_functions import causal_conv1d_bwd_function
combined = importlib.import_module('mamba_ssm.ops.triton.ssd_combined')
native_cumsum = combined._chunk_cumsum_bwd
modules = ('ssd_combined', 'ssd_chunk_state', 'ssd_chunk_scan',
'ssd_state_passing', 'ssd_bmm', 'layernorm_gated', 'k_activations')
tiles = {}
seen = set()
single_writer_kernels = {'_chunk_scan_chunk_state_bwd_dx_kernel': 64,
'_chunk_state_bwd_db_kernel': 128,
'_chunk_scan_bwd_dc_kernel': 128}
constrained = set()
for suffix in modules:
module = importlib.import_module('mamba_ssm.ops.triton.' + suffix)
for name, kernel in vars(module).items():
if not hasattr(kernel, 'configs') or not hasattr(kernel, 'cache') or id(kernel) in seen:
continue
seen.add(id(kernel))
configs = list(kernel.configs)
if not configs:
raise RuntimeError('Native autotuner has no configuration: ' + name)
# Large native candidates can exceed H100 shared memory; autotuning
# normally rejects those. Freeze the smallest available native tile
# subject to the single-writer head/state-axis constraint.
minimum_n = single_writer_kernels.get(name, 0)
if minimum_n:
constrained.add(name)
eligible = [c for c in configs if c.kwargs.get('BLOCK_SIZE_N', 0) >= minimum_n]
if not eligible:
raise RuntimeError('No native single-writer configuration: ' + name)
def tile_size(config):
dims = [v for k, v in config.kwargs.items() if k.startswith('BLOCK_SIZE') and isinstance(v, int)]
import math
return (math.prod(dims), config.num_stages, config.num_warps,
tuple(sorted(config.kwargs.items())))
chosen = min(eligible, key=tile_size)
kernel.configs = [chosen]
kernel.cache.clear()
tiles[suffix + '.' + name] = {'kwargs': dict(chosen.kwargs),
'num_warps': chosen.num_warps, 'num_stages': chosen.num_stages}
if constrained != set(single_writer_kernels):
raise RuntimeError('Pinned native single-writer kernel inventory changed')
def ordered_conv(x, weight, bias, dout, seq_idx, initial_states,
dfinal_states, dx, return_dinitial_states, silu_activation):
if any(v is not None for v in (seq_idx, initial_states, dfinal_states)) or return_dinitial_states:
raise ValueError('Ordered training convolution requires reset sequences')
# Retain the original full-batch input gradient and its memory layout.
dx_out, _, _, _ = causal_conv1d_bwd_function(
x, weight, bias, dout, None, None, None, dx, False, silu_activation)
dw = torch.zeros_like(weight, dtype=torch.float32)
db = None if bias is None else torch.zeros_like(bias, dtype=torch.float32)
# Native channels-first B=1 reduction has no cross-example atomics.
for batch in range(x.shape[0]):
_, part_w, part_b, _ = causal_conv1d_bwd_function(
x[batch:batch + 1].contiguous(), weight.float(),
None if bias is None else bias.float(), dout[batch:batch + 1].contiguous(),
None, None, None, None, False, silu_activation)
dw.add_(part_w)
if db is not None:
db.add_(part_b)
return dx_out, dw.to(weight.dtype), None if db is None else db.to(bias.dtype), None
def ordered_cumsum(ddA, ddt_out, dt, A, dt_bias=None, dt_softplus=False,
dt_limit=(0.0, float('inf')), ddt=None):
# Native cumsum backward's parameter sums use atomic_add across B/C.
# Each call below has B=C=1; accumulate those partials in lexical order.
batch, length, heads = dt.shape
chunks, width = ddA.shape[-2:]
if ddA.shape != (batch, heads, chunks, width) or ddt_out.shape != ddA.shape:
raise ValueError('Unexpected SSD chunk derivatives')
if ddt is None:
ddt = torch.empty_like(dt)
dA = torch.zeros_like(A, dtype=torch.float32)
db = None if dt_bias is None else torch.zeros_like(dt_bias, dtype=torch.float32)
for b in range(batch):
for c in range(chunks):
begin, end = c * width, min((c + 1) * width, length)
if begin >= end:
continue
_, part_a, part_b = native_cumsum(
ddA[b:b + 1, :, c:c + 1], ddt_out[b:b + 1, :, c:c + 1],
dt[b:b + 1, begin:end], A, dt_bias=dt_bias,
dt_softplus=dt_softplus, dt_limit=dt_limit,
ddt=ddt[b:b + 1, begin:end])
dA.add_(part_a)
if db is not None:
db.add_(part_b)
return ddt, dA, db
combined.causal_conv1d_bwd_function = ordered_conv
combined._chunk_cumsum_bwd = ordered_cumsum
_installed = {'policy': POLICY, 'tiles': tiles,
'convolution_reduction': 'native_B1_channels_first_fixed_FP32_sum',
'cumsum_reduction': 'native_B1_C1_fixed_FP32_sum'}
return _installed
|