From bde531d7204d68deb9d7b6aad5e5cde4633086fa Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Mon, 24 Aug 2026 02:34:37 +0900 Subject: [PATCH 1/9] fix(rocm): fall back for single-bank fast index copy --- python/freetoken/kernel/fast_index_copy.py | 44 ++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/python/freetoken/kernel/fast_index_copy.py b/python/freetoken/kernel/fast_index_copy.py index 1aaa1303..05b6a4f8 100644 --- a/python/freetoken/kernel/fast_index_copy.py +++ b/python/freetoken/kernel/fast_index_copy.py @@ -22,6 +22,42 @@ def _skip_fast_index_copy_enabled() -> bool: return os.getenv(SKIP_FAST_INDEX_COPY_ENV, "").strip().lower() in _TRUE_VALUES +def _is_rocm() -> bool: + return getattr(torch.version, "hip", None) is not None + + +def _rocm_index_copy_fallback( + dst: torch.Tensor, + dst_indices: torch.Tensor, + src: torch.Tensor, + src_indices: torch.Tensor, + num_indices: torch.Tensor | None = None, +) -> None: + """Correctness-first ROCm fallback for the CUDA-specific copy JIT. + + The native fast-index-copy header still contains NVIDIA inline PTX and CUDA + DLPack device matchers. Until that kernel has a HIP implementation, keep + ROCm functional by gathering only the requested source rows and moving that + bounded selection to the destination device. CUDA continues to use the + existing JIT unchanged. + """ + count = dst_indices.numel() if num_indices is None else int(num_indices.item()) + assert 0 <= count <= dst_indices.numel() + assert count <= src_indices.numel() + if count == 0: + return + + src_index = src_indices[:count].to(device=src.device, dtype=torch.long) + dst_index = dst_indices[:count].to(device=dst.device, dtype=torch.long) + rows = src.index_select(0, src_index) + if rows.device != dst.device: + rows = rows.to( + device=dst.device, + non_blocking=src.device.type == "cpu" and src.is_pinned(), + ) + dst.index_copy_(0, dst_index, rows) + + @lru_cache(maxsize=None) def _jit_update_flag_module() -> Module: return load_jit( @@ -114,6 +150,14 @@ def fast_index_copy_jit( dst = dst.as_strided(size=(dst.size(0), num_dst_feature), stride=(num_dst_feature, 1)) src = src.as_strided(size=(src.size(0), num_src_feature), stride=(num_src_feature, 1)) + if _is_rocm(): + if priority is not None: + raise NotImplementedError( + "ROCm fast-index-copy fallback does not implement high/normal priority scheduling" + ) + _rocm_index_copy_fallback(dst, dst_indices, src, src_indices, num_indices) + return + feature_size = dst.size(-1) * dst.element_size() num_block = num_block or DEFAULT_NUM_BLOCKS worker_threads = worker_threads or _default_worker_threads(feature_size) From 9b5c79b75d803eacef28119c606f22b80ed749b2 Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Mon, 24 Aug 2026 02:34:52 +0900 Subject: [PATCH 2/9] test(rocm): cover fast index copy fallback --- tests/kernels/test_fast_index_copy_rocm.py | 64 ++++++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 tests/kernels/test_fast_index_copy_rocm.py diff --git a/tests/kernels/test_fast_index_copy_rocm.py b/tests/kernels/test_fast_index_copy_rocm.py new file mode 100644 index 00000000..e73795a3 --- /dev/null +++ b/tests/kernels/test_fast_index_copy_rocm.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import pytest +import torch + +from freetoken.kernel import fast_index_copy as fast_copy + + +def test_rocm_fallback_copies_only_requested_rows() -> None: + src = torch.arange(20, dtype=torch.float32).reshape(5, 4).to(torch.bfloat16) + dst = torch.full((4, 4), -1, dtype=torch.bfloat16) + src_indices = torch.tensor([4, 1, 3], dtype=torch.int32) + dst_indices = torch.tensor([2, 0, 3], dtype=torch.int32) + num_indices = torch.tensor([2], dtype=torch.int64) + + fast_copy._rocm_index_copy_fallback( + dst, + dst_indices, + src, + src_indices, + num_indices, + ) + + torch.testing.assert_close(dst[2], src[4], rtol=0, atol=0) + torch.testing.assert_close(dst[0], src[1], rtol=0, atol=0) + torch.testing.assert_close(dst[1], torch.full((4,), -1, dtype=torch.bfloat16), rtol=0, atol=0) + torch.testing.assert_close(dst[3], torch.full((4,), -1, dtype=torch.bfloat16), rtol=0, atol=0) + + +def test_rocm_dispatch_does_not_build_cuda_jit(monkeypatch: pytest.MonkeyPatch) -> None: + src = torch.arange(12, dtype=torch.float32).reshape(3, 4) + dst = torch.zeros((3, 4), dtype=torch.float32) + src_indices = torch.tensor([2, 0], dtype=torch.int32) + dst_indices = torch.tensor([1, 2], dtype=torch.int32) + + monkeypatch.setattr(fast_copy, "_is_rocm", lambda: True) + + def fail_jit(**_kwargs): + raise AssertionError("ROCm dispatch must not compile the CUDA fast-index-copy JIT") + + monkeypatch.setattr(fast_copy, "_jit_fast_index_copy_module", fail_jit) + + fast_copy.fast_index_copy_jit(dst, dst_indices, src, src_indices) + + torch.testing.assert_close(dst[1], src[2], rtol=0, atol=0) + torch.testing.assert_close(dst[2], src[0], rtol=0, atol=0) + + +def test_rocm_priority_mode_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None: + src = torch.zeros((2, 4), dtype=torch.float32) + dst = torch.zeros((2, 4), dtype=torch.float32) + indices = torch.tensor([0], dtype=torch.int32) + + monkeypatch.setattr(fast_copy, "_is_rocm", lambda: True) + + with pytest.raises(NotImplementedError, match="priority scheduling"): + fast_copy.fast_index_copy_jit( + dst, + indices, + src, + indices, + priority="high", + sync_flag=torch.zeros((1,), dtype=torch.int32), + ) From a6cea4cadd2b3e2991193e5c576ee0edd241e7ee Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Mon, 24 Aug 2026 02:35:13 +0900 Subject: [PATCH 3/9] test(rocm): use explicit fast copy module import --- tests/kernels/test_fast_index_copy_rocm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/kernels/test_fast_index_copy_rocm.py b/tests/kernels/test_fast_index_copy_rocm.py index e73795a3..37f0d5f6 100644 --- a/tests/kernels/test_fast_index_copy_rocm.py +++ b/tests/kernels/test_fast_index_copy_rocm.py @@ -3,7 +3,7 @@ import pytest import torch -from freetoken.kernel import fast_index_copy as fast_copy +import freetoken.kernel.fast_index_copy as fast_copy def test_rocm_fallback_copies_only_requested_rows() -> None: From 9203d884493d60467c4ff94350aac39968f02095 Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Mon, 24 Aug 2026 03:27:34 +0900 Subject: [PATCH 4/9] fix(rocm): complete RDNA3 runtime path --- .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 486 +----------------- .../csrc/include/freetoken/hip_compat.h | 24 + .../kernel/csrc/include/freetoken/utils.cuh | 6 + .../kernel/csrc/jit/fast_index_copy.cuh | 24 + python/freetoken/kernel/fast_index_copy.py | 44 -- python/freetoken/kernel/triton/activation.py | 10 +- python/freetoken/kernel/triton/e4m3_compat.py | 4 + python/freetoken/kernel/triton/norm.py | 8 +- setup.py | 3 + tests/kernels/test_fast_index_copy_rocm.py | 64 --- 10 files changed, 79 insertions(+), 594 deletions(-) delete mode 100644 tests/kernels/test_fast_index_copy_rocm.py diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index 880e8637..ad397fdd 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -1,3 +1,6 @@ +Warning: truncated output (original token count: 25647) +Total output lines: 2150 + // CPU-compute MoE executor for the "cpu" offload backend. // // Decode ships activations to the CPU, computes the routed experts here (reading @@ -29,7 +32,7 @@ #include #include -#include +#include #include #if defined(__linux__) @@ -826,486 +829,7 @@ float dot_dsfp4_avx512(const uint8_t* packed, const uint8_t* scale, const float* // AVX2: a 32-K block is 16 bytes -> two 8-lane halves (8 even + 8 odd each). __attribute__((target("avx2,fma"))) inline __m256 dsfp4_half_avx2(const uint8_t* pk, const float* xeb, const float* xob, __m256 mag8) { - __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); - __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); - __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); - return _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); -} - -__attribute__((target("avx2,fma"))) -float dot_dsfp4_avx2(const uint8_t* packed, const uint8_t* scale, const float* xe, - const float* xo, int K, const float* e2m1, const float* e8m0) { - const __m256 mag8 = _mm256_loadu_ps(e2m1); - __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* pk = packed + (size_t)b * 16; - const float* xeb = xe + (size_t)b * 16; - const float* xob = xo + (size_t)b * 16; - const __m256 sc = _mm256_set1_ps(e8m0[scale[b]]); - acc0 = _mm256_fmadd_ps(dsfp4_half_avx2(pk, xeb, xob, mag8), sc, acc0); - acc1 = _mm256_fmadd_ps(dsfp4_half_avx2(pk + 8, xeb + 8, xob + 8, mag8), sc, acc1); - } - return hsum256(_mm256_add_ps(acc0, acc1)); -} -#endif - -dsdot_fn select_dsdot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - if (t >= ISA_AVX512) return dot_dsfp4_avx512; - if (t >= ISA_AVX2) return dot_dsfp4_avx2; -#endif - (void)t; - return dot_dsfp4_scalar; -} - -// ------------------------- mxfp4 (gpt-oss) GEMV ----------------------------- -// Transposed split-K layout: blk[Kpairs, N2] (N innermost), scl[Kpairs/16, N2] -// e8m0 per 32-K. Computes out[c] = sum_kb (E2M1[lo]*x[2kb] + E2M1[hi]*x[2kb+1]) -// * 2^(e8m0-127) for a contiguous column tile (blk/scl already offset to col 0 of -// the tile). Vectorized over N (16 columns / __m512), K stays the outer (cache- -// sequential) loop. Used by both gate_up (K=H) and down (K=I). -using mxgemv_fn = void (*)(float*, const uint8_t*, const uint8_t*, const bf16_t*, int, int, - int, const float*, const float*); - -void mxfp4_gemv_scalar(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, - int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { - for (int c = 0; c < ncol; ++c) out[c] = 0.0f; - for (int kb = 0; kb < Kpairs; ++kb) { - const uint8_t* w = blk + (size_t)kb * N2; - const uint8_t* s = scl + (size_t)(kb >> 4) * N2; - const float xl = bf16_to_f32(x[2 * kb]); - const float xh = bf16_to_f32(x[2 * kb + 1]); - for (int c = 0; c < ncol; ++c) { - const uint8_t byte = w[c]; - out[c] += (e2m1[byte & 0xF] * xl + e2m1[byte >> 4] * xh) * e8m0[s[c]]; - } - } -} - -#if CPU_MOE_X86 -__attribute__((target("avx512f"))) -void mxfp4_gemv_avx512(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, - int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { - (void)e8m0; // e8m0[c]=2^(c-127) computed via bit construction (no gather) - const __m512 lut = _mm512_loadu_ps(e2m1); - const __m512i loma = _mm512_set1_epi32(0xF); - // K-outer / N-inner: each kb cache line is read once and all live column chunks - // (up to 4 -> 64 cols) accumulate from registers, so DRAM/L2 stream the tile once. - int c0 = 0; - for (; c0 + 16 <= ncol; c0 += 64) { - const int nchunk = std::min(4, (ncol - c0) / 16); - __m512 acc[4]; - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_setzero_ps(); - for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = 32 K = one scale row - __m512 sc[4]; - for (int ci = 0; ci < nchunk; ++ci) { - __m128i sraw = _mm_loadu_si128(reinterpret_cast( - scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 16)); - sc[ci] = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu8_epi32(sraw), 23)); - } - __m512 blk_acc[4]; - for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm512_setzero_ps(); - for (int kk = 0; kk < 16; ++kk) { - const int kb = kblk + kk; - const uint8_t* wbase = blk + (size_t)kb * N2 + c0; - // The transposed layout strides K by N2 bytes; prefetch ahead so the strided - // reads are not exposed to DRAM latency (the HW streamer misses big strides). - constexpr int PFD = 8; - if (kb + PFD < Kpairs) - _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), - _MM_HINT_T0); - const __m512 xl = _mm512_set1_ps(bf16_to_f32(x[2 * kb])); - const __m512 xh = _mm512_set1_ps(bf16_to_f32(x[2 * kb + 1])); - for (int ci = 0; ci < nchunk; ++ci) { - __m512i wi = _mm512_cvtepu8_epi32( - _mm_loadu_si128(reinterpret_cast(wbase + ci * 16))); - __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); - __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); - blk_acc[ci] = _mm512_fmadd_ps(vlo, xl, blk_acc[ci]); - blk_acc[ci] = _mm512_fmadd_ps(vhi, xh, blk_acc[ci]); - } - } - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); - } - for (int ci = 0; ci < nchunk; ++ci) _mm512_storeu_ps(out + c0 + ci * 16, acc[ci]); - } - for (int c = c0; c < ncol; ++c) { // tail columns (< 16) - float o = 0.0f; - for (int kb = 0; kb < Kpairs; ++kb) { - const uint8_t byte = blk[(size_t)kb * N2 + c]; - uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; - float sc; - std::memcpy(&sc, &bits, 4); - o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + - e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; - } - out[c] = o; - } -} - -__attribute__((target("avx2,fma"))) -void mxfp4_gemv_avx2(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, - int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { - (void)e8m0; // e8m0[s]=2^(s-127) built via s<<23 (no gather) - const __m256 mag8 = _mm256_loadu_ps(e2m1); - int c0 = 0; - for (; c0 + 8 <= ncol; c0 += 32) { // up to 4 chunks of 8 = 32 cols - const int nchunk = std::min(4, (ncol - c0) / 8); - __m256 acc[4]; - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_setzero_ps(); - for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = one scale row - __m256 sc[4]; - for (int ci = 0; ci < nchunk; ++ci) { - __m128i sraw = _mm_loadl_epi64(reinterpret_cast( - scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 8)); - sc[ci] = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu8_epi32(sraw), 23)); - } - __m256 blk_acc[4]; - for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm256_setzero_ps(); - for (int kk = 0; kk < 16; ++kk) { - const int kb = kblk + kk; - const uint8_t* wbase = blk + (size_t)kb * N2 + c0; - constexpr int PFD = 8; - if (kb + PFD < Kpairs) - _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), - _MM_HINT_T0); - const __m256 xl = _mm256_set1_ps(bf16_to_f32(x[2 * kb])); - const __m256 xh = _mm256_set1_ps(bf16_to_f32(x[2 * kb + 1])); - for (int ci = 0; ci < nchunk; ++ci) { - __m256i wi = _mm256_cvtepu8_epi32( - _mm_loadl_epi64(reinterpret_cast(wbase + ci * 8))); - __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); - __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); - blk_acc[ci] = _mm256_fmadd_ps(vlo, xl, blk_acc[ci]); - blk_acc[ci] = _mm256_fmadd_ps(vhi, xh, blk_acc[ci]); - } - } - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); - } - for (int ci = 0; ci < nchunk; ++ci) _mm256_storeu_ps(out + c0 + ci * 8, acc[ci]); - } - for (int c = c0; c < ncol; ++c) { // tail columns (< 8); none when ncol%8==0 - float o = 0.0f; - for (int kb = 0; kb < Kpairs; ++kb) { - const uint8_t byte = blk[(size_t)kb * N2 + c]; - uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; - float sc; - std::memcpy(&sc, &bits, 4); - o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + - e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; - } - out[c] = o; - } -} -#endif - -mxgemv_fn select_mxgemv() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - if (t >= ISA_AVX512) return mxfp4_gemv_avx512; - if (t >= ISA_AVX2) return mxfp4_gemv_avx2; -#endif - (void)t; - return mxfp4_gemv_scalar; -} - -// Round a clamped |x|<=448 to nearest float8-e4m3 (RNE), back to fp32. Matches -// torch.float8_e4m3fn / triton .to(float8e4nv). -inline float e4m3_round(float x) { - const float sign = x < 0.0f ? -1.0f : 1.0f; - const float a = std::fabs(x); - if (a == 0.0f) return 0.0f; - if (a >= 448.0f) return sign * 448.0f; - int e; - std::frexp(a, &e); // a in [2^(e-1), 2^e) - float step = std::ldexp(1.0f, e - 4); - const float min_step = std::ldexp(1.0f, -9); // e4m3 subnormal step (2^-9) - if (step < min_step) step = min_step; - float r = std::nearbyint(a / step) * step; - if (r > 448.0f) r = 448.0f; - return sign * r; -} - -// IEEE ceil(log2(v)) for v>0 (matches dsv4 _log2_ceil / fast_round_scale). -inline int ceil_log2_pos(float v) { - uint32_t bits; - std::memcpy(&bits, &v, sizeof(bits)); - const int exp = (int)((bits >> 23) & 0xFF); - const int man = (int)(bits & 0x7FFFFF); - return exp - 127 + (man != 0 ? 1 : 0); -} - -// Split an interleaved bf16 row into fp32 even/odd halves (even[m]=src[2m]). -// bf16->fp32 is exact, so this only reorders -- done once per token/route and -// reused across every output row of the GEMV. -inline void deinterleave_bf16_f32(const bf16_t* src, float* even, float* odd, int K) { - for (int m = 0; m < K / 2; ++m) { - even[m] = bf16_to_f32(src[2 * m]); - odd[m] = bf16_to_f32(src[2 * m + 1]); - } -} - -// DeepSeek-V4 activation FP8 round-trip (bf16 in/out): per 128-block, -// s = 2^ceil(log2(max(|x|,1e-4)/448)); y = round_e4m3(clamp(x/s,+-448)) * s. -void fp8_roundtrip_bf16(const bf16_t* src, bf16_t* dst, int K) { - for (int b0 = 0; b0 < K; b0 += 128) { - const int b1 = std::min(K, b0 + 128); - float amax = 1e-4f; - for (int i = b0; i < b1; ++i) amax = std::max(amax, std::fabs(bf16_to_f32(src[i]))); - const float s = std::ldexp(1.0f, ceil_log2_pos(amax * (1.0f / 448.0f))); - const float inv_s = 1.0f / s; - for (int i = b0; i < b1; ++i) { - float q = bf16_to_f32(src[i]) * inv_s; - q = std::min(448.0f, std::max(-448.0f, q)); - dst[i] = f32_to_bf16(e4m3_round(q) * s); - } - } -} - -// --------------------------------- executor --------------------------------- - -struct CpuMoeExecutor; - -struct MoeTask { - CpuMoeExecutor* exec; - int layer_id; - int num_tokens; - const bf16_t* x; // [num_tokens, H] - const int32_t* ids; // [num_tokens, top_k] (raw expert ids; <0 = skip) - const float* w; // [num_tokens, top_k] - bf16_t* y; // [num_tokens, H] -}; - -// Output-row tiling. Small enough to give every worker independent work even at -// batch size 1; large enough to amortize the atomic work-grab. -// -// Bandwidth notes (Sapphire Rapids 8480+, 13 cores): the two passes already read -// every expert weight byte exactly once per token (each output row block is owned -// by one worker), and x stays hot in L1 across a (token,expert)'s rows -- so the -// kernel is single-read bandwidth-optimal at bs=1 (~205 GB/s vs ~55 GB/s PCIe). -// One worker per *physical* core, pinned, is the sweet spot; SMT oversubscription -// thrashes the spin-barrier. Deferred (not worth it here / for this workload): -// - AMX-bf16: a GEMM tile engine; decode is M=1 GEMV so tiles sit idle. It would -// only pay off in a grouped/batched (dedup) path. -// - expert dedup for bs>1: read each distinct expert once and GEMM its tokens. -// Helps locality+bytes when bs is large; decode batches here are tiny (<=4). -// - NUMA: a single node is assumed. Multi-socket machines would split each -// expert's K dimension per node (banks are already per-row contiguous). -constexpr int IBLK = 32; -constexpr int HBLK = 32; - -// -------------------------------- Q4_0 (W4A8) -------------------------------- -// Native GGUF Q4_0 experts (gemma4 GGUF): per-32 block = fp16 scale d + 16 packed -// bytes; byte j holds element j in its low nibble and j+16 in its high nibble, so a -// block's storage order is [lo0..lo15, hi0..hi15] and w = (nibble - 8) * d. Matches -// the reference dequant (models/gguf/dequant.py) and the packed banks the GPU offload -// path streams. -// -// llama.cpp ggml_vec_dot_q4_0_q8_0: W4A8. The activation is pre-quantized to Q8_0 -// (per-32-block int8 ``aq`` + fp32 scale ``asb``); each block unpacks its 16 bytes to -// 32 int8 weights in [-8,7] (bytes_from_nibbles_32: low nibbles -> elems 0..15, high -// -> 16..31) and runs an integer block dot -- VPDPBUSD (AVX-VNNI) or VPMADDUBSW+VPMADDWD -// (AVX2) with the ggml sign trick |w|*(sign(w)*a)=w*a, or a scalar int loop -- then -// scales the block sum by wd*xd in fp32. No fp weight dequant / shuffle chain. The GPU -// offload path (ggml_moe_a8_vec / MMVQ) is also W4A8, so cpu and hybrid stay close. -using q4dot_fn = float (*)(const uint8_t*, const int8_t*, const float*, int); - -float q4_0_dot_i8_scalar(const uint8_t* w, const int8_t* aq, const float* asb, int K) { - float acc = 0.0f; - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* blk = w + (size_t)b * 18; - uint16_t dh; - std::memcpy(&dh, blk, sizeof(dh)); - const uint8_t* q = blk + 2; // 16 nibble bytes - const int8_t* a = aq + (size_t)b * 32; - int isum = 0; - for (int j = 0; j < 16; ++j) { - isum += ((int)(q[j] & 0x0F) - 8) * (int)a[j]; // elem j - isum += ((int)(q[j] >> 4) - 8) * (int)a[16 + j]; // elem 16+j - } - acc += fp16_to_f32(dh) * asb[b] * (float)isum; - } - return acc; -} - -#if CPU_MOE_X86 -// fp16 block scale -> fp32 via HW F16C (single value in lane 0). -__attribute__((target("f16c"))) -static inline float q4_scale(uint16_t h) { - return _mm_cvtss_f32(_mm_cvtph_ps(_mm_cvtsi32_si128((int)h))); -} - -// Unpack one Q4_0 block's 16 bytes -> 32 int8 weights in [-8,7] (elems 0..15 = low -// nibbles, 16..31 = high nibbles). ``eight`` = _mm256_set1_epi8(8). -__attribute__((target("avx2"))) -static inline __m256i q4_unpack32(const uint8_t* blk, __m128i mask, __m256i eight) { - const __m128i qb = _mm_loadu_si128(reinterpret_cast(blk + 2)); - const __m128i lo = _mm_and_si128(qb, mask); - const __m128i hi = _mm_and_si128(_mm_srli_epi16(qb, 4), mask); - return _mm256_sub_epi8(_mm256_set_m128i(hi, lo), eight); -} - -// AVX2 W4A8 (llama.cpp non-VNNI mul_sum_i8_pairs): integer block dot via VPMADDUBSW + -// VPMADDWD (sign trick), scaled by wd*xd. |aw*sa| pair sums <= 8*127*2 < 32767 -> no -// int16 saturation. This is the fast path on AVX2 CPUs without AVX-VNNI (and the -// avx512-tier fallback, since the block dot is 256-bit either way). -__attribute__((target("avx2,fma,f16c"))) -float q4_0_dot_i8_avx2(const uint8_t* w, const int8_t* aq, const float* asb, int K) { - const __m128i mask = _mm_set1_epi8(0x0F); - const __m256i eight = _mm256_set1_epi8(8); - const __m256i ones16 = _mm256_set1_epi16(1); - __m256 accF = _mm256_setzero_ps(); - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* blk = w + (size_t)b * 18; - _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); - uint16_t dh; - std::memcpy(&dh, blk, sizeof(dh)); - __m256i wq = q4_unpack32(blk, mask, eight); - __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); - __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) - __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) - __m256i d32 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, sa), ones16); // 8 int32 - accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d32), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); - } - return hsum256(accF); -} - -// AVX-VNNI W4A8: one VPDPBUSD per block (the fast path on modern CPUs). -__attribute__((target("avx2,avxvnni,fma,f16c"))) -float q4_0_dot_i8_vnni(const uint8_t* w, const int8_t* aq, const float* asb, int K) { - const __m128i mask = _mm_set1_epi8(0x0F); - const __m256i eight = _mm256_set1_epi8(8); - __m256 accF = _mm256_setzero_ps(); - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* blk = w + (size_t)b * 18; - _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); - uint16_t dh; - std::memcpy(&dh, blk, sizeof(dh)); - __m256i wq = q4_unpack32(blk, mask, eight); - __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); - __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) - __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) - __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); - // All 32 elems of the block share wd*xd; distribute over di's 8 partial sums and - // reduce at the end (equivalent to scale * block_total). - accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); - } - return hsum256(accF); -} -#endif // CPU_MOE_X86 - -// All tiers are W4A8 (int8 activations pre-quantized to Q8_0). AVX-VNNI is orthogonal to -// the ISA tier (gated by cpu_has_avxvnni() / FREETOKEN_CPU_MOE_NO_VNNI), so it wins when -// present; otherwise the 256-bit VPMADDUBSW kernel covers both the avx2 and avx512 tiers. -q4dot_fn select_q4dot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - if (cpu_has_avxvnni()) return q4_0_dot_i8_vnni; - if (t >= ISA_AVX2) return q4_0_dot_i8_avx2; -#endif - (void)t; - return q4_0_dot_i8_scalar; -} - -enum WFmt { WF_BF16 = 0, WF_NVFP4 = 1, WF_MXFP4 = 2, WF_DSFP4 = 3, WF_Q4_0 = 4 }; - -// Each ctor pointer arg is the address of a CPU int64 array of length -// num_layers (one base address per layer, built by cpu_executor.py's -// _make_table), not a single flat bank. tbl_at resolves -// tbl[layer_id] once per task/pass; a null table (bank unused by this fmt, ptr -// arg 0) resolves to nullptr without dereferencing. -inline const void* tbl_at(const uint64_t* tbl, int layer_id) { - return tbl ? reinterpret_cast(tbl[layer_id]) : nullptr; -} - -struct CpuMoeExecutor { - int num_threads; - int num_layers, num_experts, top_k; - int H, I; - int act, apply_on_input; - int fmt; // WFmt - bool needs_di = false; // pre-deinterleave activations to fp32 (nvfp4/ds_fp4) - // Per-layer pointer tables (one base address per layer, see tbl_at). gate_up_tbl - // doubles as the bf16 gate_up table and the nvfp4/mxfp4/q4_0/ds_fp4 packed-gate_up - // table (down_tbl likewise for down); which reinterpretation applies is picked by - // fmt at each resolve site (see gemm1_dot/gemm2_dot/do_pass1_mxfp4/do_pass1_dsfp4). - const uint64_t* gate_up_tbl; // bf16: [E,2I,H] rows; else: packed e2m1/mxfp4-blocks - const uint64_t* down_tbl; // bf16: [E,H,I] rows; else: packed e2m1/mxfp4-blocks - const uint64_t* gu_scale_tbl; // nvfp4/mxfp4/ds_fp4: [E,2I,*] block scales - const uint64_t* gu_global_tbl; // nvfp4: [E,2I] fp16 row globals - const uint64_t* dn_scale_tbl; // nvfp4/mxfp4/ds_fp4: [E,H,*] block scales - const uint64_t* dn_global_tbl; // nvfp4: [E,H] fp16 row globals - const uint64_t* gu_bias_tbl; // mxfp4: [E,2I] bf16 biases - const uint64_t* dn_bias_tbl; // mxfp4: [E,H] bf16 biases - float swiglu_alpha; - float swiglu_limit; // +inf == no clamp - dot_fn dot; - nvdot_fn nvdot; - nvi8dot_fn nvi8dot = nullptr; // AVX-VNNI W4A8 nvfp4 dot (nullptr -> use fp32 nvdot) - bool use_vnni = false; // nvfp4 + AVX-VNNI: decode via int8 VPDPBUSD (W4A8) - bool use_q4a8 = false; // q4_0: always W4A8 (llama.cpp Q4_0 x Q8_0); int8 pre-quant - dsdot_fn dsdot; - mxgemv_fn mxgemv; - q4dot_fn q4dot; - // ds_fp4: the caller already FP8-round-tripped the input activations on the GPU - // (same reference grid), so submit() must not repeat it on the host-callback - // thread. That scalar per-element pass is single-threaded ON THE DECODE CRITICAL - // PATH (~0.3ms/layer at H=4096, every worker and the GPU waiting on it); moving - // it to a captured GPU elementwise kernel removes it while keeping the official - // W4A8 numerics bit-exact. Set via set_input_prequant (see cpu_executor.py). - bool input_prequant = false; - // Q4_0 packed-row byte strides (H/32*18 for gate_up over K=H, I/32*18 for down over K=I). - int q4_gu_row_bytes = 0, q4_dn_row_bytes = 0; - float e2m1_lut[16]; - float e4m3_lut[256]; - float e8m0_lut[256]; // mxfp4 block scale: 2^(s-127), s clamped to [0,254] - const char* isa; - - std::vector g_scratch; // [max_tokens * top_k * I] intermediate - std::vector xq_scratch; // [max_tokens * H] ds_fp4 fp8-roundtripped input - // ds_fp4 activations pre-deinterleaved to fp32 (even/odd K) for the row-major dot. - std::vector xe_scratch, xo_scratch; // [max_tokens * H/2] (input) - std::vector ge_scratch, go_scratch; // [max_tokens*top_k*I/2] (intermediate) - // AVX-VNNI W4A8: per-16-block int8 activations [even(8),odd(8)] + per-block scale. - std::vector xi8_scratch, gi8_scratch; // [max_tokens*H], [max_tokens*top_k*I] - std::vector xas_scratch, gas_scratch; // [max_tokens*H/16], [..*top_k*I/16] - std::string isa_str; - - std::vector workers; - std::mutex task_mtx; - std::condition_variable task_cv; - std::mutex sync_mtx; - std::condition_variable sync_cv; - - bool stop = false; - uint64_t cur_gen = 0; - MoeTask* cur_task = nullptr; - std::atomic submitted{0}; - std::atomic completed{0}; - - std::atomic p1_next{0}; - std::atomic p2_next{0}; - std::atomic prt_next{0}; // ds_fp4 intermediate fp8 round-trip phase - int64_t p1_total = 0, p2_total = 0, prt_total = 0; - int n_iblk = 0, n_hblk = 0; - std::atomic done_count{0}; - std::atomic bar_count{0}; - std::atomic bar_sense{0}; - - std::vector owned_tasks; // persistent task descriptors (graph-stable) - std::vector core_ids; // worker tid -> logical CPU to pin to (may be empty) - - // ---- Flag-based GPU<->CPU handshake (replaces the per-layer cudaLaunchHostFunc pair) ---- - // A tiny GPU kernel bumps ready_flags[slot] at submit; this coordinator thread busy-polls - // it, runs the slot's task on the worker pool, and sets done_flags[slot], which a GPU - // spin-wait kernel polls at sync. This removes the ~2x30-50us host-func dispatch round - // trips per MoE layer per decode step that otherwise idle the GPU (~6 ms/step on a - // 75-layer model). One slot per (layer, decode batch size) pair -- the Python side + __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cas…5647 tokens truncated…h size) pair -- the Python side // allocates slots as tasks are created. Flags live in mapped-pinned host memory (UVA: // the same pointers are used by the GPU kernels and by this thread). std::thread coord_thread; diff --git a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h index 99539d1f..00a1b540 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h +++ b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h @@ -19,6 +19,14 @@ #include #include +#ifndef CUDART_CB +#define CUDART_CB +#endif + +#ifndef __grid_constant__ +#define __grid_constant__ +#endif + // --- API name mapping (CUDA -> HIP) --- // HIP already defines most cuda* names as macros that expand to hip* equivalents // via hip_runtime.h, but a few are missing or differ in signature. Define them @@ -52,6 +60,14 @@ #define cudaHostAlloc hipHostMalloc #endif +#ifndef cudaHostAllocPortable +#define cudaHostAllocPortable hipHostMallocPortable +#endif + +#ifndef cudaHostAllocMapped +#define cudaHostAllocMapped hipHostMallocMapped +#endif + #ifndef cudaHostRegister #define cudaHostRegister hipHostRegister #endif @@ -122,6 +138,14 @@ #define cudaStream_t hipStream_t #endif +#ifndef cudaStreamSynchronize +#define cudaStreamSynchronize hipStreamSynchronize +#endif + +#ifndef cudaLaunchHostFunc +#define cudaLaunchHostFunc hipLaunchHostFunc +#endif + #ifndef dim3 // HIP already provides dim3; this is a no-op guard. #endif diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index 8e917832..f21b9585 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -1,5 +1,6 @@ #pragma once +#include #include #include @@ -115,6 +116,10 @@ public: } auto with_attr(bool use_pdl) -> LaunchKernel & { +#ifdef __HIP__ + (void)use_pdl; + m_config.numAttrs = 0; +#else if (use_pdl) { m_attr_cache.id = ::cudaLaunchAttributeProgrammaticStreamSerialization; m_attr_cache.val.programmaticStreamSerializationAllowed = 1; @@ -123,6 +128,7 @@ public: } else { m_config.numAttrs = 0; } +#endif return *this; } diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index bb83c23e..fe3f6be4 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -34,40 +34,64 @@ inline constexpr auto get_mem_package() { } __always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 { +#ifdef __HIP_PLATFORM_AMD__ + return *src; +#else uint32_t tmp; asm volatile("ld.global.L1::no_allocate.b32 %0,[%1];" : "=r"(tmp) : "l"(src)); return uint1{tmp}; +#endif } __always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 { +#ifdef __HIP_PLATFORM_AMD__ + return *src; +#else uint32_t tmp0, tmp1; asm volatile("ld.global.L1::no_allocate.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src)); return uint2{tmp0, tmp1}; +#endif } __always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 { +#ifdef __HIP_PLATFORM_AMD__ + return *src; +#else uint32_t tmp0, tmp1, tmp2, tmp3; asm volatile("ld.global.L1::no_allocate.v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src)); return uint4{tmp0, tmp1, tmp2, tmp3}; +#endif } __always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& value) { +#ifdef __HIP_PLATFORM_AMD__ + *dst = value; +#else uint32_t tmp = value.x; asm volatile("st.global.wt.b32 [%0],%1;" ::"l"(dst), "r"(tmp)); +#endif } __always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) { +#ifdef __HIP_PLATFORM_AMD__ + *dst = value; +#else uint32_t tmp0 = value.x; uint32_t tmp1 = value.y; asm volatile("st.global.wt.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1)); +#endif } __always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& value) { +#ifdef __HIP_PLATFORM_AMD__ + *dst = value; +#else uint32_t tmp0 = value.x; uint32_t tmp1 = value.y; uint32_t tmp2 = value.z; uint32_t tmp3 = value.w; asm volatile("st.global.wt.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3)); +#endif } __always_inline __device__ void wait_flag_clear(const int32_t* __restrict__ flag_ptr) { diff --git a/python/freetoken/kernel/fast_index_copy.py b/python/freetoken/kernel/fast_index_copy.py index 05b6a4f8..1aaa1303 100644 --- a/python/freetoken/kernel/fast_index_copy.py +++ b/python/freetoken/kernel/fast_index_copy.py @@ -22,42 +22,6 @@ def _skip_fast_index_copy_enabled() -> bool: return os.getenv(SKIP_FAST_INDEX_COPY_ENV, "").strip().lower() in _TRUE_VALUES -def _is_rocm() -> bool: - return getattr(torch.version, "hip", None) is not None - - -def _rocm_index_copy_fallback( - dst: torch.Tensor, - dst_indices: torch.Tensor, - src: torch.Tensor, - src_indices: torch.Tensor, - num_indices: torch.Tensor | None = None, -) -> None: - """Correctness-first ROCm fallback for the CUDA-specific copy JIT. - - The native fast-index-copy header still contains NVIDIA inline PTX and CUDA - DLPack device matchers. Until that kernel has a HIP implementation, keep - ROCm functional by gathering only the requested source rows and moving that - bounded selection to the destination device. CUDA continues to use the - existing JIT unchanged. - """ - count = dst_indices.numel() if num_indices is None else int(num_indices.item()) - assert 0 <= count <= dst_indices.numel() - assert count <= src_indices.numel() - if count == 0: - return - - src_index = src_indices[:count].to(device=src.device, dtype=torch.long) - dst_index = dst_indices[:count].to(device=dst.device, dtype=torch.long) - rows = src.index_select(0, src_index) - if rows.device != dst.device: - rows = rows.to( - device=dst.device, - non_blocking=src.device.type == "cpu" and src.is_pinned(), - ) - dst.index_copy_(0, dst_index, rows) - - @lru_cache(maxsize=None) def _jit_update_flag_module() -> Module: return load_jit( @@ -150,14 +114,6 @@ def fast_index_copy_jit( dst = dst.as_strided(size=(dst.size(0), num_dst_feature), stride=(num_dst_feature, 1)) src = src.as_strided(size=(src.size(0), num_src_feature), stride=(num_src_feature, 1)) - if _is_rocm(): - if priority is not None: - raise NotImplementedError( - "ROCm fast-index-copy fallback does not implement high/normal priority scheduling" - ) - _rocm_index_copy_fallback(dst, dst_indices, src, src_indices, num_indices) - return - feature_size = dst.size(-1) * dst.element_size() num_block = num_block or DEFAULT_NUM_BLOCKS worker_threads = worker_threads or _default_worker_threads(feature_size) diff --git a/python/freetoken/kernel/triton/activation.py b/python/freetoken/kernel/triton/activation.py index 2c38b533..0b7c945c 100644 --- a/python/freetoken/kernel/triton/activation.py +++ b/python/freetoken/kernel/triton/activation.py @@ -20,8 +20,9 @@ import triton.language as tl from triton.language.extra import libdevice from triton.language.extra.cuda import gdc_wait, gdc_launch_dependents +from triton.language import target_info -from freetoken.utils.arch import is_sm90_supported +from freetoken.utils.arch import is_rocm, is_sm90_supported SILU = 0 GELU = 1 @@ -48,6 +49,8 @@ def _pdl_supported() -> bool: @triton.jit def _fast_tanh(x): + if target_info.is_hip(): + return libdevice.tanh(x) # PTX tanh.approx.f32 — single HW op, matches flashinfer math::tanh. return tl.inline_asm_elementwise( "tanh.approx.f32 $0, $1;", "=f,f", [x], @@ -57,6 +60,8 @@ def _fast_tanh(x): @triton.jit def _fast_ex2(x): + if target_info.is_hip(): + return libdevice.exp2(x) # PTX ex2.approx.f32 — matches __expf fast path used by flashinfer silu. return tl.inline_asm_elementwise( "ex2.approx.f32 $0, $1;", "=f,f", [x], @@ -134,8 +139,9 @@ def _act_and_mul( block_d = min(triton.next_power_of_2(d), 1024 if M >= 4096 else 512) num_stages = 2 if block_d == 1024 else 3 _act_and_mul_kernel[grid]( - o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=pdl, launch_pdl=pdl, + o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=pdl, BLOCK_D=block_d, num_warps=4, num_stages=num_stages, + **({} if is_rocm() else {"launch_pdl": pdl}), ) return out diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 61d3a0e7..e52095cc 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/python/freetoken/kernel/triton/e4m3_compat.py @@ -59,6 +59,10 @@ def e4m3_native() -> bool: if _native is None: if FORCE_EMU: _native = False + elif torch.version.hip is not None: + # ROCm reports gfx1101 as capability (11, 0), which is not a CUDA + # compute capability and must not select the native fp8e4nv path. + _native = False else: native = {torch.cuda.get_device_capability(i) >= (8, 9) for i in range(torch.cuda.device_count())} diff --git a/python/freetoken/kernel/triton/norm.py b/python/freetoken/kernel/triton/norm.py index 3f95c29f..9071e1df 100644 --- a/python/freetoken/kernel/triton/norm.py +++ b/python/freetoken/kernel/triton/norm.py @@ -30,7 +30,7 @@ import triton.language as tl from triton.language.extra.cuda import gdc_launch_dependents, gdc_wait -from freetoken.utils.arch import is_sm90_supported +from freetoken.utils.arch import is_rocm, is_sm90_supported _HEUR = {"BLOCK": lambda a: triton.next_power_of_2(a["H"])} @@ -144,7 +144,8 @@ def _rmsnorm(input, weight, eps, out, gemma: bool): pdl = contig and is_sm90_supported() _rmsnorm_kernel[(A, B)]( out, input, weight, eps, H, sxa, sxb, soa, sob, - CONTIG=contig, ENABLE_PDL=pdl, launch_pdl=pdl, GEMMA=gemma, + CONTIG=contig, ENABLE_PDL=pdl, GEMMA=gemma, + **({} if is_rocm() else {"launch_pdl": pdl}), num_warps=_num_warps(A * B), num_stages=1, ) return out @@ -172,7 +173,8 @@ def _fused_add_rmsnorm(input, residual, weight, eps, gemma: bool): pdl = contig and is_sm90_supported() _fused_add_rmsnorm_kernel[(A, B)]( input, residual, weight, eps, H, sxa, sxb, sra, srb, - CONTIG=contig, ENABLE_PDL=pdl, launch_pdl=pdl, GEMMA=gemma, + CONTIG=contig, ENABLE_PDL=pdl, GEMMA=gemma, + **({} if is_rocm() else {"launch_pdl": pdl}), num_warps=_num_warps(A * B), num_stages=1, ) diff --git a/setup.py b/setup.py index 8ba0640a..698ad780 100644 --- a/setup.py +++ b/setup.py @@ -9,6 +9,7 @@ ROOT = Path(__file__).parent +KERNEL_INCLUDE = ROOT / "python" / "freetoken" / "kernel" / "csrc" / "include" def _check_toolchain() -> None: @@ -61,6 +62,8 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: runtime_lib = "cudart" extra_compile = ["-O3", "-std=c++17"] +runtime_include_dirs.append(str(KERNEL_INCLUDE)) + _check_toolchain() diff --git a/tests/kernels/test_fast_index_copy_rocm.py b/tests/kernels/test_fast_index_copy_rocm.py deleted file mode 100644 index 37f0d5f6..00000000 --- a/tests/kernels/test_fast_index_copy_rocm.py +++ /dev/null @@ -1,64 +0,0 @@ -from __future__ import annotations - -import pytest -import torch - -import freetoken.kernel.fast_index_copy as fast_copy - - -def test_rocm_fallback_copies_only_requested_rows() -> None: - src = torch.arange(20, dtype=torch.float32).reshape(5, 4).to(torch.bfloat16) - dst = torch.full((4, 4), -1, dtype=torch.bfloat16) - src_indices = torch.tensor([4, 1, 3], dtype=torch.int32) - dst_indices = torch.tensor([2, 0, 3], dtype=torch.int32) - num_indices = torch.tensor([2], dtype=torch.int64) - - fast_copy._rocm_index_copy_fallback( - dst, - dst_indices, - src, - src_indices, - num_indices, - ) - - torch.testing.assert_close(dst[2], src[4], rtol=0, atol=0) - torch.testing.assert_close(dst[0], src[1], rtol=0, atol=0) - torch.testing.assert_close(dst[1], torch.full((4,), -1, dtype=torch.bfloat16), rtol=0, atol=0) - torch.testing.assert_close(dst[3], torch.full((4,), -1, dtype=torch.bfloat16), rtol=0, atol=0) - - -def test_rocm_dispatch_does_not_build_cuda_jit(monkeypatch: pytest.MonkeyPatch) -> None: - src = torch.arange(12, dtype=torch.float32).reshape(3, 4) - dst = torch.zeros((3, 4), dtype=torch.float32) - src_indices = torch.tensor([2, 0], dtype=torch.int32) - dst_indices = torch.tensor([1, 2], dtype=torch.int32) - - monkeypatch.setattr(fast_copy, "_is_rocm", lambda: True) - - def fail_jit(**_kwargs): - raise AssertionError("ROCm dispatch must not compile the CUDA fast-index-copy JIT") - - monkeypatch.setattr(fast_copy, "_jit_fast_index_copy_module", fail_jit) - - fast_copy.fast_index_copy_jit(dst, dst_indices, src, src_indices) - - torch.testing.assert_close(dst[1], src[2], rtol=0, atol=0) - torch.testing.assert_close(dst[2], src[0], rtol=0, atol=0) - - -def test_rocm_priority_mode_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None: - src = torch.zeros((2, 4), dtype=torch.float32) - dst = torch.zeros((2, 4), dtype=torch.float32) - indices = torch.tensor([0], dtype=torch.int32) - - monkeypatch.setattr(fast_copy, "_is_rocm", lambda: True) - - with pytest.raises(NotImplementedError, match="priority scheduling"): - fast_copy.fast_index_copy_jit( - dst, - indices, - src, - indices, - priority="high", - sync_flag=torch.zeros((1,), dtype=torch.int32), - ) From 4e11ed0a1f23e17af49ba70c31cee5336fb33ab0 Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Mon, 24 Aug 2026 03:28:27 +0900 Subject: [PATCH 5/9] fix(rocm): preserve CPU MoE implementation --- .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 484 +++++++++++++++++- 1 file changed, 480 insertions(+), 4 deletions(-) diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index ad397fdd..56ab93df 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -1,6 +1,3 @@ -Warning: truncated output (original token count: 25647) -Total output lines: 2150 - // CPU-compute MoE executor for the "cpu" offload backend. // // Decode ships activations to the CPU, computes the routed experts here (reading @@ -829,7 +826,486 @@ float dot_dsfp4_avx512(const uint8_t* packed, const uint8_t* scale, const float* // AVX2: a 32-K block is 16 bytes -> two 8-lane halves (8 even + 8 odd each). __attribute__((target("avx2,fma"))) inline __m256 dsfp4_half_avx2(const uint8_t* pk, const float* xeb, const float* xob, __m256 mag8) { - __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cas…5647 tokens truncated…h size) pair -- the Python side + __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); + __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); + __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); + return _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); +} + +__attribute__((target("avx2,fma"))) +float dot_dsfp4_avx2(const uint8_t* packed, const uint8_t* scale, const float* xe, + const float* xo, int K, const float* e2m1, const float* e8m0) { + const __m256 mag8 = _mm256_loadu_ps(e2m1); + __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* pk = packed + (size_t)b * 16; + const float* xeb = xe + (size_t)b * 16; + const float* xob = xo + (size_t)b * 16; + const __m256 sc = _mm256_set1_ps(e8m0[scale[b]]); + acc0 = _mm256_fmadd_ps(dsfp4_half_avx2(pk, xeb, xob, mag8), sc, acc0); + acc1 = _mm256_fmadd_ps(dsfp4_half_avx2(pk + 8, xeb + 8, xob + 8, mag8), sc, acc1); + } + return hsum256(_mm256_add_ps(acc0, acc1)); +} +#endif + +dsdot_fn select_dsdot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (t >= ISA_AVX512) return dot_dsfp4_avx512; + if (t >= ISA_AVX2) return dot_dsfp4_avx2; +#endif + (void)t; + return dot_dsfp4_scalar; +} + +// ------------------------- mxfp4 (gpt-oss) GEMV ----------------------------- +// Transposed split-K layout: blk[Kpairs, N2] (N innermost), scl[Kpairs/16, N2] +// e8m0 per 32-K. Computes out[c] = sum_kb (E2M1[lo]*x[2kb] + E2M1[hi]*x[2kb+1]) +// * 2^(e8m0-127) for a contiguous column tile (blk/scl already offset to col 0 of +// the tile). Vectorized over N (16 columns / __m512), K stays the outer (cache- +// sequential) loop. Used by both gate_up (K=H) and down (K=I). +using mxgemv_fn = void (*)(float*, const uint8_t*, const uint8_t*, const bf16_t*, int, int, + int, const float*, const float*); + +void mxfp4_gemv_scalar(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, + int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { + for (int c = 0; c < ncol; ++c) out[c] = 0.0f; + for (int kb = 0; kb < Kpairs; ++kb) { + const uint8_t* w = blk + (size_t)kb * N2; + const uint8_t* s = scl + (size_t)(kb >> 4) * N2; + const float xl = bf16_to_f32(x[2 * kb]); + const float xh = bf16_to_f32(x[2 * kb + 1]); + for (int c = 0; c < ncol; ++c) { + const uint8_t byte = w[c]; + out[c] += (e2m1[byte & 0xF] * xl + e2m1[byte >> 4] * xh) * e8m0[s[c]]; + } + } +} + +#if CPU_MOE_X86 +__attribute__((target("avx512f"))) +void mxfp4_gemv_avx512(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, + int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { + (void)e8m0; // e8m0[c]=2^(c-127) computed via bit construction (no gather) + const __m512 lut = _mm512_loadu_ps(e2m1); + const __m512i loma = _mm512_set1_epi32(0xF); + // K-outer / N-inner: each kb cache line is read once and all live column chunks + // (up to 4 -> 64 cols) accumulate from registers, so DRAM/L2 stream the tile once. + int c0 = 0; + for (; c0 + 16 <= ncol; c0 += 64) { + const int nchunk = std::min(4, (ncol - c0) / 16); + __m512 acc[4]; + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_setzero_ps(); + for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = 32 K = one scale row + __m512 sc[4]; + for (int ci = 0; ci < nchunk; ++ci) { + __m128i sraw = _mm_loadu_si128(reinterpret_cast( + scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 16)); + sc[ci] = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu8_epi32(sraw), 23)); + } + __m512 blk_acc[4]; + for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm512_setzero_ps(); + for (int kk = 0; kk < 16; ++kk) { + const int kb = kblk + kk; + const uint8_t* wbase = blk + (size_t)kb * N2 + c0; + // The transposed layout strides K by N2 bytes; prefetch ahead so the strided + // reads are not exposed to DRAM latency (the HW streamer misses big strides). + constexpr int PFD = 8; + if (kb + PFD < Kpairs) + _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), + _MM_HINT_T0); + const __m512 xl = _mm512_set1_ps(bf16_to_f32(x[2 * kb])); + const __m512 xh = _mm512_set1_ps(bf16_to_f32(x[2 * kb + 1])); + for (int ci = 0; ci < nchunk; ++ci) { + __m512i wi = _mm512_cvtepu8_epi32( + _mm_loadu_si128(reinterpret_cast(wbase + ci * 16))); + __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); + __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); + blk_acc[ci] = _mm512_fmadd_ps(vlo, xl, blk_acc[ci]); + blk_acc[ci] = _mm512_fmadd_ps(vhi, xh, blk_acc[ci]); + } + } + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); + } + for (int ci = 0; ci < nchunk; ++ci) _mm512_storeu_ps(out + c0 + ci * 16, acc[ci]); + } + for (int c = c0; c < ncol; ++c) { // tail columns (< 16) + float o = 0.0f; + for (int kb = 0; kb < Kpairs; ++kb) { + const uint8_t byte = blk[(size_t)kb * N2 + c]; + uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; + float sc; + std::memcpy(&sc, &bits, 4); + o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + + e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; + } + out[c] = o; + } +} + +__attribute__((target("avx2,fma"))) +void mxfp4_gemv_avx2(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, + int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { + (void)e8m0; // e8m0[s]=2^(s-127) built via s<<23 (no gather) + const __m256 mag8 = _mm256_loadu_ps(e2m1); + int c0 = 0; + for (; c0 + 8 <= ncol; c0 += 32) { // up to 4 chunks of 8 = 32 cols + const int nchunk = std::min(4, (ncol - c0) / 8); + __m256 acc[4]; + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_setzero_ps(); + for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = one scale row + __m256 sc[4]; + for (int ci = 0; ci < nchunk; ++ci) { + __m128i sraw = _mm_loadl_epi64(reinterpret_cast( + scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 8)); + sc[ci] = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu8_epi32(sraw), 23)); + } + __m256 blk_acc[4]; + for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm256_setzero_ps(); + for (int kk = 0; kk < 16; ++kk) { + const int kb = kblk + kk; + const uint8_t* wbase = blk + (size_t)kb * N2 + c0; + constexpr int PFD = 8; + if (kb + PFD < Kpairs) + _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), + _MM_HINT_T0); + const __m256 xl = _mm256_set1_ps(bf16_to_f32(x[2 * kb])); + const __m256 xh = _mm256_set1_ps(bf16_to_f32(x[2 * kb + 1])); + for (int ci = 0; ci < nchunk; ++ci) { + __m256i wi = _mm256_cvtepu8_epi32( + _mm_loadl_epi64(reinterpret_cast(wbase + ci * 8))); + __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); + __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); + blk_acc[ci] = _mm256_fmadd_ps(vlo, xl, blk_acc[ci]); + blk_acc[ci] = _mm256_fmadd_ps(vhi, xh, blk_acc[ci]); + } + } + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); + } + for (int ci = 0; ci < nchunk; ++ci) _mm256_storeu_ps(out + c0 + ci * 8, acc[ci]); + } + for (int c = c0; c < ncol; ++c) { // tail columns (< 8); none when ncol%8==0 + float o = 0.0f; + for (int kb = 0; kb < Kpairs; ++kb) { + const uint8_t byte = blk[(size_t)kb * N2 + c]; + uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; + float sc; + std::memcpy(&sc, &bits, 4); + o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + + e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; + } + out[c] = o; + } +} +#endif + +mxgemv_fn select_mxgemv() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (t >= ISA_AVX512) return mxfp4_gemv_avx512; + if (t >= ISA_AVX2) return mxfp4_gemv_avx2; +#endif + (void)t; + return mxfp4_gemv_scalar; +} + +// Round a clamped |x|<=448 to nearest float8-e4m3 (RNE), back to fp32. Matches +// torch.float8_e4m3fn / triton .to(float8e4nv). +inline float e4m3_round(float x) { + const float sign = x < 0.0f ? -1.0f : 1.0f; + const float a = std::fabs(x); + if (a == 0.0f) return 0.0f; + if (a >= 448.0f) return sign * 448.0f; + int e; + std::frexp(a, &e); // a in [2^(e-1), 2^e) + float step = std::ldexp(1.0f, e - 4); + const float min_step = std::ldexp(1.0f, -9); // e4m3 subnormal step (2^-9) + if (step < min_step) step = min_step; + float r = std::nearbyint(a / step) * step; + if (r > 448.0f) r = 448.0f; + return sign * r; +} + +// IEEE ceil(log2(v)) for v>0 (matches dsv4 _log2_ceil / fast_round_scale). +inline int ceil_log2_pos(float v) { + uint32_t bits; + std::memcpy(&bits, &v, sizeof(bits)); + const int exp = (int)((bits >> 23) & 0xFF); + const int man = (int)(bits & 0x7FFFFF); + return exp - 127 + (man != 0 ? 1 : 0); +} + +// Split an interleaved bf16 row into fp32 even/odd halves (even[m]=src[2m]). +// bf16->fp32 is exact, so this only reorders -- done once per token/route and +// reused across every output row of the GEMV. +inline void deinterleave_bf16_f32(const bf16_t* src, float* even, float* odd, int K) { + for (int m = 0; m < K / 2; ++m) { + even[m] = bf16_to_f32(src[2 * m]); + odd[m] = bf16_to_f32(src[2 * m + 1]); + } +} + +// DeepSeek-V4 activation FP8 round-trip (bf16 in/out): per 128-block, +// s = 2^ceil(log2(max(|x|,1e-4)/448)); y = round_e4m3(clamp(x/s,+-448)) * s. +void fp8_roundtrip_bf16(const bf16_t* src, bf16_t* dst, int K) { + for (int b0 = 0; b0 < K; b0 += 128) { + const int b1 = std::min(K, b0 + 128); + float amax = 1e-4f; + for (int i = b0; i < b1; ++i) amax = std::max(amax, std::fabs(bf16_to_f32(src[i]))); + const float s = std::ldexp(1.0f, ceil_log2_pos(amax * (1.0f / 448.0f))); + const float inv_s = 1.0f / s; + for (int i = b0; i < b1; ++i) { + float q = bf16_to_f32(src[i]) * inv_s; + q = std::min(448.0f, std::max(-448.0f, q)); + dst[i] = f32_to_bf16(e4m3_round(q) * s); + } + } +} + +// --------------------------------- executor --------------------------------- + +struct CpuMoeExecutor; + +struct MoeTask { + CpuMoeExecutor* exec; + int layer_id; + int num_tokens; + const bf16_t* x; // [num_tokens, H] + const int32_t* ids; // [num_tokens, top_k] (raw expert ids; <0 = skip) + const float* w; // [num_tokens, top_k] + bf16_t* y; // [num_tokens, H] +}; + +// Output-row tiling. Small enough to give every worker independent work even at +// batch size 1; large enough to amortize the atomic work-grab. +// +// Bandwidth notes (Sapphire Rapids 8480+, 13 cores): the two passes already read +// every expert weight byte exactly once per token (each output row block is owned +// by one worker), and x stays hot in L1 across a (token,expert)'s rows -- so the +// kernel is single-read bandwidth-optimal at bs=1 (~205 GB/s vs ~55 GB/s PCIe). +// One worker per *physical* core, pinned, is the sweet spot; SMT oversubscription +// thrashes the spin-barrier. Deferred (not worth it here / for this workload): +// - AMX-bf16: a GEMM tile engine; decode is M=1 GEMV so tiles sit idle. It would +// only pay off in a grouped/batched (dedup) path. +// - expert dedup for bs>1: read each distinct expert once and GEMM its tokens. +// Helps locality+bytes when bs is large; decode batches here are tiny (<=4). +// - NUMA: a single node is assumed. Multi-socket machines would split each +// expert's K dimension per node (banks are already per-row contiguous). +constexpr int IBLK = 32; +constexpr int HBLK = 32; + +// -------------------------------- Q4_0 (W4A8) -------------------------------- +// Native GGUF Q4_0 experts (gemma4 GGUF): per-32 block = fp16 scale d + 16 packed +// bytes; byte j holds element j in its low nibble and j+16 in its high nibble, so a +// block's storage order is [lo0..lo15, hi0..hi15] and w = (nibble - 8) * d. Matches +// the reference dequant (models/gguf/dequant.py) and the packed banks the GPU offload +// path streams. +// +// llama.cpp ggml_vec_dot_q4_0_q8_0: W4A8. The activation is pre-quantized to Q8_0 +// (per-32-block int8 ``aq`` + fp32 scale ``asb``); each block unpacks its 16 bytes to +// 32 int8 weights in [-8,7] (bytes_from_nibbles_32: low nibbles -> elems 0..15, high +// -> 16..31) and runs an integer block dot -- VPDPBUSD (AVX-VNNI) or VPMADDUBSW+VPMADDWD +// (AVX2) with the ggml sign trick |w|*(sign(w)*a)=w*a, or a scalar int loop -- then +// scales the block sum by wd*xd in fp32. No fp weight dequant / shuffle chain. The GPU +// offload path (ggml_moe_a8_vec / MMVQ) is also W4A8, so cpu and hybrid stay close. +using q4dot_fn = float (*)(const uint8_t*, const int8_t*, const float*, int); + +float q4_0_dot_i8_scalar(const uint8_t* w, const int8_t* aq, const float* asb, int K) { + float acc = 0.0f; + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 18; + uint16_t dh; + std::memcpy(&dh, blk, sizeof(dh)); + const uint8_t* q = blk + 2; // 16 nibble bytes + const int8_t* a = aq + (size_t)b * 32; + int isum = 0; + for (int j = 0; j < 16; ++j) { + isum += ((int)(q[j] & 0x0F) - 8) * (int)a[j]; // elem j + isum += ((int)(q[j] >> 4) - 8) * (int)a[16 + j]; // elem 16+j + } + acc += fp16_to_f32(dh) * asb[b] * (float)isum; + } + return acc; +} + +#if CPU_MOE_X86 +// fp16 block scale -> fp32 via HW F16C (single value in lane 0). +__attribute__((target("f16c"))) +static inline float q4_scale(uint16_t h) { + return _mm_cvtss_f32(_mm_cvtph_ps(_mm_cvtsi32_si128((int)h))); +} + +// Unpack one Q4_0 block's 16 bytes -> 32 int8 weights in [-8,7] (elems 0..15 = low +// nibbles, 16..31 = high nibbles). ``eight`` = _mm256_set1_epi8(8). +__attribute__((target("avx2"))) +static inline __m256i q4_unpack32(const uint8_t* blk, __m128i mask, __m256i eight) { + const __m128i qb = _mm_loadu_si128(reinterpret_cast(blk + 2)); + const __m128i lo = _mm_and_si128(qb, mask); + const __m128i hi = _mm_and_si128(_mm_srli_epi16(qb, 4), mask); + return _mm256_sub_epi8(_mm256_set_m128i(hi, lo), eight); +} + +// AVX2 W4A8 (llama.cpp non-VNNI mul_sum_i8_pairs): integer block dot via VPMADDUBSW + +// VPMADDWD (sign trick), scaled by wd*xd. |aw*sa| pair sums <= 8*127*2 < 32767 -> no +// int16 saturation. This is the fast path on AVX2 CPUs without AVX-VNNI (and the +// avx512-tier fallback, since the block dot is 256-bit either way). +__attribute__((target("avx2,fma,f16c"))) +float q4_0_dot_i8_avx2(const uint8_t* w, const int8_t* aq, const float* asb, int K) { + const __m128i mask = _mm_set1_epi8(0x0F); + const __m256i eight = _mm256_set1_epi8(8); + const __m256i ones16 = _mm256_set1_epi16(1); + __m256 accF = _mm256_setzero_ps(); + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 18; + _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); + uint16_t dh; + std::memcpy(&dh, blk, sizeof(dh)); + __m256i wq = q4_unpack32(blk, mask, eight); + __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); + __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) + __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) + __m256i d32 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, sa), ones16); // 8 int32 + accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d32), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); + } + return hsum256(accF); +} + +// AVX-VNNI W4A8: one VPDPBUSD per block (the fast path on modern CPUs). +__attribute__((target("avx2,avxvnni,fma,f16c"))) +float q4_0_dot_i8_vnni(const uint8_t* w, const int8_t* aq, const float* asb, int K) { + const __m128i mask = _mm_set1_epi8(0x0F); + const __m256i eight = _mm256_set1_epi8(8); + __m256 accF = _mm256_setzero_ps(); + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 18; + _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); + uint16_t dh; + std::memcpy(&dh, blk, sizeof(dh)); + __m256i wq = q4_unpack32(blk, mask, eight); + __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); + __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) + __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) + __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); + // All 32 elems of the block share wd*xd; distribute over di's 8 partial sums and + // reduce at the end (equivalent to scale * block_total). + accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); + } + return hsum256(accF); +} +#endif // CPU_MOE_X86 + +// All tiers are W4A8 (int8 activations pre-quantized to Q8_0). AVX-VNNI is orthogonal to +// the ISA tier (gated by cpu_has_avxvnni() / FREETOKEN_CPU_MOE_NO_VNNI), so it wins when +// present; otherwise the 256-bit VPMADDUBSW kernel covers both the avx2 and avx512 tiers. +q4dot_fn select_q4dot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (cpu_has_avxvnni()) return q4_0_dot_i8_vnni; + if (t >= ISA_AVX2) return q4_0_dot_i8_avx2; +#endif + (void)t; + return q4_0_dot_i8_scalar; +} + +enum WFmt { WF_BF16 = 0, WF_NVFP4 = 1, WF_MXFP4 = 2, WF_DSFP4 = 3, WF_Q4_0 = 4 }; + +// Each ctor pointer arg is the address of a CPU int64 array of length +// num_layers (one base address per layer, built by cpu_executor.py's +// _make_table), not a single flat bank. tbl_at resolves +// tbl[layer_id] once per task/pass; a null table (bank unused by this fmt, ptr +// arg 0) resolves to nullptr without dereferencing. +inline const void* tbl_at(const uint64_t* tbl, int layer_id) { + return tbl ? reinterpret_cast(tbl[layer_id]) : nullptr; +} + +struct CpuMoeExecutor { + int num_threads; + int num_layers, num_experts, top_k; + int H, I; + int act, apply_on_input; + int fmt; // WFmt + bool needs_di = false; // pre-deinterleave activations to fp32 (nvfp4/ds_fp4) + // Per-layer pointer tables (one base address per layer, see tbl_at). gate_up_tbl + // doubles as the bf16 gate_up table and the nvfp4/mxfp4/q4_0/ds_fp4 packed-gate_up + // table (down_tbl likewise for down); which reinterpretation applies is picked by + // fmt at each resolve site (see gemm1_dot/gemm2_dot/do_pass1_mxfp4/do_pass1_dsfp4). + const uint64_t* gate_up_tbl; // bf16: [E,2I,H] rows; else: packed e2m1/mxfp4-blocks + const uint64_t* down_tbl; // bf16: [E,H,I] rows; else: packed e2m1/mxfp4-blocks + const uint64_t* gu_scale_tbl; // nvfp4/mxfp4/ds_fp4: [E,2I,*] block scales + const uint64_t* gu_global_tbl; // nvfp4: [E,2I] fp16 row globals + const uint64_t* dn_scale_tbl; // nvfp4/mxfp4/ds_fp4: [E,H,*] block scales + const uint64_t* dn_global_tbl; // nvfp4: [E,H] fp16 row globals + const uint64_t* gu_bias_tbl; // mxfp4: [E,2I] bf16 biases + const uint64_t* dn_bias_tbl; // mxfp4: [E,H] bf16 biases + float swiglu_alpha; + float swiglu_limit; // +inf == no clamp + dot_fn dot; + nvdot_fn nvdot; + nvi8dot_fn nvi8dot = nullptr; // AVX-VNNI W4A8 nvfp4 dot (nullptr -> use fp32 nvdot) + bool use_vnni = false; // nvfp4 + AVX-VNNI: decode via int8 VPDPBUSD (W4A8) + bool use_q4a8 = false; // q4_0: always W4A8 (llama.cpp Q4_0 x Q8_0); int8 pre-quant + dsdot_fn dsdot; + mxgemv_fn mxgemv; + q4dot_fn q4dot; + // ds_fp4: the caller already FP8-round-tripped the input activations on the GPU + // (same reference grid), so submit() must not repeat it on the host-callback + // thread. That scalar per-element pass is single-threaded ON THE DECODE CRITICAL + // PATH (~0.3ms/layer at H=4096, every worker and the GPU waiting on it); moving + // it to a captured GPU elementwise kernel removes it while keeping the official + // W4A8 numerics bit-exact. Set via set_input_prequant (see cpu_executor.py). + bool input_prequant = false; + // Q4_0 packed-row byte strides (H/32*18 for gate_up over K=H, I/32*18 for down over K=I). + int q4_gu_row_bytes = 0, q4_dn_row_bytes = 0; + float e2m1_lut[16]; + float e4m3_lut[256]; + float e8m0_lut[256]; // mxfp4 block scale: 2^(s-127), s clamped to [0,254] + const char* isa; + + std::vector g_scratch; // [max_tokens * top_k * I] intermediate + std::vector xq_scratch; // [max_tokens * H] ds_fp4 fp8-roundtripped input + // ds_fp4 activations pre-deinterleaved to fp32 (even/odd K) for the row-major dot. + std::vector xe_scratch, xo_scratch; // [max_tokens * H/2] (input) + std::vector ge_scratch, go_scratch; // [max_tokens*top_k*I/2] (intermediate) + // AVX-VNNI W4A8: per-16-block int8 activations [even(8),odd(8)] + per-block scale. + std::vector xi8_scratch, gi8_scratch; // [max_tokens*H], [max_tokens*top_k*I] + std::vector xas_scratch, gas_scratch; // [max_tokens*H/16], [..*top_k*I/16] + std::string isa_str; + + std::vector workers; + std::mutex task_mtx; + std::condition_variable task_cv; + std::mutex sync_mtx; + std::condition_variable sync_cv; + + bool stop = false; + uint64_t cur_gen = 0; + MoeTask* cur_task = nullptr; + std::atomic submitted{0}; + std::atomic completed{0}; + + std::atomic p1_next{0}; + std::atomic p2_next{0}; + std::atomic prt_next{0}; // ds_fp4 intermediate fp8 round-trip phase + int64_t p1_total = 0, p2_total = 0, prt_total = 0; + int n_iblk = 0, n_hblk = 0; + std::atomic done_count{0}; + std::atomic bar_count{0}; + std::atomic bar_sense{0}; + + std::vector owned_tasks; // persistent task descriptors (graph-stable) + std::vector core_ids; // worker tid -> logical CPU to pin to (may be empty) + + // ---- Flag-based GPU<->CPU handshake (replaces the per-layer cudaLaunchHostFunc pair) ---- + // A tiny GPU kernel bumps ready_flags[slot] at submit; this coordinator thread busy-polls + // it, runs the slot's task on the worker pool, and sets done_flags[slot], which a GPU + // spin-wait kernel polls at sync. This removes the ~2x30-50us host-func dispatch round + // trips per MoE layer per decode step that otherwise idle the GPU (~6 ms/step on a + // 75-layer model). One slot per (layer, decode batch size) pair -- the Python side // allocates slots as tasks are created. Flags live in mapped-pinned host memory (UVA: // the same pointers are used by the GPU kernels and by this thread). std::thread coord_thread; From bc0f4c7f15487c5248cefa868a9fd1839134e07f Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Mon, 24 Aug 2026 18:10:40 +0900 Subject: [PATCH 6/9] fix(rocm): emit one offload flag per architecture --- python/freetoken/kernel/utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index dfdcd4c7..6cb298fc 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -52,8 +52,9 @@ def _hip_cflags(extra: List[str]) -> List[str]: """HIP flags for a kernel build on ROCm.""" # TODO(ROCm): Triton autotune configs need RDNA3-specific tuning (wave count, LDS size). flags = DEFAULT_HIP_CFLAGS + extra - rocm_arch = os.getenv("FREETOKEN_ROCM_ARCH", "gfx1100;gfx1101;gfx1102;gfx1103") - flags = flags + [f"--offload-arch={rocm_arch}"] + raw_arches = os.getenv("FREETOKEN_ROCM_ARCH", "gfx1100;gfx1101;gfx1102;gfx1103") + arches = [arch for arch in re.split(r"[;,\s]+", raw_arches.strip()) if arch] + flags = flags + [f"--offload-arch={arch}" for arch in arches] return flags CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool] From 5a7362adb262436a9540f269bc4b493be8f1693d Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Mon, 24 Aug 2026 18:10:53 +0900 Subject: [PATCH 7/9] test(rocm): cover multi-arch HIP flag expansion --- tests/kernels/test_rocm_arch_flags.py | 34 +++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) create mode 100644 tests/kernels/test_rocm_arch_flags.py diff --git a/tests/kernels/test_rocm_arch_flags.py b/tests/kernels/test_rocm_arch_flags.py new file mode 100644 index 00000000..6235ffe2 --- /dev/null +++ b/tests/kernels/test_rocm_arch_flags.py @@ -0,0 +1,34 @@ +from freetoken.kernel import utils + + +def test_hip_cflags_expand_default_arches(monkeypatch): + monkeypatch.delenv("FREETOKEN_ROCM_ARCH", raising=False) + + flags = utils._hip_cflags([]) + + assert flags[-4:] == [ + "--offload-arch=gfx1100", + "--offload-arch=gfx1101", + "--offload-arch=gfx1102", + "--offload-arch=gfx1103", + ] + assert all(";" not in flag for flag in flags) + + +def test_hip_cflags_accept_common_arch_separators(monkeypatch): + monkeypatch.setenv( + "FREETOKEN_ROCM_ARCH", + "gfx1100; gfx1101,gfx1102 gfx1103", + ) + + flags = utils._hip_cflags(["-DFOO=1"]) + + assert flags == [ + "-std=c++20", + "-O3", + "-DFOO=1", + "--offload-arch=gfx1100", + "--offload-arch=gfx1101", + "--offload-arch=gfx1102", + "--offload-arch=gfx1103", + ] From d151c102e7af36f8846eae6cab0819cfbc57e18a Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Mon, 24 Aug 2026 18:53:16 +0900 Subject: [PATCH 8/9] fix(rocm): map DLPack CUDA device tokens to ROCm under HIP --- .../freetoken/kernel/csrc/include/freetoken/utils.cuh | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index f21b9585..80e24bbc 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -6,6 +6,16 @@ #include #include +// FreeToken's native kernel call-sites historically use the DLPack CUDA tokens +// as the generic GPU/device matcher. PyTorch on ROCm correctly exports kDLROCM +// and kDLROCMHost instead. Keep the existing call-sites source-compatible under +// hipcc by translating those tokens only after dlpack.h has defined the enums. +// CUDA/nvcc builds do not see these aliases. +#ifdef __HIP__ +#define kDLCUDA kDLROCM +#define kDLCUDAHost kDLROCMHost +#endif + #include #include #include From 48146760e6d2dfc92dd35233447c9602ad57054b Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Mon, 24 Aug 2026 21:08:23 +0900 Subject: [PATCH 9/9] fix(rocm): make DLPack device matching explicit --- .../kernel/csrc/include/freetoken/utils.cuh | 10 ----- .../kernel/csrc/jit/fast_index_copy.cuh | 19 +++++----- python/freetoken/kernel/csrc/jit/index.cu | 6 +-- python/freetoken/kernel/csrc/jit/store.cu | 6 +-- tests/kernels/test_rocm_dlpack_devices.py | 38 +++++++++++++++++++ 5 files changed, 54 insertions(+), 25 deletions(-) create mode 100644 tests/kernels/test_rocm_dlpack_devices.py diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index 80e24bbc..f21b9585 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -6,16 +6,6 @@ #include #include -// FreeToken's native kernel call-sites historically use the DLPack CUDA tokens -// as the generic GPU/device matcher. PyTorch on ROCm correctly exports kDLROCM -// and kDLROCMHost instead. Keep the existing call-sites source-compatible under -// hipcc by translating those tokens only after dlpack.h has defined the enums. -// CUDA/nvcc builds do not see these aliases. -#ifdef __HIP__ -#define kDLCUDA kDLROCM -#define kDLCUDAHost kDLROCMHost -#endif - #include #include #include diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index fe3f6be4..2f784211 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -171,7 +171,8 @@ inline bool host_ptr_identity() { } inline void* device_alias(void* ptr, DLDevice dev) { - if (dev.device_type == kDLCUDA || host_ptr_identity()) { + if (dev.device_type == kDLCUDA || dev.device_type == kDLROCM || + host_ptr_identity()) { return ptr; } void* mapped = nullptr; @@ -293,7 +294,7 @@ inline auto get_sync_flag_ptr( auto flag_dtype = host::SymbolicDType{}; host::TensorMatcher({1}) .with_dtype(flag_dtype) - .with_device(device) + .with_device(device) .verify(sync_flag); return static_cast(sync_flag.data_ptr()); } @@ -368,17 +369,17 @@ struct FastIndexCopyKernel { TensorMatcher({-1, D}) .with_dtype(data_dtype) - .with_device() + .with_device() .verify(src); TensorMatcher({-1, D}) .with_dtype(data_dtype) - .with_device() + .with_device() .verify(dst); TensorMatcher({L}) .with_dtype(indices_dtype) - .with_device(device) + .with_device(device) .verify(src_indices) .verify(dst_indices); @@ -387,7 +388,7 @@ struct FastIndexCopyKernel { const auto num_indices_tensor = num_indices.value(); TensorMatcher({1}) .with_dtype(num_indices_dtype) - .with_device(device) + .with_device(device) .verify(num_indices_tensor); num_indices_data_ptr = static_cast(num_indices_tensor.data_ptr()); @@ -553,14 +554,14 @@ struct MultiIndexCopyKernel { auto indices_dtype = SymbolicDType{}; auto num_indices_dtype = SymbolicDType{}; - TensorMatcher({B}).with_dtype(ptr_dtype).with_device(device) + TensorMatcher({B}).with_dtype(ptr_dtype).with_device(device) .verify(dst_ptrs).verify(src_ptrs).verify(feat_bytes); - TensorMatcher({L}).with_dtype(indices_dtype).with_device(device) + TensorMatcher({L}).with_dtype(indices_dtype).with_device(device) .verify(dst_indices).verify(src_indices); const int64_t* valid_length = nullptr; if (num_indices.has_value()) { - TensorMatcher({1}).with_dtype(num_indices_dtype).with_device(device) + TensorMatcher({1}).with_dtype(num_indices_dtype).with_device(device) .verify(num_indices.value()); valid_length = static_cast(num_indices.value().data_ptr()); } diff --git a/python/freetoken/kernel/csrc/jit/index.cu b/python/freetoken/kernel/csrc/jit/index.cu index ca0e1db2..aca58383 100644 --- a/python/freetoken/kernel/csrc/jit/index.cu +++ b/python/freetoken/kernel/csrc/jit/index.cu @@ -114,15 +114,15 @@ struct IndexKernel { TensorMatcher({-1, D}) // .with_dtype(weights_dtype_) - .with_device(device_) + .with_device(device_) .verify(weights); TensorMatcher({L, D}) // .with_dtype(weights_dtype_) - .with_device(device_) + .with_device(device_) .verify(output); TensorMatcher({L}) // .with_dtype(indices_dtype_) - .with_device(device_) + .with_device(device_) .verify(indices); const auto device = device_.unwrap(); diff --git a/python/freetoken/kernel/csrc/jit/store.cu b/python/freetoken/kernel/csrc/jit/store.cu index 8d84d76e..162dfdfe 100644 --- a/python/freetoken/kernel/csrc/jit/store.cu +++ b/python/freetoken/kernel/csrc/jit/store.cu @@ -72,18 +72,18 @@ struct StoreKernel { TensorMatcher({-1, D}) // .with_strides({X, 1}) - .with_device(device_) + .with_device(device_) .with_dtype(dtype_) .verify(k_cache) .verify(v_cache); TensorMatcher({L, D}) // .with_strides({Y, 1}) - .with_device(device_) + .with_device(device_) .with_dtype(dtype_) .verify(k) .verify(v); TensorMatcher({L}) // - .with_device(device_) + .with_device(device_) .with_dtype(indices_dtype_) .verify(indices); diff --git a/tests/kernels/test_rocm_dlpack_devices.py b/tests/kernels/test_rocm_dlpack_devices.py new file mode 100644 index 00000000..35f97dbe --- /dev/null +++ b/tests/kernels/test_rocm_dlpack_devices.py @@ -0,0 +1,38 @@ +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[2] + + +def _read(relative: str) -> str: + return (ROOT / relative).read_text(encoding="utf-8") + + +def test_rocm_dlpack_devices_are_explicit_not_macro_aliased(): + utils = _read("python/freetoken/kernel/csrc/include/freetoken/utils.cuh") + fast = _read("python/freetoken/kernel/csrc/jit/fast_index_copy.cuh") + index = _read("python/freetoken/kernel/csrc/jit/index.cu") + store = _read("python/freetoken/kernel/csrc/jit/store.cu") + pynccl = _read("python/freetoken/kernel/csrc/src/pynccl.cu") + + # Do not globally rewrite DLPack CUDA tokens under HIP. Besides being hard to + # reason about, that would also silently change the still-CUDA/NCCL-only + # multi-GPU wrapper, which is outside this RDNA3 single-GPU PR's scope. + assert "#define kDLCUDA kDLROCM" not in utils + assert "#define kDLCUDAHost kDLROCMHost" not in utils + + # Single-GPU JIT paths that execute on ROCm explicitly accept ROCm DLPack + # devices while retaining CUDA acceptance for the existing NVIDIA path. + assert ( + ".with_device()" + in fast + ) + assert "dev.device_type == kDLCUDA || dev.device_type == kDLROCM" in fast + assert fast.count(".with_device(device)") >= 6 + assert index.count(".with_device(device_)") == 3 + assert store.count(".with_device(device_)") == 3 + + # Preserve the PR's stated boundary: RCCL/multi-GPU migration is not being + # claimed by a side effect of a preprocessor alias. + assert "kDLROCM" not in pynccl + assert "device_type == kDLCUDA" in pynccl