mixle.models.sparsity_2_4 module

2:4 structured sparsity, end to end (roadmap I4): training-time mask ramp + cuSPARSELt-format export.

Two pieces, glued by ONE borrowed primitive rather than two reimplementations:

  1. TwoFourSparsityRamp – a schedulable training-time mask ramp. It does not reinvent 2:4 masking: at every ramp step it calls G2’s mixle.models.sigma_weighted_projection. sigma_weighted_block_sparse() with the literal "2:4" pattern, which is the actual masking/value -readjustment mechanism (see that module’s docstring for the alternating-projection algorithm). This module’s own job is just the RAMP: which fraction of a weight matrix’s rows are already under the hard 2:4 constraint at a given training step, growing from 0% at start_step to 100% at end_step. The ramp’s step-to-fraction map is a plain, swappable callable (schedule) precisely so a future roadmap-H3 “structure-edit schedule” controller could drive it (supply a different schedule, or mutate start_step/end_step on the fly) without touching this module – H3 itself is NOT built here, only the seam it would plug into.

  2. export_2_4_compressed() / decompress() – the actual cuSPARSELt-style compressed-matrix format: for every contiguous group of 4 weights along the input (last) axis that already satisfies the 2:4 constraint (exactly 2 nonzeros), store the 2 surviving VALUES plus a small INDEX recording which 2 of the 4 in-group positions they came from. This is the documented shape of NVIDIA’s semi-structured sparse storage (see e.g. the cuSPARSELt / Ampere structured-sparse-tensor-core docs and torch.sparse.SparseSemiStructuredTensor): compressed values at half the density, plus small per-group metadata recording nonzero positions (not full-size, since only C(4,2) = 6 patterns are possible per group). NVIDIA does not publish the exact bit-for-bit metadata layout their kernels consume (it is treated as an opaque, hardware/kernel-generation-specific detail even inside torch.sparse.SparseSemiStructuredTensor), so the exact packing implemented here (2 bits per in-group position, 2 positions packed per byte-nibble, 2 nibbles per byte) is THIS module’s own documented, round-trip-correct encoding of the publicly documented “values + position indices” shape – not a claim of bit-exact compatibility with a specific cuSPARSELt release. Values are portable float32/float64; a real cuSPARSELt handle would additionally repack this into its internal opaque compressed-matrix object via cusparseLtSpMMACompress (or, in this torch build, torch.sparse.to_sparse_semi_structured), which requires a CUDA tensor and a cuSPARSELt-capable GPU – see cusparselt_status() and the module docstring in mixle/tests/sparsity_2_4_test.py for what was actually checked/measured in THIS environment (no CUDA device here – see that check’s output).

Environment note (checked, not assumed): this environment has no CUDA device (torch.cuda.is_available() is False), so torch.backends.cusparselt reports unavailable and torch.sparse.to_sparse_semi_structured raises (CPU tensors are not supported by that call at all, CUDA or not). The compress/decompress round trip below is pure CPU numpy/torch and is exercised directly (round-trip exactness + measured byte-size compression ratio are pinned by tests); no accelerated cuSPARSELt kernel is run anywhere in this module. See cusparselt_status() for a queryable summary of what this torch build actually offers.

class TwoFourSparsityRamp(start_step, end_step, target_density=0.5, schedule=None)[source]

Bases: object

Schedulable training-time 2:4 mask ramp.

start_step/end_step: training steps between which the constrained-row fraction grows from 0 to 1 (see schedule). target_density: the density the 2:4 CONSTRAINT itself enforces once applied – structurally fixed at 0.5 (2 of every 4), kept as an explicit constructor argument (rather than hardcoded) so the acceptance criterion “2:4 model at stated … density” is documented at the call site and so a mismatched caller expectation fails loudly at construction time instead of silently. schedule(step, start_step, end_step) -> fraction in [0, 1]: the pluggable ramp-shape callable – THE seam a future roadmap-H3 structure-edit-schedule controller would drive (a non-linear cubic ramp, a warmup-then-hold schedule, one keyed off validation loss instead of step count, etc.); defaults to _linear_ramp(). H3 itself is not implemented here – this is deliberately just a callable/ parameter another component could supply.

Parameters:

  • start_step (int)

  • end_step (int)

  • target_density (float)

  • schedule (ScheduleFn | None)

fraction(step)[source]

Fraction of the weight matrix’s OUTPUT ROWS that are under the hard 2:4 constraint at step, in [0, 1], clipped in case a custom schedule over/undershoots.

Parameters:

step (int)

Return type:

float

project(weight, step, sigma=None)[source]

