Instructions to use OpenOneRec/KSA-V2-4B-base with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use OpenOneRec/KSA-V2-4B-base with Transformers:

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

pipe = pipeline("text-generation", model="OpenOneRec/KSA-V2-4B-base", trust_remote_code=True)
messages = [
    {"role": "user", "content": "Who are you?"},
]
pipe(messages)

# pip install -U transformers accelerate
# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("OpenOneRec/KSA-V2-4B-base", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("OpenOneRec/KSA-V2-4B-base", trust_remote_code=True, device_map="auto")
messages = [
    {"role": "user", "content": "Who are you?"},
]
inputs = tokenizer.apply_chat_template(
	messages,
	add_generation_prompt=True,
	tokenize=True,
	return_dict=True,
	return_tensors="pt",
).to(model.device)

outputs = model.generate(**inputs, max_new_tokens=256)
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:]))

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use OpenOneRec/KSA-V2-4B-base with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "OpenOneRec/KSA-V2-4B-base"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "OpenOneRec/KSA-V2-4B-base",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker

docker model run hf.co/OpenOneRec/KSA-V2-4B-base

SGLang

How to use OpenOneRec/KSA-V2-4B-base 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 "OpenOneRec/KSA-V2-4B-base" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "OpenOneRec/KSA-V2-4B-base",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

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 "OpenOneRec/KSA-V2-4B-base" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "OpenOneRec/KSA-V2-4B-base",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Docker Model Runner
How to use OpenOneRec/KSA-V2-4B-base with Docker Model Runner:
```
docker model run hf.co/OpenOneRec/KSA-V2-4B-base
```

KSA-V2-4B-base / elementwise_fusion.py

OpenOneRec

Release KSA-V2-4B-base weights and inference runtime

a1111f3 verified 2 days ago

Raw History Blame Contribute Delete

3.03 kB

"""BF16 inference fusion; FMA disabled, original cast points explicit.
Torch pow/mean is retained verbatim: no alternate RMS reduction is substituted.
The Q/K path materializes normalized pre-RoPE keys because meta pooling uses them.
"""
import torch
import triton as tr
import triton.language as tl
import triton.language.extra.cuda.libdevice as lib


