disco-kernel-dev

name

disco-kernel-dev

description

Expert guide for adding, optimizing, or debugging DISCO convolution kernels in torch-harmonics. Use this skill whenever the user mentions: adding a new kernel variant or GPU arch path, wiring up a kernel through the Python dispatch layer, fixing bf16/fp16/autocast issues in DISCO, profiling or benchmarking the CSR or kpacked forward/backward, propagating kernel changes to the distributed DISCO path, or writing tests that cover all dispatch branches. Also use for debugging correctness issues in the DISCO fwd/bwd path or distributed a2a collectives.

DISCO Kernel Development Guide

Architecture snapshot

disco_interface.cpp          TORCH_LIBRARY("disco_kernels") — raw op schema
  ├── forward(inp, …)        CSR sparse contraction inp → (B,C,K,H,W)
  ├── backward(inp, …)       CSR transpose contraction (B,C,K,H,W) → inp
  └── forward_kpacked(…)     WGMMA kpacked forward (SM_90a + bf16/fp16 only)

disco_optimized.py           Python dispatch layer
  ├── _disco_s2_contraction_optimized         custom_op wrapping forward
  ├── _disco_s2_transpose_contraction_optimized  custom_op wrapping backward
  ├── _disco_s2_fused_conv_optimized          custom_op: contraction + einsum
  ├── _DiscoKpackedFn(autograd.Function)      WGMMA fwd + CSR bwd (unfused)
  ├── _DiscoKpackedFusedFn(autograd.Function) WGMMA fwd + CSR bwd (fused)
  └── _maybe_kpack_psi(…)    converts CSR psi to kpacked layout at init time if required

convolution.py               DiscreteContinuousConvS2._forward() dispatch:
  _kpacked_ok = optimized_kernel and psi_kpacked_K_pad in (8,16)
                and x.dtype in (float16, bfloat16)
  fused + kpacked  →  _disco_s2_fused_conv_kpacked
  fused only       →  _disco_s2_fused_conv_optimized
  kpacked only     →  _disco_s2_contraction_kpacked
  CSR default      →  _disco_s2_contraction_optimized / torch

distributed_convolution_kernels.py   mirrors serial dispatch for a2a paths
distributed_convolution.py           builds kpacked buffers in _build_local_psi

Key file locations

Purpose	Path
PyTorch reference kernels	torch_harmonics/disco/kernels_torch/disco_torch.py
CUDA kernel headers	torch_harmonics/disco/optimized/kernels_cuda/disco_cuda.cuh
CSR forward kernel	torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_fwd.cu
CSR backward kernel	torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_bwd.cu (BC_TILE optimized)
SM_90 kpacked kernel	torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_fwd_dense_kpacked_sm90.cu
PTX helpers (WGMMA)	torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_ptx.cuh
CPU OpenMP forward kernel	torch_harmonics/disco/optimized/kernels_cpu/disco_cpu_fwd.py
CPU OpenMP backward kernel	torch_harmonics/disco/optimized/kernels_cpu/disco_cpu_bwd.py
C++ interface	torch_harmonics/disco/optimized/disco_interface.cpp
Python dispatch	torch_harmonics/disco/optimized/disco_optimized.py
Serial conv (dispatch)	torch_harmonics/disco/convolution.py
Distributed conv	torch_harmonics/distributed/distributed_convolution.py
Distributed kernels	torch_harmonics/distributed/kernels/distributed_convolution_kernels.py
Build config	setup.py (CUDA sources list)
Serial tests	tests/test_convolution.py (TestKpackedPath)
Distributed tests	tests/test_distributed_convolution.py

---

Adding a new kernel variant — checklist

Walk through every layer in order. Each layer has a hard dependency on the previous one.

1. CUDA kernel (.cu / .cuh)

• Declare the host function in disco_cuda.cuh.
• Implement in a new .cu file named disco_cuda_{direction}_{arch}.cu (direction: fwd or bwd; arch: sm90, sm100, …).
• Add a #if defined(__CUDA_ARCH_FEAT_SM{NN}_ALL) guard so the kernel body compiles to empty on other arches — the host launcher enforces the arch check at runtime via TORCH_CHECK(props.major == N, …).
• Add the new .cu to setup.py CUDA sources.

2. C++ interface (disco_interface.cpp)

