Instructions to use WaveCut/YuE2-3B-OrbitQuant-W4A4 with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use WaveCut/YuE2-3B-OrbitQuant-W4A4 with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="WaveCut/YuE2-3B-OrbitQuant-W4A4")

# pip install -U transformers accelerate
# Load model directly
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("WaveCut/YuE2-3B-OrbitQuant-W4A4", device_map="auto")

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use WaveCut/YuE2-3B-OrbitQuant-W4A4 with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "WaveCut/YuE2-3B-OrbitQuant-W4A4"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "WaveCut/YuE2-3B-OrbitQuant-W4A4",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker

docker model run hf.co/WaveCut/YuE2-3B-OrbitQuant-W4A4

SGLang

How to use WaveCut/YuE2-3B-OrbitQuant-W4A4 with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "WaveCut/YuE2-3B-OrbitQuant-W4A4" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "WaveCut/YuE2-3B-OrbitQuant-W4A4",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "WaveCut/YuE2-3B-OrbitQuant-W4A4" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "WaveCut/YuE2-3B-OrbitQuant-W4A4",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Docker Model Runner
How to use WaveCut/YuE2-3B-OrbitQuant-W4A4 with Docker Model Runner:
```
docker model run hf.co/WaveCut/YuE2-3B-OrbitQuant-W4A4
```

YuE2-3B-OrbitQuant-W4A4 / src /kernel-source /rms-quant /orbitquant_gemv_cuda /rmsquant.cu

WaveCut

Publish OrbitQuant 0.9.5 byte-pair decode and optional RMS activation fusion

b35ab1f verified about 2 months ago

Raw History Blame Contribute Delete

5.43 kB

#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <torch/torch.h>
#include <cmath>

using BF = c10::BFloat16;

// The positive RMS scale cancels in quantization except for quant_eps.
// Keep it in the returned token norm; deliberately omit BF16 intermediates.
__global__ void rmsquant_kernel(uint8_t* packed, float* norms, const BF* x,
    const BF* weight, const int32_t* permutation, const int8_t* signs,
    const float* boundaries, float rms_eps, float quant_eps) {
  constexpr int dim=2048;
  __shared__ float values[dim], raw_sum[256], weighted_sum[256], bounds[15];
  const int t=threadIdx.x, row=blockIdx.x;
  float sx=0, sy=0;
  for(int j=t;j<dim;j+=256){
    int source=permutation[j];
    float raw=float(x[row*dim+source]);
    float weighted=raw*float(weight[source]);
    values[j]=weighted*float(signs[j]);
    sx=fmaf(raw,raw,sx);sy=fmaf(weighted,weighted,sy);
  }
  raw_sum[t]=sx;weighted_sum[t]=sy;
  if(t<15)bounds[t]=boundaries[t];
  __syncthreads();
  for(int stride=128;stride>=32;stride>>=1){
    if(t<stride){raw_sum[t]+=raw_sum[t+stride];weighted_sum[t]+=weighted_sum[t+stride];}
    __syncthreads();
  }
  // Preserve the original reduction tree, but finish inside the first warp.
  if(t<32){
    float raw=raw_sum[t],weighted=weighted_sum[t];
    #pragma unroll
    for(int stride=16;stride;stride>>=1){
      raw+=__shfl_down_sync(0xffffffff,raw,stride);
      weighted+=__shfl_down_sync(0xffffffff,weighted,stride);
    }
    if(t==0){raw_sum[0]=raw;weighted_sum[0]=weighted;}
  }
  __syncthreads();
  float scale=rsqrtf(raw_sum[0]/float(dim)+rms_eps);
  float norm=sqrtf(weighted_sum[0]);
  if(t==0)norms[row]=norm*scale;
  float inv=1.f/(norm+quant_eps/scale);
  float local[8];
  #pragma unroll
  for(int i=0;i<8;i++)local[i]=values[t+256*i]*inv;
  // The first five Hadamard stages stay inside each warp.
  #pragma unroll
  for(int width=1;width<32;width<<=1){
    #pragma unroll
    for(int i=0;i<8;i++){
      float other=__shfl_xor_sync(0xffffffff,local[i],width);
      local[i]=(t&width)?other-local[i]:local[i]+other;
    }
  }
  #pragma unroll
  for(int i=0;i<8;i++)values[t+256*i]=local[i];
  __syncthreads();
  // Only three stages exchange values across warps.
  for(int width=32;width<=128;width<<=1){
    for(int b=t;b<dim/2;b+=256){
      int group=b/width,offset=b-group*width;
      int left=group*width*2+offset,right=left+width;
      float a=values[left],v=values[right];values[left]=a+v;values[right]=a-v;
    }
    __syncthreads();
  }
  #pragma unroll
  for(int i=0;i<8;i++)local[i]=values[t+256*i];
  // Dimensions256/512/1024 are eight register values owned by one thread.
  #pragma unroll
  for(int width=1;width<8;width<<=1){
    #pragma unroll
    for(int b=0;b<4;b++){
      int group=b/width,offset=b-group*width;
      int left=group*width*2+offset,right=left+width;
      float a=local[left],v=local[right];local[left]=a+v;local[right]=a-v;
    }
  }
  // All reads finished before shared storage is reused for packing.
  __syncthreads();
  #pragma unroll
  for(int i=0;i<8;i++)values[t+256*i]=local[i];
  __syncthreads();
  for(int j=t;j<dim/2;j+=256){
    uint8_t codes[2];
    #pragma unroll
    for(int k=0;k<2;k++){
      float v=values[2*j+k]*0.02209708691207961f;
      int index=v>bounds[7]?8:0;
      index+=v>bounds[index+3]?4:0;
      index+=v>bounds[index+1]?2:0;
      index+=v>bounds[index]?1:0;
      codes[k]=index;
    }
    packed[row*(dim/2)+j]=codes[0]|(codes[1]<<4);
  }
}

