causal_conv1d kernel improvement using PTO kernel & benchmark results - #601
Open
ChristosMatzoros wants to merge 19 commits into
Open
causal_conv1d kernel improvement using PTO kernel & benchmark results#601ChristosMatzoros wants to merge 19 commits into
ChristosMatzoros wants to merge 19 commits into
Conversation
Replace the depthwise causal conv1d + bias + (optional) SiLU op with a PTO-ISA tile kernel, built as its own ascendc_library (like mega_chunk_gdn). The op name and signature are unchanged, so this is a drop-in replacement: pytorch_extensions.cpp and the Python wrapper are untouched. Improvements over the previous AscendC kernel: - faster on the GDN prefill matrix (1.5-5.8x e2e on Ascend 910B2); - numerically correct at large batch, where the previous kernel raced; - generic over dim (the previous kernel only supported a fixed dim set). conv_states / varlen (query_start_loc) / has_initial_state / pad_slot_id semantics are preserved and verified against the prefill contract test (tests/python/sgl_kernel_npu/test_conv1d_prefill.py), which also covers the exact conv_states writeback. Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
…guard
- Template the PTO kernel on the filter width K (and per-K tile width MAX_W) and
compile a dedicated entry per supported width {2,3,4,5,8,16,32,64} x {fp16,bf16};
the host dispatches by weight.size(0). (Previously width-4 only.)
- Fix a state-writeback race: when the writeback grid has more tasks than cores a
core reuses its UB tiles across tasks, so the next task's load must wait for the
previous task's store (MTE3 -> MTE2). conv_states was silently wrong at large
batch (the common stateful-decode path) before this.
- Guard the GM_ADDR define with clang-format off: PointerAlignment: Right rewrote
'type* name' to 'type *name', which broke the AscendC launch codegen.
- Tidy the host launcher (compact per-width dispatch, constexpr helpers, comments).
Verified vs an fp32 reference: the prefill contract test (all widths, fp16/bf16,
dense/varlen, initial-state, bias, activation, pad-slot) and a 1024-case
width x dim sweep, all exact (conv_states bit-exact).
Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
Replace the per-width (K) template with a per-ring-size template (RS = roundUpToPow2(K)) and pass the filter width K as a runtime argument. Six compiled variants {2,4,8,16,32,64} x {fp16,bf16} now serve any width in [2,64] (e.g. 3->rs4, 5/6/7->rs8), replacing the eight per-width entries. The UB layout is sized for the worst case K==RS so every offset stays compile-time (no codegen regression).
Collapse the (RS, MAX_W) set into one X-macro list (CC1D_RS) inlined in the kernel and host; the host forward-declares the launch stubs from it instead of including 24 generated aclrtlaunch headers, and the entry-definition macros are compacted.
Tests: extend the contract test and fp32-reference sweep to non-power-of-two widths (runtime K < RS). Validated on Ascend 910B2: contract test all-pass (conv_states bit-exact), fp32 sweep 0/1664 failures; e2e and kernel-only performance unchanged vs the per-width build.
Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
Move the fp16/bf16 -> fp32 widening of weight/bias from the host (.to(at::kFloat)) into the kernel: stage the native tiles into the idle accumulator-region scratch (loads pipeline), one MTE2->V barrier, then pipelined TCVT into the resident fp32 tiles. Matches the main/PR-555 pattern (cast on device) and drops two per-call host cast launches. Numerically identical (bit-exact vs fp32 ref; 1664-case sweep + contract test pass for fp16/bf16). Faster: up to ~2.1x e2e at small batch and 1.02-1.19x kernel-only, tied at large batch. Also adds bias dtype/shape validation. Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
Unify the per-sequence row-index/length space as int32 (start, len, l0, l1, cacheIdx) and add a signed halo = K-1, so the signed<->unsigned cast churn in the index/halo/offset arithmetic disappears. Also drop the redundant bitmask casts (the & already converts a non-negative operand), consolidate the lanes ternary, and remove the dead `len` parameter from convChunk. Explicit casts 40 -> 16; the rest are load-bearing (uint64_t offset-overflow guards, the halo derivation, (int32_t)k, lanes). Bit-exact vs the fp32 reference (contract + 1664-case sweep pass, fp16/bf16) and perf-neutral (kernel-only TaskDur tied with the prior commit on all shapes). Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
Format the PTO kernel/host to the repo clang-format style (csrc/.clang-format) so `pre-commit run --all-files` (CI lint) passes. No functional change: the GM_ADDR / launch-stub regions stay under `// clang-format off` because their pointer glueing is required by the AscendC launch codegen. Rebuild confirms all 24 kernel entries still compile; whitespace-only, so numerics are unchanged. Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
Signed-off-by: Raphael Steiner <raphael.steiner@huawei.com>
Signed-off-by: raphaelsteiner <raphael.steiner@huawei.com>
Signed-off-by: raphaelsteiner <raphael.steiner@huawei.com>
Signed-off-by: raphaelsteiner <raphael.steiner@huawei.com>
Signed-off-by: raphaelsteiner <raphael.steiner@huawei.com>
Signed-off-by: raphaelsteiner <raphael.steiner@huawei.com>
Signed-off-by: raphaelsteiner <raphael.steiner@huawei.com>
Signed-off-by: raphaelsteiner <raphael.steiner@huawei.com>
Signed-off-by: raphaelsteiner <raphael.steiner@huawei.com>
…ode/num_accepted) sgl-project#592 replaced this op with the AscendC kernel and changed the registered schema (added num_accepted_tokens + run_mode, made activation_mode an int). The PTO kernel this PR restores uses the original signature — bias as an optional after the required index tensors, bool activation_mode, no run_mode — so update pytorch_extensions.cpp to match. Callers on sgl-project#592's signature (with run_mode) must adapt; the Python wrapper and other ops are untouched. Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request refactors the causal_conv1d NPU kernel and host-side implementation by replacing the legacy AscendC code with a new implementation based on the PTO tile ISA, simplifying the dispatch logic and removing several obsolete files. The code review highlights several critical correctness and robustness issues: in the kernel, negative index loading when hasInit is false can read stale history or uninitialized memory, requiring dummy loads and zeroing; on the host side, the defined MIN_CHUNK_ROWS constant is unused, leading to excessive sequence chunking overhead; there is a lack of validation for input tensor sizes (query_start_loc, cache_indices, and has_initial_state) which could cause underflows or out-of-bounds reads; and finally, the PyTorch extension wrapper needs explicit type casting for default integer and boolean tensors to prevent strict type assertion failures.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Signed-off-by: Christos Konstantinos Matzoros <christos.konstantinos.matzoros@h-partners.com>
Contributor
Author
|
All points addressed:
Correctness & perf: unit tests + 656/656 fp32-reference configs pass; a 400-shape kernel-time A/B (this branch vs the pre-fix commit) gives geomean 1.00× (0/400 shapes outside ±2%). No performance change. |
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
2 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.
causal_conv1d: PTO-ISA tiling kernel vsmain(AscendC) — benchmark resultsReimplements
causal_conv1d(Mamba/GDN depthwise causal conv1d + bias + optional SiLU) on the PTO tile ISA, and improves the launch tiling. Closes #596Head-to-head of this branch's PTO-ISA
causal_conv1dagainst the AscendC kernel onmain, on an Ascend 910B2 (48 AIV cores).main—f451a59(AscendC)causal-conv1d-tiling-onmain@13ad3a2(PTO-ISA), rebased ontomain(f451a59)speedup = main / branch(>1 ⇒ this branch is faster). One idle NPU, idle-gated before + after each run.The two layouts. Every shape is measured in both input layouts the op supports:
[batch, seq, dim]: every sequence in the batch has the same lengthseq(a rectangular, padded tensor).[cu_seqlen, dim]: variable-length sequences concatenated end-to-end into one flat 2D tensor, withquery_start_locmarking the boundaries — the real continuous-batching / serving layout. The causal conv resets its window at each sequence boundary (a sequence never reads the previous one's tokens).Summary
Geometric-mean speedup (
main / branch, >1 ⇒ branch faster), by channel count (dim):Full per-shape results
3D dense · dim=1024 — 40 shapes (kernel median 1.78×, e2e median 1.71×)
3D dense · dim=2048 — 40 shapes (kernel median 1.62×, e2e median 1.63×)
3D dense · dim=4096 — 40 shapes (kernel median 1.40×, e2e median 1.48×)
3D dense · dim=5120 — 40 shapes (kernel median 2.15×, e2e median 2.13×)
3D dense · dim=6144 — 40 shapes (kernel median 1.88×, e2e median 1.90×)
varlen / packed · dim=1024 — 40 shapes (kernel median 1.45×, e2e median 1.61×)
varlen / packed · dim=2048 — 40 shapes (kernel median 1.33×, e2e median 1.61×)
varlen / packed · dim=4096 — 40 shapes (kernel median 1.20×, e2e median 1.32×)
varlen / packed · dim=5120 — 40 shapes (kernel median 1.83×, e2e median 1.88×)
varlen / packed · dim=6144 — 40 shapes (kernel median 1.58×, e2e median 1.67×)
Reproduce these numbers
The benchmark scripts live in the
convolution_benchmarkingbranch of the fork:Take the two files from that directory — you don't need anything else from that branch:
bench_causal_conv1d_compare.py— the benchmark/compare toolcompare.md— step-by-step guide (environment setup, build, run, compare)Then, following
compare.md: buildmainandrunit, build this branch andrunit, thencomparethe two result files. Each run auto-detects the op signature (AscendC or PTO), times every shape both ways (e2e + kernel), checks correctness against an fp32 reference, and idle-gates the NPU. Example: