Skip to content

cleanup(rocm): trim dead W4A8 explore sweep from production copy #41

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
BlivionIaG wants to merge 1 commit into from
Closed

cleanup(rocm): trim dead W4A8 explore sweep from production copy #41

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 0 additions & 31 deletions csrc/rocm/moe_w4a8_rdna2.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,6 @@

#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <map>
#include <mutex>
Expand Down Expand Up @@ -321,30 +320,6 @@ bool on_gfx1030() {
return cached_ok;
}

bool w4a8_moe_debug_enabled() {
static const bool enabled = [] {
const char* v = std::getenv("VLLM_RDNA2_W4A8_MOE_DEBUG");
return v != nullptr && std::strcmp(v, "1") == 0;
}();
return enabled;
}

void w4a8_moe_log_shape(int m, int k, int n, int group, int block_m) {
static std::mutex mu;
static std::map<std::tuple<int, int, int, int, int>, bool> seen;
const std::tuple<int, int, int, int, int> key{m, k, n, group, block_m};
{
std::lock_guard<std::mutex> lock(mu);
if (!seen.emplace(key, true).second) {
return;
}
}
fprintf(stderr,
"[W4A8-MOE-DEBUG] fast path fired: M=%d K=%d N=%d G=%d bsm=%d\n", m,
k, n, group, block_m);
fflush(stderr);
}

// Stable per-(m, k, groups, device) activation workspace. Held for the life of
// the process so the device pointers captured into a graph never move; the
// dense path's per-call allocation is the suspected graph-aliasing source.
Expand Down Expand Up @@ -579,10 +554,4 @@ void moe_w4a8_gemm_rdna2(torch::Tensor a, torch::Tensor c,

TORCH_WARN_ONCE(
"RDNA2 W4A8 sdot4 MoE path active (config moe_a8_k32_ag)");
if (w4a8_moe_debug_enabled()) {
w4a8_moe_log_shape(static_cast<int>(size_m), static_cast<int>(size_k),
static_cast<int>(size_n),
static_cast<int>(group_size),
static_cast<int>(block_size_m));
}
}
78 changes: 6 additions & 72 deletions csrc/rocm/w4a8_sdot4_rdna2.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,22 +20,14 @@
// the SAME packed buffer RDNA2W4A16LinearKernel already produces
// (zero-extended nibbles + gptq_shuffle).
//
// The kernel body lives in the sibling header `w4a8_sdot4_rdna2.cuh`
// (copied verbatim from the explore-only tree at
// /tmp/pr9_branch/csrc/rocm/explore/w4a8_sdot4.cuh and renamed so the
// production C ABI and its header are co-located; the source of truth stays
// untouched per the "do not modify csrc/rocm/explore/* in place" rule).
// The kernel body lives in the sibling header `w4a8_sdot4_rdna2.cuh`.
// This TU is the production wrapper: the gemm entry owns the shape/LDS
// eligibility and its own internal W4A16 fallback, so Python never branches
// on a runtime value in the traced forward.

#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <mutex>
#include <set>
#include <tuple>

#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
Expand Down Expand Up @@ -67,10 +59,6 @@ enum Err : int {

constexpr int kMaxSplit = 16;

// Recommended config: per-(token, group) activation scales, 32-wide K step,
// 8-row M tile. This is the only A_GROUP tile among the explore sweep and the
// one the G1/G2 accuracy numbers ("per-(token, G=64)") were measured with.
constexpr int kDefaultConfigId = 8; // a8_lds_k32_ag
// Prefill-only: decode (M < kW4a8MinRows) keeps the W4A16 arms; the wired
// W4A8 config is a8_lds_k32_ag (M_TILE 8).
constexpr int kW4a8MinRows = 33;
Expand Down Expand Up @@ -101,34 +89,6 @@ int group_index(int64_t group_size) {
}
}

// Env-gated diagnostic: when VLLM_RDNA2_W4A8_DEBUG=1, log each distinct
// (m, k, n, group) shape that actually fires the W4A8 fast path, once per
// shape per process. Read the env once (thread-safe) so the traced forward
// never touches it.
bool w4a8_debug_enabled() {
static const bool enabled = [] {
const char* v = std::getenv("VLLM_RDNA2_W4A8_DEBUG");
return v != nullptr && std::strcmp(v, "1") == 0;
}();
return enabled;
}

void w4a8_log_shape(int m, int k, int n, int group) {
static std::mutex mu;
static std::set<std::tuple<int, int, int, int>> seen;
const std::tuple<int, int, int, int> key{m, k, n, group};
{
std::lock_guard<std::mutex> lock(mu);
if (!seen.insert(key).second) {
return;
}
}
fprintf(stderr,
"[W4A8-DEBUG] fast path fired: m=%d k=%d n=%d group=%d\n", m, k, n,
group);
fflush(stderr);
}

// Mirror of pick_split_k's LDS cap: some group-aligned split <= 16 must bring
// M_TILE=8 rows of K plus the per-(token, group) scales under 64 KiB of LDS.
//
Expand Down Expand Up @@ -245,34 +205,12 @@ int launch_gemm(const GemmArgs& p, hipStream_t stream) {
using LaunchFn = int (*)(const GemmArgs&, hipStream_t);
using SplitFn = int (*)(int, int, int);

struct ConfigEntry {
int id;
int m_tile;
int a_group;
LaunchFn launch[3]; // group 32, 64, 128
const LaunchFn kLaunchA8LdsK32Ag[3] = {
&launch_gemm<ex::Cfg<256, 4, 32, 8, 32, ex::ASrc::kLds, true>>,
&launch_gemm<ex::Cfg<256, 4, 32, 8, 64, ex::ASrc::kLds, true>>,
&launch_gemm<ex::Cfg<256, 4, 32, 8, 128, ex::ASrc::kLds, true>>,
};

#define W4A8_CFG(th, npt, ks, mt, g, src, ag) \
ex::Cfg<th, npt, ks, mt, g, ex::ASrc::src, (ag) != 0>
#define W4A8_ENTRY(id, name, th, npt, ks, mt, src, ag) \
{id, \
mt, \
ag, \
{&launch_gemm<W4A8_CFG(th, npt, ks, mt, 32, src, ag)>, \
&launch_gemm<W4A8_CFG(th, npt, ks, mt, 64, src, ag)>, \
&launch_gemm<W4A8_CFG(th, npt, ks, mt, 128, src, ag)>}},

const ConfigEntry kConfigs[] = {W4A8_EXPLORE_CONFIGS(W4A8_ENTRY)};

const ConfigEntry* find_config(int id) {
for (const ConfigEntry& c : kConfigs) {
if (c.id == id) {
return &c;
}
}
return nullptr;
}

// Thread t owns row (t % MT) of its block's row tile; the per-group activation
// scale variant (A_GROUP) needs no block reduction, so the launch is identical
// for both — only the a_scale layout differs (per token vs [T][K/G][MT]).
Expand Down Expand Up @@ -382,7 +320,6 @@ at::Tensor w4a8_gemm_rdna2(torch::Tensor a, torch::Tensor b_q_weight,
use_v2_format);
}

const ConfigEntry* c = find_config(kDefaultConfigId);
const int gi = group_index(group_size);
const GemmArgs p{static_cast<const int8_t*>(a_i8.data_ptr()),
static_cast<const uint32_t*>(b_q_weight.data_ptr()),
Expand All @@ -393,14 +330,11 @@ at::Tensor w4a8_gemm_rdna2(torch::Tensor a, torch::Tensor b_q_weight,
out.data_ptr(),
size_m, size_n, size_k, zero_offset,
/*split_k=*/1, /*out_f32=*/0};
if (c->launch[gi](p, stream) != 0) {
if (kLaunchA8LdsK32Ag[gi](p, stream) != 0) {
return gptq_gemm_rdna2_prefill(a, b_q_weight, b_qzeros, b_scales, b_g_idx,
use_v2_format);
}

TORCH_WARN_ONCE("RDNA2 W4A8 sdot4 path active (config a8_lds_k32_ag)");
if (w4a8_debug_enabled()) {
w4a8_log_shape(size_m, size_k, size_n, group_size);
}
return out;
}
112 changes: 0 additions & 112 deletions csrc/rocm/w4a8_sdot4_rdna2.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,17 +1,6 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Explore-only W4A8 draft for gfx1030: W uint4 x A int8 -> v_dot4_i32_i8.
//
// NOT WIRED: not in VLLM_ROCM_EXT_SRC, not registered in torch_bindings.cpp,
// not reachable from any Python dispatcher. The only consumers are
// w4a8_sdot4_capi.cu (a standalone .so for the V620 harness) and
// benchmarks/kernels/w4a8_sdot4_explore/isa_check.py.
// See docs/explore/w4a8-sdot4/DESIGN.md.
//
// Depends only on <cstddef>, <cstdint> and clang builtins, so the file
// compiles under hipcc and under plain clang with w4a8_sdot4_isa_shim.h.
//
// Contract (T = ceil(M / M_TILE) row tiles):
// w [K/8, N] uint32 GPTQ K-packed, then gptq_shuffle: the
// buffer the RDNA2 W4A16 kernels read.
Expand Down Expand Up @@ -523,79 +512,6 @@ __global__ __launch_bounds__(THREADS) void w4a8_act_quant_kernel(
}
}

// ---------------------------------------------------------------------------
// G0 peak probes: CHAINS independent accumulators per lane keep the loop
// issue-bound rather than latency-bound. The result is kept alive by a store
// that never fires in practice.
// ---------------------------------------------------------------------------
enum class Probe : int { kSdot4 = 0, kFdot2 = 1, kFmaF32 = 2 };

template <Probe P, int CHAINS>
__global__ __launch_bounds__(256) void w4a8_probe_kernel(int iters,
uint32_t seed,
uint32_t* out) {
const uint32_t t = thread_x() + block_x() * 256;
const uint32_t a = seed ^ (t * 2654435761u);
const uint32_t b = seed * 40503u + t;
uint32_t sig = 0;
if constexpr (P == Probe::kSdot4) {
int32_t acc[CHAINS];
#pragma unroll
for (int c = 0; c < CHAINS; ++c) {
acc[c] = c;
}
for (int i = 0; i < iters; ++i) {
#pragma unroll
for (int c = 0; c < CHAINS; ++c) {
acc[c] = sdot4(a + c, b, acc[c]);
}
}
#pragma unroll
for (int c = 0; c < CHAINS; ++c) {
sig ^= static_cast<uint32_t>(acc[c]);
}
} else if constexpr (P == Probe::kFdot2) {
const f16x2_t x = __builtin_bit_cast(f16x2_t, a & 0x3BFF3BFFu);
const f16x2_t y = __builtin_bit_cast(f16x2_t, b & 0x3BFF3BFFu);
float acc[CHAINS];
#pragma unroll
for (int c = 0; c < CHAINS; ++c) {
acc[c] = static_cast<float>(c);
}
for (int i = 0; i < iters; ++i) {
#pragma unroll
for (int c = 0; c < CHAINS; ++c) {
acc[c] = __builtin_amdgcn_fdot2(x, y, acc[c], /*clamp=*/false);
}
}
#pragma unroll
for (int c = 0; c < CHAINS; ++c) {
sig ^= __builtin_bit_cast(uint32_t, acc[c]);
}
} else {
const float x = static_cast<float>(a & 0xFFu) * 1e-3f;
const float y = static_cast<float>(b & 0xFFu) * 1e-3f;
float acc[CHAINS];
#pragma unroll
for (int c = 0; c < CHAINS; ++c) {
acc[c] = static_cast<float>(c);
}
for (int i = 0; i < iters; ++i) {
#pragma unroll
for (int c = 0; c < CHAINS; ++c) {
acc[c] = __builtin_fmaf(x, y, acc[c]);
}
}
#pragma unroll
for (int c = 0; c < CHAINS; ++c) {
sig ^= __builtin_bit_cast(uint32_t, acc[c]);
}
}
if (sig == 0x9E3779B9u) {
out[t] = sig;
}
}

#else // non-gfx1030 device pass: same signatures, no body.

template <class C>
Expand All @@ -607,35 +523,7 @@ template <int THREADS, int MT, bool PER_GROUP>
__global__ __launch_bounds__(THREADS) void w4a8_act_quant_kernel(
const f16_t*, int64_t, int8_t*, float*, int32_t*, int, int, int) {}

enum class Probe : int { kSdot4 = 0, kFdot2 = 1, kFmaF32 = 2 };

template <Probe P, int CHAINS>
__global__ __launch_bounds__(256) void w4a8_probe_kernel(int, uint32_t,
uint32_t*) {}

#endif // __gfx1030__ || !__HIP_DEVICE_COMPILE__

} // namespace explore_w4a8
} // namespace vllm

// Explore sweep, instantiated for group sizes 32, 64 and 128:
// X(id, name, THREADS, NPT, K_STEP, M_TILE, A_SRC, A_GROUP)
// Shared by w4a8_sdot4_capi.cu (dispatch table) and isa_check.py.
#define W4A8_EXPLORE_CONFIGS(X) \
X(0, "a16_lds_k32", 256, 4, 32, 16, kLds, 0) /* ConfigA-class */ \
X(1, "a16_lds_k16", 256, 4, 16, 16, kLds, 0) \
X(2, "a16_smem_k16", 256, 4, 16, 16, kSmem, 0) /* "LDS=0" */ \
X(3, "a8_smem_k32", 256, 4, 32, 8, kSmem, 0) \
X(4, "a8_lds_k32", 256, 4, 32, 8, kLds, 0) /* fewer VGPRs */ \
X(5, "a32n2_lds_k32", 256, 2, 32, 32, kLds, 0) /* 2x W reuse */ \
X(6, "c16_lds_k32", 128, 4, 32, 16, kLds, 0) /* ConfigC-class */ \
X(7, "a16_lds_k32_ag", 256, 4, 32, 16, kLds, 1) /* per-group A scale */ \
X(8, "a8_lds_k32_ag", 256, 4, 32, 8, kLds, 1)

