Add fused CuTe-DSL forward+backward kernels for the CSA Compressor gated pooling (THD) - #5984
Open
zkyue wants to merge 4 commits into
Open
Add fused CuTe-DSL forward+backward kernels for the CSA Compressor gated pooling (THD)#5984zkyue wants to merge 4 commits into
zkyue wants to merge 4 commits into
Conversation
The gated softmax pooling region of Compressor._forward_thd (gather-index
build -> gather -> +APE -> overlap-window transform -> fp32 softmax ->
gated weighted sum -> bf16 cast) decomposes into ~39 forward and ~51
backward kernel launches per call at compress_ratio = 4, materializes
(total_comp, 2*ratio, 1, head_dim) window intermediates that are also
saved for autograd, and its backward pays an indexing_backward plus a
radix sort for the gather.
Add one forward and one backward kernel written in the CuTe Python DSL
(nvidia-cutlass-dsl, JIT-compiled on first use) covering exactly that
region for the THD packed path, coff in {1, 2}:
* fp32 arithmetic with a single final bf16 rounding; mul/fma pinned in
PTX so results do not depend on compiler FMA contraction. dKV/dScore
are bit-identical to an fp32-intermediate eager reference; against an
fp64 oracle the fused kernels max abs error is <= the eager path
error on every tested output.
* Backward is atomic-free for dKV/dScore: every consumed input element
belongs to exactly one pooling window, so gradient stores are
disjoint; unconsumed elements keep exact zeros from the
zero-initialized grad buffers, matching autograd. dAPE is reduced
with one fp32 atomic per (k, dim) per CTA and is therefore not
bitwise run-to-run deterministic (forward, dKV, dScore are); the
backward raises under torch.use_deterministic_algorithms(True).
* Static-capacity padding rows (fixed_total_comp) are supported:
forward computes them exactly like the eager code (gather from token
0 with first-in-segment semantics), backward ignores incoming
gradients on padding rows.
* CUDA-graph capture compatible after a one-step eager warmup per
(ratio, head_dim, coff, device) configuration; a first call that
would JIT under capture raises a clear RuntimeError instead of
corrupting the capture.
* A cached-launch fast path snapshots the DSL launch state once per
(kernel, config, device, thread) and replays it with in-place
argument mutation, removing per-call host wrapper overhead; any
structural mismatch on a future nvidia-cutlass-dsl upgrade raises at
construction and the wrapper permanently falls back to the regular
launch path.
Addresses NVIDIA#5968.
Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
Additive fast path: maybe_compress_thd_fused() returns the pooled tensor for the supported configuration (THD non-pre-grouped path, compress_ratio == 4, bf16 kv/score, fp32 ape, compute capability 10.0, int32 flat offsets) and None otherwise, in which case the eager region runs unchanged. SBHD and the pre-grouped CP-prep path are untouched. The dispatch also keeps the eager path under torch.use_deterministic_algorithms(True) (dAPE fp32 atomics) and under torch.compile tracing (raw-pointer launch path is not traceable). MCORE_CSA_FUSED_COMPRESSOR=0 disables the dispatch entirely. compress_ratio == 128 is functionally supported by the explicit op API but stays on the eager path until tuned (see NVIDIA#5968). Addresses NVIDIA#5968. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
Covers numerics vs an fp32-intermediate eager reference (bitwise dKV/dScore), vs the verbatim upstream eager numerics (tolerance) and vs an fp64 oracle (fused error <= eager error) over ragged THD packs including segments shorter than ratio; coff == 1 (compress_ratio == 128) functional correctness; fixed_total_comp static-capacity padding rows (including a pack whose leading segment is shorter than ratio, and that nonzero padding-row gradients are ignored); run-to-run determinism of forward/dKV/dScore; CUDA-graph capture -> replay (including replay with new data and a smaller device-side true row count) and the loud JIT-under-capture error; dispatch gating and eager fallback (kill switch, ratio 128, non-bf16, layout, empty output, deterministic mode), plus Compressor._forward_thd-level integration. The module is marked launch_on_gb200 for the CC 10.0 CI lane and skips cleanly without CUDA, nvidia-cutlass-dsl, or compute capability 10.0. Addresses NVIDIA#5968. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
Standalone harness (not collected by pytest) behind the numbers in issue NVIDIA#5968 and the PR: --check compares fused fwd+bwd against the verbatim upstream eager region, an fp32-intermediate reference and an fp64 oracle; --bench measures CUDA-event wall clock of the replaced region; --nsys provides a cudaProfilerApi-gated loop with NVTX FWD/BWD ranges for pure-kernel-time capture. The eager reference reuses the upstream batch_of_row and Compressor._overlap_transform_thd implementations. Addresses NVIDIA#5968. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
Contributor
|
Hi @zkyue Thanks for your contribution! I think it is better to put the kernels in cudnn-frontend like other CSA/HCA fused kernels. Is it convenient for you to do so, or do you need our help to port the kernel? |
Author
|
@hxbai Sure, happy to do that — the kernels are self-contained CuTe DSL and should map cleanly onto the cudnn-frontend python API structure (we've contributed there before). Plan: I'll open a PR against cudnn-frontend porting the fwd+bwd kernels (APIBase-style wrapper, THD semantics, tests), then rework this PR into a thin dispatch that uses the cudnn-frontend implementation when available, keeping the eager fallback. Will link the kernel PR here once it's up. |
Author
|
@hxbai The cudnn-frontend kernel PR is up: NVIDIA/cudnn-frontend#427 (fwd+bwd kernels, APIBase-style wrapper, THD semantics, 40 tests; outputs bit-identical to the kernels in this PR). Once it lands I'll rework this PR into the thin dispatch that uses the cudnn-frontend implementation when available, keeping the eager fallback. |
Anerudhan
pushed a commit
to NVIDIA/cudnn-frontend
that referenced
this pull request
Jul 30, 2026
…rom Megatron-LM) (#427) * CSA: add fused Compressor forward+backward CuTe-DSL kernels Port the fused CSA/HCA Compressor gated-pooling kernels from Megatron-LM (NVIDIA/Megatron-LM#5984) into the FE-OSS Python API, per the maintainer request in NVIDIA/Megatron-LM#5984 (comment). One forward and one backward CuTe-DSL kernel fuse the THD gated-softmax pooling region (gather -> +APE -> overlap-window transform -> fp32 softmax -> gated weighted sum -> bf16 cast). Kernel math is unchanged from the Megatron original (verified bitwise on forward/dKV/dScore); the interface is reshaped to the APIBase pattern: CSACompressorForward / CSACompressorBackward classes, csa_compressor_forward_wrapper / csa_compressor_backward_wrapper, a new python/cudnn/csa package, docs, and fe_api tests (numerics vs fp32/upstream eager references and an fp64 oracle, ragged packs, static-capacity padding, run-to-run determinism, CUDA-graph capture/replay, check_support boundaries). Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> * CSA compressor: vectorize forward bf16 access (32-bit, CuTe autovec_copy) Widen every bf16 load/store in the fused Compressor forward kernel from scalar (16-bit) to paired 32-bit accesses: each thread now owns 2 adjacent head dims and moves its window slices through 32-bit universal copies (cute.autovec_copy over register fragments; cute.assume makes the alignment provable). The per-lane fp32 math is unchanged and in the same order, so the forward output stays bitwise identical to the previous kernel (re-verified against the Megatron-LM original on the same 5 THD shape-mix gate as the original port). Odd head_dims keep a scalar (vec == 1) instantiation of the same kernel; the head_dim 128/512 production shapes take vec == 2. The launch schedule moves from (rows_per_cta=4, threads=128) to one row per CTA with 64-thread column groups, which widens the sub-wave grids that limit the small shapes. Measured (nsys pure kernel time, B200, ratio=4 coff=2, THD packs of 8192-token sequences, forward kernel only): shape before after speedup 1x8192 d128 6.0 us 4.5 us 1.36x 3x8192 d128 12.6 us 10.0 us 1.26x 1x8192 d512 16.0 us 12.4 us 1.29x 3x8192 d512 42.1 us 35.0 us 1.20x Alternatives measured and rejected on the same matrix (all bitwise-equal): a PTX-unpack vec2 prototype (slower than the pure-DSL version: 4.7/10.5/12.9/35.7 us), and 64/128/256-bit per-thread vectors, which cut executed instructions but lose to register pressure and occupancy (80/147/255 regs vs 48; the 256-bit variant spills). Registers stay spill-free at 48 (previous kernel: 40). Tests: numerics shape matrix extended with an odd head_dim case to cover the scalar-layout instantiation, and a check_support rejection for head_dims beyond the launch bound. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> * CSA compressor: fold dKV/dScore zero-init into the backward kernel The backward previously required zero-initialized dKV/dScore buffers: the kernel only wrote consumed positions, and two tensor-wide bf16 fill kernels (torch.zeros_like) supplied the zeros everywhere else. Those fills cost as much DRAM traffic as the kernel's own stores. The kernel now writes every position itself: consumed positions get their gradients as before (disjoint, atomic-free stores), and each never-consumed slot class gets an exact zero from a unique natural owner, keeping all stores disjoint and dKV/dScore bitwise run-to-run deterministic: - first-half columns of each segment's last block's own tokens (coff == 2; no next block consumes them) -> that last block; - per-segment tail tokens (seqlen % ratio, both halves) -> the segment's last block; - all tokens of segments with zero output blocks (seqlen < ratio) -> CTA column bidx == 0; - tokens beyond cu_seqlens[-1] (static token-capacity padding of the gradient buffers, the CUDA-graph case) -> grid-strided across CTA columns. The backward wrapper allocates grad_kv/grad_score with torch.empty_like instead of torch.zeros_like (the fp32 dAPE fill stays: it is the atomic accumulator). When total_comp == 0 the kernel cannot launch, so the wrapper falls back to zeroed allocations to preserve autograd's exact-zero semantics; the class API documents the same caveat. Measured on the backward region (kernel + remaining fills vs kernel + three fills before), B200, ratio=4 coff=2, THD packs of 8192-token sequences; nsys = sum of kernel durations per iteration, wall = CUDA events around the wrapper call: shape nsys before -> after wall before -> after 1x8192 d128 15.5 -> 12.8 us (1.21x) 56.2 -> 47.2 us (1.19x) 3x8192 d128 29.9 -> 22.2 us (1.35x) 61.8 -> 54.6 us (1.13x) 1x8192 d512 34.2 -> 22.8 us (1.50x) 67.9 -> 57.3 us (1.19x) 3x8192 d512 90.9 -> 66.0 us (1.38x) 111.6 -> 101.8 us (1.10x) The zero-writes cost the kernel 8 registers (64 -> 72, occupancy 44% -> 41% at 3x8192 d512) and a small amount of extra store traffic, repaid several times over by dropping the two whole-tensor fills. dKV/dScore remain bitwise identical to the Megatron-LM original on the 5-shape port gate and to the fp32 eager reference on every test shape. Tests: NaN-canary suite proving the kernel fully overwrites uninitialized buffers (ragged, degenerate short segments, all-tiny packs, head_dim 512, static-capacity padded rows) bitwise against zero-initialized runs and the eager reference; a zeros-fallback test for packs where no segment reaches ratio tokens. The existing CUDA-graph replay test covers the token-capacity padding class (it fails without the grid-strided sweep). Docs: csa.md updated (buffer contract, fill count, perf tables refreshed for both this commit and the forward vectorization). Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> * Address review feedback: docs bounds, lint, warning stacklevel - docs: list the total_comp * head_dim < 2**31 and total_comp > 0 with total_tokens < ratio support boundaries already enforced in check_support - docs: run the compressor tests from test/python per repo convention - csa/__init__.py: spell out __all__ as an explicit literal (Ruff PLE0604) - compressor/__init__.py: sort __all__ (Ruff RUF022) - compressor/api.py: add stacklevel=2 to the deterministic-mode RuntimeWarning Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> * CSA compressor: rewrite graph benchmark to fwd-only + fwd+bwd total graphs Address review feedback on the PR #427 fairness benchmark (benchmark/csa/bench_csa_compressor.py). - Drop the subtraction-derived backward-graph estimate. Graph variants are now two directly measured, symmetric columns for both eager and fused: a forward-only graph (eager under no_grad) and a forward+backward total graph. - Capture the eager fwd+bwd graph with stable, pre-allocated zero .grad buffers (the supported torch CUDA-graph pattern): the captured region zeros them in place, then runs forward + autograd backward, so every replay accumulates into a zeroed buffer -- numerically identical to a single fresh backward. A per-shape graph-vs-fresh-backward cross-check is printed (all shapes bitwise-equal). - The fused total graph captures the forward wrapper immediately followed by the backward wrapper. - Emit both the per-call and graph tables from a single run; graph speedups are eager/fused of the displayed us, truncated to one decimal (never rounded up). The backward replay is reported only as ~total-fwd (a reference), never as a column. Not collected by pytest; black -l160 clean. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> * docs(csa): refresh graph table to fwd/total columns; soften wording Address review feedback on the CSA compressor performance tables (docs/fe-oss-apis/csa.md). - Replace the CUDA-graph table (which published a subtraction-derived backward column) with two directly measured columns: eager/fused forward-only graph and eager/fused forward+backward total graph. No backward-graph column; backward is noted only as ~total-fwd, never as a measured value. - Refresh both the per-call and graph tables from a single run of the benchmark harness; graph speedups are eager/fused of the displayed us, truncated to one decimal. - State that capture collapses each side's per-op launches into a single replay rather than removes launch/host overhead; say less launch-bound rather than compute-bound. - Note how the re-run reproduces the previously published per-call wall clock, and document the eager-backward capture basis (stable zero .grad buffers). Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> * bench(csa): guard speedup division, dedupe leaf construction, ruff nits Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> * CSA compressor: widen the validated envelope to coff in {1, 2} The kernels are already generic over coff — for coff=1 the window is the block's own ratio tokens (no overlap transform, always valid); only the check_support gate restricted the APIs to the production coff=2 form. Widen the gate to coff in {1, 2} (ratio stays gated at 4), update the error message and the envelope prose in api.py / docs / README, and parametrize the tests over both coff forms: numerics vs the fp32/upstream/fp64 eager references (the eager reference already implements coff=1 faithfully — it skips the overlap transform), static- capacity padding, replay determinism, empty outputs, NaN-canary zero-writes, deterministic-mode error paths, CUDA-graph capture, class-vs-wrapper equivalence, and check_support accept/reject boundaries (new coff=0 / coff=3 rejects; a new empty-segment shape covers both coff forms). No kernel changes. Addresses hxbai's coff=1 question on the PR. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> * CSA: add missing docstrings across the compressor Python files Bring the PR's Python docstring coverage over CodeRabbit's 80% pre-merge gate: module docstrings for the csa package/compressor __init__ files and one-line docstrings for the remaining undocumented helpers in api.py (cache/ctor/compile plumbing), compressor_sm100.py (launcher ctors), the test helpers, and the benchmark helpers. interrogate 1.7.0 over the PR's six CSA Python files: 69.4% -> 100%. Docstrings only — no logic, kernel, or launch-path changes. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> * CSA compressor: add dedicated ratio=128 kernels (forward+backward) with ratio dispatch The generic CSA compressor kernels keep the whole coff * ratio pooling window in per-thread registers; past ratio ~32 that hits the 255-register wall and spills kilobytes of local memory per thread, so ratio=128 needs dedicated kernels rather than a lifted gate. This adds them behind the same public API: check_support now admits ratio in {4, 128} x coff in {1, 2} (ratio=128 additionally gated to head_dim in {128, 512}, the hardware-validated set) and the wrappers and class APIs route to the matching kernel family by ratio transparently. Design: the forward streams the window with a chunked online softmax (one CTA per output row, per-column (max, denom, acc) triples merged in one fixed smem order); its launch schedule is bucketed by output-row count (small/default/large, all precompiled for CUDA-graph safety) and per bucket selects two-phase accumulation and an ex2.approx fast exp where measured faster. The backward stages each row's window through shared memory chunk-parallel, fuses the per-chunk den/S partial sums into the same pass, merges partials in a fixed chunk order, and stores gradients with a hoisted 1/den multiply — keeping kernel-side zero-writes to never-consumed dKV/dScore slots and fp32-atomic dAPE exactly as at ratio=4. ptxas (sm_100a): forward 32-51 registers, backward 48-128, 0 spill / 0 stack across all 16 shipped (config, schedule) kernels (reproducer: benchmark/csa/reg_probe_csa_compressor_r128.py). Numerics contract, per ratio family (docs/fe-oss-apis/csa.md): ratio=4 keeps its bitwise contract UNCHANGED — dKV/dScore bit-identical to the fp32-intermediate eager reference, at coff 1 and 2. ratio=128 is deterministic + faithful to that same fp32-intermediate eager reference: bitwise run-to-run determinism of out/dKV/dScore (NaN-prefill replay tested); the same values within final-bf16 rounding at the gate tolerances (differing elements <= max(1, 0.1%), max_abs <= 1.6e-2, calibrated on the gate's documented input distribution — bf16 deviations scale 2^k with 2^k inputs while the fp32 intermediates stay finite); the eager reference's non-finite propagation (finite inputs that overflow its fp32 intermediates poison both sides alike, with the stated caveat that the fused order saturates earlier near fp32 max — un-normalized chunk partials); and, on inputs whose fp32 intermediates stay finite in both evaluation orders, fp64-oracle parity (at least as close to an fp64 oracle as the eager reference). The reduction reorders and fast-exp buckets were each adopted on a measured win and are covered by that contract. Performance (nsys isolated GPU kernel time, 1x B200, vs the fp32-intermediate eager reference region on identical inputs): forward 13.2-44.4x, backward 7.4-21.9x across coff {1, 2} x head_dim {128, 512} x 8k-131k tokens; e.g. coff=2, d=128, 65k tokens: fwd 726.3 -> 16.4 us (44.4x), bwd 960.0 -> 61.2 us (15.7x). Full tables with methodology in docs/fe-oss-apis/csa.md. Validation (B200, CC 10.0): the committed 21-case contract gate (benchmark/csa/gate_csa_compressor_r128.py) passes 21/21 with shipped schedule coverage asserted (forward 10/10, backward 6/6 buckets); pytest fe_api/csa/test_CSA_compressor.py: 88 passed, 1 skipped at the default L0 level and 98 passed, 1 skipped with -m "L0 or L1" (both ratio families; the ratio=4 rows, including the coff=1 envelope, stay on their bitwise assertions). New tests cover the ratio=128 dispatch envelope (schedule selection at every nb_total bucket boundary; every shipped (config, schedule) kernel executed against the full contract at L1), exact-zero assertions on every never-consumed dKV/dScore slot class in the NaN-canary, CUDA-graph capture/replay for the new family, and grad_ape zeroing-ownership regressions. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> --------- Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Addresses #5968 (contribution invited by the maintainers in
#5968 (comment)).
The
Compressorused by the CSA/HCA experimental attention variants implements its gatedsoftmax pooling in eager PyTorch. In
_forward_thd, the region between the two projectionGEMMs and the RMSNorm — gather-index build, gather,
+ APE, overlap-window transform(
coff == 2), fp32 softmax over the window, gated weighted sum — decomposes into ~39forward and ~51 backward kernel launches per call (at
compress_ratio = 4), materializes(total_comp, 2*ratio, 1, head_dim)window intermediates (also saved for autograd), andits backward pays an
indexing_backwardplus a radix sort for the gather.This PR adds an additive fused fast path for exactly that region: one forward and one
backward kernel written in the CuTe Python DSL (
nvidia-cutlass-dsl, JIT-compiled atfirst call), keeping the existing eager implementation as the semantic reference and
fallback. Forward+backward replace the ~90 eager launches with 1 + 1 kernels (plus 3
grad-buffer zero-fills in backward).
Files
megatron/core/transformer/experimental_attention_variant/csa_fused_compressor.py— new,self-contained: import-guarded DSL imports, the two kernels, a cached-launch fast path,
the autograd wrapper, the explicit op API (
compress_thd_fused) and the dispatch helper(
maybe_compress_thd_fused).megatron/core/transformer/experimental_attention_variant/csa.py— minimal insertion in_forward_thd: try the fused path, fall into the (unchanged, re-indented) eager regionwhen it returns
None.tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_compressor.py— unit tests.tests/unit_tests/transformer/experimental_attention_variant/bench_csa_fused_compressor.py—the standalone benchmark/correctness harness behind the numbers below (not collected by pytest).
Dispatch / gating (initial)
The fast path engages only for: THD packed non-pre-grouped path,
compress_ratio == 4(
coff == 2), bf16kv/score, fp32ape, compute capability 10.0 (the onlyvalidated architecture; the kernels use no TMA/tcgen05/MMA features, so wider coverage is
likely straightforward but stays opt-in until validated), and int32 flat offsets
(
total_tokens * coff * head_dim < 2**31). Everything else — SBHD, the pre-grouped CP-prepinput path,
compress_ratio == 128, missing DSL, other devices — keeps the eagerimplementation.
MCORE_CSA_FUSED_COMPRESSOR=0disables the dispatch entirely.compress_ratio == 128(coff == 1) is functionally supported by the explicit op API andcovered by the tests, but stays on eager until tuned (currently only ~3.1× fwd / ~1.2× bwd
wall-clock there). The dispatch also keeps the eager path under
torch.use_deterministic_algorithms(True)(seedAPEbelow) and undertorch.compiletracing (the raw-pointer launch path is not traceable; eager lets the compiler fuse the
region itself).
Numerics
All arithmetic is fp32 with a single final bf16 rounding;
mul.rn.f32/fma.rn.f32arepinned in PTX so results do not depend on compiler FMA contraction. Fresh numbers from the
included harness on the ported code (B200, details below):
rounding):
dKV/dScorebit-identical on every tested shape; forward differs on<0.006% of elements with max abs difference 3.9e-3–7.8e-3.
itself rounds intermediates to bf16 (
torch.softmax(...).to(kv.dtype), bf16 multiply).Against an fp64 oracle the fused kernels' max abs error is equal to or lower than
eager's on every tested output (table in the harness output / issue).
dKV/dScore(disjoint stores; unconsumed elements keepexact zeros, matching autograd). Forward,
dKV,dScorereplay bitwise identically runto run.
dAPEis accumulated with fp32 atomics and is not bitwise run-to-rundeterministic (max run-to-run deviation 1.5e-5–6.9e-5 over 10 replays per shape in the
issue measurements, fp32); the dispatch therefore falls back to eager under
torch.use_deterministic_algorithms(True), and the explicit op raises in backward inthat mode. A deterministic per-CTA-partials reduction is possible follow-up work if
required.
CUDA graphs &
fixed_total_comp(new vs the issue)Both items the issue left open are implemented and tested here:
per
(ratio, head_dim, coff)configuration; a first call that would JIT under captureraises a clear
RuntimeErrorinstead of corrupting the capture. The unit test capturesfwd+bwd in one graph and replays with new data and a smaller device-side true row
count (static capacity, changed
cu_seqlenscontents), checking bitwisedKV/dScoreagainst the eager reference.
code (they gather the window from token 0 with first-in-segment semantics, as the eager
gather does); backward ignores incoming gradients on padding rows (they are tail
padding, not consumed downstream) — tested, including that nonzero padding-row gradients
change nothing and including a pack whose leading segment is shorter than
ratio.Performance
Measured on 1×B200 (CC 10.0, driver 590.48.01); PyTorch 2.13.0a0 (CUDA 13.3),
nvidia-cutlass-dsl4.6.1; bf16kv/score(total, 1, coff*head_dim), fp32ape;compress_ratio = 4,coff = 2; THD packs of 8192-token sequences. Both implementationsrun on identical inputs over exactly the replaced region (projection GEMMs, RMSNorm, RoPE
excluded — unchanged). Numbers re-measured on this PR's code with the included harness.
Isolated GPU kernel time (nsys, sum of kernel durations per iteration, 50 iterations
after 20 warmup; no launch/host overhead; backward includes its grad-buffer fills):
End-to-end wall clock of the same region (CUDA events, median of 100, includes launch
overhead and, for backward, autograd engine overhead — not comparable to the kernel-time
numbers above):
(The fused wall numbers go through the explicit op API of the harness and therefore
include its host-side argument validation, a few µs per call; the production dispatch
path performs its cheaper gating checks instead. They are slightly below the issue's
wall numbers, which were measured on the pre-review version of the wrapper without the
added validation; the isolated kernel times above are unchanged.)
Tests
tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_compressor.py(all skip cleanly without CUDA / DSL / CC 10.0):
dKV/dScore), vs verbatim upstream eager (tolerance),vs fp64 oracle (fused error ≤ eager error), over ragged multi-segment packs including
segments shorter than
ratio;coff == 1(compress_ratio == 128) functional correctness;fixed_total_compstatic-capacity padding (fwd semantics + ignored padding grads,including a pack with a leading segment shorter than
ratio);dKV/dScore;maybe_compress_thd_fused), including thedeterministic-mode fallback (and the explicit op raising in backward under
torch.use_deterministic_algorithms(True)), plusCompressor._forward_thd-levelintegration (fused engages, matches eager, gradients flow).
The module carries
pytestmark = pytest.mark.launch_on_gb200so the CC 10.0 CI laneexercises it; everywhere else the tests skip cleanly (no CUDA / no DSL / other archs).
Also ran the existing related suites with the fused dispatch live:
test_attention_variant_csa.py,test_csa_cp_layout_kernels.py,test_csa_cp_utils.py—136 passed, 1 skipped (needs ≥2 ranks), 0 failed.
Notes
MCORE_CSA_FUSED_COMPRESSOR=0kill switch), matching the scope proposed in [ENHANCEMENT] Fused forward+backward CUDA kernel for the CSA/HCACompressorgated pooling (THD path) #5968. Itdoes change training numerics relative to the current eager code (more accurate against
the fp64 oracle, but not bit-identical) — happy to flip this to opt-in if you prefer.
per-call host overhead; any structural mismatch on a future
nvidia-cutlass-dslupgrade raises during construction and the wrapper permanently falls back to the
regular launch path (it never launches from a partially built snapshot).
MCORE_CSA_FUSED_COMPRESSOR_FAST_LAUNCH=0disables just this optimization.dAPEreduction, tuning forcompress_ratio == 128, wider arch enablement after validation.