Skip to content

sampling.cuh may be missing hipcub includes — the ROCm 10 gate from #4754 may be unnecessary #4819

Description

@zufayu

Follow-up to #4754, which disables cpp_itfs sampling on ROCm 10 after
sampling.cuh failed to compile with no type named 'Traits' in namespace 'hipcub'.

That PR was the right call to unblock — it is a no-op on every currently
supported ROCm and carries no regression risk. This issue records a hypothesis
we did not get to test before merging, so it isn't lost.

Hypothesis

The compile failure may be a missing #include on our side rather than a
hipCUB API removal. If so, ROCm 10 can keep the fast sampling kernels and the
gate becomes unnecessary.

Evidence

sampling.cuh never includes the headers for the symbols it uses.

includes:  hipcub/block/{block_adjacent_difference, block_load,
                         block_reduce, block_scan, block_store}.hpp
uses:      hipcub::Traits  (L87)
           hipcub::Min     (L505)
           hipcub::Max     (L1063)

hipcub/util_type.hpp and hipcub/thread/thread_operators.hpp are never
included directly — all three symbols arrive transitively today. Deleting one
#include on ROCm 7.0 reproduces the exact six diagnostics from #4754, same
lines and columns. A missing include and a removed API are indistinguishable
from the error text alone.

hipcub::Traits is still present upstream. ROCm/hipCUB@develop,
hipcub/include/hipcub/backend/rocprim/util_type.hpp:839:

template<typename T>
struct Traits : NumericTraits<typename std::remove_cv<T>::type>
{};

Not deprecated — HIPCUB_CLANG_SUPPRESS_DEPRECATED_POP closes above it, and
detail::is_primitive uses it immediately below. hipCUB's release tags also
stop at rocm-7.2.4; there is no ROCm 10 tag yet.

Caveat: this is the public hipCUB repo. The toolchain in #4754 came from the
_rocm_sdk_devel pip wheel, which may ship a different hipCUB.

How to settle it

Needs a ROCm 10 environment. The gate in #4754 reads AITER_ROCM_VERSION
before torch.version.hip, so it can be bypassed without reverting anything:

  1. Add to csrc/cpp_itfs/sampling/sampling.cuh, next to the other hipcub includes:
#include <hipcub/thread/thread_operators.hpp>
#include <hipcub/util_type.hpp>
  1. Bypass the gate and force a JIT build:
AITER_ROCM_VERSION=7 python -c "
import torch
from aiter.ops import sampling  # noqa: F401  (registers torch.ops.aiter.*)
probs = torch.softmax(torch.randn(8, 4096, device='cuda'), dim=-1).contiguous()
print(torch.ops.aiter.top_p_sampling_from_probs(probs, None, None, 0.9, deterministic=True))
"
  • Compiles and runs → missing include. Land the two #include lines and
    remove the gate; ROCm 10 keeps the fast kernels.
  • Still no type named 'Traits' → the API is gone. Close this; the gate in
    [Bugfix] Disable cpp_itfs sampling on ROCm 10 #4754 is the correct long-term answer.

Either way, worth confirming — the two headers are the correct direct
dependencies regardless of the outcome, and relying on block_*.hpp to pull in
Traits/Min/Max is fragile against any hipCUB header reshuffle.

Nit, separate from the above

The gate keys off torch.version.hip, which is the ROCm torch was built
against. sampling.cuh is compiled by whichever hipcc is on PATH.
compile_template_op already resolves that with get_hip_version()
(csrc/cpp_itfs/utils.py:113) and uses it for five flag decisions
(utils.py:216-228), including the -amdgpu-coerce-illegal-types=1 that shows
up in the #4754 log. The new gate is the only place in that call path asking
torch instead of the compiler.

Not urgent — the two agree in the environments we ship — but keying it off
get_hip_version() would be consistent with the rest of the file.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions