[detorch] remove vestigial py_itfs_common.h includes (non-CK); top_k_per_row now torch-free - #4756
Open
amd-ruitang3 wants to merge 2 commits into
Open
Conversation
asm_topk_per_row_{decode,prefill}.cu each carried an
`#include "py_itfs_common.h"` that was their ONLY torch source (that header
has an unguarded `#include <torch/all.h>`), yet neither TU uses any symbol
py_itfs_common.h exports. The kernel TU (topk_per_row_kernels.cu) was already
de-torched in ROCm#4702; these two ASM sibling TUs were missed, so the whole
module_top_k_per_row closure kept pulling the torch/ATen header tree on every
JIT build.
Both files already include aiter_tensor.h, which pulls aiter_hip_common.h
(providing AITER_CHECK / HIP_CALL / HipDeviceGuard + the HIP runtime) — so the
include is simply removed, no replacement needed. The full module closure
(13 headers) is now torch-free.
Verified in-container (gfx942): module_top_k_per_row rebuilds clean;
test_topk_per_row.py (decode) all shapes match torch.topk + mb_workspace_reuse
PASS; test_topk_row_prefill.py standard + fast kernels Passed (all_close=True).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
Follow-up to the topk cleanup: remove the vestigial `#include "py_itfs_common.h"` from the other non-CK TUs that include it but use none of its symbols. These files keep their own direct torch includes (they are still torch-bound modules), so this is not an immediate compile-time win — it removes a latent regression: when these modules are later de-torched, a leftover py_itfs_common.h (which has an unguarded `#include <torch/all.h>`) would silently re-introduce the torch header tree, exactly the trap that left module_top_k_per_row half-de-torched after ROCm#4702. Files: - py_itfs_cu/asm_mha_{fwd,bwd,varlen_fwd,varlen_bwd}.cu — torch/all.h + ATen/hip/HIPContext.h remain (load-bearing; ~44-74 torch tokens each); getCurrentHIPStream here is at::hip::, so no aiter_hip_common.h symbol was used. - kernels/fused_ar_mhc_post.cu — already includes aiter_hip_common.h directly (line 8) for AITER_CHECK/HIP_CALL; ATen/c10 includes remain. Verified in-container (gfx942): module_fmha_v3_fwd (39.4s) and module_fused_ar_mhc (36.5s) both rebuild clean from a cleared build dir. asm_mha_fwd.cu is include-identical to the other three asm_mha TUs. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.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
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.
Summary
Removes every vestigial
#include "py_itfs_common.h"from the non-CK TUsthat include it but reference none of its exported symbols. That header carries
an unguarded
#include <torch/all.h>, so a stray include silently drags the wholetorch/ATen header tree into a module's JIT build.
Two kinds of benefit:
module_top_k_per_row. Its kernel TU wasde-torched in [module_topk_*] de-torch topk_per_row / topk_plain + externalize workspaces #4702, but the two ASM sibling TUs
(
asm_topk_per_row_{decode,prefill}.cu) were missed — each included onlyaiter_tensor.h+py_itfs_common.h, so the header was their sole torchsource. Removing it makes the entire module closure (13 headers) torch-free.
asm_mha_{fwd,bwd,varlen_fwd,varlen_bwd}.cuand
fused_ar_mhc_post.cukeep their own direct torch includes (they are stilltorch-bound), so this is not a compile-time win today. But it deletes the exact
trap that left topk half-de-torched: once those modules are de-torched later, a
leftover
py_itfs_common.hwould silently re-introduce<torch/all.h>.Why each removal is safe
asm_topk_per_row_{decode,prefill}.cuaiter_tensor.h(already included, line 3) pullsaiter_hip_common.h→AITER_CHECK/HIP_CALL/HipDeviceGuard+ HIP runtime; now torch-freeasm_mha_{fwd,bwd,varlen_fwd,varlen_bwd}.cutorch/all.h+ATen/hip/HIPContext.hremain (load-bearing);getCurrentHIPStreamhere isat::hip::, so no aiter_hip_common symbol was in playfused_ar_mhc_post.cu#include "aiter_hip_common.h"directly (line 8); ATen/c10 includes remainpy_itfs_common.h's torch-typed exports (torch_fp8/torch_fp4x2/t2ck/torchDTypeToStr) are used only by CK TUs that need torch at the boundary anyway;those includes are left untouched (out of scope).
Validation (in-container, gfx942 / ROCm 7.2.3)
module_top_k_per_row— fresh JIT rebuild from a cleared build dir compiles clean; recursive include closure = 0 torch/ATen/c10 headers.op_tests/test_topk_per_row.py(decode): all shapes matchtorch.topk;[mb_workspace_reuse] PASS.op_tests/test_topk_row_prefill.py(prefill): standard and fast kernelsPassed,all_close_standard/all_close_fast == True.module_fmha_v3_fwd— fresh rebuild clean (39.4s).asm_mha_fwd.cuis include-identical to the other threeasm_mhaTUs.module_fused_ar_mhc— fresh rebuild clean (36.5s).🤖 Generated with Claude Code