base_IIXIV

Instructions to use mainline777/base_IIXIV with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use mainline777/base_IIXIV with Transformers:

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

pipe = pipeline("text-generation", model="mainline777/base_IIXIV", trust_remote_code=True)
messages = [
    {"role": "user", "content": "Who are you?"},
]
pipe(messages)

# Load model directly
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("mainline777/base_IIXIV", trust_remote_code=True, device_map="auto")

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use mainline777/base_IIXIV with vLLM:

Install from pip and serve model

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

Use Docker

docker model run hf.co/mainline777/base_IIXIV

SGLang

How to use mainline777/base_IIXIV 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 "mainline777/base_IIXIV" \
    --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": "mainline777/base_IIXIV",
		"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 "mainline777/base_IIXIV" \
        --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": "mainline777/base_IIXIV",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Docker Model Runner
How to use mainline777/base_IIXIV with Docker Model Runner:
```
docker model run hf.co/mainline777/base_IIXIV
```

base_IIXIV / fla /ops /common /chunk_o.py

mainline777

Duplicate from silx-ai/Quasar-Preview

41865df about 2 months ago

Raw

History Blame Contribute Delete

25 kB

	# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang

	import torch
	import triton
	import triton.language as tl

	from fla.ops.utils import prepare_chunk_indices
	from fla.ops.utils.op import exp, exp2
	from fla.utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem

	BKV_LIST = [64, 128] if check_shared_mem() else ([32, 64] if check_shared_mem('ada') else [32])
	NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8]


	@triton.heuristics({
	'USE_G': lambda args: args['g'] is not None,
	'USE_G_GAMMA': lambda args: args['g_gamma'] is not None,
	'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
	})
	@triton.autotune(
	configs=[
	triton.Config({'BK': 128, 'BV': 128}, num_warps=8, num_stages=3),
	triton.Config({'BK': 64, 'BV': 64}, num_warps=4, num_stages=3),
	triton.Config({'BK': 32, 'BV': 32}, num_warps=2, num_stages=3),
	],
	key=['H', 'K', 'V', 'BT', 'TRANSPOSE_STATE'],
	**autotune_cache_kwargs,
	)
	@triton.jit(do_not_specialize=['T'])
	def chunk_fwd_kernel_o(
	q,
	k,
	v,
	h,
	g,
	g_gamma,
	o,
	cu_seqlens,
	chunk_indices,
	scale,
	T,
	H: tl.constexpr,
	K: tl.constexpr,
	V: tl.constexpr,
	BT: tl.constexpr,
	BK: tl.constexpr,
	BV: tl.constexpr,
	USE_G: tl.constexpr,
	USE_G_GAMMA: tl.constexpr,
	USE_EXP2: tl.constexpr,
	TRANSPOSE_STATE: tl.constexpr,
	IS_VARLEN: tl.constexpr,
	):
	i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
	i_b, i_h = i_bh // H, i_bh % H

	if IS_VARLEN:
	i_tg = i_t
	i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
	bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
	T = eos - bos
	NT = tl.cdiv(T, BT)
	else:
	NT = tl.cdiv(T, BT)
	i_tg = i_b * NT + i_t
	bos, eos = i_b * T, i_b * T + T

	# offset calculation
	q += (bos * H + i_h) * K
	k += (bos * H + i_h) * K
	v += (bos * H + i_h) * V
	o += (bos * H + i_h) * V
	h += (i_tg * H + i_h).to(tl.int64) * K*V

	b_o = tl.zeros([BT, BV], dtype=tl.float32)
	b_A = tl.zeros([BT, BT], dtype=tl.float32)

	for i_k in range(tl.cdiv(K, BK)):
	p_q = tl.make_block_ptr(q, (T, K), (HK, 1), (i_t BT, i_k * BK), (BT, BK), (1, 0))
	p_k = tl.make_block_ptr(k, (K, T), (1, HK), (i_k BK, i_t * BT), (BK, BT), (0, 1))
	if TRANSPOSE_STATE:
	p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0))
	else:
	p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
	# [BT, BK]
	b_q = tl.load(p_q, boundary_check=(0, 1))
	# [BK, BT]
	b_k = tl.load(p_k, boundary_check=(0, 1))
	b_h = tl.load(p_h, boundary_check=(0, 1))

	# [BT, BK] @ [BK, BV] -> [BT, BV]
	if TRANSPOSE_STATE:
	b_o += tl.dot(b_q, tl.trans(b_h))
	else:
	b_o += tl.dot(b_q, b_h)
	# [BT, BK] @ [BK, BT] -> [BT, BT]
	b_A += tl.dot(b_q, b_k)

	if USE_G:
	g += bos * H + i_h
	p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,))
	b_g = tl.load(p_g, boundary_check=(0,))
	if USE_EXP2:
	b_o = b_o * exp2(b_g)[:, None]
	b_A = b_A * exp2(b_g[:, None] - b_g[None, :])
	else:
	b_o = b_o * exp(b_g)[:, None]
	b_A = b_A * exp(b_g[:, None] - b_g[None, :])

	if USE_G_GAMMA:
	b_gamma = tl.load(g_gamma + i_h)
	b_g = b_gamma * (tl.arange(0, BT) + 1)
	if USE_EXP2:
	b_o = b_o * exp2(b_g)[:, None]
	b_A = b_A * exp2(b_g[:, None] - b_g[None, :])
	else:
	b_o = b_o * exp(b_g)[:, None]
	b_A = b_A * exp(b_g[:, None] - b_g[None, :])

	o_t = i_t * BT + tl.arange(0, BT)
	m_t = o_t < T
	m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t)
	b_A = tl.where(m_A, b_A, 0)

	p_v = tl.make_block_ptr(v, (T, V), (HV, 1), (i_t BT, i_v * BV), (BT, BV), (1, 0))
	p_o = tl.make_block_ptr(o, (T, V), (HV, 1), (i_t BT, i_v * BV), (BT, BV), (1, 0))

	b_v = tl.load(p_v, boundary_check=(0, 1))
	# to fix mma -> mma layout conversion
	# already solved by triton v3.2 or higher
	b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale
	tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))


	@triton.heuristics({
	'USE_G': lambda args: args['g'] is not None,
	'USE_G_GAMMA': lambda args: args['g_gamma'] is not None,
	'USE_DW': lambda args: args['dw'] is not None,
	'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
	})
	@triton.jit(do_not_specialize=['T'])
	def chunk_bwd_kernel_dqkwg(
	q,
	k,
	v,
	g,
	g_gamma,
	h,
	do,
	dh,
	dq,
	dk,
	dw,
	dv,
	dg,
	cu_seqlens,
	chunk_indices,
	scale,
	B: tl.constexpr,
	T,
	H: tl.constexpr,
	K: tl.constexpr,
	V: tl.constexpr,
	BT: tl.constexpr,
	BK: tl.constexpr,
	BV: tl.constexpr,
	USE_G: tl.constexpr,
	USE_G_GAMMA: tl.constexpr,
	USE_EXP2: tl.constexpr,
	USE_DW: tl.constexpr,
	TRANSPOSE_STATE: tl.constexpr,
	IS_VARLEN: tl.constexpr,
	):
	i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
	i_b, i_h = i_bh // H, i_bh % H

	all = B * T
	if IS_VARLEN:
	i_tg = i_t
	i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
	bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
	T = eos - bos
	NT = tl.cdiv(T, BT)
	else:
	NT = tl.cdiv(T, BT)
	i_tg = i_b * NT + i_t
	bos, eos = i_b * T, i_b * T + T

	# offset calculation
	v += (bos * H + i_h) * V
	do += (bos * H + i_h) * V
	h += (i_tg * H + i_h).to(tl.int64) * K*V
	dh += (i_tg * H + i_h).to(tl.int64) * K*V
	q += (bos * H + i_h) * K
	k += (bos * H + i_h) * K
	dq += (bos * H + i_h) * K
	dk += (bos * H + i_h) * K

	# for delta rule only
	if USE_DW:
	dw += (bos * H + i_h) * K
	dv += (bos * H + i_h) * V

	if USE_G:
	dg += i_k * all * H
	b_dg_last = tl.zeros([1], dtype=tl.float32) if USE_G else None
	if USE_G_GAMMA:
	b_gamma = tl.load(g_gamma + i_h)
	b_g = b_gamma * (tl.arange(0, BT) + 1)
	b_g_last = b_gamma * min(BT, T - i_t * BT)
	b_dq = tl.zeros([BT, BK], dtype=tl.float32)
	b_dk = tl.zeros([BT, BK], dtype=tl.float32)
	b_ds = tl.zeros([BT, BT], dtype=tl.float32)
	b_dw = tl.zeros([BT, BK], dtype=tl.float32) if USE_DW else None

	for i_v in range(tl.cdiv(V, BV)):
	p_v = tl.make_block_ptr(v, (T, V), (HV, 1), (i_t BT, i_v * BV), (BT, BV), (1, 0))
	p_do = tl.make_block_ptr(do, (T, V), (HV, 1), (i_t BT, i_v * BV), (BT, BV), (1, 0))
	if TRANSPOSE_STATE:
	p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0))
	p_dh = tl.make_block_ptr(dh, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0))
	else:
	p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1))
	p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1))
	# [BT, BV]
	b_v = tl.load(p_v, boundary_check=(0, 1))
	b_do = tl.load(p_do, boundary_check=(0, 1))
	# [BV, BK]
	b_h = tl.load(p_h, boundary_check=(0, 1))
	b_dh = tl.load(p_dh, boundary_check=(0, 1))
	if USE_G:
	b_dg_last += (tl.sum(b_h * b_dh))
	# [BT, BV] @ [BV, BT] -> [BT, BT]
	b_ds += tl.dot(b_do, tl.trans(b_v))
	# [BT, BV] @ [BV, BK] -> [BT, BK]
	b_dq += tl.dot(b_do, b_h.to(b_do.dtype))
	# [BT, BV] @ [BV, BK] -> [BT, BK]
	b_dk += tl.dot(b_v, b_dh.to(b_v.dtype))
	if USE_DW:
	p_dv = tl.make_block_ptr(dv, (T, V), (HV, 1), (i_t BT, i_v * BV), (BT, BV), (1, 0))
	b_dv = tl.load(p_dv, boundary_check=(0, 1))
	b_dw += tl.dot(b_dv.to(b_v.dtype), b_h.to(b_v.dtype))

	if USE_DW:
	p_dw = tl.make_block_ptr(dw, (T, K), (HK, 1), (i_t BT, i_k * BK), (BT, BK), (1, 0))
	tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1))

	tl.debug_barrier()
	p_q = tl.make_block_ptr(q, (T, K), (HK, 1), (i_t BT, i_k * BK), (BT, BK), (1, 0))
	p_k = tl.make_block_ptr(k, (T, K), (HK, 1), (i_t BT, i_k * BK), (BT, BK), (1, 0))
	b_q = tl.load(p_q, boundary_check=(0, 1))
	b_k = tl.load(p_k, boundary_check=(0, 1))

	p_dq = tl.make_block_ptr(dq, (T, K), (HK, 1), (i_t BT, i_k * BK), (BT, BK), (1, 0))
	p_dk = tl.make_block_ptr(dk, (T, K), (HK, 1), (i_t BT, i_k * BK), (BT, BK), (1, 0))

	o_t = i_t * BT + tl.arange(0, BT)
	m_t = o_t < T
	m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t)
	if USE_G:
	b_dg = tl.zeros([BT], dtype=tl.float32)
	g += bos * H + i_h
	dg += bos * H + i_h
	p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,))
	b_g = tl.load(p_g, boundary_check=(0,))
	b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H)
	if USE_EXP2:
	b_dg_last *= exp2(b_g_last)
	b_dq = b_dq * exp2(b_g)[:, None] * scale
	else:
	b_dg_last *= exp(b_g_last)
	b_dq = b_dq * exp(b_g)[:, None] * scale
	b_dg += tl.sum(b_dq * b_q, axis=1)

	if USE_EXP2:
	b_dk = b_dk * tl.where(m_t, exp2(-b_g + b_g_last), 0)[:, None]
	else:
	b_dk = b_dk * tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None]
	b_dg -= tl.sum(b_k * b_dk, axis=1)
	b_dg_last += tl.sum(b_dk * b_k)

	if USE_EXP2:
	b_ds = tl.where(m_A, b_ds * exp2(b_g[:, None] - b_g[None, :]), 0) * scale
	else:
	b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale
	b_ds2 = b_ds * tl.dot(b_q, tl.trans(b_k))
	b_dg += tl.sum(b_ds2, axis=1)
	b_dg -= tl.sum(b_ds2, axis=0)

	b_ds = b_ds.to(b_k.dtype)
	# [BT, BK]
	b_dq += tl.dot(b_ds, b_k)
	b_dk += tl.dot(tl.trans(b_ds), b_q)
	p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,))
	# (SY 09/21) revcumsum in a separate kernel due to strange triton compiler issue
	# b_dg = tl.dot(tl.where(o_t[:, None] <= o_t[None, :], 1., 0.), b_dg, allow_tf32=False) + b_dg_last)
	b_dg = tl.where(o_t < min(i_t * BT + BT, T) - 1, b_dg, b_dg + b_dg_last)
	tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1))
	tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
	tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,))

	elif USE_G_GAMMA:
	if USE_EXP2:
	b_dq = b_dq * exp2(b_g)[:, None] * scale
	b_dk = b_dk * tl.where(m_t, exp2(-b_g + b_g_last), 0)[:, None]
	b_ds = tl.where(m_A, b_ds * exp2(b_g[:, None] - b_g[None, :]), 0) * scale
	else:
	b_dq = b_dq * exp(b_g)[:, None] * scale
	b_dk = b_dk * tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None]
	b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale
	b_ds = b_ds.to(b_k.dtype)
	# [BT, BK]
	b_dq += tl.dot(b_ds, b_k)
	b_dk += tl.dot(tl.trans(b_ds), b_q)
	tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1))
	tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))

	else:
	b_ds = tl.where(m_A, b_ds, 0)
	b_ds = b_ds.to(b_k.dtype)
	b_dq += tl.dot(b_ds, b_k)
	b_dk += tl.dot(tl.trans(b_ds), b_q) * scale
	b_dq *= scale
	tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1))
	tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))


	@triton.heuristics({
	'USE_G': lambda args: args['g'] is not None,
	'USE_G_GAMMA': lambda args: args['g_gamma'] is not None,
	'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
	})
	@triton.jit(do_not_specialize=['T'])
	def chunk_bwd_kernel_dv(
	q,
	k,
	g,
	g_gamma,
	do,
	dv,
	dh,
	cu_seqlens,
	chunk_indices,
	scale,
	T,
	H: tl.constexpr,
	K: tl.constexpr,
	V: tl.constexpr,
	BT: tl.constexpr,
	BK: tl.constexpr,
	BV: tl.constexpr,
	USE_G: tl.constexpr,
	USE_G_GAMMA: tl.constexpr,
	USE_EXP2: tl.constexpr,
	IS_VARLEN: tl.constexpr,
	):
	i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
	i_b, i_h = i_bh // H, i_bh % H
	if IS_VARLEN:
	i_tg = i_t
	i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
	bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
	T = eos - bos
	NT = tl.cdiv(T, BT)
	else:
	NT = tl.cdiv(T, BT)
	i_tg = i_b * NT + i_t
	bos, eos = i_b * T, i_b * T + T

	b_dv = tl.zeros([BT, BV], dtype=tl.float32)

	# offset calculation
	q += (bos * H + i_h) * K
	k += (bos * H + i_h) * K
	do += (bos * H + i_h) * V
	dv += (bos * H + i_h) * V
	dh += (i_tg * H + i_h).to(tl.int64) * K*V

	b_A = tl.zeros([BT, BT], dtype=tl.float32)
	for i_k in range(tl.cdiv(K, BK)):
	p_k = tl.make_block_ptr(k, (T, K), (HK, 1), (i_t BT, i_k * BK), (BT, BK), (1, 0))
	p_q = tl.make_block_ptr(q, (K, T), (1, HK), (i_k BK, i_t * BT), (BK, BT), (0, 1))
	b_q = tl.load(p_q, boundary_check=(0, 1))
	b_k = tl.load(p_k, boundary_check=(0, 1))
	b_A += tl.dot(b_k, b_q)
	p_dh = tl.make_block_ptr(dh, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
	b_dh = tl.load(p_dh, boundary_check=(0, 1))
	b_dv += tl.dot(b_k, b_dh.to(b_k.dtype))

	o_t = i_t * BT + tl.arange(0, BT)
	m_t = o_t < T
	if USE_G:
	g += bos * H + i_h
	p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,))
	b_g = tl.load(p_g, boundary_check=(0,))
	b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H)
	if USE_G_GAMMA:
	b_gamma = tl.load(g_gamma + i_h)
	b_g = b_gamma * (tl.arange(0, BT) + 1)
	b_g_last = b_gamma * min(BT, T - i_t * BT)

	m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t)
	if USE_G or USE_G_GAMMA:
	if USE_EXP2:
	b_A = tl.where(m_A, b_A * exp2(b_g[None, :] - b_g[:, None]) * scale, 0).to(do.dtype.element_ty)
	b_dv *= tl.where(m_t, exp2(-b_g + b_g_last), 0)[:, None]
	else:
	b_A = tl.where(m_A, b_A * exp(b_g[None, :] - b_g[:, None]) * scale, 0).to(do.dtype.element_ty)
	b_dv *= tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None]
	else:
	b_A = tl.where(m_A, b_A * scale, 0).to(do.dtype.element_ty)
	p_do = tl.make_block_ptr(do, (T, V), (HV, 1), (i_t BT, i_v * BV), (BT, BV), (1, 0))
	p_dv = tl.make_block_ptr(dv, (T, V), (HV, 1), (i_t BT, i_v * BV), (BT, BV), (1, 0))
	b_do = tl.load(p_do, boundary_check=(0, 1))
	b_dv += tl.dot(b_A.to(b_do.dtype), b_do)
	tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1))


	@triton.heuristics({
	'USE_G': lambda args: args['g'] is not None,
	'USE_G_GAMMA': lambda args: args['g_gamma'] is not None,
	'USE_A': lambda args: args['A'] is not None,
	'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
	})
	@triton.jit(do_not_specialize=['T'])
	def chunk_bwd_kernel_dv_local(
	q,
	k,
	g,
	g_gamma,
	A,
	do,
	dv,
	cu_seqlens,
	chunk_indices,
	scale,
	T,
	H: tl.constexpr,
	K: tl.constexpr,
	V: tl.constexpr,
	BT: tl.constexpr,
	BK: tl.constexpr,
	BV: tl.constexpr,
	USE_G: tl.constexpr,
	USE_G_GAMMA: tl.constexpr,
	USE_EXP2: tl.constexpr,
	USE_A: tl.constexpr,
	IS_VARLEN: tl.constexpr,
	):
	i_t, i_bh = tl.program_id(0), tl.program_id(1)
	i_b, i_h = i_bh // H, i_bh % H
	if IS_VARLEN:
	i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
	bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
	T = eos - bos
	else:
	bos, eos = i_b * T, i_b * T + T

	# offset calculation
	q += (bos * H + i_h) * K
	k += (bos * H + i_h) * K
	do += (bos * H + i_h) * V
	dv += (bos * H + i_h) * V

	if USE_A:
	p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (BT, T), (1, HBT), (0, i_t BT), (BT, BT), (0, 1))
	b_A = tl.load(p_A, boundary_check=(0, 1))
	else:
	if USE_G:
	g += bos * H + i_h
	p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,))
	b_g = tl.load(p_g, boundary_check=(0,))
	if USE_G_GAMMA:
	b_gamma = tl.load(g_gamma + i_h)
	b_g = b_gamma * (tl.arange(0, BT) + 1)

	b_A = tl.zeros([BT, BT], dtype=tl.float32)
	for i_k in range(tl.cdiv(K, BK)):
	p_k = tl.make_block_ptr(k, (T, K), (HK, 1), (i_t BT, i_k * BK), (BT, BK), (1, 0))
	p_q = tl.make_block_ptr(q, (K, T), (1, HK), (i_k BK, i_t * BT), (BK, BT), (0, 1))

	b_k = tl.load(p_k, boundary_check=(0, 1))
	b_q = tl.load(p_q, boundary_check=(0, 1))
	b_A += tl.dot(b_k, b_q) * scale
	if USE_G or USE_G_GAMMA:
	if USE_EXP2:
	b_A *= exp2(b_g[None, :] - b_g[:, None])
	else:
	b_A *= exp(b_g[None, :] - b_g[:, None])

	o_t = i_t * BT + tl.arange(0, BT)
	m_t = o_t < T
	m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t)
	b_A = tl.where(m_A, b_A, 0).to(do.dtype.element_ty)

	for i_v in range(tl.cdiv(V, BV)):
	p_do = tl.make_block_ptr(do, (T, V), (HV, 1), (i_t BT, i_v * BV), (BT, BV), (1, 0))
	p_dv = tl.make_block_ptr(dv, (T, V), (HV, 1), (i_t BT, i_v * BV), (BT, BV), (1, 0))
	b_do = tl.load(p_do, boundary_check=(0, 1))
	b_dv = tl.dot(b_A.to(b_do.dtype), b_do)
	tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1))


	def chunk_fwd_o(
	q: torch.Tensor,
	k: torch.Tensor,
	v: torch.Tensor,
	h: torch.Tensor,
	g: torch.Tensor \| None = None,
	g_gamma: torch.Tensor \| None = None,
	scale: float \| None = None,
	cu_seqlens: torch.LongTensor \| None = None,
	chunk_size: int = 64,
	chunk_indices: torch.LongTensor \| None = None,
	use_exp2: bool = False,
	transpose_state_layout: bool = False,
	) -> torch.Tensor:
	B, T, H, K, V = *q.shape, v.shape[-1]
	BT = chunk_size
	if chunk_indices is None and cu_seqlens is not None:
	chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
	NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
	if scale is None:
	scale = k.shape[-1] ** -0.5

	o = torch.empty_like(v)
	def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H)
	chunk_fwd_kernel_o[grid](
	q=q,
	k=k,
	v=v,
	h=h,
	g=g,
	g_gamma=g_gamma,
	o=o,
	cu_seqlens=cu_seqlens,
	chunk_indices=chunk_indices,
	scale=scale,
	T=T,
	H=H,
	K=K,
	V=V,
	BT=BT,
	USE_EXP2=use_exp2,
	TRANSPOSE_STATE=transpose_state_layout,
	)
	return o


	def chunk_bwd_dv(
	q: torch.Tensor,
	k: torch.Tensor,
	do: torch.Tensor,
	dh: torch.Tensor,
	g: torch.Tensor \| None = None,
	g_gamma: torch.Tensor \| None = None,
	scale: float \| None = None,
	cu_seqlens: torch.LongTensor \| None = None,
	chunk_size: int = 64,
	chunk_indices: torch.LongTensor \| None = None,
	use_exp2: bool = False,
	) -> torch.Tensor:
	B, T, H, K, V = *k.shape, do.shape[-1]
	BT = chunk_size
	if chunk_indices is None and cu_seqlens is not None:
	chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
	# H100 can have larger block size
	if check_shared_mem('hopper', k.device.index):
	CONST_TILING = 128
	elif check_shared_mem('ada', k.device.index):
	CONST_TILING = 64
	else:
	CONST_TILING = 32
	BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING)
	BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING)
	NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
	NV = triton.cdiv(V, BV)
	if scale is None:
	scale = k.shape[-1] ** -0.5

	dv = torch.empty_like(do)
	grid = (NV, NT, B * H)
	chunk_bwd_kernel_dv[grid](
	q=q,
	k=k,
	g=g,
	g_gamma=g_gamma,
	do=do,
	dv=dv,
	dh=dh,
	cu_seqlens=cu_seqlens,
	chunk_indices=chunk_indices,
	scale=scale,
	T=T,
	H=H,
	K=K,
	V=V,
	BT=BT,
	BK=BK,
	BV=BV,
	USE_EXP2=use_exp2,
	# H200 Triton autotune can wedge while timing this backward kernel.
	# Use the same conservative fixed launch across ranks.
	num_warps=4,
	num_stages=3,
	)
	return dv


	def chunk_bwd_dv_local(
	q: torch.Tensor,
	k: torch.Tensor,
	do: torch.Tensor,
	g: torch.Tensor \| None = None,
	g_gamma: torch.Tensor \| None = None,
	A: torch.Tensor \| None = None,
	scale: float = None,
	cu_seqlens: torch.LongTensor \| None = None,
	chunk_size: int = 64,
	chunk_indices: torch.LongTensor \| None = None,
	use_exp2: bool = False,
	) -> torch.Tensor:
	B, T, H, K, V = *k.shape, do.shape[-1]
	BT = chunk_size
	if chunk_indices is None and cu_seqlens is not None:
	chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
	# H100 can have larger block size
	if check_shared_mem('hopper', k.device.index):
	CONST_TILING = 128
	elif check_shared_mem('ada', k.device.index):
	CONST_TILING = 64
	else:
	CONST_TILING = 32
	BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING)
	BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING)
	NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)

	dv = torch.empty_like(do)
	grid = (NT, B * H)
	chunk_bwd_kernel_dv_local[grid](
	q=q,
	k=k,
	g=g,
	g_gamma=g_gamma,
	A=A,
	do=do,
	dv=dv,
	cu_seqlens=cu_seqlens,
	chunk_indices=chunk_indices,
	scale=scale,
	T=T,
	H=H,
	K=K,
	V=V,
	BT=BT,
	BK=BK,
	BV=BV,
	USE_EXP2=use_exp2,
	# Keep the local DV variant deterministic too; otherwise it can hit
	# the same H200 autotune stall after the global DV kernel is fixed.
	num_warps=4,
	num_stages=3,
	)
	return dv


	def chunk_bwd_dqkwg(
	q: torch.Tensor,
	k: torch.Tensor,
	v: torch.Tensor,
	do: torch.Tensor,
	h: torch.Tensor,
	dh: torch.Tensor,
	w: torch.Tensor \| None = None,
	g: torch.Tensor \| None = None,
	g_gamma: torch.Tensor \| None = None,
	dv: torch.Tensor \| None = None,
	scale: float \| None = None,
	cu_seqlens: torch.LongTensor \| None = None,
	chunk_size: int = 64,
	chunk_indices: torch.LongTensor \| None = None,
	use_exp2: bool = False,
	transpose_state_layout: bool = False,
	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

	B, T, H, K, V = *k.shape, v.shape[-1]
	BT = chunk_size
	if chunk_indices is None and cu_seqlens is not None:
	chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
	NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)

	if check_shared_mem('hopper', k.device.index):
	CONST_TILING = 128
	elif check_shared_mem('ada', k.device.index):
	CONST_TILING = 64
	else:
	CONST_TILING = 32
	BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING)
	BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING)
	NK = triton.cdiv(K, BK)
	dq = torch.empty_like(q)
	dk = torch.empty_like(k)
	dg = torch.empty(NK, *g.shape, dtype=torch.float32, device=g.device) if g is not None else None
	dw = torch.empty_like(w) if w is not None else None

	grid = (NK, NT, B * H)
	# H200 Triton autotune can stall for this backward kernel while benchmarking
	# configs. Use a fixed conservative launch; avoiding autotune matters more
	# than chasing a small per-kernel speedup during long 20B training runs.
	chunk_bwd_kernel_dqkwg[grid](
	q=q,
	k=k,
	v=v,
	g=g,
	g_gamma=g_gamma,
	h=h,
	do=do,
	dh=dh,
	dw=dw,
	dq=dq,
	dk=dk,
	dv=dv,
	dg=dg,
	cu_seqlens=cu_seqlens,
	chunk_indices=chunk_indices,
	scale=scale,
	B=B,
	T=T,
	H=H,
	K=K,
	V=V,
	BT=BT,
	BK=BK,
	BV=BV,
	USE_EXP2=use_exp2,
	TRANSPOSE_STATE=transpose_state_layout,
	num_warps=4,
	num_stages=3,
	)

	if dg is not None:
	dg = dg.sum(0)
	return dq, dk, dw, dg