Attention¶

`decoder_attention`	Compute paged-KV attention for the decode stage.
`encoder_attention`	Variable-length BF16 attention forward for prefill with GQA and optional causal/custom mask.
`get_qk_token`	Compute actual_cu_seq_k, seq_lens_kv, and return max/total token statistics for subsequent attention kernels.
`write_cache_kv_decoder`	Process the current decode token fused QKV, generate Q output for decode attention, and append K/V to paged KV cache.
`write_cache_kv_encoder`	Split fused QKV, apply RoPE and optional RMSNorm to Q/K, output contiguous Q/K/V, and simultaneously write K/V to paged KV cache.

decoder_attention¶

decoder_attention( q_input, cu_seq_q, seq_lens_kv, seq_len_encoder, seq_len_decoder, cache_k, cache_v, block_tables, attn_out, q_dequant_scale=None, kv_dequant_scale=None, head_num=0, kv_head_num=0, head_dim=128, max_seq_k=0, mtp_steps=0, cache_quant_type='none', )¶

Compute paged-KV attention for the decode stage.

Reads the current-step query together with the paged K/V cache, attends over each sample’s visible KV range located through block_tables, and writes the attention result into attn_out. Supports grouped-query attention (GQA), multi-token prediction (MTP), and both a BF16 cache and a block-wise FP8 cache.

Parameters:

q_input (Tensor) – Query, shape [batch, kv_head_num, 64, head_dim]. Dtype is bfloat16 when cache_quant_type is “none” and float8_e4m3fn when it is “block_wise_fp8”. The 64 rows hold the GQA-folded query heads padded to a fixed width.
cu_seq_q (Tensor) – Prefix-sum of decode query tokens per sample, shape [batch + 1], dtype int32.
seq_lens_kv (Tensor) – Total visible KV length per sample, shape [batch], dtype int32.
seq_len_encoder (Tensor) – Encoder/prefill length per sample, shape [batch], dtype int32.
seq_len_decoder (Tensor) – Decode prefix length per sample, shape [batch], dtype int32.
cache_k (Tensor) – Paged K cache, shape [num_blocks, kv_head_num, block_size, head_dim]. Dtype is bfloat16 for the “none” cache and uint8 (packed FP8) for the “block_wise_fp8” cache.
cache_v (Tensor) – Paged V cache, same shape and dtype as cache_k.
block_tables (Tensor) – Mapping from each sample to its physical cache blocks, shape [batch, max_blocks_per_seq], dtype int32.
attn_out (Tensor) – Output buffer written in place, shape [q_token_num, head_num * head_dim], dtype bfloat16.
q_dequant_scale (Tensor, optional) – Query dequantization scale used only by the FP8 path, shape [batch, kv_head_num], dtype float32. Ignored when cache_quant_type is “none”. Default: None.
kv_dequant_scale (Tensor, optional) – Block-wise K/V dequantization scale used only by the FP8 path, shape [num_blocks, kv_head_num, 4], dtype float32. Ignored when cache_quant_type is “none”. Default: None.
head_num (int) – Number of query heads.
kv_head_num (int) – Number of key/value heads.
head_dim (int) – Per-head dimension. Default: 128.
max_seq_k (int) – Maximum visible KV length across the current batch.
mtp_steps (int, optional) – Number of extra multi-token-prediction tokens; 0 means ordinary single-token decode. Default: 0.
cache_quant_type (str, optional) – KV cache quantization mode. Supported values are “none” (BF16 cache) and “block_wise_fp8” (block-wise FP8 cache). Default: “none”.

Returns:

None. The attention result is written in place into attn_out.

Example

import torch
from halops import (
    write_cache_kv_encoder,
    write_cache_kv_decoder,
    decoder_attention,
)

head_num, kv_head_num, head_dim = 14, 1, 128
batch, seq_len_enc, block_size = 1, 128, 128
max_seq_len = 4096
num_blocks = batch * max_seq_len // block_size
total_kv = seq_len_enc + 1  # prefill tokens + 1 decode token

# Shared paged KV cache and block table.
cache_k = torch.zeros(num_blocks, kv_head_num, block_size, head_dim,
                      dtype=torch.bfloat16, device="cuda")