void rmsquant(torch::Tensor &packed, torch::Tensor &norms, torch::Tensor const &x,
    torch::Tensor const &weight, torch::Tensor const &permutation,
    torch::Tensor const &signs, torch::Tensor const &boundaries,
    double rms_eps, double quant_eps) {
  TORCH_CHECK(x.is_cuda() && x.is_contiguous() && x.scalar_type()==torch::kBFloat16
    && x.dim()==2 && x.size(1)==2048,"requires contiguous CUDA BF16 [rows,2048]");
  for(auto* t:{&weight,&permutation,&signs,&boundaries})
    TORCH_CHECK(t->device()==x.device() && t->is_contiguous() && t->dim()==1,"invalid constant device/layout");
  TORCH_CHECK(weight.scalar_type()==torch::kBFloat16 && weight.numel()==2048,"invalid RMS weight");
  TORCH_CHECK(permutation.scalar_type()==torch::kInt && permutation.numel()==2048,"invalid permutation");
  TORCH_CHECK(signs.scalar_type()==torch::kChar && signs.numel()==2048,"invalid signs");
  TORCH_CHECK(boundaries.scalar_type()==torch::kFloat && boundaries.numel()==15,"invalid boundaries");
  TORCH_CHECK(packed.device()==x.device() && packed.is_contiguous() && packed.scalar_type()==torch::kUInt8
    && packed.dim()==2 && packed.size(0)==x.size(0) && packed.size(1)==1024,"invalid packed output");
  TORCH_CHECK(norms.device()==x.device() && norms.is_contiguous() && norms.scalar_type()==torch::kFloat
    && norms.dim()==1 && norms.numel()==x.size(0),"invalid norm output");
  TORCH_CHECK(std::isfinite(rms_eps) && rms_eps>0 && std::isfinite(quant_eps) && quant_eps>0,"epsilons must be finite and positive");
  const at::cuda::OptionalCUDAGuard guard(device_of(x));
  if(!x.size(0))return;
  rmsquant_kernel<<<x.size(0),256,0,at::cuda::getCurrentCUDAStream()>>>(packed.data_ptr<uint8_t>(),norms.data_ptr<float>(),
    x.data_ptr<BF>(),weight.data_ptr<BF>(),permutation.data_ptr<int32_t>(),signs.data_ptr<int8_t>(),
    boundaries.data_ptr<float>(),rms_eps,quant_eps);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

	#include <ATen/cuda/CUDAContext.h>
	#include <c10/cuda/CUDAGuard.h>
	#include <c10/cuda/CUDAException.h>
	#include <torch/torch.h>
	#include <cmath>

	using BF = c10::BFloat16;

	// The positive RMS scale cancels in quantization except for quant_eps.
	// Keep it in the returned token norm; deliberately omit BF16 intermediates.
	__global__ void rmsquant_kernel(uint8_t* packed, float* norms, const BF* x,
	const BF* weight, const int32_t* permutation, const int8_t* signs,
	const float* boundaries, float rms_eps, float quant_eps) {
	constexpr int dim=2048;
	__shared__ float values[dim], raw_sum[256], weighted_sum[256], bounds[15];
	const int t=threadIdx.x, row=blockIdx.x;
	float sx=0, sy=0;
	for(int j=t;j<dim;j+=256){
	int source=permutation[j];
	float raw=float(x[row*dim+source]);
	float weighted=raw*float(weight[source]);
	values[j]=weighted*float(signs[j]);
	sx=fmaf(raw,raw,sx);sy=fmaf(weighted,weighted,sy);
	}
	raw_sum[t]=sx;weighted_sum[t]=sy;
	if(t<15)bounds[t]=boundaries[t];
	__syncthreads();
	for(int stride=128;stride>=32;stride>>=1){
	if(t<stride){raw_sum[t]+=raw_sum[t+stride];weighted_sum[t]+=weighted_sum[t+stride];}
	__syncthreads();
	}
	// Preserve the original reduction tree, but finish inside the first warp.
	if(t<32){
	float raw=raw_sum[t],weighted=weighted_sum[t];
	#pragma unroll
	for(int stride=16;stride;stride>>=1){
	raw+=__shfl_down_sync(0xffffffff,raw,stride);
	weighted+=__shfl_down_sync(0xffffffff,weighted,stride);
	}
	if(t==0){raw_sum[0]=raw;weighted_sum[0]=weighted;}
	}
	__syncthreads();
	float scale=rsqrtf(raw_sum[0]/float(dim)+rms_eps);
	float norm=sqrtf(weighted_sum[0]);
	if(t==0)norms[row]=norm*scale;
	float inv=1.f/(norm+quant_eps/scale);
	float local[8];
	#pragma unroll
	for(int i=0;i<8;i++)local[i]=values[t+256i]inv;
	// The first five Hadamard stages stay inside each warp.
	#pragma unroll
	for(int width=1;width<32;width<<=1){
	#pragma unroll
	for(int i=0;i<8;i++){
	float other=__shfl_xor_sync(0xffffffff,local[i],width);
	local[i]=(t&width)?other-local[i]:local[i]+other;
	}
	}
	#pragma unroll
	for(int i=0;i<8;i++)values[t+256*i]=local[i];
	__syncthreads();
	// Only three stages exchange values across warps.
	for(int width=32;width<=128;width<<=1){
	for(int b=t;b<dim/2;b+=256){
	int group=b/width,offset=b-group*width;
	int left=groupwidth2+offset,right=left+width;
	float a=values[left],v=values[right];values[left]=a+v;values[right]=a-v;
	}
	__syncthreads();
	}
	#pragma unroll
	for(int i=0;i<8;i++)local[i]=values[t+256*i];
	// Dimensions256/512/1024 are eight register values owned by one thread.
	#pragma unroll
	for(int width=1;width<8;width<<=1){
	#pragma unroll
	for(int b=0;b<4;b++){
	int group=b/width,offset=b-group*width;
	int left=groupwidth2+offset,right=left+width;
	float a=local[left],v=local[right];local[left]=a+v;local[right]=a-v;
	}
	}
	// All reads finished before shared storage is reused for packing.
	__syncthreads();
	#pragma unroll
	for(int i=0;i<8;i++)values[t+256*i]=local[i];
	__syncthreads();
	for(int j=t;j<dim/2;j+=256){
	uint8_t codes[2];
	#pragma unroll
	for(int k=0;k<2;k++){
	float v=values[2j+k]0.02209708691207961f;
	int index=v>bounds[7]?8:0;
	index+=v>bounds[index+3]?4:0;
	index+=v>bounds[index+1]?2:0;
	index+=v>bounds[index]?1:0;
	codes[k]=index;
	}
	packed[row*(dim/2)+j]=codes[0]\|(codes[1]<<4);
	}
	}

	void rmsquant(torch::Tensor &packed, torch::Tensor &norms, torch::Tensor const &x,
	torch::Tensor const &weight, torch::Tensor const &permutation,
	torch::Tensor const &signs, torch::Tensor const &boundaries,
	double rms_eps, double quant_eps) {
	TORCH_CHECK(x.is_cuda() && x.is_contiguous() && x.scalar_type()==torch::kBFloat16
	&& x.dim()==2 && x.size(1)==2048,"requires contiguous CUDA BF16 [rows,2048]");
	for(auto* t:{&weight,&permutation,&signs,&boundaries})
	TORCH_CHECK(t->device()==x.device() && t->is_contiguous() && t->dim()==1,"invalid constant device/layout");
	TORCH_CHECK(weight.scalar_type()==torch::kBFloat16 && weight.numel()==2048,"invalid RMS weight");
	TORCH_CHECK(permutation.scalar_type()==torch::kInt && permutation.numel()==2048,"invalid permutation");
	TORCH_CHECK(signs.scalar_type()==torch::kChar && signs.numel()==2048,"invalid signs");
	TORCH_CHECK(boundaries.scalar_type()==torch::kFloat && boundaries.numel()==15,"invalid boundaries");
	TORCH_CHECK(packed.device()==x.device() && packed.is_contiguous() && packed.scalar_type()==torch::kUInt8
	&& packed.dim()==2 && packed.size(0)==x.size(0) && packed.size(1)==1024,"invalid packed output");
	TORCH_CHECK(norms.device()==x.device() && norms.is_contiguous() && norms.scalar_type()==torch::kFloat
	&& norms.dim()==1 && norms.numel()==x.size(0),"invalid norm output");
	TORCH_CHECK(std::isfinite(rms_eps) && rms_eps>0 && std::isfinite(quant_eps) && quant_eps>0,"epsilons must be finite and positive");
	const at::cuda::OptionalCUDAGuard guard(device_of(x));
	if(!x.size(0))return;
	rmsquant_kernel<<<x.size(0),256,0,at::cuda::getCurrentCUDAStream()>>>(packed.data_ptr<uint8_t>(),norms.data_ptr<float>(),
	x.data_ptr<BF>(),weight.data_ptr<BF>(),permutation.data_ptr<int32_t>(),signs.data_ptr<int8_t>(),
	boundaries.data_ptr<float>(),rms_eps,quant_eps);
	C10_CUDA_KERNEL_LAUNCH_CHECK();
	}