Skip to content

Feat/gemma4 int4 group32 #499

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

Draft
hzjane wants to merge 2 commits into
base: upgrade/vllm-xpu-v0.21.0
Choose a base branch
from
Draft

Feat/gemma4 int4 group32 #499

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
be42dbf
[esimd][int4] GEMV: support group_size=32 (q4_0) alongside 128
hzjane Jun 26, 2026
35e7115
[esimd][int4] GEMM (pgrp): support group_size=32 (q4_0) alongside 128
hzjane Jun 26, 2026
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
10 changes: 8 additions & 2 deletions vllm/custom-esimd-kernels-vllm/csrc/xpu/esimd_kernel.sycl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -160,10 +160,14 @@ at::Tensor esimd_gemv_int4(
at::Tensor output) {
int64_t N = weight.size(0);
int64_t K = weight.size(1) * 2; // packed: K/2 bytes → K elements
// Infer group size from scale shape [N, K/group_size]: 128 (legacy GGML)
// or 32 (vLLM 0.21 sym_int4 / q4_0 32-element blocks).
int64_t n_groups = weight_scale.size(weight_scale.dim() - 1);
int group_size = (n_groups > 0) ? (int)(K / n_groups) : 128;
EXTRACT_PTR(p_in, input); EXTRACT_PTR(p_w, weight);
EXTRACT_PTR(p_sc, weight_scale); EXTRACT_PTR(p_out, output);
auto& dpcpp_queue = get_device_queue(input);
GEMV_int4_host(p_in, p_w, p_sc, p_out, N, K, dpcpp_queue);
GEMV_int4_host(p_in, p_w, p_sc, p_out, N, K, dpcpp_queue, group_size);
return output;
}