• Register the new op schema in TORCH_LIBRARY("disco_kernels", m).
• Implement the CUDA dispatch in TORCH_LIBRARY_IMPL("disco_kernels", CUDA, m).

3. Python dispatch layer (disco_optimized.py)

Every op that participates in autograd needs all four of these:

What	How
Fake kernel (shape inference)	@torch.library.register_fake("disco_kernels::op_name")
AutocastCUDA handler	@torch.library.impl("disco_kernels::op_name", "AutocastCUDA") — cast float inputs to torch.get_autocast_dtype("cuda"), call .default inside autocast(enabled=False)
Backward	torch.library.register_autograd(…) for simple ops; torch.autograd.Function subclass when forward and backward use different kernel paths
Public wrapper	thin Python function that calls op.apply(…) or op.default(…)

When to use autograd.Function vs register_autograd: Use autograd.Function whenever forward and backward run different kernels — the canonical case is WGMMA forward + CSR backward or when composing ops with other PyTorch kernels, for example the fused kernel variants. register_autograd assumes the backward is structurally parallel to the forward; mixing kernel types breaks that assumption.

AutocastCUDA pattern (copy-paste template):

@torch.library.impl("disco_kernels::my_op", "AutocastCUDA")
def _(inp, ...):
    cast_dtype = torch.get_autocast_dtype("cuda")
    with torch.amp.autocast("cuda", enabled=False):
        return my_op(inp.to(cast_dtype), ...)

4. Serial conv dispatch (convolution.py)

• Add _kpacked_ok (or equivalent arch gate) before the dispatch branch.
• The 4-way dispatch pattern (fused×kpacked) should be kept symmetric:

  fused + new_path → new fused variant
  fused only       → existing fused CSR variant
  new_path only    → new unfused variant
  default          → CSR
  

5. Distributed path (distributed_convolution.py + _kernels.py)

The distributed path must mirror the serial dispatch exactly:

1. In _build_local_psi: call the same psi-preparation helpers as the serial _build_local_psi (e.g. pack_psi_dense + _maybe_kpack_psi for kpacked). Set self.{feature}_K_pad = None unconditionally before the if optimized_kernel: block so forward always has the attribute.
2. In distributed_convolution_kernels.py: add the same _feature_ok guard and dispatch branches as the serial path.
3. In forward: thread the new buffers through via getattr(self, "buf", None).

---

Backward direction — gather vs scatter

This is the key architectural insight for DISCO; get it wrong and you silently lose performance or correctness.

• CSR forward (disco_kernels::forward): inp → K-expanded. Gather direction — each output pixel reads from a bounded neighbourhood. Input-pixel-parallel, no atomics.
• CSR backward (disco_kernels::backward): K-expanded grad → inp grad. Also gather direction — each input pixel accumulates from its neighbourhood. Input-pixel-parallel, no atomics. This is the correct backward for any convolution with overlapping support sets, for the same reason cuDNN uses implicit GEMM not col2im scatter.
• WGMMA forward_kpacked: accelerated CSR forward using Tensor Cores, restricted to SM_90a + bf16/fp16 + K_PAD ∈ {8, 16}.
• WGMMA backward (retired): scatter direction — output → input — causes massive atomicAdd contention when support sets overlap. Do NOT reintroduce. This is why it is currently not in the active code base.

The _DiscoKpackedFn / _DiscoKpackedFusedFn autograd.Function classes exist precisely to pair WGMMA forward with CSR backward.

---

BC_TILE optimization (CSR backward)

The CSR backward had poor FMA utilisation (L1/TEX bound, ~12% FMA) because each CTA covered one channel and redundantly loaded the psi index arrays. BC_TILE amortises the index loads: one CTA processes BC_TILE channels, loading psi indices once.

• BC_TILE is selected at runtime: BC >= 8 → 8, BC >= 4 → 4, else 1.
• Non-divisible BC is handled by ceiling division + if (bc >= BC_total) continue guards; invalid slots get zero-filled registers (harmless FMAs).
• cudaFuncSetAttribute(cudaFuncAttributeMaxDynamicSharedMemorySize) is required for BC_TILE=8 (49152 bytes, exactly at the default carveout limit).

---

kpacked layout

_maybe_kpack_psi converts pack_psi_dense output to the K-packed format:

