| """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) |
| |
| |
| |
| 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') |
| |
| 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) |
| |
| 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): |
| |
| |
| 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 |
|
|
