[module_mla_reduce] detorch reduce.cu + delete dead no_redundant Python port - #4796
Open
amd-ruitang3 wants to merge 1 commit into
Open
[module_mla_reduce] detorch reduce.cu + delete dead no_redundant Python port#4796amd-ruitang3 wants to merge 1 commit into
amd-ruitang3 wants to merge 1 commit into
Conversation
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
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.
Clean-rebuild wall for module_mla_reduce (2 TUs): ~243s (torch-included) → ~222–231s de-torched (~5–9%; device codegen of the 1269-line opus/CK kernel dominates, so torch-header removal is a modest slice)
1.
csrc/kernels/mla/reduce.cu— remove torch dependencyThe Python side (
develop=True), the C++ signature (mla_reduce.husesaiter_tensor_t) and the pybind (rocm_ops.hppis torch-free) were already de-torched in #4729; onlyreduce.custill pulled torch. Fixed the two remaining torch deps:TORCH_WARN(torch macro, host-side occupancy warning) →fprintf(stderr, ...). Same trigger frequency as before (the original usedTORCH_WARN, not_ONCE), and it only fires in the degraded high-splits regime — no kernel/GPU perf impact.std::optionalin themla_reduce_v1signature was satisfied only transitively by the torch headers; after dropping<ATen/hip/HIPContext.h>/<ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h>/<torch/python.h>it needs an explicit#include <optional>(+<cstdio>for the fprintf).2.
aiter/ops/attention.py— delete dead codeget_mla_metadata_v1_no_redundantwas reimplemented as a pure-Python port in #4729 but has zero in-repo callers and no tests — dead code. Removed it (and its now-unusedimport heapq).Why it matters (latent build break)
mla_reduce_v1is@compile_ops(develop=True), whose production build setstorch_exclude=True— this strips pytorch's-isystem .../torch/includefrom the compile. The originalreduce.cufails under that path:So this isn't only a build-time win — the de-torched TU is what makes the module actually buildable in the torch-excluded develop path.
Verification (in-container, gfx942 / MI308X)
op_tests/test_mla_reduce.py— ALL PASSED (all splits / head / dv combos)black+ruff check— cleanmodule_mla_reduce(2 TUs): ~243s (torch-included) → ~222–231s de-torched (~5–9%; device codegen of the 1269-line opus/CK kernel dominates, so torch-header removal is a modest slice)🤖 Generated with Claude Code