Skip to content

[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
ROCm:mainfrom
amd-ruitang3:detorch/topk-per-row-drop-py-itfs-common
Open

[detorch] remove vestigial py_itfs_common.h includes (non-CK); top_k_per_row now torch-free#4756
amd-ruitang3 wants to merge 2 commits into
ROCm:mainfrom
amd-ruitang3:detorch/topk-per-row-drop-py-itfs-common

Conversation

@amd-ruitang3

@amd-ruitang3 amd-ruitang3 commented Aug 14, 2026

Copy link
Copy Markdown
Contributor

Summary

Removes every vestigial #include "py_itfs_common.h" from the non-CK TUs
that 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 whole
torch/ATen header tree into a module's JIT build.

Two kinds of benefit:

  1. Immediate torch-free win — module_top_k_per_row. Its kernel TU was
    de-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 only
    aiter_tensor.h + py_itfs_common.h, so the header was their sole torch
    source. Removing it makes the entire module closure (13 headers) torch-free.
  2. Latent-regression removal — the rest. asm_mha_{fwd,bwd,varlen_fwd,varlen_bwd}.cu
    and fused_ar_mhc_post.cu keep their own direct torch includes (they are still
    torch-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.h would silently re-introduce <torch/all.h>.

Why each removal is safe

file py_itfs_common symbols used how it still gets what it needs
asm_topk_per_row_{decode,prefill}.cu none aiter_tensor.h (already included, line 3) pulls aiter_hip_common.hAITER_CHECK/HIP_CALL/HipDeviceGuard + HIP runtime; now torch-free
asm_mha_{fwd,bwd,varlen_fwd,varlen_bwd}.cu none torch/all.h + ATen/hip/HIPContext.h remain (load-bearing); getCurrentHIPStream here is at::hip::, so no aiter_hip_common symbol was in play
fused_ar_mhc_post.cu none already #include "aiter_hip_common.h" directly (line 8); ATen/c10 includes remain

py_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 match torch.topk; [mb_workspace_reuse] PASS.
    • op_tests/test_topk_row_prefill.py (prefill): standard and fast kernels Passed, all_close_standard/all_close_fast == True.
  • module_fmha_v3_fwd — fresh rebuild clean (39.4s). asm_mha_fwd.cu is include-identical to the other three asm_mha TUs.
  • module_fused_ar_mhc — fresh rebuild clean (36.5s).

🤖 Generated with Claude Code

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>
@amd-ruitang3
amd-ruitang3 requested a review from a team August 14, 2026 08:58
@github-actions

Copy link
Copy Markdown
Contributor

🏷️ CI Guide

Runs automatically on every PR:

  • ✅ Pre-checks (submodule verification, code formatting)
  • ✅ Aiter op tests (gfx942 + gfx950)
  • ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label Tests
ci:gfx1250-ffm-triton Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4756 --add-label <label>

@amd-ruitang3 amd-ruitang3 changed the title [detorch] top_k_per_row: drop vestigial py_itfs_common.h include (module now torch-free) [detorch] drop vestigial py_itfs_common.h include (module now torch-free) Aug 14, 2026
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>
@amd-ruitang3 amd-ruitang3 changed the title [detorch] drop vestigial py_itfs_common.h include (module now torch-free) [detorch] remove vestigial py_itfs_common.h includes (non-CK); top_k_per_row now torch-free Aug 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant