Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
12 changes: 4 additions & 8 deletions csrc/rocm/w4a8_sdot4_rdna2.cu
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,7 @@
// 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.
Expand Down Expand Up @@ -67,9 +63,9 @@ 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.
// Production config: per-(token, group) activation scales, 32-wide K step,
// 8-row M tile. G1/G2 accuracy numbers ("per-(token, G=64)") were measured
// on this tile.
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).
Expand Down
114 changes: 5 additions & 109 deletions csrc/rocm/w4a8_sdot4_rdna2.cuh
Original file line number Diff line number Diff line change
@@ -1,16 +1,8 @@
// 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.
// gfx1030 W4A8 primitives (uint4 weights x int8 activations, v_dot4_i32_i8).
// Built for gfx1030; the dense GEMM launches a8_lds_k32_ag (config 8) only.
//
// Contract (T = ceil(M / M_TILE) row tiles):
// w [K/8, N] uint32 GPTQ K-packed, then gptq_shuffle: the
Expand Down Expand Up @@ -523,79 +515,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 +526,12 @@ 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:
// Dense config launched by w4a8_gemm_rdna2, instantiated at 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);
#define W4A8_EXPLORE_CONFIGS(X) X(8, "a8_lds_k32_ag", 256, 4, 32, 8, kLds, 1)