Apply the ramp to a weight matrix (d_out x d_in, d_in a multiple of 4): the first n_constrained_rows(step, d_out) rows are hard-projected onto the 2:4 pattern via G2’s sigma_weighted_block_sparse() (the real masking/value-adjustment mechanism – this is not a magnitude-only reimplementation); the remaining rows are left untouched (still fully dense at this point in the ramp). Returns a NEW array/tensor of the same type and shape as weight; does not mutate weight in place (callers doing in-place training updates copy the result back themselves, see apply_ramp_to_linear_()).

sigma: optional d_in x d_in covariance to weight the projection by (e.g. a real propagated law from G1/moment_propagation, wired the same way G2’s own acceptance test wires it). Defaults to the identity, which reduces the Sigma-weighted objective to plain Frobenius reconstruction – i.e. plain magnitude-based 2:4 value selection when no activation-covariance estimate is available, which is the common case during ordinary token-level training where no data-free law has been propagated for this weight.

Rows are constrained from index 0 upward (a fixed, deterministic ordering) rather than by magnitude or another heuristic – keeps the ramp’s mask-growth trajectory the same across calls for the same step, which is what the ramp-correctness test below checks against.

Parameters:
Return type:

Any

apply_(linear, step, sigma=None)[source]

In-place convenience: project linear.weight.data through project() and copy the result back. This is the “mask-then-continue-training” step a caller runs after every optimizer step (see train_with_ramp below) – gradients keep flowing through the dense parameter between projections (straight-through), and the projection is re-applied (with a re-selected pattern, per G2) every call, which is what lets already-constrained rows keep adapting their surviving VALUES as training continues, not just their support.

Parameters:
Return type:

None

class Compressed2to4(values, indices, shape)[source]

Bases: object

A 2:4 semi-structured-sparse compressed matrix, cuSPARSELt-shaped: the 2 surviving values per group of 4 (values, half the density of the dense matrix) plus a small per-group index recording which 2 of the 4 in-group positions they occupy (indices). See the module docstring for exactly how indices is packed (2 bits per in-group position, 2 groups’ worth of index nibbles per byte) and why that specific bit layout is this module’s own documented encoding rather than a claim of bit-for-bit compatibility with a specific NVIDIA driver’s opaque compressed-matrix object.

values: float array, shape (d_out, d_in // 4, 2) – the 2 surviving values per group, in ASCENDING in-group-position order (position of values[..., 0] <= position of values[..., 1]; this ordering convention is what lets indices unambiguously address them on decompress). indices: uint8 array, shape (d_out, ceil(d_in // 4 / 2)) – packed 2-bits-per-position metadata, 2 groups per byte (low nibble = group 2k’s two 2-bit positions, high nibble = group 2k+1’s, when present). shape: the original dense (d_out, d_in) shape, needed to reconstruct decompress’s output (the last packed index byte may cover an odd group count).

Parameters:
export_2_4_compressed(weight_2_4_masked)[source]

Convert an already-2:4-masked dense weight matrix into the compressed (values + indices) storage format cuSPARSELt-style semi-structured sparse GEMM kernels consume. Precondition: every contiguous group of 4 entries along the last axis already has AT MOST 2 nonzeros (checked; raises if violated – this function does not itself select the 2:4 pattern, that’s TwoFourSparsityRamp.project() / G2’s sigma_weighted_block_sparse, this is purely the storage-format conversion of an already- constrained matrix, matching how a real cuSPARSELt workflow separates “prune/select the pattern” from “compress for the kernel”).

Parameters:

weight_2_4_masked (Any)

Return type:

Compressed2to4

decompress(compressed)[source]

Exact inverse of export_2_4_compressed(): reconstruct the dense (2:4-sparse) matrix from values/indices. Round-trips EXACTLY (bit-for-bit on the surviving values, genuine zeros elsewhere) – pinned directly by the compress/decompress test.

Parameters:

compressed (Compressed2to4)

Return type:

ndarray

cusparselt_status()[source]

Honest, queryable snapshot of what THIS torch build/environment actually offers for real cuSPARSELt-accelerated 2:4 sparse GEMM – used by the tests to decide (and clearly LABEL) whether an “inference speedup” number is a real measurement or a theoretical bound. Never raises: every field degrades to False/None if the relevant torch API is missing entirely (older torch builds).

Return type:

dict[str, Any]

is_2_4_sparse(weight)[source]

True iff every contiguous group of 4 entries along the last axis has AT MOST 2 nonzeros (the 2:4 structural constraint the export format below requires as a precondition).

Parameters:

weight (Any)

Return type:

bool