// Explicit instantiation for the ISA check (the .so instantiates through its
// launches instead).
#define W4A8_EXPLORE_INSTANTIATE_GEMM(...) \
template __global__ void vllm::explore_w4a8::w4a8_gemm_kernel<__VA_ARGS__>( \
const int8_t*, const uint32_t*, const uint32_t*, \
const vllm::explore_w4a8::f16_t*, const float*, const int32_t*, void*, \
int, int, int, int, int, int, int);
1 change: 0 additions & 1 deletion docs/rdna2/w4a8-moe-wiring.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,6 @@ default is unchanged. Full logs: `bench_results/2026-09-29_w4a8-moe/`.
| Env | Effect |
|---|---|
| `VLLM_RDNA2_W4A8_SDOT4=1` | Opt in to the MoE W4A8 path. Default OFF. |
| `VLLM_RDNA2_W4A8_MOE_DEBUG=1` | Per-shape `[W4A8-MOE-DEBUG]` log. |
| (anything else) | No change vs `rdna_extras`. |

`VLLM_RDNA2_W4A8_SDOT4` already exists in `vllm/envs.py` for the dense path;
Expand Down
13 changes: 6 additions & 7 deletions docs/rdna2/w4a8-wiring.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,19 +4,18 @@ Branch: `w4a8-wiring` (= `rdna_extras` + the production wrapper around the
explore kernel from PR #9). Default OFF — the dispatcher takes the existing
`gptq_gemm_rdna2_prefill` path unless `VLLM_RDNA2_W4A8_SDOT4=1`.

The kernel body (`w4a8_gemm_kernel`, `w4a8_act_quant_kernel`, `Cfg`,
`W4A8_EXPLORE_CONFIGS`, etc.) is copied verbatim from
`/tmp/pr9_branch/csrc/rocm/explore/w4a8_sdot4.cuh` into the sibling
`csrc/rocm/w4a8_sdot4_rdna2.cuh`. The explore source of truth stays
untouched per the "do not modify `csrc/rocm/explore/*` in place" rule — the
fork never had an `explore/` directory and this branch does not create one.
The kernel body (`w4a8_gemm_kernel`, `w4a8_act_quant_kernel`, `Cfg`)
lives in `csrc/rocm/w4a8_sdot4_rdna2.cuh`. Production instantiates only
`a8_lds_k32_ag` (group sizes 32, 64, and 128). The explore source of truth
stays untouched per the "do not modify `csrc/rocm/explore/*` in place" rule —
the fork never had an `explore/` directory and this branch does not create one.

## What changed

| File | Purpose |
|---|---|
| `csrc/rocm/w4a8_sdot4_rdna2.cu` | Production TU. `w4a8_gemm_rdna2` is a self-contained drop-in for `gptq_gemm_rdna2_prefill` with the same signature: allocates the int8 A buffer, A scales, A group sums, and the output internally; tries the `a8_lds_k32_ag` W4A8 fast path; falls back to `gptq_gemm_rdna2_prefill` whenever the shape/LDS budget is not eligible. `w4a8_act_quant_rdna2` is kept for tests (returns an empty tensor on ineligibility). One-time `TORCH_WARN_ONCE` marker when the fast path fires. |
| `csrc/rocm/w4a8_sdot4_rdna2.cuh` | Verbatim copy of the explore header (renamed). `W4A8_EXPLORE_CONFIGS` and the kernel templates are unchanged. |
| `csrc/rocm/w4a8_sdot4_rdna2.cuh` | Production GEMM and activation-quant kernels. Live launch is `a8_lds_k32_ag`. |
| `csrc/rocm/ops.h` | Declarations for the two ops. `w4a8_gemm_rdna2` mirrors `gptq_gemm_rdna2_prefill`'s signature; `w4a8_act_quant_rdna2` keeps the standalone-test signature. |
| `csrc/rocm/torch_bindings.cpp` | Registers both ops with the new signatures; the C++ Meta stubs are gone (the fake path is `register_fake` in `_custom_ops.py`). |
| `vllm/_custom_ops.py` | Wrappers `w4a8_act_quant_rdna2` / `w4a8_gemm_rdna2` with `register_fake` paths that match the new contract. |
Expand Down
Loading