Implementation:Turboderp org Exllamav2 Ext QAttn H
| Knowledge Sources | |
|---|---|
| Domains | Attention, Quantization, C_Extension |
| Last Updated | 2026-02-15 00:00 GMT |
Overview
C++ header declaring the API for ExLlamaV2's fused quantized multi-head attention module, including construction, two-phase forward pass, LoRA integration, and tensor-parallel attention variants.
Description
ext_qattn.h defines the interface for the quantized attention subsystem that fuses layer normalization, Q/K/V projection, RoPE application, and output projection into optimized CUDA operations.
make_q_attn is the factory function that constructs a quantized attention module handle. It accepts 32 parameters including:
- layernorm and layernorm_bias -- Pre-attention normalization weights and optional bias.
- layernorm_is_rms and headnorm_is_rms -- Flags selecting RMS vs standard layer norm.
- q_q_proj, q_k_proj, q_v_proj, q_o_proj -- Opaque handles to quantized weight matrices for query, key, value, and output projections.
- temp_state, temp_dq -- Scratch buffers for intermediate states and dequantized weights.
- Architecture parameters: hidden_size, num_heads, num_kv_heads, head_dim, max_seq_len.
- rope_style and sincos_size -- RoPE configuration.
- q_norm and k_norm -- Optional per-head query/key normalization weights.
- post_layernorm and post_layernorm_bias -- Optional post-attention normalization.
- residual_fp32 -- Whether to maintain residual connections in FP32 precision.
- use_graphs -- Enable CUDA graph capture for the attention module.
q_attn_forward_1 performs the first phase of attention: applies layer norm, projects Q/K/V through quantized matrices, and applies RoPE. It writes to temporary q_temp, k_temp, v_temp tensors.
q_attn_forward_2 performs the second phase: takes the attention output and projects it through the output matrix, applying any post-normalization.
q_attn_set_loras attaches LoRA adapter weight pairs (A and B matrices) for each of the four projections (Q, K, V, O).
MHAFwdKVCacheFunc is a typedef for the Flash Attention forward function signature with KV cache support. set_flash_attn_func initializes the Flash Attention function pointer.
tp_attn_forward_paged_ and tp_attn_forward_ are tensor-parallel attention variants that distribute computation across multiple devices, supporting paged KV cache and standard contiguous cache respectively.
Usage
This API is used by the Python ExLlamaV2Attention layer class. The module is constructed once during model loading via make_q_attn, and then q_attn_forward_1 and q_attn_forward_2 are called on each inference step. The two-phase design allows the KV cache update and attention score computation to happen between the phases.
Code Reference
Source Location
- Repository: Turboderp_org_Exllamav2
- File: exllamav2/exllamav2_ext/ext_qattn.h
- Lines: 1-167
Signature
uintptr_t make_q_attn(
torch::Tensor layernorm,
torch::Tensor layernorm_bias,
bool layernorm_is_rms,
bool headnorm_is_rms,
float norm_epsilon,
uintptr_t q_q_proj,
uintptr_t q_k_proj,
uintptr_t q_v_proj,
uintptr_t q_o_proj,
torch::Tensor temp_state,
torch::Tensor temp_dq,
int max_rows,
int hidden_size,
int num_heads,
int num_kv_heads,
int head_dim,
int max_seq_len,
bool has_residual,
int rope_style,
int sincos_size,
torch::Tensor q_norm,
torch::Tensor k_norm,
torch::Tensor post_layernorm,
torch::Tensor post_layernorm_bias,
bool residual_fp32,
bool use_graphs
);
void q_attn_forward_1(
uintptr_t q_attn,
torch::Tensor x,
int batch_size,
int q_len,
int past_len,
torch::Tensor past_lens,
torch::Tensor q_temp,
torch::Tensor k_temp,
torch::Tensor v_temp,
torch::Tensor sin,
torch::Tensor cos,
const std::vector<uintptr_t>& loras,
torch::Tensor loras_temp
);
void q_attn_forward_2(
uintptr_t q_attn,
torch::Tensor x,
torch::Tensor attn_output,
int batch_size,
int q_len,
const std::vector<uintptr_t>& loras,
torch::Tensor loras_temp
);
int q_attn_set_loras(
uintptr_t q_attn,
std::unordered_map<uintptr_t, torch::Tensor>& q_proj_lora_a,
std::unordered_map<uintptr_t, torch::Tensor>& q_proj_lora_b,
std::unordered_map<uintptr_t, torch::Tensor>& k_proj_lora_a,
std::unordered_map<uintptr_t, torch::Tensor>& k_proj_lora_b,
std::unordered_map<uintptr_t, torch::Tensor>& v_proj_lora_a,
std::unordered_map<uintptr_t, torch::Tensor>& v_proj_lora_b,
std::unordered_map<uintptr_t, torch::Tensor>& o_proj_lora_a,
std::unordered_map<uintptr_t, torch::Tensor>& o_proj_lora_b
);
typedef std::vector<at::Tensor> (*MHAFwdKVCacheFunc)(...);
void set_flash_attn_func();
void tp_attn_forward_paged_(uintptr_t tp_context, ...);
void tp_attn_forward_(uintptr_t tp_context, ...);
Import
from exllamav2 import exllamav2_ext as ext_c
# Construct quantized attention module
handle = ext_c.make_q_attn(
layernorm, layernorm_bias, layernorm_is_rms, headnorm_is_rms,
norm_epsilon, q_q_proj, q_k_proj, q_v_proj, q_o_proj,
temp_state, temp_dq, max_rows, hidden_size,
num_heads, num_kv_heads, head_dim, max_seq_len,
has_residual, rope_style, sincos_size,
q_norm, k_norm, post_layernorm, post_layernorm_bias,
residual_fp32, use_graphs
)
I/O Contract
| Function | Key Parameters | Output | Description |
|---|---|---|---|
| make_q_attn | 32 config params including norm weights, quantized projections, architecture dims | uintptr_t handle | Constructs fused attention module; returns opaque handle |
| q_attn_forward_1 | q_attn handle, input x, batch/seq dims, sin/cos for RoPE | q_temp, k_temp, v_temp written in-place | Phase 1: norm + Q/K/V projection + RoPE |
| q_attn_forward_2 | q_attn handle, residual x, attn_output | x modified in-place (residual added) | Phase 2: output projection + residual connection |
| q_attn_set_loras | q_attn handle, 8 LoRA weight maps (A+B for Q/K/V/O) | int (number of LoRAs set) | Attaches LoRA adapters to all four projections |
| tp_attn_forward_paged_ | tp_context, per-device tensors, paged block_index | hidden_states modified in-place | Tensor-parallel paged attention across devices |
| tp_attn_forward_ | tp_context, per-device tensors, past_len_tp | hidden_states modified in-place | Tensor-parallel contiguous-cache attention |
Usage Examples
from exllamav2 import exllamav2_ext as ext_c
# Phase 1: Compute Q, K, V with RoPE
ext_c.q_attn_forward_1(
q_attn_handle, x,
batch_size, q_len, past_len, past_lens,
q_temp, k_temp, v_temp,
sin, cos,
loras, loras_temp
)
# (KV cache update and attention computation happens here)
# Phase 2: Output projection and residual
ext_c.q_attn_forward_2(
q_attn_handle, x, attn_output,
batch_size, q_len,
loras, loras_temp
)
Related Pages
- Turboderp_org_Exllamav2_Ext_QMLP_H -- Quantized MLP module with similar construction pattern
- Turboderp_org_Exllamav2_Ext_Norm -- Normalization operations used within attention
- Turboderp_org_Exllamav2_Ext_RoPE -- RoPE positional encoding applied during forward_1
- Turboderp_org_Exllamav2_Ext_TP_H -- Tensor parallelism context for tp_attn_forward variants