Expand All @@ -175,6 +179,8 @@ at::Tensor esimd_gemv_int4_fused2(
at::Tensor w1, at::Tensor s1, at::Tensor o1) {
int64_t K = w0.size(1) * 2; // packed: K/2 bytes → K elements
int64_t N0 = w0.size(0), N1 = w1.size(0);
int64_t ng0 = s0.size(s0.dim() - 1);
int group_size = (ng0 > 0) ? (int)(K / ng0) : 128;
EXTRACT_PTR(p_in, input);
EXTRACT_PTR(pw0, w0); EXTRACT_PTR(ps0, s0); EXTRACT_PTR(po0, o0);
EXTRACT_PTR(pw1, w1); EXTRACT_PTR(ps1, s1); EXTRACT_PTR(po1, o1);
Expand All @@ -183,7 +189,7 @@ at::Tensor esimd_gemv_int4_fused2(
uint8_t* sp[2] = {ps0, ps1};
uint8_t* op[2] = {po0, po1};
uint32_t ns[2] = {(uint32_t)N0, (uint32_t)N1};
GEMV_int4_fused_host<2>(p_in, wp, sp, op, ns, (uint32_t)K, dpcpp_queue);
GEMV_int4_fused_host<2>(p_in, wp, sp, op, ns, (uint32_t)K, dpcpp_queue, group_size);
return o0;
}

Expand Down
23 changes: 16 additions & 7 deletions vllm/custom-esimd-kernels-vllm/csrc/xpu/esimd_kernel_gemm.sycl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,8 @@ at::Tensor esimd_gemm_fp8_pert(

// ---- INT4 GEMM per-group scale: DPAS kernel for M>=2 ----
// input: [M, K] fp16, weight: [N, K/2] uint8 packed int4,
// weight_scale: [N, K/128] fp16 per-group, output: [M, N] fp16 pre-allocated.
// weight_scale: [N, K/group] fp16 per-group (group=128 GGML or 32 q4_0),
// output: [M, N] fp16 pre-allocated.
// Complementary to esimd_gemv_int4 (M=1); for M>=2 use this.
at::Tensor esimd_gemm_int4_pgrp(
at::Tensor input, at::Tensor weight, at::Tensor weight_scale,
Expand All @@ -66,12 +67,20 @@ at::Tensor esimd_gemm_int4_pgrp(
"esimd_gemm_int4_pgrp: output must be fp16");
TORCH_CHECK(N % 16 == 0,
"esimd_gemm_int4_pgrp: N must be a multiple of 16 (DPAS N tile), got N=", N);
TORCH_CHECK(K % INT4_GEMM_GROUP_SIZE == 0,
"esimd_gemm_int4_pgrp: K must be a multiple of 128 (group_size), got K=", K);
TORCH_CHECK(weight_scale.size(0) == N && weight_scale.size(1) == K / INT4_GEMM_GROUP_SIZE,
TORCH_CHECK(K % 128 == 0,
"esimd_gemm_int4_pgrp: K must be a multiple of 128 (DPAS K_LOAD), got K=", K);
// Infer group size from scale shape [N, K/group_size]: 128 (legacy GGML)
// or 32 (vLLM 0.21 sym_int4 / q4_0 32-element blocks).
int64_t n_groups = weight_scale.size(weight_scale.dim() - 1);
int group_size = (n_groups > 0) ? (int)(K / n_groups) : INT4_GEMM_GROUP_SIZE;
TORCH_CHECK(group_size == 128 || group_size == 32,
"esimd_gemm_int4_pgrp: unsupported group_size=", group_size,
" (expected 128 or 32); scale shape [", weight_scale.size(0),
", ", weight_scale.size(weight_scale.dim()-1), "], K=", K);
TORCH_CHECK(weight_scale.size(0) == N && (int64_t)group_size * n_groups == K,
"esimd_gemm_int4_pgrp: weight_scale shape mismatch; expected [", N,
", ", K / INT4_GEMM_GROUP_SIZE, "], got [", weight_scale.size(0),
", ", weight_scale.size(1), "]");
", ", K / group_size, "], got [", weight_scale.size(0),
", ", weight_scale.size(weight_scale.dim()-1), "]");
TORCH_CHECK(output.size(0) == M && output.size(1) == N,
"esimd_gemm_int4_pgrp: output shape mismatch; expected [", M, ", ", N, "]");

Expand All @@ -83,6 +92,6 @@ at::Tensor esimd_gemm_int4_pgrp(
reinterpret_cast<const uint8_t*>(p_w),
reinterpret_cast<const fp16*>(p_sc),
reinterpret_cast<fp16*>(p_out),
(uint32_t)M, (uint32_t)N, (uint32_t)K, dpcpp_queue);
(uint32_t)M, (uint32_t)N, (uint32_t)K, dpcpp_queue, group_size);
return output;
}
95 changes: 60 additions & 35 deletions vllm/custom-esimd-kernels-vllm/csrc/xpu/esimd_kernels/int4_GEMM.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,10 @@ namespace xesimd = sycl::ext::intel::experimental::esimd;
using fp16 = sycl::half;

// ============================================================================
// INT4 GEMM via DPAS (XMX matrix engine), per-group scale (group_size = 128).
// INT4 GEMM via DPAS (XMX matrix engine), per-group scale.
// Supports GROUP_SIZE = 128 (legacy GGML) and 32 (vLLM 0.21 sym_int4 / q4_0).
// GROUP_SIZE is a template param (default 128); host dispatches on the scale
// last dim, so the op signature is unchanged and group128 callers are intact.
//
// Mirrors the FP8 DPAS V9 kernel (fp8_GEMM_pert.h:FP8_GEMM_DPAS_V9) and
// adapts three things that are INT4-specific:
Expand Down Expand Up @@ -42,8 +45,10 @@ using fp16 = sycl::half;
//
// Requirements:
// N % 16 == 0
// K % (K_THREADS * 128) == 0 (each K-thread owns whole scale groups)
// input fp16, weight uint8 [N, K/2], weight_scale fp16 [N, K/128],
// K % (K_THREADS * K_LOAD) == 0 (each K-thread owns whole 128-wide loads;
// K_LOAD=128 is a multiple of both 128 and 32, so this also keeps each
// thread owning whole scale groups for either GROUP_SIZE)
// input fp16, weight uint8 [N, K/2], weight_scale fp16 [N, K/GROUP_SIZE],
// output fp16 [M, N] pre-allocated.
// ============================================================================

Expand Down Expand Up @@ -78,19 +83,24 @@ int4_pair_to_vnni_scaled(simd<uint32_t, 16> byte_u32,
return interleaved.template bit_cast_view<uint32_t>();
}

template<int K_THREADS, int M_TILES>
template<int K_THREADS, int M_TILES, int GROUP_SIZE = INT4_GEMM_GROUP_SIZE>
struct GEMM_int4_pgrp_kernel {
const fp16* input; // [M, K]
const uint8_t* weight; // [N, K/2] packed int4
const fp16* scale; // [N, K/GROUP_SIZE]
const fp16* scale; // [N, K/GROUP_SIZE] (per-group scales)
fp16* output; // [M, N]
int M, N, K;
int n_groups; // K / GROUP_SIZE

void operator()(sycl::nd_item<1> item) const SYCL_ESIMD_KERNEL {
constexpr int N_TILE = 16;
constexpr int M_TILE = 8;
constexpr int K_LOAD = INT4_GEMM_GROUP_SIZE; // 128
// K_LOAD stays 128 for DPAS efficiency regardless of GROUP_SIZE.
constexpr int K_LOAD = 128;
// How many scale groups a single K_LOAD (=128 K) spans: 1 for
// GROUP_SIZE=128, 4 for GROUP_SIZE=32. K_SUB(=16) divides GROUP_SIZE
// for both, so every K_SUB falls entirely inside one scale group.
constexpr int SUBS_PER_GROUP = GROUP_SIZE / 16; // 8 (g128) or 2 (g32)
constexpr int K_SUB = 16;
constexpr int SUBS_PER_KLOAD = K_LOAD / K_SUB; // 8
constexpr int SLM_PER_THREAD = M_TILES * 128 * 4;
Expand Down Expand Up @@ -131,18 +141,27 @@ struct GEMM_int4_pgrp_kernel {
const fp16* s_base = scale + (size_t)n_start * n_groups;

for (int k_base = k_start; k_base < k_end; k_base += K_LOAD) {
int group_idx = k_base / INT4_GEMM_GROUP_SIZE;
simd<fp16, N_TILE> scales_f16;
#pragma unroll
for (int n = 0; n < N_TILE; n++) {
scales_f16[n] = s_base[n * n_groups + group_idx];
}
simd<fp16, N_TILE> scale_m8_f16 = scales_f16 * fp16(-8.0f);
// Base group index of this K_LOAD. For GROUP_SIZE=128 this is the
// single group; for 32 it is the first of SUBS_PER_KLOAD/SUBS_PER_GROUP
// (=4) groups covered by the 128-wide load.
int base_group = k_base / GROUP_SIZE;

#pragma unroll
for (int sub = 0; sub < SUBS_PER_KLOAD; sub++) {
int k_sub = k_base + sub * K_SUB;

// Which scale group this K_SUB belongs to. Each K_SUB (=16)
// lies entirely in one group (K_SUB divides GROUP_SIZE). For
// GROUP_SIZE=128, SUBS_PER_GROUP=8 so all subs share base_group;
// for 32, SUBS_PER_GROUP=2 so the group advances every 2 subs.
int group_idx = base_group + sub / SUBS_PER_GROUP;
simd<fp16, N_TILE> scales_f16;
#pragma unroll
for (int n = 0; n < N_TILE; n++) {
scales_f16[n] = s_base[n * n_groups + group_idx];
}
simd<fp16, N_TILE> scale_m8_f16 = scales_f16 * fp16(-8.0f);

// set_x is a uint32-element offset; one uint32 covers 8 int4
// K-elements, so stepping by k_sub K means set_x(k_sub / 8).
payB_t.set_x((uint32_t)(k_sub / 8));
Expand Down Expand Up @@ -269,17 +288,17 @@ struct GEMM_int4_pgrp_kernel {
}
};

template<int K_THREADS, int M_TILES>
template<int K_THREADS, int M_TILES, int GROUP_SIZE = INT4_GEMM_GROUP_SIZE>
inline void gemm_int4_pgrp_host_impl(
const fp16* input, const uint8_t* weight, const fp16* scale,
fp16* output, uint32_t M, uint32_t N, uint32_t K,
sycl::queue& q) {
int n_groups = (int)K / INT4_GEMM_GROUP_SIZE;
int n_groups = (int)K / GROUP_SIZE;
int num_wg = ((int)N + 15) / 16;
q.submit([&](sycl::handler& h) {
h.parallel_for(
sycl::nd_range<1>({(size_t)(num_wg * K_THREADS)}, {(size_t)K_THREADS}),
GEMM_int4_pgrp_kernel<K_THREADS, M_TILES>{
GEMM_int4_pgrp_kernel<K_THREADS, M_TILES, GROUP_SIZE>{
input, weight, scale, output,
(int)M, (int)N, (int)K, n_groups});
});
Expand All @@ -291,30 +310,36 @@ inline void gemm_int4_pgrp_host_impl(
inline void GEMM_int4_pgrp_host(
const fp16* input, const uint8_t* weight, const fp16* scale,
fp16* output, uint32_t M, uint32_t N, uint32_t K,
sycl::queue& q) {
sycl::queue& q, int group_size = INT4_GEMM_GROUP_SIZE) {
int m_tiles = ((int)M + 7) / 8;
int n_wgs = ((int)N + 15) / 16;
int k_threads = std::max(1, std::min(4, 640 / std::max(n_wgs, 1)));
while (k_threads > 1 && ((int)K % (k_threads * INT4_GEMM_GROUP_SIZE) != 0)) k_threads--;
// Each K-thread must own whole 128-wide K_LOADs (K_LOAD=128 is a multiple
// of both group sizes), so the per-thread K span must be a multiple of 128.
// This is independent of group_size.
while (k_threads > 1 && ((int)K % (k_threads * 128) != 0)) k_threads--;
if (k_threads == 3) k_threads = 2;

#define DISPATCH(KT, MT) gemm_int4_pgrp_host_impl<KT, MT>( \
#define DISPATCH(KT, MT, G) gemm_int4_pgrp_host_impl<KT, MT, G>( \
input, weight, scale, output, M, N, K, q)
if (k_threads >= 4) {
if (m_tiles <= 1) DISPATCH(4, 1);
else if (m_tiles <= 2) DISPATCH(4, 2);
else if (m_tiles <= 4) DISPATCH(4, 4);
else DISPATCH(4, 8);
} else if (k_threads >= 2) {
if (m_tiles <= 1) DISPATCH(2, 1);
else if (m_tiles <= 2) DISPATCH(2, 2);
else if (m_tiles <= 4) DISPATCH(2, 4);
else DISPATCH(2, 8);
} else {
if (m_tiles <= 1) DISPATCH(1, 1);
else if (m_tiles <= 2) DISPATCH(1, 2);
else if (m_tiles <= 4) DISPATCH(1, 4);
else DISPATCH(1, 8);
}
#define DISPATCH_MT(KT, G) \
do { \
if (m_tiles <= 1) DISPATCH(KT, 1, G); \
else if (m_tiles <= 2) DISPATCH(KT, 2, G); \
else if (m_tiles <= 4) DISPATCH(KT, 4, G); \
else DISPATCH(KT, 8, G); \
} while (0)
#define DISPATCH_KT(G) \
do { \
if (k_threads >= 4) DISPATCH_MT(4, G); \
else if (k_threads >= 2) DISPATCH_MT(2, G); \
else DISPATCH_MT(1, G); \
} while (0)

if (group_size == 32) { DISPATCH_KT(32); }
else { DISPATCH_KT(128); }

#undef DISPATCH_KT
#undef DISPATCH_MT
#undef DISPATCH
}
Loading