Repository navigation
Add chunk kda fwd - #11
Merged
Merged
Conversation
Signed-off-by: zhaozx-cn <zhaozx2116@163.com>
The vendored catlass copy under op_kernel/catlass/ was imported wholesale.
Its two hub headers pull in every implementation they know about, so the
chunk_kda_fwd translation unit was parsing ~26k lines of templates that no
KDA code path ever instantiates.
Auditing every Catlass symbol referenced from arch22/, arch35/, kernel_utils/
and tla/ shows the KDA kernels use exactly:
- Gemm::Block::BlockMmadTla with MmadPingpong (catlass) and with
MmadPingpongTlaMulti / MmadPingpongTlaPreloadAL1B (kernel_utils/block/,
included directly by their users)
- Epilogue::Block::BlockEpilogue with EpilogueAtlasKdaFwdHVnew / Update,
both specialized under arch22|arch35/fwd_h/epilogue/block/
Everything else -- 39 BlockMmad implementations, 20 catlass epilogues, the
AtlasA2-only BlockMmad/BlockGemm/BlockMmadAiv/BlockMmadSparseTla primaries,
the block swizzles (included by both FwdH kernels but never referenced), and
the helpers reachable only from those files -- is dead.
So kda_block_mmad.hpp and kda_block_epilogue.hpp now keep their primary
templates and include only what is selected, and the 70 orphaned files are
removed.
Preprocessed lines entering the compiler for the A5 build drop from 40943 to
28267 (116 files to 83); the A3 build drops to 23720. Verified for both
CATLASS_ARCH configurations that no include dangles and that no surviving
file references a removed symbol.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…alize
ChunkKdaFwdPostWuKernel and ChunkKdaFwdFinalizeKernel are templated on the
beta element type, but their beta_ tensor is only ever bound in Init() and is
never read -- in either the arch35 or the generic implementation:
$ grep -n "beta_" arch35/chunk_kda_fwd_finalize_arch35.h
255: beta_.SetGlobalBuffer((__gm__ BETA_T *)beta);
1430: GlobalTensor<BETA_T> beta_;
So the parameter bought nothing and cost a full second instantiation of each
kernel. Removing it also frees RunGenericBackEnd, whose only use of BETA_T was
forwarding it to Finalize; that one matters twice over because the function
has two overloads (with and without a TPipe) and each inlines the whole
FwdH + Finalize body.
Instantiations in the A5 translation unit:
arch35 PostWu (1959 lines) 2 -> 1
arch35 Finalize (1475 lines) 2 -> 1
RunGenericBackEnd 4 -> 2
The beta GM_ADDR arguments stay in the signatures so no call site changes,
and RunFrontEnd keeps BETA_T for now because Prepare still needs it.
No behavioural change: the removed state was write-only.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…template
Prepare was templated on BETA_T, so the 4336-line arch35 kernel was
instantiated twice over a parameter that only selects one vector cast:
if constexpr (IsSameType<BETA_T, float>::value) {
Adds(betaFp32, betaTyped, 0.0f, tileRows);
} else {
Cast(betaFp32, betaTyped, RoundMode::CAST_NONE, tileRows);
}
Bind beta under both element types and pick between them with a runtime flag
taken from tiling.betaDataType -- the same field the kernel entry was already
dispatching on. Three call shapes are affected: the five (arch35) / four
(generic) LoadAsFloatRow sites become LoadBetaAsFloatRow, the two staged
conversions become ConvertStagedBeta, and the fused-path staging copy branches
on the flag. Both branches now live in the binary at those sites, which costs
roughly 50 lines against a whole 4336-line instantiation.
With Prepare independent of the beta dtype, BETA_T drops out of RunFrontEnd,
RunGeneric and arch35::Run too, so the entry point's three
`if (betaDataType == 0) ... else ...` blocks collapse to single calls.
Instantiations of arch35 Prepare in the A5 translation unit: 8 -> 4
(SAFE_GATE x COMPILE_BT/K/V remain). Combined with the previous commit, total
weighted instantiated body lines drop from ~41.5k to ~20.8k.
The UB staging slot is now always sized for fp32. Checked that this changes
nothing observable: with K=128 both the old bf16 (3076 B/row) and fp32
(3080 B/row) layouts already saturate KDA_GATE_TILE_ROWS=16, so the tile row
count is unchanged; the bf16-beta path grows by 64 bytes of UB (49216 ->
49280) into the layout the fp32-beta path already used, leaving 14.7 KB of
headroom under the 192 KB budget.
The surrounding PipeBarrier / SetFlag / WaitFlag sequences are untouched.
Note for review: the staging path (ConvertStagedBeta and the fused copy) sits
inside `if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 &&
COMPILE_V == 128)`, which safe_gate=False callers never reach. Accuracy runs
covering that combination with bf16 beta are the ones that matter here.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
build_cmake_modules() unconditionally wipes the build directory, so every
invocation is a cold build -- including one where nothing changed. That is
the dominant cost while iterating, because the AscendC kernel targets
(chunk_kda_fwd above all) do not benefit from the ccache launcher configured
for the host compiler.
KEEP_BUILD_DIR=1 reuses the existing CMake cache and object files. Default
behaviour is unchanged, so nothing about CI or a release build moves; drop
the variable whenever the CMake configuration itself changes.
Also replace the hardcoded -j 16 with nproc (overridable via BUILD_JOBS).
bash build.sh -a kernels Ascend950PR_9599 # as before
KEEP_BUILD_DIR=1 bash build.sh -a kernels Ascend950PR_9599 # incremental
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Scaling -j to nproc was wrong: build parallelism here is bounded by RAM, not
by cores. The CCE compiler needs roughly 1-2 GB per job and the AscendC kernel
translation units are the hungry ones, so on a 28-core host with 15 GiB of RAM
(9 GiB actually available) -j 28 wants ~40 GB and drives the machine into swap
thrashing or the OOM killer. It hung a build at the no_workspace_kernel target.
The change also bought very little, since the long pole is a single serial
translation unit that extra jobs cannot help.
Restore the historical default and keep BUILD_JOBS for hosts that can take
more (or, as here, need less):
BUILD_JOBS=8 bash build.sh -a kernels Ascend950PR_9599
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
There was a problem hiding this comment.
🟡 Changes recommended
There are compile-/correctness-critical issues in newly added headers (missing required standard includes) and a vector-layout bug that computes a boolean size instead of a rounded-up length.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR adds a new fused chunk KDA forward operator for NPU, integrating a direct-launch AscendC kernel plus the necessary PyTorch extension registration and build-system wiring so it can be invoked from Python.
Changes:
- Added a Python convenience wrapper for
torch.ops.npu.chunk_kda_fwd. - Registered the new operator in the PyTorch extension and exposed a C++ API entry in the NPU ops header.
- Added a dedicated AscendC kernel target and supporting Catlass/TLA utility headers for the new kernel.
File summaries
| File | Description |
|---|---|
| python/sgl_kernel_npu/sgl_kernel_npu/fla/chunk_kda_fwd.py | Python wrapper for the new NPU op |
| include/sgl_kenel_npu_ops.h | Declares C++ entrypoint for chunk_kda_fwd |
| csrc/pytorch_extensions.cpp | Registers chunk_kda_fwd with Torch dispatcher |
| csrc/CMakeLists.txt | Builds/links chunk_kda_fwd host+kernel targets |
| csrc/chunk_kda_fwd/op_kernel/tla/type_traits.hpp | Adds TLA type traits utilities |
| csrc/chunk_kda_fwd/op_kernel/tla/tuple.hpp | Adds TLA tuple implementation |
| csrc/chunk_kda_fwd/op_kernel/tla/numeric/math.hpp | Adds TLA numeric helpers |
| csrc/chunk_kda_fwd/op_kernel/tla/numeric/integral_constant.hpp | Adds TLA integral constant utilities |
| csrc/chunk_kda_fwd/op_kernel/tla/numeric/integer_sequence.hpp | Adds TLA integer sequence utilities |
| csrc/chunk_kda_fwd/op_kernel/tla/int_tuple.hpp | Adds TLA int-tuple utilities (apply/transform/etc.) |
| csrc/chunk_kda_fwd/op_kernel/kernel_utils/vector/regbase.hpp | Adds vector register helper utilities |
| csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp | Kernel entrypoint + dispatch logic |
| csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_varlen.h | Varlen chunk-range resolution helper |
| csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_tiling_data.h | Direct-launch tiling-data struct definition |
| csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_h_runtime_tiling.h | Runtime tiling struct for fwd_h path |
| csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_vector.hpp | Catlass vector layout helpers for kernel |
| csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_tensor.hpp | Catlass tensor layout helpers |
| csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_layout.hpp | Catlass layout umbrella header |
| csrc/chunk_kda_fwd/op_kernel/catlass/kda_numeric_size.hpp | Numeric size utilities for Catlass |
| csrc/chunk_kda_fwd/op_kernel/catlass/kda_matrix_coord.hpp | Matrix coordinate helpers |
| csrc/chunk_kda_fwd/op_kernel/catlass/kda_gemm_coord.hpp | GEMM coordinate helpers |
| csrc/chunk_kda_fwd/op_kernel/catlass/kda_debug.hpp | Debug/check helpers for Catlass |
| csrc/chunk_kda_fwd/op_kernel/catlass/kda_coord.hpp | Core coordinate container utilities |
| csrc/chunk_kda_fwd/op_kernel/catlass/kda_catlass.hpp | Catlass core constants/macros |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_copy_tla.hpp | GEMM tile copy TLA interfaces |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_ub_to_gm.hpp | GEMM UB->GM copy selector header |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_l0b.hpp | GEMM L1->L0B copy selector header |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_l0a.hpp | GEMM L1->L0A copy selector header |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_fp.hpp | GEMM L1->FixPipe copy selector header |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_bt.hpp | GEMM L1->BiasTable copy selector header |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l0c_to_gm.hpp | GEMM L0C->GM copy selector header |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_gm_to_ub.hpp | GEMM GM->UB copy selector header |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_gm_to_l1.hpp | GEMM GM->L1 copy selector header |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_int8_to_fp16.hpp | Cast selector header (int8->fp16) |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_int4_to_int8.hpp | Cast selector header (int4->int8) |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_fp8_to_fp16.hpp | Cast selector header (fp8->fp16) |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/copy_l0c_to_ub.hpp | L0C->UB copy selector header |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_cast_int8_to_fp16.hpp | AtlasA2 cast implementation |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_ub_to_gm.hpp | AtlasA2 UB->GM copy implementation |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_fp.hpp | AtlasA2 L1->FP copy implementation |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_bt.hpp | AtlasA2 L1->BT copy implementation |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_ub.hpp | AtlasA2 GM->UB copy implementation |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_cast_int4_to_int8.hpp | AtlasA2 int4->int8 cast implementation |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_bt.hpp | Ascend950 L1->BT copy implementation |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_dst.hpp | Ascend950 L0C->Dst copy/quant plumbing |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/copy_l0c_to_ub.hpp | Ascend950 L0C->UB copy implementation |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_type.hpp | Catlass GEMM type definitions |
| csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad.hpp | Block MMAD interface + pingpong include |
| csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_epilogue_tile_copy.hpp | Epilogue tile-copy plumbing |
| csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_ub_to_gm.hpp | Epilogue UB->GM copy implementations |
| csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_ub_to_gm_tla.hpp | Epilogue UB->GM TLA implementations |
| csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_gm_to_ub.hpp | Epilogue GM->UB copy implementations |
| csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_gm_to_ub_tla.hpp | Epilogue GM->UB TLA implementations |
| csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/copy_ub_to_l1_tla.hpp | Epilogue UB->L1 TLA implementation (Ascend950) |
| csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue.hpp | Epilogue block interface |
| csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_tag_to_layout.hpp | Tag-to-layout mapping utilities |
| csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_macros.hpp | Catlass macros for device/host annotations |
| csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_dependent_false.hpp | Dependent-false helper for static_assert |
| csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_alignment.hpp | Alignment/rounding helpers |
| csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_resource.hpp | Resource buffers + pipe wrapper |
| csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_local_tensor_buffer.hpp | Local tensor buffer abstractions |
| csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_cross_core_sync.hpp | Cross-core barrier/flag helpers |
| csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_arch.hpp | Arch constants/types (AtlasA2/Ascend950) |
| csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/kda_fwd_h_epilogue_policies.hpp | Arch35 epilogue policy types |
| csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_impl.h | Arch35 dispatch implementation wrapper |
| csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/kda_fwd_h_epilogue_policies.hpp | Arch22 epilogue policy types |
| csrc/chunk_kda_fwd/op_host/arch35/chunk_kda_fwd_tiling_impl.h | Host tiling option selection for arch35 |
| build.sh | Adds KEEP_BUILD_DIR/BUILD_JOBS build controls |
Review details
- Files reviewed: 52/109 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+14
to
+16
| #include <tuple> | ||
|
|
||
| #define TLA_REQUIRES(...) typename std::enable_if<(__VA_ARGS__)>::type * = nullptr |
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
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.
No description provided.