@tr.jit
def norm_tail(X,V,W,Y,N:tl.constexpr,D:tl.constexpr,EPS:tl.constexpr,BLOCK:tl.constexpr):
 i=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK)
 x=tl.load(X+i,i<N,0).to(tl.float32)
 v=tl.load(V+i//D,i<N,0)
 scale=lib.rsqrt(v+EPS)
 n=(x*scale).to(tl.bfloat16).to(tl.float32)
 w=tl.load(W+i%D,i<N,0).to(tl.float32)
 tl.store(Y+i,n*w,i<N)

_norm_launchers={}
_rope_launchers={}

def tail(x,variance,weight,eps,fast=False):
 y=torch.empty_like(x);n=x.numel()
 key=(x.device.index,n,x.shape[-1],eps)
 if fast and key in _norm_launchers:
  _norm_launchers[key](x,variance,weight,y)
 else:
  compiled=norm_tail[(tr.cdiv(n,256),)](x,variance,weight,y,n,x.shape[-1],eps,256,enable_fp_fusion=False)
  if fast:_norm_launchers[key]=compiled[(tr.cdiv(n,256),1,1)]
 return y

def rms(x,weight,eps):
 # EXACT original FP32 pow + torch.mean dispatch and shape.
 f=x.to(torch.float32);variance=f.pow(2).mean(-1,keepdim=True)
 return tail(x,variance,weight,eps,fast=True)


@tr.jit
def rope_kernel(X,C,S,Y,N:tl.constexpr,H:tl.constexpr,T:tl.constexpr,D:tl.constexpr,ST0:tl.constexpr,ST1:tl.constexpr,ST2:tl.constexpr,ST3:tl.constexpr,SY0:tl.constexpr,SY1:tl.constexpr,SY2:tl.constexpr,SY3:tl.constexpr,BLOCK:tl.constexpr):
 i=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK);d=i%D;t=i//D%T;h=i//(D*T)%H;b=i//(D*T*H)
 ptr=b*ST0+h*ST1+t*ST2+d*ST3
 rd=tl.where(d<D//2,d+D//2,d-D//2)
 x=tl.load(X+ptr,i<N,0).to(tl.float32)
 rot=tl.load(X+b*ST0+h*ST1+t*ST2+rd*ST3,i<N,0).to(tl.float32)
 rot=tl.where(d<D//2,-rot,rot)
 c=tl.load(C+(b*T+t)*D+d,i<N,0).to(tl.float32);s=tl.load(S+(b*T+t)*D+d,i<N,0).to(tl.float32)
 a=(x*c).to(tl.bfloat16).to(tl.float32);z=(rot*s).to(tl.bfloat16).to(tl.float32)
 tl.store(Y+b*SY0+h*SY1+t*SY2+d*SY3,a+z,i<N)

def rope_one(x,cos,sin):
 b,h,t,d=x.shape
 # TensorIterator canonicalizes dense degenerate dimensions (e.g. T=1),
 # while preserving the physical order of non-contiguous dense inputs.
 y=torch.empty(x.shape,device=x.device,dtype=x.dtype) if x.is_contiguous() else torch.empty_like(x)
 n=x.numel()
 key=(x.device.index,n,h,t,d,x.stride(),y.stride())
 if key in _rope_launchers:
  _rope_launchers[key](x,cos,sin,y)
 else:
  compiled=rope_kernel[(tr.cdiv(n,256),)](x,cos,sin,y,n,h,t,d,*x.stride(),*y.stride(),256,enable_fp_fusion=False)
  _rope_launchers[key]=compiled[(tr.cdiv(n,256),1,1)]
 return y

def rope(q,k,cos,sin,position_ids=None,unsqueeze_dim=1):
 assert unsqueeze_dim==1 and cos.is_contiguous() and sin.is_contiguous()
 return rope_one(q,cos,sin),rope_one(k,cos,sin)




def enabled(x):
    return x.is_cuda and x.device.index == torch.cuda.current_device() and x.dtype == torch.bfloat16 and x.ndim >= 2 and x.shape[0] == 1 and not torch.is_grad_enabled()

	"""BF16 inference fusion; FMA disabled, original cast points explicit.

	Torch pow/mean is retained verbatim: no alternate RMS reduction is substituted.
	The Q/K path materializes normalized pre-RoPE keys because meta pooling uses them.
	"""
	import torch
	import triton as tr
	import triton.language as tl
	import triton.language.extra.cuda.libdevice as lib


	@tr.jit
	def norm_tail(X,V,W,Y,N:tl.constexpr,D:tl.constexpr,EPS:tl.constexpr,BLOCK:tl.constexpr):
	i=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK)
	x=tl.load(X+i,i<N,0).to(tl.float32)
	v=tl.load(V+i//D,i<N,0)
	scale=lib.rsqrt(v+EPS)
	n=(x*scale).to(tl.bfloat16).to(tl.float32)
	w=tl.load(W+i%D,i<N,0).to(tl.float32)
	tl.store(Y+i,n*w,i<N)

	_norm_launchers={}
	_rope_launchers={}

	def tail(x,variance,weight,eps,fast=False):
	y=torch.empty_like(x);n=x.numel()
	key=(x.device.index,n,x.shape[-1],eps)
	if fast and key in _norm_launchers:
	_norm_launchers[key](x,variance,weight,y)
	else:
	compiled=norm_tail[(tr.cdiv(n,256),)](x,variance,weight,y,n,x.shape[-1],eps,256,enable_fp_fusion=False)
	if fast:_norm_launchers[key]=compiled[(tr.cdiv(n,256),1,1)]
	return y

	def rms(x,weight,eps):
	# EXACT original FP32 pow + torch.mean dispatch and shape.
	f=x.to(torch.float32);variance=f.pow(2).mean(-1,keepdim=True)
	return tail(x,variance,weight,eps,fast=True)


	@tr.jit
	def rope_kernel(X,C,S,Y,N:tl.constexpr,H:tl.constexpr,T:tl.constexpr,D:tl.constexpr,ST0:tl.constexpr,ST1:tl.constexpr,ST2:tl.constexpr,ST3:tl.constexpr,SY0:tl.constexpr,SY1:tl.constexpr,SY2:tl.constexpr,SY3:tl.constexpr,BLOCK:tl.constexpr):
	i=tl.program_id(0)BLOCK+tl.arange(0,BLOCK);d=i%D;t=i//D%T;h=i//(DT)%H;b=i//(DTH)
	ptr=bST0+hST1+tST2+dST3
	rd=tl.where(d<D//2,d+D//2,d-D//2)
	x=tl.load(X+ptr,i<N,0).to(tl.float32)
	rot=tl.load(X+bST0+hST1+tST2+rdST3,i<N,0).to(tl.float32)
	rot=tl.where(d<D//2,-rot,rot)
	c=tl.load(C+(bT+t)D+d,i<N,0).to(tl.float32);s=tl.load(S+(bT+t)D+d,i<N,0).to(tl.float32)
	a=(xc).to(tl.bfloat16).to(tl.float32);z=(rots).to(tl.bfloat16).to(tl.float32)
	tl.store(Y+bSY0+hSY1+tSY2+dSY3,a+z,i<N)

	def rope_one(x,cos,sin):
	b,h,t,d=x.shape
	# TensorIterator canonicalizes dense degenerate dimensions (e.g. T=1),
	# while preserving the physical order of non-contiguous dense inputs.
	y=torch.empty(x.shape,device=x.device,dtype=x.dtype) if x.is_contiguous() else torch.empty_like(x)
	n=x.numel()
	key=(x.device.index,n,h,t,d,x.stride(),y.stride())
	if key in _rope_launchers:
	_rope_launchers[key](x,cos,sin,y)
	else:
	compiled=rope_kernel[(tr.cdiv(n,256),)](x,cos,sin,y,n,h,t,d,x.stride(),y.stride(),256,enable_fp_fusion=False)
	_rope_launchers[key]=compiled[(tr.cdiv(n,256),1,1)]
	return y

	def rope(q,k,cos,sin,position_ids=None,unsqueeze_dim=1):
	assert unsqueeze_dim==1 and cos.is_contiguous() and sin.is_contiguous()
	return rope_one(q,cos,sin),rope_one(k,cos,sin)




	def enabled(x):
	return x.is_cuda and x.device.index == torch.cuda.current_device() and x.dtype == torch.bfloat16 and x.ndim >= 2 and x.shape[0] == 1 and not torch.is_grad_enabled()