mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2026-05-25 12:04:27 +08:00
Merge branch 'main' into fix-oom-test-entrypoints
This commit is contained in:
commit
5e3ab94a00
@ -666,7 +666,7 @@ steps:
|
|||||||
# Quantization
|
# Quantization
|
||||||
- pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'
|
- pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'
|
||||||
- pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py
|
- pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py
|
||||||
# - pytest -v -s tests/kernels/quantization/test_silu_nvfp4_quant_fusion.py
|
- pytest -v -s tests/kernels/quantization/test_silu_nvfp4_quant_fusion.py
|
||||||
- pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py
|
- pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py
|
||||||
- pytest -v -s tests/kernels/quantization/test_flashinfer_scaled_mm.py
|
- pytest -v -s tests/kernels/quantization/test_flashinfer_scaled_mm.py
|
||||||
- pytest -v -s tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py
|
- pytest -v -s tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py
|
||||||
@ -676,7 +676,7 @@ steps:
|
|||||||
- pytest -v -s tests/compile/test_fusion_all_reduce.py
|
- pytest -v -s tests/compile/test_fusion_all_reduce.py
|
||||||
- pytest -v -s tests/compile/test_fusion_attn.py::test_attention_quant_pattern
|
- pytest -v -s tests/compile/test_fusion_attn.py::test_attention_quant_pattern
|
||||||
- pytest -v -s tests/kernels/moe/test_flashinfer.py
|
- pytest -v -s tests/kernels/moe/test_flashinfer.py
|
||||||
# - pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
|
- pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
|
||||||
|
|
||||||
##### 1 GPU test #####
|
##### 1 GPU test #####
|
||||||
##### multi gpus test #####
|
##### multi gpus test #####
|
||||||
|
|||||||
@ -403,7 +403,7 @@ class RandomDataset(BenchmarkDataset):
|
|||||||
# [6880, 6881] -> ['Ġcalls', 'here'] ->
|
# [6880, 6881] -> ['Ġcalls', 'here'] ->
|
||||||
# [1650, 939, 486] -> ['Ġcall', 'sh', 'ere']
|
# [1650, 939, 486] -> ['Ġcall', 'sh', 'ere']
|
||||||
# To avoid uncontrolled change of the prompt length,
|
# To avoid uncontrolled change of the prompt length,
|
||||||
# the encoded sequence is truncated before being decode again.
|
# the encoded sequence is truncated before being decoded again.
|
||||||
total_input_len = prefix_len + int(input_lens[i])
|
total_input_len = prefix_len + int(input_lens[i])
|
||||||
re_encoded_sequence = tokenizer.encode(prompt, add_special_tokens=False)[
|
re_encoded_sequence = tokenizer.encode(prompt, add_special_tokens=False)[
|
||||||
:total_input_len
|
:total_input_len
|
||||||
|
|||||||
@ -637,7 +637,7 @@ def bench_optype(
|
|||||||
# Clear LoRA optimization hash-maps.
|
# Clear LoRA optimization hash-maps.
|
||||||
_LORA_A_PTR_DICT.clear()
|
_LORA_A_PTR_DICT.clear()
|
||||||
_LORA_B_PTR_DICT.clear()
|
_LORA_B_PTR_DICT.clear()
|
||||||
# Run bench function so that _LORA_A_PTR_DICT and _LORA_B_PTR_DICT are setup
|
# Run bench function so that _LORA_A_PTR_DICT and _LORA_B_PTR_DICT are set up
|
||||||
for kwargs in kwargs_list:
|
for kwargs in kwargs_list:
|
||||||
op_type.bench_fn()(**kwargs)
|
op_type.bench_fn()(**kwargs)
|
||||||
torch.cuda.synchronize()
|
torch.cuda.synchronize()
|
||||||
|
|||||||
@ -962,7 +962,7 @@ async def main_mp(
|
|||||||
|
|
||||||
# At this point all the clients finished,
|
# At this point all the clients finished,
|
||||||
# collect results (TTFT, TPOT, etc.) from all the clients.
|
# collect results (TTFT, TPOT, etc.) from all the clients.
|
||||||
# This needs to happens before calling join on the clients
|
# This needs to happen before calling join on the clients
|
||||||
# (result_queue should be emptied).
|
# (result_queue should be emptied).
|
||||||
while not result_queue.empty():
|
while not result_queue.empty():
|
||||||
client_metrics.append(result_queue.get())
|
client_metrics.append(result_queue.get())
|
||||||
|
|||||||
@ -52,15 +52,6 @@
|
|||||||
#define VLLM_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \
|
#define VLLM_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \
|
||||||
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__))
|
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__))
|
||||||
|
|
||||||
#define AT_DISPATCH_BYTE_CASE(enum_type, ...) \
|
|
||||||
AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, byte_t, __VA_ARGS__)
|
|
||||||
|
|
||||||
#define VLLM_DISPATCH_CASE_BYTE_TYPES(...) \
|
|
||||||
AT_DISPATCH_BYTE_CASE(at::ScalarType::Byte, __VA_ARGS__)
|
|
||||||
|
|
||||||
#define VLLM_DISPATCH_BYTE_TYPES(TYPE, NAME, ...) \
|
|
||||||
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_BYTE_TYPES(__VA_ARGS__))
|
|
||||||
|
|
||||||
#define VLLM_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \
|
#define VLLM_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \
|
||||||
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__))
|
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__))
|
||||||
|
|
||||||
|
|||||||
@ -130,8 +130,7 @@ void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
|
|||||||
void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
|
void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
|
||||||
torch::Tensor& scale);
|
torch::Tensor& scale);
|
||||||
|
|
||||||
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
#ifndef USE_ROCM
|
||||||
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
|
|
||||||
void silu_and_mul_nvfp4_quant(torch::Tensor& out,
|
void silu_and_mul_nvfp4_quant(torch::Tensor& out,
|
||||||
torch::Tensor& output_block_scale,
|
torch::Tensor& output_block_scale,
|
||||||
torch::Tensor& input,
|
torch::Tensor& input,
|
||||||
|
|||||||
@ -26,164 +26,17 @@
|
|||||||
#include "dispatch_utils.h"
|
#include "dispatch_utils.h"
|
||||||
|
|
||||||
#include "cuda_utils.h"
|
#include "cuda_utils.h"
|
||||||
|
#include "nvfp4_utils.cuh"
|
||||||
|
|
||||||
namespace vllm {
|
namespace vllm {
|
||||||
|
|
||||||
// Get type2 from type or vice versa (applied to half and bfloat16)
|
|
||||||
template <typename T>
|
|
||||||
struct TypeConverter {
|
|
||||||
using Type = half2;
|
|
||||||
}; // keep for generality
|
|
||||||
|
|
||||||
template <>
|
|
||||||
struct TypeConverter<half2> {
|
|
||||||
using Type = c10::Half;
|
|
||||||
};
|
|
||||||
|
|
||||||
template <>
|
|
||||||
struct TypeConverter<c10::Half> {
|
|
||||||
using Type = half2;
|
|
||||||
};
|
|
||||||
|
|
||||||
template <>
|
|
||||||
struct TypeConverter<__nv_bfloat162> {
|
|
||||||
using Type = c10::BFloat16;
|
|
||||||
};
|
|
||||||
|
|
||||||
template <>
|
|
||||||
struct TypeConverter<c10::BFloat16> {
|
|
||||||
using Type = __nv_bfloat162;
|
|
||||||
};
|
|
||||||
|
|
||||||
#define ELTS_PER_THREAD 8
|
|
||||||
|
|
||||||
constexpr int CVT_FP4_ELTS_PER_THREAD = 8;
|
|
||||||
constexpr int CVT_FP4_SF_VEC_SIZE = 16;
|
|
||||||
|
|
||||||
// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
|
|
||||||
inline __device__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) {
|
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
|
||||||
uint32_t val;
|
|
||||||
asm volatile(
|
|
||||||
"{\n"
|
|
||||||
".reg .b8 byte0;\n"
|
|
||||||
".reg .b8 byte1;\n"
|
|
||||||
".reg .b8 byte2;\n"
|
|
||||||
".reg .b8 byte3;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
|
|
||||||
"mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
|
|
||||||
"}"
|
|
||||||
: "=r"(val)
|
|
||||||
: "f"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]),
|
|
||||||
"f"(array[4]), "f"(array[5]), "f"(array[6]), "f"(array[7]));
|
|
||||||
return val;
|
|
||||||
#else
|
|
||||||
return 0;
|
|
||||||
#endif
|
|
||||||
}
|
|
||||||
|
|
||||||
// Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t).
|
|
||||||
inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) {
|
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
|
||||||
uint32_t val;
|
|
||||||
asm volatile(
|
|
||||||
"{\n"
|
|
||||||
".reg .b8 byte0;\n"
|
|
||||||
".reg .b8 byte1;\n"
|
|
||||||
".reg .b8 byte2;\n"
|
|
||||||
".reg .b8 byte3;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
|
|
||||||
"mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
|
|
||||||
"}"
|
|
||||||
: "=r"(val)
|
|
||||||
: "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y),
|
|
||||||
"f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y));
|
|
||||||
return val;
|
|
||||||
#else
|
|
||||||
return 0;
|
|
||||||
#endif
|
|
||||||
}
|
|
||||||
|
|
||||||
// Fast reciprocal.
|
|
||||||
inline __device__ float reciprocal_approximate_ftz(float a) {
|
|
||||||
float b;
|
|
||||||
asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
|
|
||||||
return b;
|
|
||||||
}
|
|
||||||
|
|
||||||
template <class SFType, int CVT_FP4_NUM_THREADS_PER_SF>
|
|
||||||
__device__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(int rowIdx, int colIdx,
|
|
||||||
int numCols,
|
|
||||||
SFType* SFout) {
|
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
|
||||||
static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 ||
|
|
||||||
CVT_FP4_NUM_THREADS_PER_SF == 2);
|
|
||||||
|
|
||||||
// One pair of threads write one SF to global memory.
|
|
||||||
// TODO: stage through smem for packed STG.32
|
|
||||||
// is it better than STG.8 from 4 threads ?
|
|
||||||
if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF == 0) {
|
|
||||||
// SF vector index (16 elements share one SF in the K dimension).
|
|
||||||
int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF;
|
|
||||||
int32_t mIdx = rowIdx;
|
|
||||||
|
|
||||||
// SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)]
|
|
||||||
// --> index [mTileIdx, kTileIdx, outerMIdx, innerMIdx, innerKIdx]
|
|
||||||
|
|
||||||
int32_t mTileIdx = mIdx / (32 * 4);
|
|
||||||
// SF vector size 16.
|
|
||||||
int factor = CVT_FP4_SF_VEC_SIZE * 4;
|
|
||||||
int32_t numKTiles = (numCols + factor - 1) / factor;
|
|
||||||
int64_t mTileStride = numKTiles * 32 * 4 * 4;
|
|
||||||
|
|
||||||
int32_t kTileIdx = (kIdx / 4);
|
|
||||||
int64_t kTileStride = 32 * 4 * 4;
|
|
||||||
|
|
||||||
// M tile layout [32, 4] is column-major.
|
|
||||||
int32_t outerMIdx = (mIdx % 32);
|
|
||||||
int64_t outerMStride = 4 * 4;
|
|
||||||
|
|
||||||
int32_t innerMIdx = (mIdx % (32 * 4)) / 32;
|
|
||||||
int64_t innerMStride = 4;
|
|
||||||
|
|
||||||
int32_t innerKIdx = (kIdx % 4);
|
|
||||||
int64_t innerKStride = 1;
|
|
||||||
|
|
||||||
// Compute the global offset.
|
|
||||||
int64_t SFOffset = mTileIdx * mTileStride + kTileIdx * kTileStride +
|
|
||||||
outerMIdx * outerMStride + innerMIdx * innerMStride +
|
|
||||||
innerKIdx * innerKStride;
|
|
||||||
|
|
||||||
return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
|
|
||||||
}
|
|
||||||
#endif
|
|
||||||
return nullptr;
|
|
||||||
}
|
|
||||||
|
|
||||||
// Define a 16 bytes packed data type.
|
|
||||||
template <class Type>
|
|
||||||
struct PackedVec {
|
|
||||||
typename TypeConverter<Type>::Type elts[4];
|
|
||||||
};
|
|
||||||
|
|
||||||
template <>
|
|
||||||
struct PackedVec<__nv_fp8_e4m3> {
|
|
||||||
__nv_fp8x2_e4m3 elts[8];
|
|
||||||
};
|
|
||||||
|
|
||||||
template <class Type>
|
template <class Type>
|
||||||
__inline__ __device__ PackedVec<Type> compute_silu(PackedVec<Type>& vec,
|
__inline__ __device__ PackedVec<Type> compute_silu(PackedVec<Type>& vec,
|
||||||
PackedVec<Type>& vec2) {
|
PackedVec<Type>& vec2) {
|
||||||
PackedVec<Type> result;
|
PackedVec<Type> result;
|
||||||
#pragma unroll
|
#pragma unroll
|
||||||
for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; ++i) {
|
for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; ++i) {
|
||||||
if constexpr (std::is_same_v<Type, c10::Half>) {
|
if constexpr (std::is_same_v<Type, half>) {
|
||||||
half2 val(0.5f, 0.5f);
|
half2 val(0.5f, 0.5f);
|
||||||
half2 t0 = __hmul2(vec.elts[i], val);
|
half2 t0 = __hmul2(vec.elts[i], val);
|
||||||
half2 t1 = __hfma2(h2tanh(t0), val, val);
|
half2 t1 = __hfma2(h2tanh(t0), val, val);
|
||||||
@ -206,13 +59,12 @@ __device__ uint32_t silu_and_cvt_warp_fp16_to_fp4(PackedVec<Type>& vec,
|
|||||||
PackedVec<Type>& vec2,
|
PackedVec<Type>& vec2,
|
||||||
float SFScaleVal,
|
float SFScaleVal,
|
||||||
uint8_t* SFout) {
|
uint8_t* SFout) {
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
|
||||||
PackedVec<Type> out_silu = compute_silu(vec, vec2);
|
PackedVec<Type> out_silu = compute_silu(vec, vec2);
|
||||||
// Get absolute maximum values among the local 8 values.
|
// Get absolute maximum values among the local 8 values.
|
||||||
auto localMax = __habs2(out_silu.elts[0]);
|
auto localMax = __habs2(out_silu.elts[0]);
|
||||||
|
|
||||||
// Local maximum value.
|
// Local maximum value.
|
||||||
#pragma unroll
|
#pragma unroll
|
||||||
for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
|
for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
|
||||||
localMax = __hmax2(localMax, __habs2(out_silu.elts[i]));
|
localMax = __hmax2(localMax, __habs2(out_silu.elts[i]));
|
||||||
}
|
}
|
||||||
@ -259,9 +111,9 @@ __device__ uint32_t silu_and_cvt_warp_fp16_to_fp4(PackedVec<Type>& vec,
|
|||||||
// Convert the input to float.
|
// Convert the input to float.
|
||||||
float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
|
float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
|
||||||
|
|
||||||
#pragma unroll
|
#pragma unroll
|
||||||
for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
|
for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
|
||||||
if constexpr (std::is_same_v<Type, c10::Half>) {
|
if constexpr (std::is_same_v<Type, half>) {
|
||||||
fp2Vals[i] = __half22float2(out_silu.elts[i]);
|
fp2Vals[i] = __half22float2(out_silu.elts[i]);
|
||||||
} else {
|
} else {
|
||||||
fp2Vals[i] = __bfloat1622float2(out_silu.elts[i]);
|
fp2Vals[i] = __bfloat1622float2(out_silu.elts[i]);
|
||||||
@ -275,22 +127,14 @@ __device__ uint32_t silu_and_cvt_warp_fp16_to_fp4(PackedVec<Type>& vec,
|
|||||||
|
|
||||||
// Write the e2m1 values to global memory.
|
// Write the e2m1 values to global memory.
|
||||||
return e2m1Vec;
|
return e2m1Vec;
|
||||||
#else
|
|
||||||
return 0;
|
|
||||||
#endif
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Use UE4M3 by default.
|
// Use UE4M3 by default.
|
||||||
template <class Type, bool UE8M0_SF = false>
|
template <class Type, bool UE8M0_SF = false>
|
||||||
__global__ void
|
__global__ void __launch_bounds__(1024, 4)
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
silu_and_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
|
||||||
__launch_bounds__(1024, 4) silu_and_cvt_fp16_to_fp4(
|
float const* SFScale, uint32_t* out,
|
||||||
#else
|
uint32_t* SFout) {
|
||||||
silu_and_cvt_fp16_to_fp4(
|
|
||||||
#endif
|
|
||||||
int32_t numRows, int32_t numCols, Type const* in, float const* SFScale,
|
|
||||||
uint32_t* out, uint32_t* SFout) {
|
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
|
||||||
using PackedVec = PackedVec<Type>;
|
using PackedVec = PackedVec<Type>;
|
||||||
static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
|
static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
|
||||||
(CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
|
(CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
|
||||||
@ -328,22 +172,25 @@ silu_and_cvt_fp16_to_fp4(
|
|||||||
in_vec, in_vec2, SFScaleVal, sf_out);
|
in_vec, in_vec2, SFScaleVal, sf_out);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
#endif
|
|
||||||
}
|
}
|
||||||
|
|
||||||
} // namespace vllm
|
} // namespace vllm
|
||||||
|
|
||||||
void silu_and_mul_nvfp4_quant(torch::Tensor& output, // [..., d]
|
void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output, // [..., d]
|
||||||
torch::Tensor& output_sf,
|
torch::Tensor& output_sf,
|
||||||
torch::Tensor& input, // [..., 2 * d]
|
torch::Tensor& input, // [..., 2 * d]
|
||||||
torch::Tensor& input_sf) {
|
torch::Tensor& input_sf) {
|
||||||
TORCH_CHECK(input.dtype() == torch::kFloat16 ||
|
|
||||||
input.dtype() == torch::kBFloat16);
|
|
||||||
int32_t m = input.size(0);
|
int32_t m = input.size(0);
|
||||||
int32_t n = input.size(1) / 2;
|
int32_t n = input.size(1) / 2;
|
||||||
|
|
||||||
TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
|
TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
|
||||||
|
TORCH_CHECK(input.scalar_type() == at::ScalarType::Half ||
|
||||||
|
input.scalar_type() == at::ScalarType::BFloat16,
|
||||||
|
"Unsupported input data type for quantize_to_fp4.");
|
||||||
|
|
||||||
int multiProcessorCount =
|
int multiProcessorCount =
|
||||||
get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
|
get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
|
||||||
|
|
||||||
auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr());
|
auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr());
|
||||||
auto sf_out = static_cast<int32_t*>(output_sf.data_ptr());
|
auto sf_out = static_cast<int32_t*>(output_sf.data_ptr());
|
||||||
auto output_ptr = static_cast<int64_t*>(output.data_ptr());
|
auto output_ptr = static_cast<int64_t*>(output.data_ptr());
|
||||||
@ -352,17 +199,14 @@ void silu_and_mul_nvfp4_quant(torch::Tensor& output, // [..., d]
|
|||||||
dim3 block(std::min(int(n / ELTS_PER_THREAD), 1024));
|
dim3 block(std::min(int(n / ELTS_PER_THREAD), 1024));
|
||||||
int const numBlocksPerSM = 2048 / block.x;
|
int const numBlocksPerSM = 2048 / block.x;
|
||||||
dim3 grid(std::min(int(m), multiProcessorCount * numBlocksPerSM));
|
dim3 grid(std::min(int(m), multiProcessorCount * numBlocksPerSM));
|
||||||
|
|
||||||
VLLM_DISPATCH_HALF_TYPES(
|
VLLM_DISPATCH_HALF_TYPES(
|
||||||
input.scalar_type(), "act_and_mul_quant_kernel", [&] {
|
input.scalar_type(), "silu_and_mul_nvfp4_quant_kernel", [&] {
|
||||||
auto input_ptr = reinterpret_cast<scalar_t const*>(input.data_ptr());
|
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
|
||||||
VLLM_DISPATCH_BYTE_TYPES(
|
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
|
||||||
output.scalar_type(), "fused_act_and_mul_quant_kernel_nvfp4_type",
|
vllm::silu_and_cvt_fp16_to_fp4<cuda_type><<<grid, block, 0, stream>>>(
|
||||||
[&] {
|
m, n, input_ptr, input_sf_ptr,
|
||||||
vllm::silu_and_cvt_fp16_to_fp4<scalar_t>
|
reinterpret_cast<uint32_t*>(output_ptr),
|
||||||
<<<grid, block, 0, stream>>>(
|
reinterpret_cast<uint32_t*>(sf_out));
|
||||||
m, n, input_ptr, input_sf_ptr,
|
|
||||||
reinterpret_cast<uint32_t*>(output_ptr),
|
|
||||||
reinterpret_cast<uint32_t*>(sf_out));
|
|
||||||
});
|
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|||||||
@ -1,3 +1,19 @@
|
|||||||
|
/*
|
||||||
|
* Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
|
||||||
|
*
|
||||||
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
* you may not use this file except in compliance with the License.
|
||||||
|
* You may obtain a copy of the License at
|
||||||
|
*
|
||||||
|
* http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
*
|
||||||
|
* Unless required by applicable law or agreed to in writing, software
|
||||||
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
* See the License for the specific language governing permissions and
|
||||||
|
* limitations under the License.
|
||||||
|
*/
|
||||||
|
|
||||||
#include <torch/all.h>
|
#include <torch/all.h>
|
||||||
#include <cutlass/arch/arch.h>
|
#include <cutlass/arch/arch.h>
|
||||||
|
|
||||||
|
|||||||
@ -1,247 +1,42 @@
|
|||||||
|
/*
|
||||||
|
* Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
|
||||||
|
*
|
||||||
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
* you may not use this file except in compliance with the License.
|
||||||
|
* You may obtain a copy of the License at
|
||||||
|
*
|
||||||
|
* http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
*
|
||||||
|
* Unless required by applicable law or agreed to in writing, software
|
||||||
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
* See the License for the specific language governing permissions and
|
||||||
|
* limitations under the License.
|
||||||
|
*/
|
||||||
|
|
||||||
#include <torch/all.h>
|
#include <torch/all.h>
|
||||||
|
|
||||||
|
#include <cuda_runtime_api.h>
|
||||||
|
#include <cuda_runtime.h>
|
||||||
|
|
||||||
#include <ATen/cuda/CUDAContext.h>
|
#include <ATen/cuda/CUDAContext.h>
|
||||||
#include <c10/cuda/CUDAGuard.h>
|
#include <c10/cuda/CUDAGuard.h>
|
||||||
|
|
||||||
#include <cuda_runtime.h>
|
|
||||||
#include <cuda_fp8.h>
|
#include <cuda_fp8.h>
|
||||||
|
#include "dispatch_utils.h"
|
||||||
|
|
||||||
template <typename T>
|
#include "nvfp4_utils.cuh"
|
||||||
struct TypeConverter {
|
|
||||||
using Type = half2;
|
|
||||||
}; // keep for generality
|
|
||||||
|
|
||||||
template <>
|
namespace vllm {
|
||||||
struct TypeConverter<half2> {
|
|
||||||
using Type = half;
|
|
||||||
};
|
|
||||||
|
|
||||||
template <>
|
|
||||||
struct TypeConverter<half> {
|
|
||||||
using Type = half2;
|
|
||||||
};
|
|
||||||
|
|
||||||
template <>
|
|
||||||
struct TypeConverter<__nv_bfloat162> {
|
|
||||||
using Type = __nv_bfloat16;
|
|
||||||
};
|
|
||||||
|
|
||||||
template <>
|
|
||||||
struct TypeConverter<__nv_bfloat16> {
|
|
||||||
using Type = __nv_bfloat162;
|
|
||||||
};
|
|
||||||
|
|
||||||
#define ELTS_PER_THREAD 8
|
|
||||||
|
|
||||||
constexpr int CVT_FP4_ELTS_PER_THREAD = 8;
|
|
||||||
constexpr int CVT_FP4_SF_VEC_SIZE = 16;
|
|
||||||
|
|
||||||
// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
|
|
||||||
inline __device__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) {
|
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
|
||||||
uint32_t val;
|
|
||||||
asm volatile(
|
|
||||||
"{\n"
|
|
||||||
".reg .b8 byte0;\n"
|
|
||||||
".reg .b8 byte1;\n"
|
|
||||||
".reg .b8 byte2;\n"
|
|
||||||
".reg .b8 byte3;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
|
|
||||||
"mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
|
|
||||||
"}"
|
|
||||||
: "=r"(val)
|
|
||||||
: "f"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]),
|
|
||||||
"f"(array[4]), "f"(array[5]), "f"(array[6]), "f"(array[7]));
|
|
||||||
return val;
|
|
||||||
#else
|
|
||||||
return 0;
|
|
||||||
#endif
|
|
||||||
}
|
|
||||||
|
|
||||||
// Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t).
|
|
||||||
inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) {
|
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
|
||||||
uint32_t val;
|
|
||||||
asm volatile(
|
|
||||||
"{\n"
|
|
||||||
".reg .b8 byte0;\n"
|
|
||||||
".reg .b8 byte1;\n"
|
|
||||||
".reg .b8 byte2;\n"
|
|
||||||
".reg .b8 byte3;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
|
|
||||||
"mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
|
|
||||||
"}"
|
|
||||||
: "=r"(val)
|
|
||||||
: "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y),
|
|
||||||
"f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y));
|
|
||||||
return val;
|
|
||||||
#else
|
|
||||||
return 0;
|
|
||||||
#endif
|
|
||||||
}
|
|
||||||
|
|
||||||
// Fast reciprocal.
|
|
||||||
inline __device__ float reciprocal_approximate_ftz(float a) {
|
|
||||||
float b;
|
|
||||||
asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
|
|
||||||
return b;
|
|
||||||
}
|
|
||||||
|
|
||||||
template <class SFType, int CVT_FP4_NUM_THREADS_PER_SF>
|
|
||||||
__device__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(int rowIdx, int colIdx,
|
|
||||||
int numCols,
|
|
||||||
SFType* SFout) {
|
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
|
||||||
static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 ||
|
|
||||||
CVT_FP4_NUM_THREADS_PER_SF == 2);
|
|
||||||
|
|
||||||
// One pair of threads write one SF to global memory.
|
|
||||||
// TODO: stage through smem for packed STG.32
|
|
||||||
// is it better than STG.8 from 4 threads ?
|
|
||||||
if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF == 0) {
|
|
||||||
// SF vector index (16 elements share one SF in the K dimension).
|
|
||||||
int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF;
|
|
||||||
int32_t mIdx = rowIdx;
|
|
||||||
|
|
||||||
// SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)]
|
|
||||||
// --> index [mTileIdx, kTileIdx, outerMIdx, innerMIdx, innerKIdx]
|
|
||||||
|
|
||||||
int32_t mTileIdx = mIdx / (32 * 4);
|
|
||||||
// SF vector size 16.
|
|
||||||
int factor = CVT_FP4_SF_VEC_SIZE * 4;
|
|
||||||
int32_t numKTiles = (numCols + factor - 1) / factor;
|
|
||||||
int64_t mTileStride = numKTiles * 32 * 4 * 4;
|
|
||||||
|
|
||||||
int32_t kTileIdx = (kIdx / 4);
|
|
||||||
int64_t kTileStride = 32 * 4 * 4;
|
|
||||||
|
|
||||||
// M tile layout [32, 4] is column-major.
|
|
||||||
int32_t outerMIdx = (mIdx % 32);
|
|
||||||
int64_t outerMStride = 4 * 4;
|
|
||||||
|
|
||||||
int32_t innerMIdx = (mIdx % (32 * 4)) / 32;
|
|
||||||
int64_t innerMStride = 4;
|
|
||||||
|
|
||||||
int32_t innerKIdx = (kIdx % 4);
|
|
||||||
int64_t innerKStride = 1;
|
|
||||||
|
|
||||||
// Compute the global offset.
|
|
||||||
int64_t SFOffset = mTileIdx * mTileStride + kTileIdx * kTileStride +
|
|
||||||
outerMIdx * outerMStride + innerMIdx * innerMStride +
|
|
||||||
innerKIdx * innerKStride;
|
|
||||||
|
|
||||||
return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
|
|
||||||
}
|
|
||||||
#endif
|
|
||||||
return nullptr;
|
|
||||||
}
|
|
||||||
|
|
||||||
// Define a 16 bytes packed data type.
|
|
||||||
template <class Type>
|
|
||||||
struct PackedVec {
|
|
||||||
typename TypeConverter<Type>::Type elts[4];
|
|
||||||
};
|
|
||||||
|
|
||||||
template <>
|
|
||||||
struct PackedVec<__nv_fp8_e4m3> {
|
|
||||||
__nv_fp8x2_e4m3 elts[8];
|
|
||||||
};
|
|
||||||
|
|
||||||
// Quantizes the provided PackedVec into the uint32_t output
|
|
||||||
template <class Type, bool UE8M0_SF = false>
|
|
||||||
__device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
|
|
||||||
uint8_t* SFout) {
|
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
|
||||||
// Get absolute maximum values among the local 8 values.
|
|
||||||
auto localMax = __habs2(vec.elts[0]);
|
|
||||||
|
|
||||||
// Local maximum value.
|
|
||||||
#pragma unroll
|
|
||||||
for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
|
|
||||||
localMax = __hmax2(localMax, __habs2(vec.elts[i]));
|
|
||||||
}
|
|
||||||
|
|
||||||
// Get the absolute maximum among all 16 values (two threads).
|
|
||||||
localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax);
|
|
||||||
// Get the final absolute maximum values.
|
|
||||||
float vecMax = float(__hmax(localMax.x, localMax.y));
|
|
||||||
|
|
||||||
// Get the SF (max value of the vector / max value of e2m1).
|
|
||||||
// maximum value of e2m1 = 6.0.
|
|
||||||
// TODO: use half as compute data type.
|
|
||||||
float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
|
|
||||||
// 8 bits representation of the SF.
|
|
||||||
uint8_t fp8SFVal;
|
|
||||||
// Write the SF to global memory (STG.8).
|
|
||||||
if constexpr (UE8M0_SF) {
|
|
||||||
// Extract the 8 exponent bits from float32.
|
|
||||||
// float 32bits = 1 sign bit + 8 exponent bits + 23 mantissa bits.
|
|
||||||
uint32_t tmp = reinterpret_cast<uint32_t&>(SFValue) >> 23;
|
|
||||||
fp8SFVal = tmp & 0xff;
|
|
||||||
// Convert back to fp32.
|
|
||||||
reinterpret_cast<uint32_t&>(SFValue) = tmp << 23;
|
|
||||||
} else {
|
|
||||||
// Here SFValue is always positive, so E4M3 is the same as UE4M3.
|
|
||||||
__nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue);
|
|
||||||
reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp;
|
|
||||||
// Convert back to fp32.
|
|
||||||
SFValue = float(tmp);
|
|
||||||
}
|
|
||||||
// Get the output scale.
|
|
||||||
// Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) *
|
|
||||||
// reciprocal(SFScaleVal))
|
|
||||||
float outputScale =
|
|
||||||
SFValue != 0 ? reciprocal_approximate_ftz(
|
|
||||||
SFValue * reciprocal_approximate_ftz(SFScaleVal))
|
|
||||||
: 0.0f;
|
|
||||||
|
|
||||||
if (SFout) {
|
|
||||||
// Write the SF to global memory (STG.8).
|
|
||||||
*SFout = fp8SFVal;
|
|
||||||
}
|
|
||||||
|
|
||||||
// Convert the input to float.
|
|
||||||
float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
|
|
||||||
|
|
||||||
#pragma unroll
|
|
||||||
for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
|
|
||||||
if constexpr (std::is_same_v<Type, half>) {
|
|
||||||
fp2Vals[i] = __half22float2(vec.elts[i]);
|
|
||||||
} else {
|
|
||||||
fp2Vals[i] = __bfloat1622float2(vec.elts[i]);
|
|
||||||
}
|
|
||||||
fp2Vals[i].x *= outputScale;
|
|
||||||
fp2Vals[i].y *= outputScale;
|
|
||||||
}
|
|
||||||
|
|
||||||
// Convert to e2m1 values.
|
|
||||||
uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals);
|
|
||||||
|
|
||||||
// Write the e2m1 values to global memory.
|
|
||||||
return e2m1Vec;
|
|
||||||
#else
|
|
||||||
return 0;
|
|
||||||
#endif
|
|
||||||
}
|
|
||||||
|
|
||||||
// Use UE4M3 by default.
|
// Use UE4M3 by default.
|
||||||
template <class Type, bool UE8M0_SF = false, bool SMALL_NUM_EXPERTS = false>
|
template <class Type, bool UE8M0_SF = false, bool SMALL_NUM_EXPERTS = false>
|
||||||
__global__ void
|
__global__ void __launch_bounds__(512, 4)
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
|
||||||
__launch_bounds__(512, 4) cvt_fp16_to_fp4(
|
float const* SFScale, uint32_t* out, uint32_t* SFout,
|
||||||
#else
|
uint32_t* input_offset_by_experts,
|
||||||
cvt_fp16_to_fp4(
|
uint32_t* output_scale_offset_by_experts, int n_experts,
|
||||||
#endif
|
bool low_latency) {
|
||||||
int32_t numRows, int32_t numCols, Type const* in, float const* SFScale,
|
|
||||||
uint32_t* out, uint32_t* SFout, uint32_t* input_offset_by_experts,
|
|
||||||
uint32_t* output_scale_offset_by_experts, int n_experts, bool low_latency) {
|
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
|
||||||
using PackedVec = PackedVec<Type>;
|
using PackedVec = PackedVec<Type>;
|
||||||
static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
|
static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
|
||||||
(CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
|
(CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
|
||||||
@ -299,8 +94,8 @@ cvt_fp16_to_fp4(
|
|||||||
&input_offset_by_experts[chunk_start + 12]));
|
&input_offset_by_experts[chunk_start + 12]));
|
||||||
local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]);
|
local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]);
|
||||||
|
|
||||||
// Check against the 16 loaded offsets
|
// Check against the 16 loaded offsets
|
||||||
#pragma unroll
|
#pragma unroll
|
||||||
for (int i = 0; i < 16; i++) {
|
for (int i = 0; i < 16; i++) {
|
||||||
if (rowIdx >= local_offsets[i] && rowIdx < local_offsets[i + 1]) {
|
if (rowIdx >= local_offsets[i] && rowIdx < local_offsets[i + 1]) {
|
||||||
rowIdx_in_expert = rowIdx - local_offsets[i];
|
rowIdx_in_expert = rowIdx - local_offsets[i];
|
||||||
@ -330,21 +125,15 @@ cvt_fp16_to_fp4(
|
|||||||
|
|
||||||
out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
|
out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
|
||||||
}
|
}
|
||||||
#endif
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Kernel for LARGE_M_TOPK = true (large m_topk optimized version)
|
// Kernel for LARGE_M_TOPK = true (large m_topk optimized version)
|
||||||
template <class Type, bool UE8M0_SF = false, bool SMALL_NUM_EXPERTS = false>
|
template <class Type, bool UE8M0_SF = false, bool SMALL_NUM_EXPERTS = false>
|
||||||
__global__ void
|
__global__ void __launch_bounds__(1024, 4)
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
|
||||||
__launch_bounds__(1024, 4) cvt_fp16_to_fp4(
|
float const* SFScale, uint32_t* out, uint32_t* SFout,
|
||||||
#else
|
uint32_t* input_offset_by_experts,
|
||||||
cvt_fp16_to_fp4(
|
uint32_t* output_scale_offset_by_experts, int n_experts) {
|
||||||
#endif
|
|
||||||
int32_t numRows, int32_t numCols, Type const* in, float const* SFScale,
|
|
||||||
uint32_t* out, uint32_t* SFout, uint32_t* input_offset_by_experts,
|
|
||||||
uint32_t* output_scale_offset_by_experts, int n_experts) {
|
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
|
||||||
using PackedVec = PackedVec<Type>;
|
using PackedVec = PackedVec<Type>;
|
||||||
static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
|
static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
|
||||||
(CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
|
(CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
|
||||||
@ -425,7 +214,6 @@ cvt_fp16_to_fp4(
|
|||||||
|
|
||||||
out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
|
out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
|
||||||
}
|
}
|
||||||
#endif
|
|
||||||
}
|
}
|
||||||
|
|
||||||
template <typename T>
|
template <typename T>
|
||||||
@ -501,6 +289,8 @@ void quant_impl(void* output, void* output_scale, void* input,
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
} // namespace vllm
|
||||||
|
|
||||||
/*Quantization entry for fp4 experts quantization*/
|
/*Quantization entry for fp4 experts quantization*/
|
||||||
#define CHECK_TH_CUDA(x, m) TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
|
#define CHECK_TH_CUDA(x, m) TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
|
||||||
#define CHECK_CONTIGUOUS(x, m) \
|
#define CHECK_CONTIGUOUS(x, m) \
|
||||||
@ -560,23 +350,17 @@ void scaled_fp4_experts_quant_sm100a(
|
|||||||
// 4 means 4 fp8 values are packed into one int32
|
// 4 means 4 fp8 values are packed into one int32
|
||||||
TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
|
TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
|
||||||
|
|
||||||
auto in_dtype = input.dtype();
|
|
||||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
|
||||||
const cudaStream_t stream =
|
const cudaStream_t stream =
|
||||||
at::cuda::getCurrentCUDAStream(input.get_device());
|
at::cuda::getCurrentCUDAStream(input.get_device());
|
||||||
if (in_dtype == at::ScalarType::Half) {
|
|
||||||
quant_impl<half>(output.data_ptr(), output_scale.data_ptr(),
|
VLLM_DISPATCH_HALF_TYPES(
|
||||||
input.data_ptr(), input_global_scale.data_ptr(),
|
input.scalar_type(), "nvfp4_experts_quant_kernel", [&] {
|
||||||
input_offset_by_experts.data_ptr(),
|
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
|
||||||
output_scale_offset_by_experts.data_ptr(), m_topk, k,
|
vllm::quant_impl<cuda_type>(
|
||||||
n_experts, stream);
|
output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
|
||||||
} else if (in_dtype == at::ScalarType::BFloat16) {
|
input_global_scale.data_ptr(), input_offset_by_experts.data_ptr(),
|
||||||
quant_impl<__nv_bfloat16>(output.data_ptr(), output_scale.data_ptr(),
|
output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
|
||||||
input.data_ptr(), input_global_scale.data_ptr(),
|
stream);
|
||||||
input_offset_by_experts.data_ptr(),
|
});
|
||||||
output_scale_offset_by_experts.data_ptr(), m_topk,
|
|
||||||
k, n_experts, stream);
|
|
||||||
} else {
|
|
||||||
TORCH_CHECK(false, "Expected input data type to be half or bfloat16");
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
@ -32,6 +32,14 @@ void scaled_fp4_experts_quant_sm100a(
|
|||||||
torch::Tensor const& output_scale_offset_by_experts);
|
torch::Tensor const& output_scale_offset_by_experts);
|
||||||
#endif
|
#endif
|
||||||
|
|
||||||
|
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
||||||
|
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
|
||||||
|
void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output,
|
||||||
|
torch::Tensor& output_sf,
|
||||||
|
torch::Tensor& input,
|
||||||
|
torch::Tensor& input_sf);
|
||||||
|
#endif
|
||||||
|
|
||||||
void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
|
void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
|
||||||
torch::Tensor& output_sf, torch::Tensor const& input_sf) {
|
torch::Tensor& output_sf, torch::Tensor const& input_sf) {
|
||||||
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
||||||
@ -54,3 +62,13 @@ void scaled_fp4_experts_quant(
|
|||||||
TORCH_CHECK_NOT_IMPLEMENTED(false,
|
TORCH_CHECK_NOT_IMPLEMENTED(false,
|
||||||
"No compiled nvfp4 experts quantization kernel");
|
"No compiled nvfp4 experts quantization kernel");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
void silu_and_mul_nvfp4_quant(torch::Tensor& output, torch::Tensor& output_sf,
|
||||||
|
torch::Tensor& input, torch::Tensor& input_sf) {
|
||||||
|
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
||||||
|
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
|
||||||
|
return silu_and_mul_nvfp4_quant_sm1xxa(output, output_sf, input, input_sf);
|
||||||
|
#endif
|
||||||
|
TORCH_CHECK_NOT_IMPLEMENTED(
|
||||||
|
false, "No compiled silu_and_mul nvfp4 quantization kernel");
|
||||||
|
}
|
||||||
|
|||||||
@ -23,245 +23,18 @@
|
|||||||
#include <c10/cuda/CUDAGuard.h>
|
#include <c10/cuda/CUDAGuard.h>
|
||||||
|
|
||||||
#include <cuda_fp8.h>
|
#include <cuda_fp8.h>
|
||||||
|
#include "dispatch_utils.h"
|
||||||
|
|
||||||
#include "cuda_utils.h"
|
#include "cuda_utils.h"
|
||||||
|
#include "nvfp4_utils.cuh"
|
||||||
|
|
||||||
// Get type2 from type or vice versa (applied to half and bfloat16)
|
namespace vllm {
|
||||||
template <typename T>
|
|
||||||
struct TypeConverter {
|
|
||||||
using Type = half2;
|
|
||||||
}; // keep for generality
|
|
||||||
|
|
||||||
template <>
|
|
||||||
struct TypeConverter<half2> {
|
|
||||||
using Type = half;
|
|
||||||
};
|
|
||||||
|
|
||||||
template <>
|
|
||||||
struct TypeConverter<half> {
|
|
||||||
using Type = half2;
|
|
||||||
};
|
|
||||||
|
|
||||||
template <>
|
|
||||||
struct TypeConverter<__nv_bfloat162> {
|
|
||||||
using Type = __nv_bfloat16;
|
|
||||||
};
|
|
||||||
|
|
||||||
template <>
|
|
||||||
struct TypeConverter<__nv_bfloat16> {
|
|
||||||
using Type = __nv_bfloat162;
|
|
||||||
};
|
|
||||||
|
|
||||||
#define ELTS_PER_THREAD 8
|
|
||||||
|
|
||||||
constexpr int CVT_FP4_ELTS_PER_THREAD = 8;
|
|
||||||
constexpr int CVT_FP4_SF_VEC_SIZE = 16;
|
|
||||||
|
|
||||||
// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
|
|
||||||
inline __device__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) {
|
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
|
||||||
uint32_t val;
|
|
||||||
asm volatile(
|
|
||||||
"{\n"
|
|
||||||
".reg .b8 byte0;\n"
|
|
||||||
".reg .b8 byte1;\n"
|
|
||||||
".reg .b8 byte2;\n"
|
|
||||||
".reg .b8 byte3;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
|
|
||||||
"mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
|
|
||||||
"}"
|
|
||||||
: "=r"(val)
|
|
||||||
: "f"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]),
|
|
||||||
"f"(array[4]), "f"(array[5]), "f"(array[6]), "f"(array[7]));
|
|
||||||
return val;
|
|
||||||
#else
|
|
||||||
return 0;
|
|
||||||
#endif
|
|
||||||
}
|
|
||||||
|
|
||||||
// Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t).
|
|
||||||
inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) {
|
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
|
||||||
uint32_t val;
|
|
||||||
asm volatile(
|
|
||||||
"{\n"
|
|
||||||
".reg .b8 byte0;\n"
|
|
||||||
".reg .b8 byte1;\n"
|
|
||||||
".reg .b8 byte2;\n"
|
|
||||||
".reg .b8 byte3;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
|
|
||||||
"cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
|
|
||||||
"mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
|
|
||||||
"}"
|
|
||||||
: "=r"(val)
|
|
||||||
: "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y),
|
|
||||||
"f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y));
|
|
||||||
return val;
|
|
||||||
#else
|
|
||||||
return 0;
|
|
||||||
#endif
|
|
||||||
}
|
|
||||||
|
|
||||||
// Fast reciprocal.
|
|
||||||
inline __device__ float reciprocal_approximate_ftz(float a) {
|
|
||||||
float b;
|
|
||||||
asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
|
|
||||||
return b;
|
|
||||||
}
|
|
||||||
|
|
||||||
template <class SFType, int CVT_FP4_NUM_THREADS_PER_SF>
|
|
||||||
__device__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(int rowIdx, int colIdx,
|
|
||||||
int numCols,
|
|
||||||
SFType* SFout) {
|
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
|
||||||
static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 ||
|
|
||||||
CVT_FP4_NUM_THREADS_PER_SF == 2);
|
|
||||||
|
|
||||||
// One pair of threads write one SF to global memory.
|
|
||||||
// TODO: stage through smem for packed STG.32
|
|
||||||
// is it better than STG.8 from 4 threads ?
|
|
||||||
if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF == 0) {
|
|
||||||
// SF vector index (16 elements share one SF in the K dimension).
|
|
||||||
int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF;
|
|
||||||
int32_t mIdx = rowIdx;
|
|
||||||
|
|
||||||
// SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)]
|
|
||||||
// --> index [mTileIdx, kTileIdx, outerMIdx, innerMIdx, innerKIdx]
|
|
||||||
|
|
||||||
int32_t mTileIdx = mIdx / (32 * 4);
|
|
||||||
// SF vector size 16.
|
|
||||||
int factor = CVT_FP4_SF_VEC_SIZE * 4;
|
|
||||||
int32_t numKTiles = (numCols + factor - 1) / factor;
|
|
||||||
int64_t mTileStride = numKTiles * 32 * 4 * 4;
|
|
||||||
|
|
||||||
int32_t kTileIdx = (kIdx / 4);
|
|
||||||
int64_t kTileStride = 32 * 4 * 4;
|
|
||||||
|
|
||||||
// M tile layout [32, 4] is column-major.
|
|
||||||
int32_t outerMIdx = (mIdx % 32);
|
|
||||||
int64_t outerMStride = 4 * 4;
|
|
||||||
|
|
||||||
int32_t innerMIdx = (mIdx % (32 * 4)) / 32;
|
|
||||||
int64_t innerMStride = 4;
|
|
||||||
|
|
||||||
int32_t innerKIdx = (kIdx % 4);
|
|
||||||
int64_t innerKStride = 1;
|
|
||||||
|
|
||||||
// Compute the global offset.
|
|
||||||
int64_t SFOffset = mTileIdx * mTileStride + kTileIdx * kTileStride +
|
|
||||||
outerMIdx * outerMStride + innerMIdx * innerMStride +
|
|
||||||
innerKIdx * innerKStride;
|
|
||||||
|
|
||||||
return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
|
|
||||||
}
|
|
||||||
#endif
|
|
||||||
return nullptr;
|
|
||||||
}
|
|
||||||
|
|
||||||
// Define a 16 bytes packed data type.
|
|
||||||
template <class Type>
|
|
||||||
struct PackedVec {
|
|
||||||
typename TypeConverter<Type>::Type elts[4];
|
|
||||||
};
|
|
||||||
|
|
||||||
template <>
|
|
||||||
struct PackedVec<__nv_fp8_e4m3> {
|
|
||||||
__nv_fp8x2_e4m3 elts[8];
|
|
||||||
};
|
|
||||||
|
|
||||||
// Quantizes the provided PackedVec into the uint32_t output
|
|
||||||
template <class Type, bool UE8M0_SF = false>
|
|
||||||
__device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
|
|
||||||
uint8_t* SFout) {
|
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
|
||||||
// Get absolute maximum values among the local 8 values.
|
|
||||||
auto localMax = __habs2(vec.elts[0]);
|
|
||||||
|
|
||||||
// Local maximum value.
|
|
||||||
#pragma unroll
|
|
||||||
for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
|
|
||||||
localMax = __hmax2(localMax, __habs2(vec.elts[i]));
|
|
||||||
}
|
|
||||||
|
|
||||||
// Get the absolute maximum among all 16 values (two threads).
|
|
||||||
localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax);
|
|
||||||
// Get the final absolute maximum values.
|
|
||||||
float vecMax = float(__hmax(localMax.x, localMax.y));
|
|
||||||
|
|
||||||
// Get the SF (max value of the vector / max value of e2m1).
|
|
||||||
// maximum value of e2m1 = 6.0.
|
|
||||||
// TODO: use half as compute data type.
|
|
||||||
float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
|
|
||||||
// 8 bits representation of the SF.
|
|
||||||
uint8_t fp8SFVal;
|
|
||||||
// Write the SF to global memory (STG.8).
|
|
||||||
if constexpr (UE8M0_SF) {
|
|
||||||
// Extract the 8 exponent bits from float32.
|
|
||||||
// float 32bits = 1 sign bit + 8 exponent bits + 23 mantissa bits.
|
|
||||||
uint32_t tmp = reinterpret_cast<uint32_t&>(SFValue) >> 23;
|
|
||||||
fp8SFVal = tmp & 0xff;
|
|
||||||
// Convert back to fp32.
|
|
||||||
reinterpret_cast<uint32_t&>(SFValue) = tmp << 23;
|
|
||||||
} else {
|
|
||||||
// Here SFValue is always positive, so E4M3 is the same as UE4M3.
|
|
||||||
__nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue);
|
|
||||||
reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp;
|
|
||||||
// Convert back to fp32.
|
|
||||||
SFValue = float(tmp);
|
|
||||||
}
|
|
||||||
// Get the output scale.
|
|
||||||
// Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) *
|
|
||||||
// reciprocal(SFScaleVal))
|
|
||||||
float outputScale =
|
|
||||||
SFValue != 0 ? reciprocal_approximate_ftz(
|
|
||||||
SFValue * reciprocal_approximate_ftz(SFScaleVal))
|
|
||||||
: 0.0f;
|
|
||||||
|
|
||||||
if (SFout) {
|
|
||||||
// Write the SF to global memory (STG.8).
|
|
||||||
*SFout = fp8SFVal;
|
|
||||||
}
|
|
||||||
|
|
||||||
// Convert the input to float.
|
|
||||||
float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
|
|
||||||
|
|
||||||
#pragma unroll
|
|
||||||
for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
|
|
||||||
if constexpr (std::is_same_v<Type, half>) {
|
|
||||||
fp2Vals[i] = __half22float2(vec.elts[i]);
|
|
||||||
} else {
|
|
||||||
fp2Vals[i] = __bfloat1622float2(vec.elts[i]);
|
|
||||||
}
|
|
||||||
fp2Vals[i].x *= outputScale;
|
|
||||||
fp2Vals[i].y *= outputScale;
|
|
||||||
}
|
|
||||||
|
|
||||||
// Convert to e2m1 values.
|
|
||||||
uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals);
|
|
||||||
|
|
||||||
// Write the e2m1 values to global memory.
|
|
||||||
return e2m1Vec;
|
|
||||||
#else
|
|
||||||
return 0;
|
|
||||||
#endif
|
|
||||||
}
|
|
||||||
|
|
||||||
// Use UE4M3 by default.
|
// Use UE4M3 by default.
|
||||||
template <class Type, bool UE8M0_SF = false>
|
template <class Type, bool UE8M0_SF = false>
|
||||||
__global__ void
|
__global__ void __launch_bounds__(512, 4)
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
|
||||||
__launch_bounds__(512, 4) cvt_fp16_to_fp4(
|
float const* SFScale, uint32_t* out, uint32_t* SFout) {
|
||||||
#else
|
|
||||||
cvt_fp16_to_fp4(
|
|
||||||
#endif
|
|
||||||
int32_t numRows, int32_t numCols, Type const* in, float const* SFScale,
|
|
||||||
uint32_t* out, uint32_t* SFout) {
|
|
||||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
|
||||||
using PackedVec = PackedVec<Type>;
|
using PackedVec = PackedVec<Type>;
|
||||||
static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
|
static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
|
||||||
(CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
|
(CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
|
||||||
@ -293,7 +66,6 @@ cvt_fp16_to_fp4(
|
|||||||
cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
|
cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
#endif
|
|
||||||
}
|
}
|
||||||
|
|
||||||
template <typename T>
|
template <typename T>
|
||||||
@ -332,6 +104,8 @@ template void invokeFP4Quantization(int m, int n, __nv_bfloat16 const* input,
|
|||||||
int multiProcessorCount,
|
int multiProcessorCount,
|
||||||
cudaStream_t stream);
|
cudaStream_t stream);
|
||||||
|
|
||||||
|
} // namespace vllm
|
||||||
|
|
||||||
void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
|
void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
|
||||||
torch::Tensor const& input,
|
torch::Tensor const& input,
|
||||||
torch::Tensor const& output_sf,
|
torch::Tensor const& output_sf,
|
||||||
@ -340,6 +114,9 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
|
|||||||
int32_t n = input.size(1);
|
int32_t n = input.size(1);
|
||||||
|
|
||||||
TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
|
TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
|
||||||
|
TORCH_CHECK(input.scalar_type() == at::ScalarType::Half ||
|
||||||
|
input.scalar_type() == at::ScalarType::BFloat16,
|
||||||
|
"Unsupported input data type for quantize_to_fp4.");
|
||||||
|
|
||||||
int multiProcessorCount =
|
int multiProcessorCount =
|
||||||
get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
|
get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
|
||||||
@ -353,24 +130,10 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
|
|||||||
// We don't support e8m0 scales at this moment.
|
// We don't support e8m0 scales at this moment.
|
||||||
bool useUE8M0 = false;
|
bool useUE8M0 = false;
|
||||||
|
|
||||||
switch (input.scalar_type()) {
|
VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_kernel", [&] {
|
||||||
case torch::kHalf: {
|
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
|
||||||
auto input_ptr = reinterpret_cast<half const*>(input.data_ptr());
|
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
|
||||||
invokeFP4Quantization(m, n, input_ptr, input_sf_ptr, output_ptr, sf_out,
|
vllm::invokeFP4Quantization(m, n, input_ptr, input_sf_ptr, output_ptr,
|
||||||
useUE8M0, multiProcessorCount, stream);
|
sf_out, useUE8M0, multiProcessorCount, stream);
|
||||||
break;
|
});
|
||||||
}
|
|
||||||
case torch::kBFloat16: {
|
|
||||||
auto input_ptr = reinterpret_cast<__nv_bfloat16 const*>(input.data_ptr());
|
|
||||||
invokeFP4Quantization(m, n, input_ptr, input_sf_ptr, output_ptr, sf_out,
|
|
||||||
useUE8M0, multiProcessorCount, stream);
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
default: {
|
|
||||||
std::cerr << "Observing: " << input.scalar_type()
|
|
||||||
<< " for the input datatype which is invalid";
|
|
||||||
throw std::runtime_error(
|
|
||||||
"Unsupported input data type for quantize_to_fp4.");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
251
csrc/quantization/fp4/nvfp4_utils.cuh
Normal file
251
csrc/quantization/fp4/nvfp4_utils.cuh
Normal file
@ -0,0 +1,251 @@
|
|||||||
|
/*
|
||||||
|
* Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
|
||||||
|
*
|
||||||
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
* you may not use this file except in compliance with the License.
|
||||||
|
* You may obtain a copy of the License at
|
||||||
|
*
|
||||||
|
* http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
*
|
||||||
|
* Unless required by applicable law or agreed to in writing, software
|
||||||
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
* See the License for the specific language governing permissions and
|
||||||
|
* limitations under the License.
|
||||||
|
*/
|
||||||
|
|
||||||
|
#pragma once
|
||||||
|
|
||||||
|
#include <cuda_runtime.h>
|
||||||
|
#include <cuda_fp8.h>
|
||||||
|
|
||||||
|
#define ELTS_PER_THREAD 8
|
||||||
|
|
||||||
|
constexpr int CVT_FP4_ELTS_PER_THREAD = 8;
|
||||||
|
constexpr int CVT_FP4_SF_VEC_SIZE = 16;
|
||||||
|
|
||||||
|
namespace vllm {
|
||||||
|
|
||||||
|
// Convert PyTorch cpp type to CUDA type
|
||||||
|
template <typename T>
|
||||||
|
struct CUDATypeConverter {
|
||||||
|
using Type = T;
|
||||||
|
};
|
||||||
|
|
||||||
|
template <>
|
||||||
|
struct CUDATypeConverter<at::Half> {
|
||||||
|
using Type = half;
|
||||||
|
};
|
||||||
|
|
||||||
|
template <>
|
||||||
|
struct CUDATypeConverter<at::BFloat16> {
|
||||||
|
using Type = __nv_bfloat16;
|
||||||
|
};
|
||||||
|
|
||||||
|
// Get type2 from type or vice versa (applied to half and bfloat16)
|
||||||
|
template <typename T>
|
||||||
|
struct TypeConverter {
|
||||||
|
using Type = half2;
|
||||||
|
}; // keep for generality
|
||||||
|
|
||||||
|
template <>
|
||||||
|
struct TypeConverter<half2> {
|
||||||
|
using Type = half;
|
||||||
|
};
|
||||||
|
|
||||||
|
template <>
|
||||||
|
struct TypeConverter<half> {
|
||||||
|
using Type = half2;
|
||||||
|
};
|
||||||
|
|
||||||
|
template <>
|
||||||
|
struct TypeConverter<__nv_bfloat162> {
|
||||||
|
using Type = __nv_bfloat16;
|
||||||
|
};
|
||||||
|
|
||||||
|
template <>
|
||||||
|
struct TypeConverter<__nv_bfloat16> {
|
||||||
|
using Type = __nv_bfloat162;
|
||||||
|
};
|
||||||
|
|
||||||
|
// Define a 16 bytes packed data type.
|
||||||
|
template <class Type>
|
||||||
|
struct PackedVec {
|
||||||
|
typename TypeConverter<Type>::Type elts[4];
|
||||||
|
};
|
||||||
|
|
||||||
|
template <>
|
||||||
|
struct PackedVec<__nv_fp8_e4m3> {
|
||||||
|
__nv_fp8x2_e4m3 elts[8];
|
||||||
|
};
|
||||||
|
|
||||||
|
// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
|
||||||
|
inline __device__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) {
|
||||||
|
uint32_t val;
|
||||||
|
asm volatile(
|
||||||
|
"{\n"
|
||||||
|
".reg .b8 byte0;\n"
|
||||||
|
".reg .b8 byte1;\n"
|
||||||
|
".reg .b8 byte2;\n"
|
||||||
|
".reg .b8 byte3;\n"
|
||||||
|
"cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
|
||||||
|
"cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
|
||||||
|
"cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
|
||||||
|
"cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
|
||||||
|
"mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
|
||||||
|
"}"
|
||||||
|
: "=r"(val)
|
||||||
|
: "f"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]),
|
||||||
|
"f"(array[4]), "f"(array[5]), "f"(array[6]), "f"(array[7]));
|
||||||
|
return val;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t).
|
||||||
|
inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) {
|
||||||
|
uint32_t val;
|
||||||
|
asm volatile(
|
||||||
|
"{\n"
|
||||||
|
".reg .b8 byte0;\n"
|
||||||
|
".reg .b8 byte1;\n"
|
||||||
|
".reg .b8 byte2;\n"
|
||||||
|
".reg .b8 byte3;\n"
|
||||||
|
"cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
|
||||||
|
"cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
|
||||||
|
"cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
|
||||||
|
"cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
|
||||||
|
"mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
|
||||||
|
"}"
|
||||||
|
: "=r"(val)
|
||||||
|
: "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y),
|
||||||
|
"f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y));
|
||||||
|
return val;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Fast reciprocal.
|
||||||
|
inline __device__ float reciprocal_approximate_ftz(float a) {
|
||||||
|
float b;
|
||||||
|
asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
|
||||||
|
return b;
|
||||||
|
}
|
||||||
|
|
||||||
|
template <class SFType, int CVT_FP4_NUM_THREADS_PER_SF>
|
||||||
|
__device__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(int rowIdx, int colIdx,
|
||||||
|
int numCols,
|
||||||
|
SFType* SFout) {
|
||||||
|
static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 ||
|
||||||
|
CVT_FP4_NUM_THREADS_PER_SF == 2);
|
||||||
|
|
||||||
|
// One pair of threads write one SF to global memory.
|
||||||
|
// TODO: stage through smem for packed STG.32
|
||||||
|
// is it better than STG.8 from 4 threads ?
|
||||||
|
if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF == 0) {
|
||||||
|
// SF vector index (16 elements share one SF in the K dimension).
|
||||||
|
int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF;
|
||||||
|
int32_t mIdx = rowIdx;
|
||||||
|
|
||||||
|
// SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)]
|
||||||
|
// --> index [mTileIdx, kTileIdx, outerMIdx, innerMIdx, innerKIdx]
|
||||||
|
|
||||||
|
int32_t mTileIdx = mIdx / (32 * 4);
|
||||||
|
// SF vector size 16.
|
||||||
|
int factor = CVT_FP4_SF_VEC_SIZE * 4;
|
||||||
|
int32_t numKTiles = (numCols + factor - 1) / factor;
|
||||||
|
int64_t mTileStride = numKTiles * 32 * 4 * 4;
|
||||||
|
|
||||||
|
int32_t kTileIdx = (kIdx / 4);
|
||||||
|
int64_t kTileStride = 32 * 4 * 4;
|
||||||
|
|
||||||
|
// M tile layout [32, 4] is column-major.
|
||||||
|
int32_t outerMIdx = (mIdx % 32);
|
||||||
|
int64_t outerMStride = 4 * 4;
|
||||||
|
|
||||||
|
int32_t innerMIdx = (mIdx % (32 * 4)) / 32;
|
||||||
|
int64_t innerMStride = 4;
|
||||||
|
|
||||||
|
int32_t innerKIdx = (kIdx % 4);
|
||||||
|
int64_t innerKStride = 1;
|
||||||
|
|
||||||
|
// Compute the global offset.
|
||||||
|
int64_t SFOffset = mTileIdx * mTileStride + kTileIdx * kTileStride +
|
||||||
|
outerMIdx * outerMStride + innerMIdx * innerMStride +
|
||||||
|
innerKIdx * innerKStride;
|
||||||
|
|
||||||
|
return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
|
||||||
|
}
|
||||||
|
return nullptr;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Quantizes the provided PackedVec into the uint32_t output
|
||||||
|
template <class Type, bool UE8M0_SF = false>
|
||||||
|
__device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
|
||||||
|
uint8_t* SFout) {
|
||||||
|
// Get absolute maximum values among the local 8 values.
|
||||||
|
auto localMax = __habs2(vec.elts[0]);
|
||||||
|
|
||||||
|
// Local maximum value.
|
||||||
|
#pragma unroll
|
||||||
|
for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
|
||||||
|
localMax = __hmax2(localMax, __habs2(vec.elts[i]));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get the absolute maximum among all 16 values (two threads).
|
||||||
|
localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax);
|
||||||
|
// Get the final absolute maximum values.
|
||||||
|
float vecMax = float(__hmax(localMax.x, localMax.y));
|
||||||
|
|
||||||
|
// Get the SF (max value of the vector / max value of e2m1).
|
||||||
|
// maximum value of e2m1 = 6.0.
|
||||||
|
// TODO: use half as compute data type.
|
||||||
|
float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
|
||||||
|
// 8 bits representation of the SF.
|
||||||
|
uint8_t fp8SFVal;
|
||||||
|
// Write the SF to global memory (STG.8).
|
||||||
|
if constexpr (UE8M0_SF) {
|
||||||
|
// Extract the 8 exponent bits from float32.
|
||||||
|
// float 32bits = 1 sign bit + 8 exponent bits + 23 mantissa bits.
|
||||||
|
uint32_t tmp = reinterpret_cast<uint32_t&>(SFValue) >> 23;
|
||||||
|
fp8SFVal = tmp & 0xff;
|
||||||
|
// Convert back to fp32.
|
||||||
|
reinterpret_cast<uint32_t&>(SFValue) = tmp << 23;
|
||||||
|
} else {
|
||||||
|
// Here SFValue is always positive, so E4M3 is the same as UE4M3.
|
||||||
|
__nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue);
|
||||||
|
reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp;
|
||||||
|
// Convert back to fp32.
|
||||||
|
SFValue = float(tmp);
|
||||||
|
}
|
||||||
|
// Get the output scale.
|
||||||
|
// Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) *
|
||||||
|
// reciprocal(SFScaleVal))
|
||||||
|
float outputScale =
|
||||||
|
SFValue != 0 ? reciprocal_approximate_ftz(
|
||||||
|
SFValue * reciprocal_approximate_ftz(SFScaleVal))
|
||||||
|
: 0.0f;
|
||||||
|
|
||||||
|
if (SFout) {
|
||||||
|
// Write the SF to global memory (STG.8).
|
||||||
|
*SFout = fp8SFVal;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Convert the input to float.
|
||||||
|
float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
|
||||||
|
|
||||||
|
#pragma unroll
|
||||||
|
for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
|
||||||
|
if constexpr (std::is_same_v<Type, half>) {
|
||||||
|
fp2Vals[i] = __half22float2(vec.elts[i]);
|
||||||
|
} else {
|
||||||
|
fp2Vals[i] = __bfloat1622float2(vec.elts[i]);
|
||||||
|
}
|
||||||
|
fp2Vals[i].x *= outputScale;
|
||||||
|
fp2Vals[i].y *= outputScale;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Convert to e2m1 values.
|
||||||
|
uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals);
|
||||||
|
|
||||||
|
// Write the e2m1 values to global memory.
|
||||||
|
return e2m1Vec;
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace vllm
|
||||||
@ -115,8 +115,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
|||||||
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
|
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
|
||||||
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
|
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
|
||||||
|
|
||||||
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
#ifndef USE_ROCM
|
||||||
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
|
|
||||||
ops.def(
|
ops.def(
|
||||||
"silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, "
|
"silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, "
|
||||||
"Tensor input, Tensor input_global_scale) -> ()");
|
"Tensor input, Tensor input_global_scale) -> ()");
|
||||||
|
|||||||
@ -375,7 +375,7 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
|
|||||||
# Install FlashInfer from source
|
# Install FlashInfer from source
|
||||||
ARG FLASHINFER_GIT_REPO="https://github.com/flashinfer-ai/flashinfer.git"
|
ARG FLASHINFER_GIT_REPO="https://github.com/flashinfer-ai/flashinfer.git"
|
||||||
# Keep this in sync with "flashinfer" extra in setup.py
|
# Keep this in sync with "flashinfer" extra in setup.py
|
||||||
ARG FLASHINFER_GIT_REF="v0.2.14.post1"
|
ARG FLASHINFER_GIT_REF="v0.3.0"
|
||||||
# Flag to control whether to compile FlashInfer AOT kernels
|
# Flag to control whether to compile FlashInfer AOT kernels
|
||||||
# Set to "true" to enable AOT compilation:
|
# Set to "true" to enable AOT compilation:
|
||||||
# docker build --build-arg FLASHINFER_AOT_COMPILE=true ...
|
# docker build --build-arg FLASHINFER_AOT_COMPILE=true ...
|
||||||
|
|||||||
@ -117,7 +117,7 @@ def run_gemma3n(question: str, audio_count: int) -> ModelRequestData:
|
|||||||
|
|
||||||
# Granite Speech
|
# Granite Speech
|
||||||
def run_granite_speech(question: str, audio_count: int) -> ModelRequestData:
|
def run_granite_speech(question: str, audio_count: int) -> ModelRequestData:
|
||||||
# NOTE - the setting in this example are somehat different than what is
|
# NOTE - the setting in this example are somewhat different from what is
|
||||||
# optimal for granite speech, and it is generally recommended to use beam
|
# optimal for granite speech, and it is generally recommended to use beam
|
||||||
# search. Check the model README for suggested settings.
|
# search. Check the model README for suggested settings.
|
||||||
# https://huggingface.co/ibm-granite/granite-speech-3.3-8b
|
# https://huggingface.co/ibm-granite/granite-speech-3.3-8b
|
||||||
|
|||||||
@ -10,6 +10,7 @@ wheel
|
|||||||
jinja2>=3.1.6
|
jinja2>=3.1.6
|
||||||
datasets # for benchmark scripts
|
datasets # for benchmark scripts
|
||||||
numba == 0.60.0 # v0.61 doesn't support Python 3.9. Required for N-gram speculative decoding
|
numba == 0.60.0 # v0.61 doesn't support Python 3.9. Required for N-gram speculative decoding
|
||||||
|
nixl==0.3.0 # for PD disaggregation
|
||||||
--extra-index-url=https://download.pytorch.org/whl/xpu
|
--extra-index-url=https://download.pytorch.org/whl/xpu
|
||||||
torch==2.8.0+xpu
|
torch==2.8.0+xpu
|
||||||
torchaudio
|
torchaudio
|
||||||
|
|||||||
2
setup.py
2
setup.py
@ -694,7 +694,7 @@ setup(
|
|||||||
"mistral_common[audio]"], # Required for audio processing
|
"mistral_common[audio]"], # Required for audio processing
|
||||||
"video": [], # Kept for backwards compatibility
|
"video": [], # Kept for backwards compatibility
|
||||||
# FlashInfer should be updated together with the Dockerfile
|
# FlashInfer should be updated together with the Dockerfile
|
||||||
"flashinfer": ["flashinfer-python==0.2.14.post1"],
|
"flashinfer": ["flashinfer-python==0.3.0"],
|
||||||
# Optional deps for AMD FP4 quantization support
|
# Optional deps for AMD FP4 quantization support
|
||||||
"petit-kernel": ["petit-kernel"],
|
"petit-kernel": ["petit-kernel"],
|
||||||
},
|
},
|
||||||
|
|||||||
@ -1,13 +1,16 @@
|
|||||||
# SPDX-License-Identifier: Apache-2.0
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
from contextlib import suppress
|
from contextlib import suppress
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from typing import Any, Optional
|
from typing import TYPE_CHECKING, Any, Optional
|
||||||
from unittest.mock import MagicMock
|
from unittest.mock import MagicMock
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
import pytest_asyncio
|
||||||
|
|
||||||
from vllm.config import MultiModalConfig
|
from vllm.config import MultiModalConfig
|
||||||
from vllm.engine.multiprocessing.client import MQLLMEngineClient
|
from vllm.engine.multiprocessing.client import MQLLMEngineClient
|
||||||
@ -17,6 +20,164 @@ from vllm.entrypoints.openai.serving_models import (BaseModelPath,
|
|||||||
OpenAIServingModels)
|
OpenAIServingModels)
|
||||||
from vllm.transformers_utils.tokenizer import get_tokenizer
|
from vllm.transformers_utils.tokenizer import get_tokenizer
|
||||||
|
|
||||||
|
from ...utils import RemoteOpenAIServer
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from openai import OpenAI
|
||||||
|
|
||||||
|
GPT_OSS_MODEL_NAME = "openai/gpt-oss-20b"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(scope="module")
|
||||||
|
def monkeypatch_module():
|
||||||
|
from _pytest.monkeypatch import MonkeyPatch
|
||||||
|
mpatch = MonkeyPatch()
|
||||||
|
yield mpatch
|
||||||
|
mpatch.undo()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(scope="module")
|
||||||
|
def gptoss_server(monkeypatch_module: pytest.MonkeyPatch):
|
||||||
|
with monkeypatch_module.context() as m:
|
||||||
|
m.setenv("VLLM_ATTENTION_BACKEND", "TRITON_ATTN_VLLM_V1")
|
||||||
|
args = [
|
||||||
|
"--enforce-eager",
|
||||||
|
"--max-model-len",
|
||||||
|
"8192",
|
||||||
|
"--tool-call-parser",
|
||||||
|
"openai",
|
||||||
|
"--reasoning-parser",
|
||||||
|
"openai_gptoss",
|
||||||
|
"--enable-auto-tool-choice",
|
||||||
|
]
|
||||||
|
with RemoteOpenAIServer(GPT_OSS_MODEL_NAME, args) as remote_server:
|
||||||
|
yield remote_server
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
async def gptoss_client(gptoss_server):
|
||||||
|
async with gptoss_server.get_async_client() as async_client:
|
||||||
|
yield async_client
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_gpt_oss_chat_tool_call_streaming(gptoss_client: OpenAI):
|
||||||
|
tools = [{
|
||||||
|
"type": "function",
|
||||||
|
"function": {
|
||||||
|
"name": "get_current_weather",
|
||||||
|
"description": "Get the current weather in a given location",
|
||||||
|
"parameters": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"city": {
|
||||||
|
"type": "string"
|
||||||
|
},
|
||||||
|
"state": {
|
||||||
|
"type": "string"
|
||||||
|
},
|
||||||
|
"unit": {
|
||||||
|
"type": "string",
|
||||||
|
"enum": ["celsius", "fahrenheit"],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": ["city", "state", "unit"],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}]
|
||||||
|
|
||||||
|
messages = [
|
||||||
|
{
|
||||||
|
"role": "user",
|
||||||
|
"content": "What is the weather in Dallas, TX?"
|
||||||
|
},
|
||||||
|
]
|
||||||
|
|
||||||
|
stream = await gptoss_client.chat.completions.create(
|
||||||
|
model=GPT_OSS_MODEL_NAME, messages=messages, tools=tools, stream=True)
|
||||||
|
|
||||||
|
name = None
|
||||||
|
args_buf = ""
|
||||||
|
async for chunk in stream:
|
||||||
|
delta = chunk.choices[0].delta
|
||||||
|
if delta.tool_calls:
|
||||||
|
tc = delta.tool_calls[0]
|
||||||
|
if tc.function and tc.function.name:
|
||||||
|
name = tc.function.name
|
||||||
|
if tc.function and tc.function.arguments:
|
||||||
|
args_buf += tc.function.arguments
|
||||||
|
|
||||||
|
assert name is not None
|
||||||
|
assert len(args_buf) > 0
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_gpt_oss_multi_turn_chat(gptoss_client: OpenAI):
|
||||||
|
tools = [{
|
||||||
|
"type": "function",
|
||||||
|
"function": {
|
||||||
|
"name": "get_current_weather",
|
||||||
|
"description": "Get the current weather in a given location",
|
||||||
|
"parameters": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"city": {
|
||||||
|
"type": "string"
|
||||||
|
},
|
||||||
|
"state": {
|
||||||
|
"type": "string"
|
||||||
|
},
|
||||||
|
"unit": {
|
||||||
|
"type": "string",
|
||||||
|
"enum": ["celsius", "fahrenheit"],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": ["city", "state", "unit"],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}]
|
||||||
|
|
||||||
|
messages = [
|
||||||
|
{
|
||||||
|
"role": "system",
|
||||||
|
"content": "you are a helpful assistant"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"role": "user",
|
||||||
|
"content": "What is the weather in Dallas, TX?"
|
||||||
|
},
|
||||||
|
]
|
||||||
|
|
||||||
|
first = await gptoss_client.chat.completions.create(
|
||||||
|
model=GPT_OSS_MODEL_NAME,
|
||||||
|
messages=messages,
|
||||||
|
tools=tools,
|
||||||
|
temperature=0.0,
|
||||||
|
)
|
||||||
|
first_msg = first.choices[0].message
|
||||||
|
assert first_msg.tool_calls is not None and len(first_msg.tool_calls) > 0
|
||||||
|
tc = first_msg.tool_calls[0]
|
||||||
|
assert tc.function is not None and tc.function.name == "get_current_weather"
|
||||||
|
args1 = tc.function.arguments
|
||||||
|
assert args1 is not None and len(args1) > 0
|
||||||
|
|
||||||
|
messages.append({"role": "assistant", "content": args1})
|
||||||
|
messages.append({
|
||||||
|
"role": "user",
|
||||||
|
"content": "Now convert to celsius and return JSON only"
|
||||||
|
})
|
||||||
|
|
||||||
|
second = await gptoss_client.chat.completions.create(
|
||||||
|
model=GPT_OSS_MODEL_NAME,
|
||||||
|
messages=messages,
|
||||||
|
tools=tools,
|
||||||
|
temperature=0.0,
|
||||||
|
)
|
||||||
|
second_msg = second.choices[0].message
|
||||||
|
assert (second_msg.content is not None and len(second_msg.content) > 0) or \
|
||||||
|
(second_msg.tool_calls is not None and len(second_msg.tool_calls) > 0) # noqa: E501
|
||||||
|
|
||||||
|
|
||||||
MODEL_NAME = "openai-community/gpt2"
|
MODEL_NAME = "openai-community/gpt2"
|
||||||
CHAT_TEMPLATE = "Dummy chat template for testing {}"
|
CHAT_TEMPLATE = "Dummy chat template for testing {}"
|
||||||
BASE_MODEL_PATHS = [BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME)]
|
BASE_MODEL_PATHS = [BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME)]
|
||||||
|
|||||||
@ -8,8 +8,7 @@ from vllm.model_executor.layers.activation import SiluAndMul
|
|||||||
from vllm.platforms import current_platform
|
from vllm.platforms import current_platform
|
||||||
from vllm.scalar_type import scalar_types
|
from vllm.scalar_type import scalar_types
|
||||||
|
|
||||||
if not (current_platform.has_device_capability(100)
|
if not current_platform.has_device_capability(100):
|
||||||
and hasattr(torch.ops._C, "silu_and_mul_nvfp4_quant")):
|
|
||||||
pytest.skip(reason="Nvfp4 Requires compute capability of 10 or above.",
|
pytest.skip(reason="Nvfp4 Requires compute capability of 10 or above.",
|
||||||
allow_module_level=True)
|
allow_module_level=True)
|
||||||
|
|
||||||
|
|||||||
@ -60,9 +60,9 @@ DEVICES = ([
|
|||||||
# prefill stage(True) or decode stage(False)
|
# prefill stage(True) or decode stage(False)
|
||||||
STAGES = [True, False]
|
STAGES = [True, False]
|
||||||
|
|
||||||
NUM_RANDOM_SEEDS = 6
|
NUM_RANDOM_SEEDS = 2
|
||||||
|
|
||||||
VOCAB_PARALLEL_EMBEDDING_TEST_NUM_RANDOM_SEEDS = 128
|
VOCAB_PARALLEL_EMBEDDING_TEST_NUM_RANDOM_SEEDS = 2
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture(autouse=True)
|
@pytest.fixture(autouse=True)
|
||||||
|
|||||||
@ -250,7 +250,7 @@ def build_video_inputs_from_test_info(
|
|||||||
|
|
||||||
def apply_image_size_scaling(image, size: Union[float, tuple[int, int]],
|
def apply_image_size_scaling(image, size: Union[float, tuple[int, int]],
|
||||||
size_type: SizeType):
|
size_type: SizeType):
|
||||||
"""Applies a size scaler to one image; this can be a an image size factor,
|
"""Applies a size scaler to one image; this can be an image size factor,
|
||||||
which scales the image while maintaining the aspect ratio"""
|
which scales the image while maintaining the aspect ratio"""
|
||||||
# Special case for embeddings; if it's a tensor, it's only valid if we
|
# Special case for embeddings; if it's a tensor, it's only valid if we
|
||||||
# are considering size factors at constant scale, i.e., we just clone
|
# are considering size factors at constant scale, i.e., we just clone
|
||||||
|
|||||||
@ -42,7 +42,7 @@ def get_filtered_test_settings(
|
|||||||
else:
|
else:
|
||||||
assert test_info.prompt_formatter is not None
|
assert test_info.prompt_formatter is not None
|
||||||
|
|
||||||
# Everything looks okay; keep if this is has correct proc handling
|
# Everything looks okay; keep if this is correct proc handling
|
||||||
if (test_info.distributed_executor_backend
|
if (test_info.distributed_executor_backend
|
||||||
is not None) == new_proc_per_test:
|
is not None) == new_proc_per_test:
|
||||||
matching_tests[test_name] = test_info
|
matching_tests[test_name] = test_info
|
||||||
|
|||||||
147
tests/tool_use/test_openai_tool_parser.py
Normal file
147
tests/tool_use/test_openai_tool_parser.py
Normal file
@ -0,0 +1,147 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
|
||||||
|
import json
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from openai_harmony import (Conversation, DeveloperContent,
|
||||||
|
HarmonyEncodingName, Message, Role, SystemContent,
|
||||||
|
load_harmony_encoding)
|
||||||
|
|
||||||
|
from vllm.entrypoints.openai.protocol import FunctionCall, ToolCall
|
||||||
|
from vllm.entrypoints.openai.tool_parsers import OpenAIToolParser
|
||||||
|
from vllm.transformers_utils.tokenizer import get_tokenizer
|
||||||
|
|
||||||
|
MODEL = "gpt2"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(scope="module")
|
||||||
|
def openai_tokenizer():
|
||||||
|
# The parser does not use the tokenizer, but the constructor requires it.
|
||||||
|
return get_tokenizer(MODEL)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def openai_tool_parser(openai_tokenizer):
|
||||||
|
return OpenAIToolParser(openai_tokenizer)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(scope="module")
|
||||||
|
def harmony_encoding():
|
||||||
|
return load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
|
||||||
|
|
||||||
|
|
||||||
|
def assert_tool_calls(
|
||||||
|
actual_tool_calls: list[ToolCall],
|
||||||
|
expected_tool_calls: list[ToolCall],
|
||||||
|
):
|
||||||
|
assert len(actual_tool_calls) == len(expected_tool_calls)
|
||||||
|
|
||||||
|
for actual_tool_call, expected_tool_call in zip(actual_tool_calls,
|
||||||
|
expected_tool_calls):
|
||||||
|
assert isinstance(actual_tool_call.id, str)
|
||||||
|
assert len(actual_tool_call.id) > 16 # Default from protocol.py
|
||||||
|
assert actual_tool_call.type == "function"
|
||||||
|
assert actual_tool_call.function == expected_tool_call.function
|
||||||
|
|
||||||
|
|
||||||
|
def test_extract_tool_calls_no_tools(openai_tool_parser, harmony_encoding):
|
||||||
|
convo = Conversation.from_messages([
|
||||||
|
Message.from_role_and_content(
|
||||||
|
Role.SYSTEM,
|
||||||
|
SystemContent.new(),
|
||||||
|
),
|
||||||
|
Message.from_role_and_content(
|
||||||
|
Role.DEVELOPER,
|
||||||
|
DeveloperContent.new().with_instructions("Talk like a pirate!")),
|
||||||
|
Message.from_role_and_content(Role.USER, "Arrr, how be you?"),
|
||||||
|
Message.from_role_and_content(Role.ASSISTANT,
|
||||||
|
"This is a test").with_channel("final")
|
||||||
|
])
|
||||||
|
token_ids = harmony_encoding.render_conversation_for_completion(
|
||||||
|
convo, Role.ASSISTANT)
|
||||||
|
extracted_info = openai_tool_parser.extract_tool_calls(
|
||||||
|
"",
|
||||||
|
request=None,
|
||||||
|
token_ids=token_ids,
|
||||||
|
)
|
||||||
|
assert not extracted_info.tools_called
|
||||||
|
assert extracted_info.tool_calls == []
|
||||||
|
assert extracted_info.content == "This is a test"
|
||||||
|
|
||||||
|
|
||||||
|
def test_extract_tool_calls_single_tool(openai_tool_parser, harmony_encoding):
|
||||||
|
convo = Conversation.from_messages([
|
||||||
|
Message.from_role_and_content(Role.USER,
|
||||||
|
"What is the weather in Tokyo?"),
|
||||||
|
Message.from_role_and_content(
|
||||||
|
Role.ASSISTANT,
|
||||||
|
'User asks: "What is the weather in Tokyo?" We need to use get_current_weather tool.', # noqa: E501
|
||||||
|
).with_channel("analysis"),
|
||||||
|
Message.from_role_and_content(
|
||||||
|
Role.ASSISTANT,
|
||||||
|
'{"location": "Tokyo"}').with_channel("commentary").with_recipient(
|
||||||
|
"functions.get_current_weather").with_content_type("json"),
|
||||||
|
])
|
||||||
|
token_ids = harmony_encoding.render_conversation_for_completion(
|
||||||
|
convo, Role.ASSISTANT)
|
||||||
|
|
||||||
|
extracted_info = openai_tool_parser.extract_tool_calls(
|
||||||
|
"",
|
||||||
|
request=None,
|
||||||
|
token_ids=token_ids,
|
||||||
|
)
|
||||||
|
assert extracted_info.tools_called
|
||||||
|
expected_tool_calls = [
|
||||||
|
ToolCall(function=FunctionCall(
|
||||||
|
name="get_current_weather",
|
||||||
|
arguments=json.dumps({"location": "Tokyo"}),
|
||||||
|
))
|
||||||
|
]
|
||||||
|
assert_tool_calls(extracted_info.tool_calls, expected_tool_calls)
|
||||||
|
assert extracted_info.content is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_extract_tool_calls_multiple_tools(
|
||||||
|
openai_tool_parser,
|
||||||
|
harmony_encoding,
|
||||||
|
):
|
||||||
|
convo = Conversation.from_messages([
|
||||||
|
Message.from_role_and_content(
|
||||||
|
Role.USER, "What is the weather in Tokyo based on where I'm at?"),
|
||||||
|
Message.from_role_and_content(
|
||||||
|
Role.ASSISTANT,
|
||||||
|
'User asks: "What is the weather in Tokyo?" based on their location. We need to use get_current_weather tool and get_user_location tool.', # noqa: E501
|
||||||
|
).with_channel("analysis"),
|
||||||
|
Message.from_role_and_content(
|
||||||
|
Role.ASSISTANT,
|
||||||
|
'{"location": "Tokyo"}').with_channel("commentary").with_recipient(
|
||||||
|
"functions.get_current_weather").with_content_type("json"),
|
||||||
|
Message.from_role_and_content(
|
||||||
|
Role.ASSISTANT,
|
||||||
|
'{"location": "Tokyo"}').with_channel("commentary").with_recipient(
|
||||||
|
"functions.get_user_location").with_content_type("json"),
|
||||||
|
])
|
||||||
|
token_ids = harmony_encoding.render_conversation_for_completion(
|
||||||
|
convo,
|
||||||
|
Role.ASSISTANT,
|
||||||
|
)
|
||||||
|
|
||||||
|
extracted_info = openai_tool_parser.extract_tool_calls(
|
||||||
|
"",
|
||||||
|
request=None,
|
||||||
|
token_ids=token_ids,
|
||||||
|
)
|
||||||
|
assert extracted_info.tools_called
|
||||||
|
expected_tool_calls = [
|
||||||
|
ToolCall(function=FunctionCall(
|
||||||
|
name="get_current_weather",
|
||||||
|
arguments=json.dumps({"location": "Tokyo"}),
|
||||||
|
)),
|
||||||
|
ToolCall(function=FunctionCall(
|
||||||
|
name="get_user_location",
|
||||||
|
arguments=json.dumps({"location": "Tokyo"}),
|
||||||
|
))
|
||||||
|
]
|
||||||
|
assert_tool_calls(extracted_info.tool_calls, expected_tool_calls)
|
||||||
|
assert extracted_info.content is None
|
||||||
@ -393,7 +393,7 @@ class MockLoggingStatLogger(LoggingStatLogger):
|
|||||||
async def test_customize_loggers(monkeypatch):
|
async def test_customize_loggers(monkeypatch):
|
||||||
"""Test that we can customize the loggers.
|
"""Test that we can customize the loggers.
|
||||||
If a customized logger is provided at the init, it should
|
If a customized logger is provided at the init, it should
|
||||||
be used directly.
|
be added to the default loggers.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
with monkeypatch.context() as m, ExitStack() as after:
|
with monkeypatch.context() as m, ExitStack() as after:
|
||||||
@ -410,7 +410,8 @@ async def test_customize_loggers(monkeypatch):
|
|||||||
|
|
||||||
stat_loggers = engine.logger_manager.per_engine_logger_dict
|
stat_loggers = engine.logger_manager.per_engine_logger_dict
|
||||||
assert len(stat_loggers) == 1
|
assert len(stat_loggers) == 1
|
||||||
assert len(stat_loggers[0]) == 1
|
assert len(
|
||||||
|
stat_loggers[0]) == 2 # LoggingStatLogger + MockLoggingStatLogger
|
||||||
stat_loggers[0][0].log.assert_called_once()
|
stat_loggers[0][0].log.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
83
tests/v1/metrics/test_engine_logger_apis.py
Normal file
83
tests/v1/metrics/test_engine_logger_apis.py
Normal file
@ -0,0 +1,83 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
import copy
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from vllm.v1.engine.async_llm import AsyncEngineArgs, AsyncLLM
|
||||||
|
from vllm.v1.metrics.ray_wrappers import RayPrometheusStatLogger
|
||||||
|
|
||||||
|
|
||||||
|
class DummyStatLogger:
|
||||||
|
"""
|
||||||
|
A dummy stat logger for testing purposes.
|
||||||
|
Implements the minimal interface expected by StatLoggerManager.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, vllm_config, engine_idx):
|
||||||
|
self.vllm_config = vllm_config
|
||||||
|
self.engine_idx = engine_idx
|
||||||
|
self.recorded = []
|
||||||
|
self.logged = False
|
||||||
|
self.engine_initialized = False
|
||||||
|
|
||||||
|
def record(self, scheduler_stats, iteration_stats, engine_idx):
|
||||||
|
self.recorded.append((scheduler_stats, iteration_stats, engine_idx))
|
||||||
|
|
||||||
|
def log(self):
|
||||||
|
self.logged = True
|
||||||
|
|
||||||
|
def log_engine_initialized(self):
|
||||||
|
self.engine_initialized = True
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def log_stats_enabled_engine_args():
|
||||||
|
"""
|
||||||
|
Shared fixture providing common AsyncEngineArgs configuration
|
||||||
|
used across multiple tests.
|
||||||
|
"""
|
||||||
|
return AsyncEngineArgs(
|
||||||
|
model="distilbert/distilgpt2",
|
||||||
|
dtype="half",
|
||||||
|
disable_log_stats=False,
|
||||||
|
enforce_eager=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_llm_replace_default_loggers(
|
||||||
|
log_stats_enabled_engine_args):
|
||||||
|
"""
|
||||||
|
RayPrometheusStatLogger should replace the default PrometheusStatLogger
|
||||||
|
"""
|
||||||
|
|
||||||
|
engine = AsyncLLM.from_engine_args(log_stats_enabled_engine_args,
|
||||||
|
stat_loggers=[RayPrometheusStatLogger])
|
||||||
|
assert isinstance(engine.logger_manager.prometheus_logger,
|
||||||
|
RayPrometheusStatLogger)
|
||||||
|
engine.shutdown()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_llm_add_to_default_loggers(log_stats_enabled_engine_args):
|
||||||
|
"""
|
||||||
|
It's still possible to use custom stat loggers exclusively by passing
|
||||||
|
disable_log_stats=True in addition to a list of custom stat loggers.
|
||||||
|
"""
|
||||||
|
# Create engine_args with disable_log_stats=True for this test
|
||||||
|
disabled_log_engine_args = copy.deepcopy(log_stats_enabled_engine_args)
|
||||||
|
disabled_log_engine_args.disable_log_stats = True
|
||||||
|
|
||||||
|
# Disable default loggers; pass custom stat logger to the constructor
|
||||||
|
engine = AsyncLLM.from_engine_args(disabled_log_engine_args,
|
||||||
|
stat_loggers=[DummyStatLogger])
|
||||||
|
|
||||||
|
assert len(engine.logger_manager.per_engine_logger_dict[0]) == 1
|
||||||
|
assert isinstance(engine.logger_manager.per_engine_logger_dict[0][0],
|
||||||
|
DummyStatLogger)
|
||||||
|
|
||||||
|
# log_stats is still True, since custom stat loggers are used
|
||||||
|
assert engine.log_stats
|
||||||
|
|
||||||
|
engine.shutdown()
|
||||||
@ -822,7 +822,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[T], Generic[T]):
|
|||||||
and context_lens_tensor is not None \
|
and context_lens_tensor is not None \
|
||||||
and context_lens_tensor[:self.num_prefills].max() > 0:
|
and context_lens_tensor[:self.num_prefills].max() > 0:
|
||||||
|
|
||||||
# NOTE: it is recommend you read the `Chunked Prefill` section in
|
# NOTE: it is recommended you read the `Chunked Prefill` section in
|
||||||
# the comment at the top of the file before trying to understand
|
# the comment at the top of the file before trying to understand
|
||||||
# the following code
|
# the following code
|
||||||
|
|
||||||
|
|||||||
@ -6,7 +6,7 @@ KV cache helper for store.
|
|||||||
from collections import defaultdict
|
from collections import defaultdict
|
||||||
from collections.abc import Sequence
|
from collections.abc import Sequence
|
||||||
from concurrent.futures import CancelledError, Future
|
from concurrent.futures import CancelledError, Future
|
||||||
from typing import Optional, cast
|
from typing import Literal, Optional, Union, cast
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
@ -196,3 +196,51 @@ class KVOutputAggregator:
|
|||||||
output_future.add_done_callback(make_callback(i))
|
output_future.add_done_callback(make_callback(i))
|
||||||
|
|
||||||
return result_future
|
return result_future
|
||||||
|
|
||||||
|
|
||||||
|
def _make_src_and_dst_indices(
|
||||||
|
src_block_ids: list[int],
|
||||||
|
dst_block_ids: list[int],
|
||||||
|
src_device: Union[torch.device, str],
|
||||||
|
dst_device: Union[torch.device, str],
|
||||||
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||||
|
src_indices = torch.tensor(src_block_ids,
|
||||||
|
device=src_device,
|
||||||
|
dtype=torch.int64)
|
||||||
|
dst_indices = torch.tensor(dst_block_ids,
|
||||||
|
device=dst_device,
|
||||||
|
dtype=torch.int64)
|
||||||
|
return src_indices, dst_indices
|
||||||
|
|
||||||
|
|
||||||
|
def copy_kv_blocks(
|
||||||
|
src_kv_caches: dict[str, torch.Tensor],
|
||||||
|
dst_kv_caches: dict[str, torch.Tensor],
|
||||||
|
src_block_ids: list[int],
|
||||||
|
dst_block_ids: list[int],
|
||||||
|
direction: Literal["h2d", "d2h"],
|
||||||
|
) -> None:
|
||||||
|
"""Copy kv blocks between different buffers."""
|
||||||
|
if not src_kv_caches or not dst_kv_caches or \
|
||||||
|
not src_block_ids or not dst_block_ids or \
|
||||||
|
len(src_block_ids) != len(dst_block_ids):
|
||||||
|
return
|
||||||
|
|
||||||
|
src_device = next(iter(src_kv_caches.values())).device
|
||||||
|
dst_device = next(iter(dst_kv_caches.values())).device
|
||||||
|
|
||||||
|
src_indices, dst_indices = _make_src_and_dst_indices(
|
||||||
|
src_block_ids=src_block_ids,
|
||||||
|
dst_block_ids=dst_block_ids,
|
||||||
|
src_device=src_device,
|
||||||
|
dst_device=dst_device)
|
||||||
|
|
||||||
|
from vllm.platforms import current_platform
|
||||||
|
if direction == "h2d":
|
||||||
|
copy_fn = current_platform.insert_blocks_to_device
|
||||||
|
else:
|
||||||
|
copy_fn = current_platform.swap_out_blocks_to_host
|
||||||
|
for layer_name in src_kv_caches:
|
||||||
|
src_tensor = src_kv_caches[layer_name]
|
||||||
|
dst_tensor = dst_kv_caches[layer_name]
|
||||||
|
copy_fn(src_tensor, dst_tensor, src_indices, dst_indices)
|
||||||
|
|||||||
@ -61,6 +61,7 @@ except ImportError:
|
|||||||
_NIXL_SUPPORTED_XPUS = {
|
_NIXL_SUPPORTED_XPUS = {
|
||||||
"cuda": ("cuda", ),
|
"cuda": ("cuda", ),
|
||||||
"tpu": ("cpu", ),
|
"tpu": ("cpu", ),
|
||||||
|
"xpu": ("cpu", ),
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@ -717,7 +717,7 @@ class AsyncLLMEngine(EngineClient):
|
|||||||
# Stop the execute model loop in parallel workers until there
|
# Stop the execute model loop in parallel workers until there
|
||||||
# are more requests to process. This avoids waiting
|
# are more requests to process. This avoids waiting
|
||||||
# indefinitely in torch.distributed ops which may otherwise
|
# indefinitely in torch.distributed ops which may otherwise
|
||||||
# timeout, and unblocks the RPC thread in the workers so that
|
# time out, and unblocks the RPC thread in the workers so that
|
||||||
# they can process any other queued control plane messages,
|
# they can process any other queued control plane messages,
|
||||||
# such as add/remove lora adapters.
|
# such as add/remove lora adapters.
|
||||||
await engine.engine.stop_remote_worker_execution_loop_async()
|
await engine.engine.stop_remote_worker_execution_loop_async()
|
||||||
|
|||||||
@ -270,7 +270,7 @@ class MQLLMEngineClient(EngineClient):
|
|||||||
queue.put_nowait(request_output)
|
queue.put_nowait(request_output)
|
||||||
|
|
||||||
async def setup(self):
|
async def setup(self):
|
||||||
"""Setup the client before it starts sending server requests."""
|
"""Set up the client before it starts sending server requests."""
|
||||||
|
|
||||||
# Start output_loop
|
# Start output_loop
|
||||||
if self.output_loop is None:
|
if self.output_loop is None:
|
||||||
|
|||||||
@ -1,5 +1,8 @@
|
|||||||
# SPDX-License-Identifier: Apache-2.0
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
import datetime
|
import datetime
|
||||||
import json
|
import json
|
||||||
from collections.abc import Iterable, Sequence
|
from collections.abc import Iterable, Sequence
|
||||||
@ -18,7 +21,8 @@ from openai_harmony import (Author, Conversation, DeveloperContent,
|
|||||||
Role, StreamableParser, SystemContent, TextContent,
|
Role, StreamableParser, SystemContent, TextContent,
|
||||||
ToolDescription, load_harmony_encoding)
|
ToolDescription, load_harmony_encoding)
|
||||||
|
|
||||||
from vllm.entrypoints.openai.protocol import ResponseInputOutputItem
|
from vllm.entrypoints.openai.protocol import (ChatCompletionToolsParam,
|
||||||
|
ResponseInputOutputItem)
|
||||||
from vllm.utils import random_uuid
|
from vllm.utils import random_uuid
|
||||||
|
|
||||||
REASONING_EFFORT = {
|
REASONING_EFFORT = {
|
||||||
@ -63,13 +67,29 @@ def get_system_message(
|
|||||||
return sys_msg
|
return sys_msg
|
||||||
|
|
||||||
|
|
||||||
def get_developer_message(instructions: Optional[str] = None,
|
def create_tool_definition(tool: Union[ChatCompletionToolsParam, Tool]):
|
||||||
tools: Optional[list[Tool]] = None) -> Message:
|
if isinstance(tool, ChatCompletionToolsParam):
|
||||||
|
return ToolDescription.new(
|
||||||
|
name=tool.function.name,
|
||||||
|
description=tool.function.description,
|
||||||
|
parameters=tool.function.parameters,
|
||||||
|
)
|
||||||
|
return ToolDescription.new(
|
||||||
|
name=tool.name,
|
||||||
|
description=tool.description,
|
||||||
|
parameters=tool.parameters,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def get_developer_message(
|
||||||
|
instructions: Optional[str] = None,
|
||||||
|
tools: Optional[list[Union[Tool, ChatCompletionToolsParam]]] = None,
|
||||||
|
) -> Message:
|
||||||
dev_msg_content = DeveloperContent.new()
|
dev_msg_content = DeveloperContent.new()
|
||||||
if instructions is not None:
|
if instructions is not None:
|
||||||
dev_msg_content = dev_msg_content.with_instructions(instructions)
|
dev_msg_content = dev_msg_content.with_instructions(instructions)
|
||||||
if tools is not None:
|
if tools is not None:
|
||||||
function_tools = []
|
function_tools: list[Union[Tool, ChatCompletionToolsParam]] = []
|
||||||
for tool in tools:
|
for tool in tools:
|
||||||
if tool.type in ("web_search_preview", "code_interpreter"):
|
if tool.type in ("web_search_preview", "code_interpreter"):
|
||||||
# These are built-in tools that are added to the system message.
|
# These are built-in tools that are added to the system message.
|
||||||
@ -80,11 +100,7 @@ def get_developer_message(instructions: Optional[str] = None,
|
|||||||
raise ValueError(f"tool type {tool.type} not supported")
|
raise ValueError(f"tool type {tool.type} not supported")
|
||||||
if function_tools:
|
if function_tools:
|
||||||
function_tool_descriptions = [
|
function_tool_descriptions = [
|
||||||
ToolDescription.new(
|
create_tool_definition(tool) for tool in function_tools
|
||||||
name=tool.name,
|
|
||||||
description=tool.description,
|
|
||||||
parameters=tool.parameters,
|
|
||||||
) for tool in function_tools
|
|
||||||
]
|
]
|
||||||
dev_msg_content = dev_msg_content.with_function_tools(
|
dev_msg_content = dev_msg_content.with_function_tools(
|
||||||
function_tool_descriptions)
|
function_tool_descriptions)
|
||||||
@ -148,16 +164,46 @@ def parse_response_input(
|
|||||||
return msg
|
return msg
|
||||||
|
|
||||||
|
|
||||||
def parse_chat_input(chat_msg) -> Message:
|
def parse_chat_input(chat_msg) -> list[Message]:
|
||||||
role = chat_msg["role"]
|
if not isinstance(chat_msg, dict):
|
||||||
content = chat_msg["content"]
|
# Handle Pydantic models
|
||||||
|
chat_msg = chat_msg.model_dump(exclude_none=True)
|
||||||
|
|
||||||
|
role = chat_msg.get("role")
|
||||||
|
|
||||||
|
# Assistant message with tool calls
|
||||||
|
tool_calls = chat_msg.get("tool_calls")
|
||||||
|
if role == "assistant" and tool_calls:
|
||||||
|
msgs: list[Message] = []
|
||||||
|
for call in tool_calls:
|
||||||
|
func = call.get("function", {})
|
||||||
|
name = func.get("name", "")
|
||||||
|
arguments = func.get("arguments", "") or ""
|
||||||
|
msg = Message.from_role_and_content(Role.ASSISTANT, arguments)
|
||||||
|
msg = msg.with_channel("commentary")
|
||||||
|
msg = msg.with_recipient(f"functions.{name}")
|
||||||
|
msg = msg.with_content_type("json")
|
||||||
|
msgs.append(msg)
|
||||||
|
return msgs
|
||||||
|
|
||||||
|
# Tool role message (tool output)
|
||||||
|
if role == "tool":
|
||||||
|
name = chat_msg.get("name", "")
|
||||||
|
content = chat_msg.get("content", "") or ""
|
||||||
|
msg = Message.from_author_and_content(
|
||||||
|
Author.new(Role.TOOL, f"functions.{name}"),
|
||||||
|
content).with_channel("commentary")
|
||||||
|
return [msg]
|
||||||
|
|
||||||
|
# Default: user/assistant/system messages with content
|
||||||
|
content = chat_msg.get("content", "")
|
||||||
if isinstance(content, str):
|
if isinstance(content, str):
|
||||||
contents = [TextContent(text=content)]
|
contents = [TextContent(text=content)]
|
||||||
else:
|
else:
|
||||||
# TODO: Support refusal.
|
# TODO: Support refusal.
|
||||||
contents = [TextContent(text=c.get("text", "")) for c in content]
|
contents = [TextContent(text=c.get("text", "")) for c in content]
|
||||||
msg = Message.from_role_and_contents(role, contents)
|
msg = Message.from_role_and_contents(role, contents)
|
||||||
return msg
|
return [msg]
|
||||||
|
|
||||||
|
|
||||||
def render_for_completion(messages: list[Message]) -> list[int]:
|
def render_for_completion(messages: list[Message]) -> list[int]:
|
||||||
|
|||||||
@ -6,7 +6,7 @@ import json
|
|||||||
import time
|
import time
|
||||||
from collections.abc import AsyncGenerator, AsyncIterator
|
from collections.abc import AsyncGenerator, AsyncIterator
|
||||||
from collections.abc import Sequence as GenericSequence
|
from collections.abc import Sequence as GenericSequence
|
||||||
from typing import Callable, Final, Optional, Union
|
from typing import TYPE_CHECKING, Callable, Final, Optional, Union
|
||||||
|
|
||||||
import jinja2
|
import jinja2
|
||||||
import partial_json_parser
|
import partial_json_parser
|
||||||
@ -489,6 +489,8 @@ class OpenAIServingChat(OpenAIServing):
|
|||||||
get_streamable_parser_for_assistant()
|
get_streamable_parser_for_assistant()
|
||||||
for _ in range(num_choices)
|
for _ in range(num_choices)
|
||||||
]
|
]
|
||||||
|
harmony_tools_streamed = [False] * num_choices
|
||||||
|
tools_streamed = [False] * num_choices
|
||||||
|
|
||||||
if isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam):
|
if isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam):
|
||||||
tool_choice_function_name = request.tool_choice.function.name
|
tool_choice_function_name = request.tool_choice.function.name
|
||||||
@ -662,13 +664,11 @@ class OpenAIServingChat(OpenAIServing):
|
|||||||
|
|
||||||
if self.use_harmony:
|
if self.use_harmony:
|
||||||
harmony_parser = harmony_parsers[i]
|
harmony_parser = harmony_parsers[i]
|
||||||
|
prev_recipient = harmony_parser.current_recipient
|
||||||
for token_id in output.token_ids:
|
for token_id in output.token_ids:
|
||||||
harmony_parser.process(token_id)
|
harmony_parser.process(token_id)
|
||||||
is_reasoning = \
|
cur_channel = harmony_parser.current_channel
|
||||||
harmony_parser.current_channel == "analysis"
|
cur_recipient = harmony_parser.current_recipient
|
||||||
if not request.include_reasoning and is_reasoning:
|
|
||||||
# Skip the reasoning content.
|
|
||||||
continue
|
|
||||||
delta_text = harmony_parser.last_content_delta or ""
|
delta_text = harmony_parser.last_content_delta or ""
|
||||||
else:
|
else:
|
||||||
delta_text = output.text
|
delta_text = output.text
|
||||||
@ -681,8 +681,7 @@ class OpenAIServingChat(OpenAIServing):
|
|||||||
delta_message: Optional[DeltaMessage]
|
delta_message: Optional[DeltaMessage]
|
||||||
|
|
||||||
# just update previous_texts and previous_token_ids
|
# just update previous_texts and previous_token_ids
|
||||||
if ((tool_choice_auto or self.reasoning_parser)
|
if tool_choice_auto or self.reasoning_parser:
|
||||||
and not self.use_harmony):
|
|
||||||
assert previous_texts is not None
|
assert previous_texts is not None
|
||||||
assert all_previous_token_ids is not None
|
assert all_previous_token_ids is not None
|
||||||
previous_text = previous_texts[i]
|
previous_text = previous_texts[i]
|
||||||
@ -696,11 +695,54 @@ class OpenAIServingChat(OpenAIServing):
|
|||||||
current_token_ids = as_list(output.token_ids)
|
current_token_ids = as_list(output.token_ids)
|
||||||
|
|
||||||
if self.use_harmony:
|
if self.use_harmony:
|
||||||
if is_reasoning:
|
if cur_channel == "final":
|
||||||
delta_message = DeltaMessage(
|
|
||||||
reasoning_content=delta_text)
|
|
||||||
else:
|
|
||||||
delta_message = DeltaMessage(content=delta_text)
|
delta_message = DeltaMessage(content=delta_text)
|
||||||
|
elif cur_channel == "analysis":
|
||||||
|
if request.include_reasoning:
|
||||||
|
delta_message = DeltaMessage(
|
||||||
|
reasoning_content=delta_text)
|
||||||
|
else:
|
||||||
|
delta_message = None
|
||||||
|
elif (cur_channel == "commentary" and cur_recipient
|
||||||
|
and cur_recipient.startswith("functions.")):
|
||||||
|
# Count completed tool calls to determine index
|
||||||
|
base_index = 0
|
||||||
|
for msg in harmony_parser.messages:
|
||||||
|
if (msg.channel == "commentary"
|
||||||
|
and msg.recipient
|
||||||
|
and msg.recipient.startswith(
|
||||||
|
"functions.")):
|
||||||
|
base_index += 1
|
||||||
|
|
||||||
|
if prev_recipient != cur_recipient:
|
||||||
|
tool_name = cur_recipient.split(
|
||||||
|
"functions.", 1)[1]
|
||||||
|
delta_message = DeltaMessage(tool_calls=[
|
||||||
|
DeltaToolCall(
|
||||||
|
id=make_tool_call_id(),
|
||||||
|
type="function",
|
||||||
|
function=DeltaFunctionCall(
|
||||||
|
name=tool_name,
|
||||||
|
arguments="",
|
||||||
|
),
|
||||||
|
index=base_index,
|
||||||
|
)
|
||||||
|
])
|
||||||
|
elif delta_text:
|
||||||
|
delta_message = DeltaMessage(tool_calls=[
|
||||||
|
DeltaToolCall(
|
||||||
|
index=base_index,
|
||||||
|
function=DeltaFunctionCall(
|
||||||
|
arguments=delta_text),
|
||||||
|
)
|
||||||
|
])
|
||||||
|
else:
|
||||||
|
delta_message = None
|
||||||
|
|
||||||
|
if delta_message is not None:
|
||||||
|
harmony_tools_streamed[i] = True
|
||||||
|
else:
|
||||||
|
delta_message = None
|
||||||
# handle streaming deltas for tools with named tool_choice
|
# handle streaming deltas for tools with named tool_choice
|
||||||
elif tool_choice_function_name:
|
elif tool_choice_function_name:
|
||||||
if (self.reasoning_parser and not reasoning_end_arr[i]
|
if (self.reasoning_parser and not reasoning_end_arr[i]
|
||||||
@ -758,6 +800,7 @@ class OpenAIServingChat(OpenAIServing):
|
|||||||
delta_message = DeltaMessage(tool_calls=[
|
delta_message = DeltaMessage(tool_calls=[
|
||||||
delta_tool_call,
|
delta_tool_call,
|
||||||
])
|
])
|
||||||
|
tools_streamed[i] = True
|
||||||
|
|
||||||
elif request.tool_choice == "required":
|
elif request.tool_choice == "required":
|
||||||
assert previous_texts is not None
|
assert previous_texts is not None
|
||||||
@ -783,6 +826,7 @@ class OpenAIServingChat(OpenAIServing):
|
|||||||
if (delta_message and delta_message.tool_calls and
|
if (delta_message and delta_message.tool_calls and
|
||||||
delta_message.tool_calls[0].id is not None):
|
delta_message.tool_calls[0].id is not None):
|
||||||
history_tool_call_cnt += 1
|
history_tool_call_cnt += 1
|
||||||
|
tools_streamed[i] = True
|
||||||
|
|
||||||
# update the previous values for the next iteration
|
# update the previous values for the next iteration
|
||||||
previous_texts[i] = current_text
|
previous_texts[i] = current_text
|
||||||
@ -859,6 +903,8 @@ class OpenAIServingChat(OpenAIServing):
|
|||||||
current_token_ids=current_token_ids,
|
current_token_ids=current_token_ids,
|
||||||
delta_token_ids=delta_token_ids,
|
delta_token_ids=delta_token_ids,
|
||||||
request=request))
|
request=request))
|
||||||
|
if delta_message and delta_message.tool_calls:
|
||||||
|
tools_streamed[i] = True
|
||||||
# when only tool calls
|
# when only tool calls
|
||||||
elif tool_choice_auto:
|
elif tool_choice_auto:
|
||||||
assert tool_parser is not None
|
assert tool_parser is not None
|
||||||
@ -871,6 +917,8 @@ class OpenAIServingChat(OpenAIServing):
|
|||||||
current_token_ids=current_token_ids,
|
current_token_ids=current_token_ids,
|
||||||
delta_token_ids=output.token_ids,
|
delta_token_ids=output.token_ids,
|
||||||
request=request))
|
request=request))
|
||||||
|
if delta_message and delta_message.tool_calls:
|
||||||
|
tools_streamed[i] = True
|
||||||
|
|
||||||
# when only reasoning
|
# when only reasoning
|
||||||
elif self.reasoning_parser:
|
elif self.reasoning_parser:
|
||||||
@ -907,7 +955,10 @@ class OpenAIServingChat(OpenAIServing):
|
|||||||
# wasn't ready to send a token, then
|
# wasn't ready to send a token, then
|
||||||
# get the next token without streaming a chunk
|
# get the next token without streaming a chunk
|
||||||
if delta_message is None:
|
if delta_message is None:
|
||||||
continue
|
if output.finish_reason is None:
|
||||||
|
continue
|
||||||
|
else:
|
||||||
|
delta_message = DeltaMessage()
|
||||||
|
|
||||||
# Log streaming delta if output logging is enabled
|
# Log streaming delta if output logging is enabled
|
||||||
if self.enable_log_outputs and self.request_logger:
|
if self.enable_log_outputs and self.request_logger:
|
||||||
@ -993,12 +1044,18 @@ class OpenAIServingChat(OpenAIServing):
|
|||||||
])
|
])
|
||||||
|
|
||||||
# Send the finish response for each request.n only once
|
# Send the finish response for each request.n only once
|
||||||
|
if auto_tools_called or tools_streamed[i] or (
|
||||||
|
self.use_harmony
|
||||||
|
and harmony_tools_streamed[i]):
|
||||||
|
finish_reason_ = "tool_calls"
|
||||||
|
else:
|
||||||
|
finish_reason_ = output.finish_reason \
|
||||||
|
if output.finish_reason else "stop"
|
||||||
choice_data = ChatCompletionResponseStreamChoice(
|
choice_data = ChatCompletionResponseStreamChoice(
|
||||||
index=i,
|
index=i,
|
||||||
delta=delta_message,
|
delta=delta_message,
|
||||||
logprobs=logprobs,
|
logprobs=logprobs,
|
||||||
finish_reason=output.finish_reason
|
finish_reason=finish_reason_,
|
||||||
if not auto_tools_called else "tool_calls",
|
|
||||||
stop_reason=output.stop_reason,
|
stop_reason=output.stop_reason,
|
||||||
token_ids=(as_list(output.token_ids)
|
token_ids=(as_list(output.token_ids)
|
||||||
if request.return_token_ids else None))
|
if request.return_token_ids else None))
|
||||||
@ -1131,31 +1188,32 @@ class OpenAIServingChat(OpenAIServing):
|
|||||||
logprobs = None
|
logprobs = None
|
||||||
|
|
||||||
if self.use_harmony:
|
if self.use_harmony:
|
||||||
reasoning_content, final_content, is_tool_call = (
|
if TYPE_CHECKING:
|
||||||
parse_chat_output(token_ids))
|
assert self.tool_parser is not None
|
||||||
if not request.include_reasoning:
|
tool_parser = self.tool_parser(tokenizer)
|
||||||
reasoning_content = None
|
# NOTE: We use token_ids for openai tool parser
|
||||||
|
tool_call_info = tool_parser.extract_tool_calls(
|
||||||
if is_tool_call:
|
"",
|
||||||
# TODO(woosuk): Implement tool call for gpt-oss.
|
request=request,
|
||||||
# For now, only Responses API supports tool call for
|
token_ids=token_ids, # type: ignore
|
||||||
# gpt-oss.
|
)
|
||||||
raise NotImplementedError(
|
reasoning_content, content = None, tool_call_info.content
|
||||||
"Tool call in Chat Completion API is not supported "
|
if request.include_reasoning:
|
||||||
"for gpt-oss yet. Please use Responses API instead.")
|
reasoning_content, content, _ = parse_chat_output(
|
||||||
else:
|
token_ids)
|
||||||
# Normal message
|
message = ChatMessage(
|
||||||
message = ChatMessage(
|
role=role,
|
||||||
role=role,
|
reasoning_content=reasoning_content,
|
||||||
reasoning_content=reasoning_content,
|
content=content,
|
||||||
content=final_content,
|
tool_calls=tool_call_info.tool_calls,
|
||||||
)
|
)
|
||||||
|
|
||||||
choice_data = ChatCompletionResponseChoice(
|
choice_data = ChatCompletionResponseChoice(
|
||||||
index=output.index,
|
index=output.index,
|
||||||
message=message,
|
message=message,
|
||||||
logprobs=logprobs,
|
logprobs=logprobs,
|
||||||
finish_reason="tool_calls" if is_tool_call else
|
finish_reason="tool_calls"
|
||||||
|
if tool_call_info.tools_called else
|
||||||
output.finish_reason if output.finish_reason else "stop",
|
output.finish_reason if output.finish_reason else "stop",
|
||||||
stop_reason=output.stop_reason,
|
stop_reason=output.stop_reason,
|
||||||
)
|
)
|
||||||
@ -1419,9 +1477,10 @@ class OpenAIServingChat(OpenAIServing):
|
|||||||
step_top_logprobs = top_logprobs[i]
|
step_top_logprobs = top_logprobs[i]
|
||||||
if step_top_logprobs is None or step_top_logprobs.get(
|
if step_top_logprobs is None or step_top_logprobs.get(
|
||||||
token_id) is None:
|
token_id) is None:
|
||||||
token = tokenizer.decode(token_id)
|
|
||||||
if should_return_as_token_id:
|
if should_return_as_token_id:
|
||||||
token = f"token_id:{token_id}"
|
token = f"token_id:{token_id}"
|
||||||
|
else:
|
||||||
|
token = tokenizer.decode(token_id)
|
||||||
|
|
||||||
logprobs_content.append(
|
logprobs_content.append(
|
||||||
ChatCompletionLogProbsContent(
|
ChatCompletionLogProbsContent(
|
||||||
@ -1503,12 +1562,12 @@ class OpenAIServingChat(OpenAIServing):
|
|||||||
messages.append(sys_msg)
|
messages.append(sys_msg)
|
||||||
|
|
||||||
# Add developer message.
|
# Add developer message.
|
||||||
dev_msg = get_developer_message()
|
dev_msg = get_developer_message(tools=request.tools)
|
||||||
messages.append(dev_msg)
|
messages.append(dev_msg)
|
||||||
|
|
||||||
# Add user message.
|
# Add user message.
|
||||||
for chat_msg in request.messages:
|
for chat_msg in request.messages:
|
||||||
messages.append(parse_chat_input(chat_msg))
|
messages.extend(parse_chat_input(chat_msg))
|
||||||
|
|
||||||
# Render prompt token ids.
|
# Render prompt token ids.
|
||||||
prompt_token_ids = render_for_completion(messages)
|
prompt_token_ids = render_for_completion(messages)
|
||||||
|
|||||||
@ -16,6 +16,7 @@ from .llama4_pythonic_tool_parser import Llama4PythonicToolParser
|
|||||||
from .llama_tool_parser import Llama3JsonToolParser
|
from .llama_tool_parser import Llama3JsonToolParser
|
||||||
from .minimax_tool_parser import MinimaxToolParser
|
from .minimax_tool_parser import MinimaxToolParser
|
||||||
from .mistral_tool_parser import MistralToolParser
|
from .mistral_tool_parser import MistralToolParser
|
||||||
|
from .openai_tool_parser import OpenAIToolParser
|
||||||
from .phi4mini_tool_parser import Phi4MiniJsonToolParser
|
from .phi4mini_tool_parser import Phi4MiniJsonToolParser
|
||||||
from .pythonic_tool_parser import PythonicToolParser
|
from .pythonic_tool_parser import PythonicToolParser
|
||||||
from .qwen3coder_tool_parser import Qwen3CoderToolParser
|
from .qwen3coder_tool_parser import Qwen3CoderToolParser
|
||||||
@ -46,4 +47,5 @@ __all__ = [
|
|||||||
"Qwen3CoderToolParser",
|
"Qwen3CoderToolParser",
|
||||||
"SeedOssToolParser",
|
"SeedOssToolParser",
|
||||||
"Step3ToolParser",
|
"Step3ToolParser",
|
||||||
|
"OpenAIToolParser",
|
||||||
]
|
]
|
||||||
|
|||||||
73
vllm/entrypoints/openai/tool_parsers/openai_tool_parser.py
Normal file
73
vllm/entrypoints/openai/tool_parsers/openai_tool_parser.py
Normal file
@ -0,0 +1,73 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from collections.abc import Sequence
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from vllm.entrypoints.harmony_utils import parse_output_into_messages
|
||||||
|
from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
|
||||||
|
DeltaMessage,
|
||||||
|
ExtractedToolCallInformation,
|
||||||
|
FunctionCall, ToolCall)
|
||||||
|
from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import (
|
||||||
|
ToolParser, ToolParserManager)
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from vllm.transformers_utils.tokenizer import AnyTokenizer
|
||||||
|
|
||||||
|
|
||||||
|
@ToolParserManager.register_module("openai")
|
||||||
|
class OpenAIToolParser(ToolParser):
|
||||||
|
|
||||||
|
def __init__(self, tokenizer: AnyTokenizer):
|
||||||
|
super().__init__(tokenizer)
|
||||||
|
|
||||||
|
def extract_tool_calls(
|
||||||
|
self,
|
||||||
|
model_output: str,
|
||||||
|
request: ChatCompletionRequest,
|
||||||
|
token_ids: Sequence[int] | None = None,
|
||||||
|
) -> ExtractedToolCallInformation:
|
||||||
|
if token_ids is None:
|
||||||
|
raise NotImplementedError(
|
||||||
|
"OpenAIToolParser requires token IDs and does not support text-based extraction." # noqa: E501
|
||||||
|
)
|
||||||
|
|
||||||
|
parser = parse_output_into_messages(token_ids)
|
||||||
|
tool_calls = []
|
||||||
|
final_content = None
|
||||||
|
|
||||||
|
if len(parser.messages) > 0:
|
||||||
|
for msg in parser.messages:
|
||||||
|
if msg.recipient and msg.recipient.startswith("functions."):
|
||||||
|
tool_calls.append(
|
||||||
|
ToolCall(
|
||||||
|
type="function",
|
||||||
|
function=FunctionCall(
|
||||||
|
name=msg.recipient.split("functions.")[1],
|
||||||
|
arguments=msg.content[0].text,
|
||||||
|
),
|
||||||
|
))
|
||||||
|
elif msg.channel == "final":
|
||||||
|
final_content = msg.content[0].text
|
||||||
|
|
||||||
|
return ExtractedToolCallInformation(
|
||||||
|
tools_called=len(tool_calls) > 0,
|
||||||
|
tool_calls=tool_calls,
|
||||||
|
content=final_content,
|
||||||
|
)
|
||||||
|
|
||||||
|
def extract_tool_calls_streaming(
|
||||||
|
self,
|
||||||
|
previous_text: str,
|
||||||
|
current_text: str,
|
||||||
|
delta_text: str,
|
||||||
|
previous_token_ids: Sequence[int],
|
||||||
|
current_token_ids: Sequence[int],
|
||||||
|
delta_token_ids: Sequence[int],
|
||||||
|
request: ChatCompletionRequest,
|
||||||
|
) -> DeltaMessage | None:
|
||||||
|
raise NotImplementedError(
|
||||||
|
"Not being used, manual parsing in serving_chat.py" # noqa: E501
|
||||||
|
)
|
||||||
@ -0,0 +1,146 @@
|
|||||||
|
{
|
||||||
|
"1": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"2": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"4": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"8": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"16": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"24": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"32": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"48": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"64": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"96": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"128": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"256": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"512": {
|
||||||
|
"BLOCK_SIZE_M": 32,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 16,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"1024": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"1536": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"2048": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"3072": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"4096": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 16,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,146 @@
|
|||||||
|
{
|
||||||
|
"1": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"2": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"4": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"8": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"16": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 8,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"24": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"32": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"48": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"64": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"96": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"128": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"256": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"512": {
|
||||||
|
"BLOCK_SIZE_M": 32,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 32,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"1024": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 16,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"1536": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"2048": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"3072": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 16,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"4096": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,146 @@
|
|||||||
|
{
|
||||||
|
"1": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"2": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"4": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"8": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 16,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"16": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 16,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"24": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 16,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"32": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"48": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"64": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"96": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"128": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"256": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"512": {
|
||||||
|
"BLOCK_SIZE_M": 32,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"1024": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 16,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"1536": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 16,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"2048": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 16,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"3072": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 16,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"4096": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 16,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,146 @@
|
|||||||
|
{
|
||||||
|
"1": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"2": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"4": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 8,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"8": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"16": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"24": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"32": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"48": {
|
||||||
|
"BLOCK_SIZE_M": 32,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"64": {
|
||||||
|
"BLOCK_SIZE_M": 32,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"96": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"128": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"256": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"512": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"1024": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"1536": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"2048": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"3072": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"4096": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,146 @@
|
|||||||
|
{
|
||||||
|
"1": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"2": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"4": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 8,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"8": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"16": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"24": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"32": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"48": {
|
||||||
|
"BLOCK_SIZE_M": 32,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"64": {
|
||||||
|
"BLOCK_SIZE_M": 32,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"96": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"128": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"256": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"512": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"1024": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"1536": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"2048": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"3072": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"4096": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,146 @@
|
|||||||
|
{
|
||||||
|
"1": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"2": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 16,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"4": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 64,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"8": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"16": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"24": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"32": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"48": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 64,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"64": {
|
||||||
|
"BLOCK_SIZE_M": 32,
|
||||||
|
"BLOCK_SIZE_N": 256,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 64,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"96": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"128": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"256": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 16,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"512": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 32,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"1024": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 32,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"1536": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 64,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"2048": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 64,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"3072": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 32,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"4096": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 256,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 32,
|
||||||
|
"num_warps": 8,
|
||||||
|
"num_stages": 4
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,146 @@
|
|||||||
|
{
|
||||||
|
"1": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"2": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 16,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"4": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"8": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"16": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"24": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"32": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"48": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"64": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 32,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"96": {
|
||||||
|
"BLOCK_SIZE_M": 32,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 16,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"128": {
|
||||||
|
"BLOCK_SIZE_M": 32,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 16,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"256": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"512": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 32,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"1024": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 32,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"1536": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"2048": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"3072": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"4096": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,146 @@
|
|||||||
|
{
|
||||||
|
"1": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"2": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"4": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"8": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"16": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"24": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"32": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"48": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"64": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"96": {
|
||||||
|
"BLOCK_SIZE_M": 32,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"128": {
|
||||||
|
"BLOCK_SIZE_M": 32,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"256": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"512": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"1024": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"1536": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"2048": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"3072": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"4096": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,146 @@
|
|||||||
|
{
|
||||||
|
"1": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"2": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 32,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"4": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 64,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"8": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 32,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"16": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"24": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"32": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"48": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"64": {
|
||||||
|
"BLOCK_SIZE_M": 16,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 32,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"96": {
|
||||||
|
"BLOCK_SIZE_M": 32,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"128": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 256,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 8,
|
||||||
|
"num_stages": 5
|
||||||
|
},
|
||||||
|
"256": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 1,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 3
|
||||||
|
},
|
||||||
|
"512": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 16,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"1024": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 32,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"1536": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 32,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"2048": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 128,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 32,
|
||||||
|
"num_warps": 4,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"3072": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 256,
|
||||||
|
"BLOCK_SIZE_K": 128,
|
||||||
|
"GROUP_SIZE_M": 64,
|
||||||
|
"num_warps": 8,
|
||||||
|
"num_stages": 4
|
||||||
|
},
|
||||||
|
"4096": {
|
||||||
|
"BLOCK_SIZE_M": 64,
|
||||||
|
"BLOCK_SIZE_N": 256,
|
||||||
|
"BLOCK_SIZE_K": 256,
|
||||||
|
"GROUP_SIZE_M": 32,
|
||||||
|
"num_warps": 8,
|
||||||
|
"num_stages": 4
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -57,13 +57,14 @@ def _valid_deep_gemm(hidden_states: torch.Tensor, w1: torch.Tensor,
|
|||||||
if not _valid_deep_gemm_shape(M, N, K):
|
if not _valid_deep_gemm_shape(M, N, K):
|
||||||
logger.debug_once(
|
logger.debug_once(
|
||||||
"DeepGemm disabled due to unaligned problem size. "
|
"DeepGemm disabled due to unaligned problem size. "
|
||||||
"M: %s, N: %s, K: %s. M should >= align size "
|
"M: %s, N: %s, K: %s. M should >= %s "
|
||||||
"and N and K must be multiples of %s."
|
"and N and K must be multiples of %s. "
|
||||||
"This is not an error and we will fall back to triton.",
|
"This is not an error and we will fall back to triton.",
|
||||||
M,
|
M,
|
||||||
N,
|
N,
|
||||||
K,
|
K,
|
||||||
align,
|
align,
|
||||||
|
align,
|
||||||
)
|
)
|
||||||
return False
|
return False
|
||||||
elif N <= 512:
|
elif N <= 512:
|
||||||
|
|||||||
@ -19,7 +19,7 @@ def awq_dequantize_kernel(
|
|||||||
num_rows, # input num rows in qweight
|
num_rows, # input num rows in qweight
|
||||||
BLOCK_SIZE_X: tl.constexpr,
|
BLOCK_SIZE_X: tl.constexpr,
|
||||||
BLOCK_SIZE_Y: tl.constexpr):
|
BLOCK_SIZE_Y: tl.constexpr):
|
||||||
# Setup the pids.
|
# Set up the pids.
|
||||||
pid_x = tl.program_id(axis=0)
|
pid_x = tl.program_id(axis=0)
|
||||||
pid_y = tl.program_id(axis=1)
|
pid_y = tl.program_id(axis=1)
|
||||||
|
|
||||||
|
|||||||
@ -128,7 +128,7 @@ class QuantizationConfig(ABC):
|
|||||||
@staticmethod
|
@staticmethod
|
||||||
def get_from_keys_or(config: dict[str, Any], keys: list[str],
|
def get_from_keys_or(config: dict[str, Any], keys: list[str],
|
||||||
default: Any) -> Any:
|
default: Any) -> Any:
|
||||||
"""Get a optional value from the model's quantization config."""
|
"""Get an optional value from the model's quantization config."""
|
||||||
try:
|
try:
|
||||||
return QuantizationConfig.get_from_keys(config, keys)
|
return QuantizationConfig.get_from_keys(config, keys)
|
||||||
except ValueError:
|
except ValueError:
|
||||||
|
|||||||
@ -256,7 +256,7 @@ class GptOssForCausalLMConfig(VerifyAndUpdateConfig):
|
|||||||
def verify_and_update_config(vllm_config: "VllmConfig") -> None:
|
def verify_and_update_config(vllm_config: "VllmConfig") -> None:
|
||||||
decoding_config = vllm_config.decoding_config
|
decoding_config = vllm_config.decoding_config
|
||||||
if decoding_config.reasoning_backend == "":
|
if decoding_config.reasoning_backend == "":
|
||||||
decoding_config.reasoning_backend = "GptOss"
|
decoding_config.reasoning_backend = "openai_gptoss"
|
||||||
|
|
||||||
# Increase the max capture size from 512 to 1024 for performance.
|
# Increase the max capture size from 512 to 1024 for performance.
|
||||||
# NOTE(woosuk): This will increase the number of CUDA graphs
|
# NOTE(woosuk): This will increase the number of CUDA graphs
|
||||||
|
|||||||
@ -200,6 +200,32 @@ class TpuPlatform(Platform):
|
|||||||
model_config: "ModelConfig") -> bool:
|
model_config: "ModelConfig") -> bool:
|
||||||
return True
|
return True
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
@torch.compile(backend="openxla")
|
||||||
|
def insert_blocks_to_device(
|
||||||
|
cls,
|
||||||
|
src_cache: torch.Tensor,
|
||||||
|
dst_cache: torch.Tensor,
|
||||||
|
src_block_indices: torch.Tensor,
|
||||||
|
dst_block_indices: torch.Tensor,
|
||||||
|
) -> None:
|
||||||
|
torch.ops.xla.dynamo_set_buffer_donor_(dst_cache, True)
|
||||||
|
dst_cache[dst_block_indices] = src_cache[src_block_indices].to(
|
||||||
|
dst_cache.device)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
@torch.compile(backend="openxla")
|
||||||
|
def swap_out_blocks_to_host(
|
||||||
|
cls,
|
||||||
|
src_cache: torch.Tensor,
|
||||||
|
dst_cache: torch.Tensor,
|
||||||
|
src_block_indices: torch.Tensor,
|
||||||
|
dst_block_indices: torch.Tensor,
|
||||||
|
) -> None:
|
||||||
|
""" tpu blocks to cpu blocks"""
|
||||||
|
torch.ops.xla.dynamo_set_buffer_donor_(src_cache, True)
|
||||||
|
dst_cache[dst_block_indices] = src_cache[src_block_indices].cpu()
|
||||||
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
from tpu_commons.platforms import TpuPlatform as TpuCommonsPlatform
|
from tpu_commons.platforms import TpuPlatform as TpuCommonsPlatform
|
||||||
|
|||||||
@ -164,6 +164,13 @@ class XPUPlatform(Platform):
|
|||||||
vllm_config.scheduler_config.max_model_len,
|
vllm_config.scheduler_config.max_model_len,
|
||||||
DEFAULT_MAX_NUM_BATCHED_TOKENS)
|
DEFAULT_MAX_NUM_BATCHED_TOKENS)
|
||||||
|
|
||||||
|
if (envs.VLLM_KV_CACHE_LAYOUT is None
|
||||||
|
or envs.VLLM_KV_CACHE_LAYOUT != "NHD"):
|
||||||
|
os.environ["VLLM_KV_CACHE_LAYOUT"] = "NHD"
|
||||||
|
logger.info(
|
||||||
|
"Setting VLLM_KV_CACHE_LAYOUT to 'NHD' for XPU; "
|
||||||
|
"only NHD layout is supported by XPU attention kernels.")
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def is_pin_memory_available(cls):
|
def is_pin_memory_available(cls):
|
||||||
return True
|
return True
|
||||||
@ -210,3 +217,27 @@ class XPUPlatform(Platform):
|
|||||||
@classmethod
|
@classmethod
|
||||||
def opaque_attention_op(cls) -> bool:
|
def opaque_attention_op(cls) -> bool:
|
||||||
return True
|
return True
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def insert_blocks_to_device(
|
||||||
|
cls,
|
||||||
|
src_cache: torch.Tensor,
|
||||||
|
dst_cache: torch.Tensor,
|
||||||
|
src_block_indices: torch.Tensor,
|
||||||
|
dst_block_indices: torch.Tensor,
|
||||||
|
) -> None:
|
||||||
|
"""Copy blocks from src_cache to dst_cache on XPU."""
|
||||||
|
_src_cache = src_cache[:, src_block_indices]
|
||||||
|
dst_cache[:, dst_block_indices] = _src_cache.to(dst_cache.device)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def swap_out_blocks_to_host(
|
||||||
|
cls,
|
||||||
|
src_cache: torch.Tensor,
|
||||||
|
dst_cache: torch.Tensor,
|
||||||
|
src_block_indices: torch.Tensor,
|
||||||
|
dst_block_indices: torch.Tensor,
|
||||||
|
) -> None:
|
||||||
|
"""Copy blocks from XPU to host (CPU)."""
|
||||||
|
_src_cache = src_cache[:, src_block_indices]
|
||||||
|
dst_cache[:, dst_block_indices] = _src_cache.cpu()
|
||||||
|
|||||||
@ -6,6 +6,7 @@ from typing import Optional, Union
|
|||||||
|
|
||||||
from transformers import PreTrainedTokenizerBase
|
from transformers import PreTrainedTokenizerBase
|
||||||
|
|
||||||
|
from vllm.entrypoints.harmony_utils import parse_chat_output
|
||||||
from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
|
from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
|
||||||
DeltaMessage)
|
DeltaMessage)
|
||||||
from vllm.logger import init_logger
|
from vllm.logger import init_logger
|
||||||
@ -14,7 +15,7 @@ from vllm.reasoning import ReasoningParser, ReasoningParserManager
|
|||||||
logger = init_logger(__name__)
|
logger = init_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@ReasoningParserManager.register_module("GptOss")
|
@ReasoningParserManager.register_module("openai_gptoss")
|
||||||
class GptOssReasoningParser(ReasoningParser):
|
class GptOssReasoningParser(ReasoningParser):
|
||||||
"""
|
"""
|
||||||
Reasoning parser for GptOss model.
|
Reasoning parser for GptOss model.
|
||||||
@ -39,9 +40,10 @@ class GptOssReasoningParser(ReasoningParser):
|
|||||||
return False
|
return False
|
||||||
|
|
||||||
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
|
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
|
||||||
raise RuntimeError(
|
_, content, _ = parse_chat_output(input_ids)
|
||||||
"GptOss model uses harmony to extract reasoning content. This "
|
if content is None:
|
||||||
"function should not be called.")
|
return []
|
||||||
|
return self.model_tokenizer.encode(content)
|
||||||
|
|
||||||
def extract_reasoning_content_streaming(
|
def extract_reasoning_content_streaming(
|
||||||
self,
|
self,
|
||||||
@ -52,13 +54,34 @@ class GptOssReasoningParser(ReasoningParser):
|
|||||||
current_token_ids: Sequence[int],
|
current_token_ids: Sequence[int],
|
||||||
delta_token_ids: Sequence[int],
|
delta_token_ids: Sequence[int],
|
||||||
) -> Union[DeltaMessage, None]:
|
) -> Union[DeltaMessage, None]:
|
||||||
raise RuntimeError(
|
prev_reasoning, prev_content, _ = parse_chat_output(
|
||||||
"GptOss model uses harmony to extract reasoning content. This "
|
list(previous_token_ids))
|
||||||
"function should not be called.")
|
cur_reasoning, cur_content, _ = parse_chat_output(
|
||||||
|
list(current_token_ids))
|
||||||
|
reasoning_delta = None
|
||||||
|
content_delta = None
|
||||||
|
if cur_reasoning is not None:
|
||||||
|
prev_r = prev_reasoning or ""
|
||||||
|
if cur_reasoning.startswith(prev_r):
|
||||||
|
reasoning_delta = cur_reasoning[len(prev_r):] or None
|
||||||
|
else:
|
||||||
|
reasoning_delta = cur_reasoning
|
||||||
|
if cur_content is not None:
|
||||||
|
prev_c = prev_content or ""
|
||||||
|
if cur_content.startswith(prev_c):
|
||||||
|
content_delta = cur_content[len(prev_c):] or None
|
||||||
|
else:
|
||||||
|
content_delta = cur_content
|
||||||
|
if reasoning_delta is None and content_delta is None:
|
||||||
|
return None
|
||||||
|
return DeltaMessage(reasoning_content=reasoning_delta,
|
||||||
|
content=content_delta)
|
||||||
|
|
||||||
def extract_reasoning_content(
|
def extract_reasoning_content(
|
||||||
self, model_output: str, request: ChatCompletionRequest
|
self,
|
||||||
|
model_output: str,
|
||||||
|
request: ChatCompletionRequest,
|
||||||
) -> tuple[Optional[str], Optional[str]]:
|
) -> tuple[Optional[str], Optional[str]]:
|
||||||
raise RuntimeError(
|
raise NotImplementedError(
|
||||||
"GptOss model uses harmony to extract reasoning content. This "
|
"gpt-oss has a special branch for parsing reasoning in non-streaming mode. This method shouldn't be used." # noqa: E501
|
||||||
"function should not be called.")
|
)
|
||||||
|
|||||||
@ -401,7 +401,7 @@ M = TypeVar("M", bound=MLACommonMetadata)
|
|||||||
|
|
||||||
|
|
||||||
def use_flashinfer_prefill() -> bool:
|
def use_flashinfer_prefill() -> bool:
|
||||||
# For blackwell default to flashinfer prefill if its available since
|
# For blackwell default to flashinfer prefill if it's available since
|
||||||
# it is faster than FA2.
|
# it is faster than FA2.
|
||||||
return (flashinfer_available and not envs.VLLM_USE_CUDNN_PREFILL
|
return (flashinfer_available and not envs.VLLM_USE_CUDNN_PREFILL
|
||||||
and current_platform.is_device_capability(100))
|
and current_platform.is_device_capability(100))
|
||||||
@ -1018,7 +1018,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
|
|||||||
return layer.weight
|
return layer.weight
|
||||||
|
|
||||||
# we currently do not have quantized bmm's which are needed for
|
# we currently do not have quantized bmm's which are needed for
|
||||||
# `W_UV` and `W_UK_T`, we we just store fp16/bf16 copies and perform
|
# `W_UV` and `W_UK_T`, we just store fp16/bf16 copies and perform
|
||||||
# the bmm's in 16-bit, the extra memory overhead of this is fairly low
|
# the bmm's in 16-bit, the extra memory overhead of this is fairly low
|
||||||
kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj).T
|
kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj).T
|
||||||
assert kv_b_proj_weight.shape == (
|
assert kv_b_proj_weight.shape == (
|
||||||
|
|||||||
@ -98,7 +98,12 @@ class AsyncLLM(EngineClient):
|
|||||||
self.model_config = vllm_config.model_config
|
self.model_config = vllm_config.model_config
|
||||||
self.vllm_config = vllm_config
|
self.vllm_config = vllm_config
|
||||||
self.log_requests = log_requests
|
self.log_requests = log_requests
|
||||||
self.log_stats = log_stats
|
|
||||||
|
self.log_stats = log_stats or (stat_loggers is not None)
|
||||||
|
if not log_stats and stat_loggers is not None:
|
||||||
|
logger.info(
|
||||||
|
"AsyncLLM created with log_stats=False and non-empty custom "
|
||||||
|
"logger list; enabling logging without default stat loggers")
|
||||||
|
|
||||||
if self.model_config.skip_tokenizer_init:
|
if self.model_config.skip_tokenizer_init:
|
||||||
self.tokenizer = None
|
self.tokenizer = None
|
||||||
@ -137,6 +142,7 @@ class AsyncLLM(EngineClient):
|
|||||||
vllm_config=vllm_config,
|
vllm_config=vllm_config,
|
||||||
engine_idxs=self.engine_core.engine_ranks_managed,
|
engine_idxs=self.engine_core.engine_ranks_managed,
|
||||||
custom_stat_loggers=stat_loggers,
|
custom_stat_loggers=stat_loggers,
|
||||||
|
enable_default_loggers=log_stats,
|
||||||
)
|
)
|
||||||
self.logger_manager.log_engine_initialized()
|
self.logger_manager.log_engine_initialized()
|
||||||
|
|
||||||
|
|||||||
@ -651,16 +651,16 @@ class StatLoggerManager:
|
|||||||
vllm_config: VllmConfig,
|
vllm_config: VllmConfig,
|
||||||
engine_idxs: Optional[list[int]] = None,
|
engine_idxs: Optional[list[int]] = None,
|
||||||
custom_stat_loggers: Optional[list[StatLoggerFactory]] = None,
|
custom_stat_loggers: Optional[list[StatLoggerFactory]] = None,
|
||||||
|
enable_default_loggers: bool = True,
|
||||||
):
|
):
|
||||||
self.engine_idxs = engine_idxs if engine_idxs else [0]
|
self.engine_idxs = engine_idxs if engine_idxs else [0]
|
||||||
|
|
||||||
factories: list[StatLoggerFactory]
|
factories: list[StatLoggerFactory] = []
|
||||||
if custom_stat_loggers is not None:
|
if custom_stat_loggers is not None:
|
||||||
factories = custom_stat_loggers
|
factories.extend(custom_stat_loggers)
|
||||||
else:
|
|
||||||
factories = []
|
if enable_default_loggers and logger.isEnabledFor(logging.INFO):
|
||||||
if logger.isEnabledFor(logging.INFO):
|
factories.append(LoggingStatLogger)
|
||||||
factories.append(LoggingStatLogger)
|
|
||||||
|
|
||||||
# engine_idx: StatLogger
|
# engine_idx: StatLogger
|
||||||
self.per_engine_logger_dict: dict[int, list[StatLoggerBase]] = {}
|
self.per_engine_logger_dict: dict[int, list[StatLoggerBase]] = {}
|
||||||
|
|||||||
@ -28,6 +28,7 @@ from vllm.config import (CompilationLevel, CUDAGraphMode, VllmConfig,
|
|||||||
from vllm.distributed.eplb.eplb_state import EplbState
|
from vllm.distributed.eplb.eplb_state import EplbState
|
||||||
from vllm.distributed.kv_transfer import (get_kv_transfer_group,
|
from vllm.distributed.kv_transfer import (get_kv_transfer_group,
|
||||||
has_kv_transfer_group)
|
has_kv_transfer_group)
|
||||||
|
from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks
|
||||||
from vllm.distributed.parallel_state import (
|
from vllm.distributed.parallel_state import (
|
||||||
get_pp_group, get_tp_group, graph_capture, is_global_first_rank,
|
get_pp_group, get_tp_group, graph_capture, is_global_first_rank,
|
||||||
prepare_communication_buffer_for_model)
|
prepare_communication_buffer_for_model)
|
||||||
@ -3139,6 +3140,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
|
|||||||
|
|
||||||
if has_kv_transfer_group():
|
if has_kv_transfer_group():
|
||||||
get_kv_transfer_group().register_kv_caches(kv_caches)
|
get_kv_transfer_group().register_kv_caches(kv_caches)
|
||||||
|
if self.device.type == 'xpu':
|
||||||
|
get_kv_transfer_group().set_host_xfer_buffer_ops(
|
||||||
|
copy_kv_blocks)
|
||||||
|
|
||||||
def may_add_encoder_only_layers_to_kv_cache_config(self) -> None:
|
def may_add_encoder_only_layers_to_kv_cache_config(self) -> None:
|
||||||
"""
|
"""
|
||||||
|
|||||||
@ -3,7 +3,7 @@
|
|||||||
import bisect
|
import bisect
|
||||||
import gc
|
import gc
|
||||||
import time
|
import time
|
||||||
from typing import TYPE_CHECKING, Any, Literal, Optional, Union, cast
|
from typing import TYPE_CHECKING, Any, Optional, cast
|
||||||
from unittest.mock import patch
|
from unittest.mock import patch
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
@ -23,6 +23,7 @@ from vllm.config import (ParallelConfig, VllmConfig,
|
|||||||
get_layers_from_vllm_config, update_config)
|
get_layers_from_vllm_config, update_config)
|
||||||
from vllm.distributed.kv_transfer import (get_kv_transfer_group,
|
from vllm.distributed.kv_transfer import (get_kv_transfer_group,
|
||||||
has_kv_transfer_group)
|
has_kv_transfer_group)
|
||||||
|
from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks
|
||||||
from vllm.forward_context import set_forward_context
|
from vllm.forward_context import set_forward_context
|
||||||
from vllm.logger import init_logger
|
from vllm.logger import init_logger
|
||||||
from vllm.lora.layers import BaseLayerWithLoRA
|
from vllm.lora.layers import BaseLayerWithLoRA
|
||||||
@ -1887,75 +1888,6 @@ def _get_padded_token_len(paddings: list[int], x: int) -> int:
|
|||||||
return paddings[index]
|
return paddings[index]
|
||||||
|
|
||||||
|
|
||||||
def _make_src_and_dst_indices(
|
|
||||||
src_block_ids: list[int],
|
|
||||||
dst_block_ids: list[int],
|
|
||||||
src_device: Union[torch.device, str],
|
|
||||||
dst_device: Union[torch.device, str],
|
|
||||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
||||||
src_indices = torch.tensor(src_block_ids,
|
|
||||||
device=src_device,
|
|
||||||
dtype=torch.int64)
|
|
||||||
dst_indices = torch.tensor(dst_block_ids,
|
|
||||||
device=dst_device,
|
|
||||||
dtype=torch.int64)
|
|
||||||
return src_indices, dst_indices
|
|
||||||
|
|
||||||
|
|
||||||
@torch.compile(backend="openxla")
|
|
||||||
def _insert_blocks_to_tpu(
|
|
||||||
cpu_cache: torch.Tensor,
|
|
||||||
tpu_cache: torch.Tensor,
|
|
||||||
cpu_block_indices: torch.Tensor,
|
|
||||||
tpu_block_indices: torch.Tensor,
|
|
||||||
) -> None:
|
|
||||||
torch.ops.xla.dynamo_set_buffer_donor_(tpu_cache, True)
|
|
||||||
tpu_cache[tpu_block_indices] = cpu_cache[cpu_block_indices].to(
|
|
||||||
tpu_cache.device)
|
|
||||||
|
|
||||||
|
|
||||||
@torch.compile(backend="openxla")
|
|
||||||
def _swap_out_tpu_blocks(
|
|
||||||
tpu_cache: torch.Tensor,
|
|
||||||
cpu_cache: torch.Tensor,
|
|
||||||
tpu_block_indices: torch.Tensor,
|
|
||||||
cpu_block_indices: torch.Tensor,
|
|
||||||
) -> None:
|
|
||||||
""" tpu blocks to cpu blocks"""
|
|
||||||
torch.ops.xla.dynamo_set_buffer_donor_(tpu_cache, True)
|
|
||||||
cpu_cache[cpu_block_indices] = tpu_cache[tpu_block_indices].cpu()
|
|
||||||
|
|
||||||
|
|
||||||
def copy_kv_blocks(
|
|
||||||
src_kv_caches: dict[str, torch.Tensor],
|
|
||||||
dst_kv_caches: dict[str, torch.Tensor],
|
|
||||||
src_block_ids: list[int],
|
|
||||||
dst_block_ids: list[int],
|
|
||||||
direction: Literal["h2d", "d2h"],
|
|
||||||
) -> None:
|
|
||||||
"""Copy kv blocks between different buffers."""
|
|
||||||
if not src_kv_caches or not dst_kv_caches or \
|
|
||||||
not src_block_ids or not dst_block_ids or \
|
|
||||||
len(src_block_ids) != len(dst_block_ids):
|
|
||||||
return
|
|
||||||
|
|
||||||
src_device = next(iter(src_kv_caches.values())).device
|
|
||||||
dst_device = next(iter(dst_kv_caches.values())).device
|
|
||||||
|
|
||||||
src_indices, dst_indices = _make_src_and_dst_indices(
|
|
||||||
src_block_ids=src_block_ids,
|
|
||||||
dst_block_ids=dst_block_ids,
|
|
||||||
src_device=src_device,
|
|
||||||
dst_device=dst_device)
|
|
||||||
|
|
||||||
_copy_fn = _insert_blocks_to_tpu if direction == "h2d" else \
|
|
||||||
_swap_out_tpu_blocks
|
|
||||||
for layer_name in src_kv_caches:
|
|
||||||
src_tensor = src_kv_caches[layer_name]
|
|
||||||
dst_tensor = dst_kv_caches[layer_name]
|
|
||||||
_copy_fn(src_tensor, dst_tensor, src_indices, dst_indices)
|
|
||||||
|
|
||||||
|
|
||||||
def _get_padded_num_kv_cache_update_slices(num_tokens: int, max_num_reqs: int,
|
def _get_padded_num_kv_cache_update_slices(num_tokens: int, max_num_reqs: int,
|
||||||
page_size: int) -> int:
|
page_size: int) -> int:
|
||||||
"""Calculates the padded number of KV cache update slices to avoid
|
"""Calculates the padded number of KV cache update slices to avoid
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user