Fused depthwise causal conv1d + bias + SiLU#189
Merged
Conversation
Per-channel (depthwise) K=4 causal short convolution fused with bias and SiLU on the PTO tile ISA, for Mamba/GDN-style short conv. fp16 and bf16 I/O, fp32 accumulate. A 2-D (channel-tile x L-chunk) work grid fills all AIV cores across batch/seq/dim, with a causal K-1 halo per chunk. Validated against a torch depthwise reference (168 shape cases) and benchmarked vs aclnnConvolution and within a chunked GDN layer. Files under examples/jit_cpp/conv1d/: - conv1d_dw_pto.cpp PTO kernel (single-seq + batched fp16/bf16 entries) - jit_util_conv1d_dw.py bisheng JIT compile + ctypes binding - run_conv1d_dw.py correctness + perf vs aclnnConvolution - sweep_conv1d_dw.py full L x W shape sweep (grid, GB/s, %peak) - prof_conv.py msprof profiling driver - run_gdn_conv_integration.py GDN conv-stage before/after integration Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
…pute) Prefetch the next input row so its TLOAD (MTE2) overlaps the current row's vector compute: a prologue load of x[hstart] plus a per-iteration (cast-current -> prefetch-next x[j+1] into the same xin_h -> compute) reorder. Zero extra UB tiles/events (single EVENT_ID0, flag accounting balanced). Numerics identical (168 cases pass, worst max-abs-error 7.8e-3); ~5-14% faster in the latency-bound regime (device-event timing), no regression on bandwidth-bound shapes. Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
No functional change; satisfies the clang-format pre-commit/CI hook. Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
Add the test infrastructure to match the JIT-example convention used by swiglu / layernorm / matmul_swizzle: - test_conv1d_dw.py: parametrized correctness vs an fp32 reference across GDN and general shapes, fp16 and bf16, activation on/off, plus the single-sequence fp16 entry. - conftest.py: --npu option and npu_device / setup_npu_device fixtures. - README.md: kernel description and suggested workflow. Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
These are not part of the self-contained example: no sibling example ships a prof_* helper, and run_gdn_conv_integration.py depends on the external pto_kernels package. Keep the example to kernel + jit_util + run/sweep + test + README, matching swiglu / layernorm / matmul_swizzle. Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
Apply black and fix ruff lint (unused import, f-string without placeholders) on the example's Python helpers so the Lint CI job passes. Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
Signed-off-by: Raphael Steiner <raphael.steiner@huawei.com>
Signed-off-by: Raphael Steiner <raphael.steiner@huawei.com>
- Ping-pong the input load across two UB slots (xin_h[2]) with independent V<->MTE2 handshakes (EVENT_ID0 / reused EVENT_ID3) so the next row's load overlaps the current row's cast and compute. - siluTile: derive the element type from TileT::DType instead of taking a separate template parameter. Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
Signed-off-by: Raphael Steiner <raphael.steiner@huawei.com>
…ling examples Make the depthwise causal conv1d+bias+SiLU example generic in filter width K, tile width MAX_W and per-core UB size via -DCONV_K / -DCONV_MAX_W / -DCONV_UB_BYTES (defaults K=4 / MAX_W=3072 / 192 KiB, unchanged). The accumulator ring is sized RS=roundUpPow2(K) and indexed with & (RS-1); a static_assert rejects any K/MAX_W/dtype combination that overflows the UB. Remove the K=4-only demo entry points and rename the I/O dtype template parameter IoT -> InOutT. Build the JIT example with the CANN-toolkit pto headers and dav-c220-vec arch (matching examples/jit_cpp/fast_hadamard/jit_util_common.py) instead of an external pto-isa path; tests recompile per (K, MAX_W) via -D flags. Verified: 152/152 pytest cases; default-config output bit-identical to the prior kernel with <=1% timing across 13 shapes; 168/168 correctness in the sweep. Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
…fp16+bf16 multi-K tests
Replace the -DCONV_K / -DCONV_MAX_W / -DCONV_UB_BYTES macros with call-site
constexpr template arguments (the production entries use K=4 / MAX_W=3072) plus a
namespace UB_BYTES_PER_CORE constant that the UB static_assert checks against.
Derive the channel-tile lane width from the I/O element size (256 B / sizeof)
instead of a hardcoded 128, hoist the grid tuning knobs into a commented block,
and promote the ring mask and the UB-offset tables to constexpr. Add
static_assert(K <= accumRingSize).
Rename identifiers to be self-documenting: num_wt -> numChanTiles, col_w ->
channelTileWidth, lchunks -> sequenceChunkCount, RS -> accumRingSize, W ->
channels, processUnit -> processWorkUnit, IoT/AccT -> IoElemType/AccumElemType,
and the loop/offset locals to match.
The generalized-K test no longer recompiles via -D macros; it compiles a
per-(K, MAX_W) wrapper that #includes the kernel and instantiates
runConvSiluBatched at that K for both fp16 and bf16 I/O. Coverage: K in
{2,3,4,5,8} x {fp16,bf16} x 5 shapes x 2 seeds.
Behavior-neutral at K=4 (output bit-identical to the previous kernel); 202/202
tests pass.
Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
Wrap the over-long call_test_kernel_fp16/bf16 selection so the file passes black (the pinned pre-commit / CI lint). No behaviour change; 202/202 tests still pass. Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
Rename the depthwise causal conv1d example (directory, files, and kernel entry symbols) from conv1d_dw to causal_conv1d, matching the Mamba / sgl-kernel-npu op name (the directory was already `conv1d`, and depthwise is implied for this op). No behaviour change: the device code is identical modulo the rename, and all 202 correctness cases still pass. Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
zouzias
reviewed
Jun 24, 2026
The general-K/MAX_W test built kernel entry points from a C++ source string embedded in test_causal_conv1d.py (_WRAPPER_SRC). Move it out of Python: the production entry points now read K/MAX_W from CAUSAL_CONV_K / CAUSAL_CONV_MAX_W (default 4 / 3072), and the test recompiles causal_conv1d_pto.cpp with -DCAUSAL_CONV_K / -DCAUSAL_CONV_MAX_W, reusing the existing call_kernel_batched / call_kernel_batched_bf16 launchers. Same -D pattern as matmul_swizzle and scan. No C++ remains in the Python file; a plain build (no -D flags) is byte-for-byte unchanged. 202/202 tests pass on 910B2. Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
Collaborator
|
zouzias
approved these changes
Jun 24, 2026
zouzias
left a comment
zouzias
left a comment
Collaborator
There was a problem hiding this comment.
python run_causal_conv1d.py gives
== correctness (vs torch depthwise conv1d + bias + SiLU = aclnnConvolution) ==
168 cases passed (over 84 shapes x 2 magnitudes)
WORST max-abs-error = 7.812e-03 at L=16 W=256 scale=8.0
== performance sweep ==
[native] npu_fused_causal_conv1d UNAVAILABLE on this box: RuntimeError('aclnnFusedCausalConv1d or aclnnFusedCausalConv1dGetWorkspaceSize not in libopapi.so, or libopapi.so not found.\n[ERROR] 2026-0
[native] (it is a 910_95-only op, K fixed=3; this box is Ascend910B2)
L W ours_us aclnn_us native_us ours GB/s vs aclnn
128 2048 119.7 180.2 n/a 8.8 1.51x
256 2048 120.5 186.3 n/a 17.4 1.55x
384 2048 179.1 192.3 n/a 17.6 1.07x
512 2048 183.9 183.1 n/a 22.8 1.00x
1024 2048 237.0 211.2 n/a 35.4 0.89x
256 4096 150.4 257.0 n/a 27.9 1.71x
256 8192 150.8 441.8 n/a 55.6 2.93x
512 4096 188.0 279.7 n/a 44.6 1.49x
2048 2048 237.6 248.9 n/a 70.6 1.05x
4096 2048 466.3 431.7 n/a 72.0 0.93x
Rework run_causal_conv1d.py perf(): sweep batched (B, seq, dim) shapes with NPU device-event timing (100-iter average) and time aclnnConvolution end-to-end (incl. the layout conversion a real pipeline pays), replacing the wall-clock / batch-1 / hoisted-aclnn version whose numbers were dominated by the ~200us host-dispatch floor. Add a silu() helper; drop the now-unused time import. Satisfy pre-commit + prospector: clang-format the CAUSAL_CONV macro comment; bare '# noqa' for the _as_parameter_ / torch_npu lines (matching the sibling jit_cpp examples); '# pylint: disable=unused-argument' on the npu_device session fixtures. Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
Collaborator
Author
|
The old Reworked to a batched (B, seq, dim) sweep with NPU device-event timing (100-iter average, adaptive) and aclnn timed end-to-end, on an idle/exclusively-monitored NPU. With that, PTO is ~1.3–1.5× over aclnn at dim=2048 and a steady ~2.0–2.2× at dim=6144, scaling flat to 6.4B elements. Results running in an isolated NPU running the new
|
zouzias
approved these changes
Jun 25, 2026
zouzias
left a comment
zouzias
left a comment
Collaborator
There was a problem hiding this comment.
............................................................................................................................................................................ [ 85%]
.............................. [100%]
================================================================================= warnings summary =================================================================================
../../../venv/lib/python3.10/site-packages/torch/jit/_script.py:362: 14 warnings
/home/zouzias/github-repos/pto-kernels/venv/lib/python3.10/site-packages/torch/jit/_script.py:362: DeprecationWarning: `torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.
warnings.warn(
-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
202 passed, 14 warnings in 7.26s
This was referenced Jun 27, 2026
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.
Add fused depthwise causal conv1d + bias + SiLU — PTO example
Adds a PTO-ISA example for the depthwise causal conv1d + per-channel bias + (optional) SiLU — the Mamba/GDN short convolution. Per-channel K-tap filter, fp16/bf16 I/O, fp32 accumulate. Built and run exactly like the sibling
jit_cppexamples (CANN-toolkit PTO headers,--cce-aicore-arch=dav-c220-vec, sharedconftest/build flags).Adds (7 files under
examples/jit_cpp/causal_conv1d/): kernelcausal_conv1d_pto.cpp,jit_util_causal_conv1d.py,test_causal_conv1d.py+conftest.py,run_causal_conv1d.py,README.md,.gitignore. No existing files touched.Design: input-stationary K-tap scatter with a power-of-two accumulator ring (
& (RS-1)), double-buffered input load, generic in K/MAX_W (chosen at the call site asconstexprtemplate params), UB-budgetstatic_assert.Correctness:
pytest test_causal_conv1d.py→ 202 pass — K=4 (batched + single-seq, fp16/bf16, act on/off, edge/GDN shapes) and generic K ∈ {2,3,4,5,8} × {fp16,bf16} vs an fp32 reference.Baselines (all on Ascend 910B2):
aclnnConvolution(torch grouped depthwise conv); the current sglcausal_conv1d(the existing AscendC kernel); and PR-555 = sgl-project/sgl-kernel-npu#555, the community "add gdn custom conv1d" PR. In the tables below, PTO = this kernel's approach (thecausal_conv1d_ptoop in sgl-kernel-npu); PR-555 and sgl are those two baselines.Performance: ~2–7× vs
aclnnConvolution. The same approach productionized in sgl-kernel-npu (causal_conv1d_pto) runs 1.1–1.9× over PR-555 and 2.6–20× kernel-time over the current sgl kernel on the GDN matrix (dims 2048/6144), and is correct where the current sgl kernel is numerically wrong at large batch. Full per-shape data (idle-gated):Kernel-only — msprof TaskDur, f16, µs/forward (16 GDN shapes)
PTO/PR 1.12–1.86×; PTO/sgl 2.6–20× (sgl's single kernel scales with seq → 7–20× at B1). bf16 tracks f16 within ~2%.
Large batch + larger seq — e2e device-event, fp16, min-of-3 (45 GDN shapes, sorted by B·seq·dim)
At small/floor-bound shapes e2e is pinned at the shared ~210µs torch-op host floor (identical for both ops) — the kernel-only table above shows the true gap there. The current sgl kernel is numerically WRONG in most large-batch cases (batch race); PTO is correct throughout.
Run:
pytest test_causal_conv1d.py -q --npu npu:0