pack_idx    [Ho, NBR_PAD, 2]      int64   (shared across all K — verify with torch.equal)
pack_val    [Ho, NBR_PAD, K_PAD]  fp32    (permuted + zero-padded to K_PAD = ceil(K/8)*8)
pack_count  [Ho]                  int64

Returns None if the per-K support sets differ (layout mismatch; CSR fallback activates). K_PAD must be 8 or 16 for the WGMMA kernel; store as psi_kpacked_K_pad on the module.

---

Test coverage requirements

Every kernel variant needs tests in both the enabled and disabled states. The monkeypatch pattern forces the fallback path without needing a non-Hopper machine:

conv.psi_kpacked_K_pad = 24   # ineligible → _kpacked_ok = False → CSR fallback

Serial tests (TestKpackedPath in test_convolution.py)

Test	Gate	What it checks
test_kpacked_forward_activates_on_sm90	SM_90	kpacked path selected, output dtype preserved
test_kpacked_fused_matches_unfused	SM_90	fused==unfused output + grad
test_kpacked_bwd_bc_tile_boundaries	SM_90	BC_TILE=1/4/8 correctness vs fp32
test_kpacked_disabled_for_unsupported_k_pad	none	K_PAD=24 → no crash (CSR fallback)
test_kpacked_disabled_fused_fallback	none	fused=True + K_PAD=24 → CSR fused path, fwd+bwd match
test_kpacked_opcheck	SM_90	PT2 opcheck contract

Distributed tests (TestDistributedDiscreteContinuousConvolution)

The parameterised AMP rows (dtype=float16/bfloat16) exercise kpacked-enabled on Hopper. Three dedicated methods cover the fallback:

Test	Covers
test_kpacked_fallback_bf16_unfused	fused=False + K_PAD=24
test_kpacked_fallback_bf16_fused	fused=True + K_PAD=24
test_kpacked_fallback_fp16_unfused	fp16 + K_PAD=24

---

Profiling

Serial Benchmarks

The code base currently does not have a benchmark. For writing profiling scripts for serial kernels, stick to a minimal implementation. Aim at running the kernel in question in isolation, comparing to existing kernels when possible. Incorporate all 3 precisions (fp32, bf16, fp16) into the benchmark and also compare the results of the kernel against its fp32 variant. Relevant shape combinations are:

CONFIGS = {
    "self_256x360x720":  dict(in_channels=256, out_channels=256, in_shape=(360, 720), out_shape=(360, 720), grid_in="legendre-gauss", grid_out="legendre-gauss", theta_cutoff=0.017, kernel_shape=(3,3), basis_type="harmonic", basis_norm_mode="mean"),
    "self_512x360x720":  dict(in_channels=512, out_channels=512, in_shape=(360, 720), out_shape=(360, 720), grid_in="legendre-gauss", grid_out="legendre-gauss", theta_cutoff=0.017,
    kernel_shape=(3,3), basis_type="harmonic", basis_norm_mode="mean"),
    "down_73x721x1440":  dict(in_channels=80, out_channels=512,, in_shape=(721, 1440), out_shape=(360, 720), grid_in="equiangular", grid_out="legendre-gauss", theta_cutoff=0.017,
    kernel_shape=(3,3), basis_type="harmonic", basis_norm_mode="mean"),
}

The specific type of grid does not affect kernel performance but might affect accuracy. For detailed profiling, run ncu

# Capture light sections (avoid --set full which hangs on multi-replay kernels)
ncu --kernel-name disco_bwd_blk_k \
    --section SpeedOfLight \
    --section MemoryWorkloadAnalysis \
    --section ComputeWorkloadAnalysis \
    -o profiles/my_profile \
    python performance/disco/my_benchmark.py

Key metrics to watch:

• FMA utilisation (ComputeWorkloadAnalysis) — target near-saturation; low FMA on an L1-bound kernel means the index load is the bottleneck → BC_TILE fixes this.
• L1/TEX hit rate (MemoryWorkloadAnalysis) — high is good for psi values that fit in L2.
• DRAM utilisation — should be low for typical model sizes (psi fits in L2).

Distributed Benchmarks

Write a small benchmark using mock communicators, imitating the actual communication pattern. Allow for splitting in latitude (h) and longitude(w). Only measure the serial custom kernel performance on the individual shards for the various precisions. This yields a floor for the expected kernel execution time.