cache_v = torch.zeros_like(cache_k)
block_tables = torch.arange(
    num_blocks, dtype=torch.int32, device="cuda"
).reshape(batch, max_seq_len // block_size)
# Identity RoPE table (cos=1, sin=0) so Q/K pass through unchanged.
rotary_embs = torch.zeros(2, max_seq_len, head_dim // 2,
                          dtype=torch.float32, device="cuda")
rotary_embs[0] = 1

# Step 1: prefill. Write the encoder Q/K/V and populate the KV cache.
qkv_enc = torch.randn(batch * seq_len_enc, head_num + 2 * kv_head_num,
                      head_dim, dtype=torch.bfloat16, device="cuda")
cu_seq_enc = torch.tensor([0, seq_len_enc], dtype=torch.int32,
                          device="cuda")
seq_len_enc_t = torch.tensor([seq_len_enc], dtype=torch.int32,
                             device="cuda")
seq_len_dec_t = torch.zeros(batch, dtype=torch.int32, device="cuda")
q_enc = torch.zeros(batch * seq_len_enc, head_num, head_dim,
                    dtype=torch.bfloat16, device="cuda")
k_enc = torch.zeros(batch * seq_len_enc, kv_head_num, head_dim,
                    dtype=torch.bfloat16, device="cuda")
v_enc = torch.zeros_like(k_enc)
write_cache_kv_encoder(
    qkv_enc, cu_seq_enc, cu_seq_enc, rotary_embs, seq_len_enc_t,
    seq_len_dec_t, cache_k, cache_v, block_tables, q_enc, k_enc, v_enc,
    head_num=head_num, kv_head_num=kv_head_num, head_dim=head_dim,
    max_seq_len_q=seq_len_enc, max_seq_len_k=seq_len_enc,
    max_input_len=max_seq_len, cache_quant_type="none",
)

# Step 2: decode. Write the current token's K/V and build the folded
# q_output of shape [batch, kv_head_num, 64, head_dim].
qkv_dec = torch.randn(batch, head_num + 2 * kv_head_num, head_dim,
                      dtype=torch.bfloat16, device="cuda")
cu_seq_q = torch.tensor([0, 1], dtype=torch.int32, device="cuda")
seq_len_enc_t2 = torch.zeros(batch, dtype=torch.int32, device="cuda")
seq_len_dec_t2 = torch.tensor([seq_len_enc], dtype=torch.int32,
                              device="cuda")
q_output, _ = write_cache_kv_decoder(
    qkv_dec, cu_seq_q, seq_len_enc_t2, seq_len_dec_t2, rotary_embs,
    cache_k, cache_v, block_tables,
    head_num=head_num, kv_head_num=kv_head_num, head_dim=head_dim,
    max_input_len=max_seq_len, cache_quant_type="none",
)

# Step 3: decode attention over the paged cache.
seq_lens_kv = torch.tensor([total_kv], dtype=torch.int32, device="cuda")
attn_out = torch.zeros(batch, head_num * head_dim,
                       dtype=torch.bfloat16, device="cuda")
decoder_attention(
    q_output, cu_seq_q, seq_lens_kv, seq_len_enc_t2, seq_len_dec_t2,
    cache_k, cache_v, block_tables, attn_out,
    head_num=head_num, kv_head_num=kv_head_num, head_dim=head_dim,
    max_seq_k=total_kv, cache_quant_type="none",
)

encoder_attention¶

encoder_attention( q_input, k_input, v_input, cu_seq_q, cu_seq_k, attn_out, mask=None, )¶

Variable-length BF16 attention forward for prefill with GQA and optional causal/custom mask.

Parameters:

q_input (torch.Tensor) – [total_q_tokens, head_num, head_dim] BF16.
k_input (torch.Tensor) – [total_k_tokens, kv_head_num, head_dim] BF16.
v_input (torch.Tensor) – [total_k_tokens, kv_head_num, head_dim] BF16.
cu_seq_q (torch.Tensor) – [batch_size+1] int32, Q token prefix-sum.
cu_seq_k (torch.Tensor) – [batch_size+1] int32, K token prefix-sum.
attn_out (torch.Tensor) – [total_q_tokens, head_num, head_dim] BF16, output tensor allocated by caller. Written in-place.
mask (torch.Tensor | None) – Optional int32 mask; when passed, uses CustomMask template branch. When None, the kernel applies the default causal mask (bottom-right aligned) per sequence.

Returns:

None. The result is written in-place to attn_out.

Example

import torch
from halops import encoder_attention

head_num, kv_head_num, head_dim = 56, 4, 128
seq_len = 1024
cu_seq_q = torch.tensor([0, seq_len], dtype=torch.int32, device="cuda")
cu_seq_k = torch.tensor([0, seq_len], dtype=torch.int32, device="cuda")
q = torch.randn(seq_len, head_num, head_dim,
                dtype=torch.bfloat16, device="cuda")
k = torch.randn(seq_len, kv_head_num, head_dim,
                dtype=torch.bfloat16, device="cuda")
v = torch.randn(seq_len, kv_head_num, head_dim,
                dtype=torch.bfloat16, device="cuda")
attn_out = torch.empty(seq_len, head_num, head_dim,
                       dtype=torch.bfloat16, device="cuda")
encoder_attention(q, k, v, cu_seq_q, cu_seq_k, attn_out)

get_qk_token¶

get_qk_token( seq_lens_encoder, seq_lens_decoder, cu_seq_q, cu_seq_k, actual_cu_seq_k, seq_lens_kv, kv_head_num=0, )¶

Compute actual_cu_seq_k, seq_lens_kv, and return max/total token statistics for subsequent attention kernels.

Parameters:

seq_lens_encoder (torch.Tensor) – [batch_size] int32, encoder token count per sample.
seq_lens_decoder (torch.Tensor) – [batch_size] int32, existing decode prefix length per sample.
cu_seq_q (torch.Tensor) – [batch_size+1] int32, Q token prefix-sum.
cu_seq_k (torch.Tensor) – [batch_size+1] int32, reserved for interface compatibility.
actual_cu_seq_k (torch.Tensor) – [batch_size+1] int32 output, actual K/V prefix-sum.
seq_lens_kv (torch.Tensor) – [max_batch_size] int32 output, per-sample KV length for decode.
kv_head_num (int) – Number of KV heads.

Returns:

A list containing a single CPU tensor with contents [max_enc, max_dec, total_q, total_k].

Return type:

List[torch.Tensor]

Example

import torch
from halops import get_qk_token

kv_head_num = 1
batch, max_batch = 2, 8
seq_lens_encoder = torch.tensor(
    [128, 0], dtype=torch.int32, device="cuda"
)
seq_lens_decoder = torch.tensor(
    [0, 512], dtype=torch.int32, device="cuda"
)
cu_seq_q = torch.tensor([0, 128, 129], dtype=torch.int32, device="cuda")
cu_seq_k = torch.zeros(batch + 1, dtype=torch.int32, device="cuda")
actual_cu_seq_k = torch.zeros(
    batch + 1, dtype=torch.int32, device="cuda"
)
seq_lens_kv = torch.zeros(max_batch, dtype=torch.int32, device="cuda")
(max_token_num,) = get_qk_token(
    seq_lens_encoder, seq_lens_decoder, cu_seq_q, cu_seq_k,
    actual_cu_seq_k, seq_lens_kv, kv_head_num,
)
max_token_num.tolist()  # [max_enc, max_dec, total_q, total_k]

[128, 513, 129, 129]

write_cache_kv_decoder¶

write_cache_kv_decoder( qkv_input, cu_seq_q, seq_len_encoder, seq_len_decoder, rotary_embs, cache_k, cache_v, block_tables, kv_dequant_scale=None, q_norm_gamma=None, k_norm_gamma=None, head_num=0, kv_head_num=0, head_dim=128, max_input_len=0, cache_quant_type='none', )¶

Process the current decode token fused QKV, generate Q output for decode attention, and append K/V to paged KV cache.

Parameters:

qkv_input (torch.Tensor) – [q_token_num, head_num + 2*kv_head_num, head_dim] BF16.
cu_seq_q (torch.Tensor) – [batch_size+1] int32, decode Q token prefix-sum.
seq_len_encoder (torch.Tensor) – [batch_size] int32; samples with >0 skip decoder-only cache write.
seq_len_decoder (torch.Tensor) – [batch_size] int32, existing decode prefix length before write.
rotary_embs (torch.Tensor) – RoPE cos/sin table.
cache_k (torch.Tensor) – paged K cache, BF16 or FP8/uint8.
cache_v (torch.Tensor) – paged V cache, BF16 or FP8/uint8.
block_tables (torch.Tensor) – [batch_size, max_blocks_per_seq] int32.
kv_dequant_scale (torch.Tensor | None) – Required for FP8 cache path, K/V block-wise dequant scale.
q_norm_gamma (torch.Tensor | None) – Optional Q RMSNorm gamma.
k_norm_gamma (torch.Tensor | None) – Optional K RMSNorm gamma.
head_num (int) – Number of Q heads.
kv_head_num (int) – Number of KV heads.
head_dim (int) – Per-head dimension (primarily targets 128).
max_input_len (int) – Maximum length of the RoPE table.
cache_quant_type (str) – "none" for BF16 cache; "block_wise_fp8" for FP8 cache.

Returns:

A list [q_output, q_dequant_scale] where –

q_output: [batch, kv_head_num, 64, head_dim], BF16 for BF16 path or float8_e4m3fn for FP8 path.
q_dequant_scale: [batch, kv_head_num] float32. In the FP8 path it holds the Q dequantization scale. In the BF16 path it is still allocated (never None) but left uninitialized and must not be used.

Return type:

List[torch.Tensor]

Example

import torch
from halops import write_cache_kv_decoder

head_num, kv_head_num, head_dim = 14, 1, 128
batch, prefix_len, block_size = 1, 128, 128
max_input_len = 4096
num_blocks = batch * max_input_len // block_size
# one new token per sequence for the current decode step
qkv = torch.randn(batch, head_num + 2 * kv_head_num, head_dim,
                  dtype=torch.bfloat16, device="cuda")
cu_seq_q = torch.tensor([0, 1], dtype=torch.int32, device="cuda")
seq_len_encoder = torch.zeros(batch, dtype=torch.int32, device="cuda")
seq_len_decoder = torch.tensor(
    [prefix_len], dtype=torch.int32, device="cuda"
)
rotary_embs = torch.zeros(2, max_input_len, head_dim // 2,
                          dtype=torch.float32, device="cuda")
rotary_embs[0] = 1  # identity RoPE
cache_k = torch.zeros(num_blocks, kv_head_num, block_size, head_dim,
                      dtype=torch.bfloat16, device="cuda")
cache_v = torch.zeros_like(cache_k)
block_tables = torch.arange(
    num_blocks, dtype=torch.int32, device="cuda"
).reshape(batch, max_input_len // block_size)
q_output, q_dequant_scale = write_cache_kv_decoder(
    qkv, cu_seq_q, seq_len_encoder, seq_len_decoder, rotary_embs,
    cache_k, cache_v, block_tables,
    head_num=head_num, kv_head_num=kv_head_num, head_dim=head_dim,
    max_input_len=max_input_len, cache_quant_type="none",
)
q_output.shape

torch.Size([1, 1, 64, 128])

write_cache_kv_encoder¶

write_cache_kv_encoder( qkv_input, cu_seq_q, cu_seq_k, rotary_embs, seq_len_encoder, seq_len_decoder, cache_k, cache_v, block_tables, q_input, k_input, v_input, kv_dequant_scale=None, q_norm_gamma=None, k_norm_gamma=None, head_num=0, kv_head_num=0, head_dim=128, max_seq_len_q=0, max_seq_len_k=0, max_input_len=0, cache_quant_type='none', )¶

Split fused QKV, apply RoPE and optional RMSNorm to Q/K, output contiguous Q/K/V, and simultaneously write K/V to paged KV cache.

Supports BF16 cache and block-wise FP8 cache.

Parameters:

qkv_input (torch.Tensor) – [total_q_tokens, head_num + 2*kv_head_num, head_dim] BF16.
cu_seq_q (torch.Tensor) – [batch_size+1] int32, Q token prefix-sum.
cu_seq_k (torch.Tensor) – [batch_size+1] int32, K/V token prefix-sum.
rotary_embs (torch.Tensor) – RoPE cos/sin table, first half is cos, second half is sin.
seq_len_encoder (torch.Tensor) – [batch_size] int32, encoder/prefill token count.
seq_len_decoder (torch.Tensor) – [batch_size] int32, existing decode prefix length.
cache_k (torch.Tensor) – paged K cache, BF16 or FP8/uint8.
cache_v (torch.Tensor) – paged V cache, BF16 or FP8/uint8.
block_tables (torch.Tensor) – [batch_size, max_blocks_per_seq] int32.
q_input (torch.Tensor) – [total_q_tokens, head_num, head_dim] BF16 output, Q after RoPE.
k_input (torch.Tensor) – [total_k_tokens, kv_head_num, head_dim] BF16 output.
v_input (torch.Tensor) – [total_k_tokens, kv_head_num, head_dim] BF16 output.
kv_dequant_scale (torch.Tensor | None) – Required for FP8 cache path, K/V block-wise dequant scale.
q_norm_gamma (torch.Tensor | None) – Optional Q RMSNorm gamma.
k_norm_gamma (torch.Tensor | None) – Optional K RMSNorm gamma.
head_num (int) – Number of Q heads.
kv_head_num (int) – Number of KV heads.
head_dim (int) – Per-head dimension (primarily targets 128).
max_seq_len_q (int) – Max Q length in current batch.
max_seq_len_k (int) – Max K length in current batch.
max_input_len (int) – Max RoPE table length.
cache_quant_type (str) – “none” for BF16 cache; “block_wise_fp8” for FP8 cache.

Returns:

None. Writes the RoPE-applied Q/K/V into q_input/k_input/v_input and appends K/V into the paged cache_k/cache_v in-place.

Example

import torch
from halops import write_cache_kv_encoder

head_num, kv_head_num, head_dim = 14, 1, 128
batch, seq_len, block_size = 1, 128, 128
max_seq_len = 4096
num_blocks = batch * max_seq_len // block_size
total_q = batch * seq_len
qkv = torch.randn(total_q, head_num + 2 * kv_head_num, head_dim,
                  dtype=torch.bfloat16, device="cuda")
cu_seq_q = torch.tensor([0, seq_len], dtype=torch.int32, device="cuda")
cu_seq_k = cu_seq_q.clone()
# identity RoPE table: cos=1, sin=0. The BF16 encoder applies
# partial NeoX RoPE with rope_dim = head_dim // 4.
rotary_embs = torch.zeros(2, max_seq_len, head_dim // 4,
                          dtype=torch.float32, device="cuda")
rotary_embs[0] = 1
seq_len_encoder = torch.tensor([seq_len], dtype=torch.int32, device="cuda")
seq_len_decoder = torch.zeros(batch, dtype=torch.int32, device="cuda")
cache_k = torch.zeros(num_blocks, kv_head_num, block_size, head_dim,
                      dtype=torch.bfloat16, device="cuda")
cache_v = torch.zeros_like(cache_k)
block_tables = torch.arange(
    num_blocks, dtype=torch.int32, device="cuda"
).reshape(batch, max_seq_len // block_size)
q_out = torch.zeros(total_q, head_num, head_dim,
                    dtype=torch.bfloat16, device="cuda")
k_out = torch.zeros(total_q, kv_head_num, head_dim,
                    dtype=torch.bfloat16, device="cuda")
v_out = torch.zeros_like(k_out)
write_cache_kv_encoder(
    qkv, cu_seq_q, cu_seq_k, rotary_embs, seq_len_encoder,
    seq_len_decoder, cache_k, cache_v, block_tables,
    q_out, k_out, v_out,
    head_num=head_num, kv_head_num=kv_head_num, head_dim=head_dim,
    max_seq_len_q=seq_len, max_seq_len_k=seq_len,
    max_input_len=max_seq_len, cache_quant_type="none",
)