Qwen.AttentionBlock.ixx File Reference

Qwen 3.8 full-attention decoder block (inference: prefill + decode). More...

#include <memory>
#include <vector>
#include <string>
#include <format>
#include <sstream>
#include <stdexcept>
#include <optional>
#include <type_traits>
import Dnn.Quantization.KvCache.Policy;
import Dnn.Quantization.Weight.PrecisionPlan;
import Dnn.Quantization.Weight.Policies;
import Serialization.ModelArchive;
import Dnn.Components.AttentionOutputGate;
import Dnn.CompositeComponent;
import Dnn.ComponentType;
import Dnn.Component;
import Compute.CpuMemoryResource;
import Dnn.ActivationType;
import Dnn.Components.QwenConfig;
import Dnn.TensorDataType;
import Compute.Device;
import Dnn.Components.ITransformerBlock;
import Dnn.TensorTypes;
import Dnn.TensorDataTypeTraits;
import Dnn.Tensor;
import Dnn.ITensor;
import Compute.DeviceAllocation;
import Compute.Observation;
import Compute.DeviceTypeTraits;
import Compute.ExecutionContext;
import Compute.DeviceId;
import Compute.DeviceType;
import Compute.ExecutionContextFactory;
import Dnn.Components.Residual;
import Compute.IExecutionContext;
import Dnn.TensorOps;
import Dnn.Components.Rope;
import Compute.GqaState;
import Dnn.Components.RmsNorm;
import Serialization.Mode;
import Dnn.Components.Swiglu;
import Dnn.Components.Gqa;
import Dnn.Components.Linear;

Classes
struct	Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvPolicy >::BlockBuildContexts
	The per-child build contexts and split-scratch geometry this block implies. More...
class	Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvPolicy >
	One Qwen 3.8 full-attention decoder block. More...
struct	Mila::Dnn::QwenAttentionBlockWorkspace< TDeviceType, TPrecision >
	Transformer-owned shared activation workspace for QwenAttentionBlock (pooling). More...

Namespaces
namespace	Mila
	Mila main API namespace.

Functions
template<DeviceType TDeviceType, TensorDataType TPrecision> requires PrecisionSupportedOnDevice<TPrecision, TDeviceType>
QwenAttentionBlockWorkspace< TDeviceType, TPrecision >	Mila::Dnn::makeQwenAttentionBlockWorkspace (const QwenConfig &config, DeviceId device, dim_t B, dim_t prefill_chunk, const std::string &name_prefix)
	Allocate the shared attention-block workspace.

Detailed Description

Qwen 3.8 full-attention decoder block (inference: prefill + decode).

One of the two block kinds in the 3:1 interleave. The other, QwenDeltaNetBlock, is a separate class rather than a flag on this one: Gemma's two kinds are ONE mixer at two geometries, so a bool kGlobal fits it, while Qwen's two kinds are different mixers and share no projection shape (Specifications/Qwen3.8.md section 8, Phase 2 revision).

Standard pre-norm decoder, with one delta that has no precedent in the tree:

• attn_output_gate. The query projection is DOUBLE width and emits [query | gate]; swish(gate) scales the attention output elementwise before o_proj (section 2). The gate half never reaches attention – it bypasses RoPE and the KV cache entirely and rejoins after the attention output is computed.

And four that are ordinary once stated: head_dim 256 is decoupled from the 5120 residual stream (so o_proj is non-square), rotary width is 64 of 256 (partial rotary), query and key carry a per-head RMSNorm before RoPE (q_norm/k_norm, [head_dim]), and the FFN is SwiGLU over a fused gate+up projection.

PER-ROLE PRECISION. The block's third parameter is a PLAN, not a policy: its four Linears carry four independently-chosen weight formats, which is the whole reason this model fits a 12 GiB card (section 4 – no uniform 4-bit path exists). A bare policy still works and means "uniform", because PrecisionPlanFor lifts it.

using QwenAttentionLayer = QwenAttentionBlock<DeviceType::Cuda, TensorDataType::BF16, QwenPrecisionPlan>;

A plan that omits a role this block builds fails the requires-clause below, naming the block – rather than silently falling back to some default precision.

Inference-only, like the Gemma blocks: implements ITransformerBlock's prefill/decode, not Component forward/backward.

WEIGHT LAYOUT CONTRACT. fc_qkv_proj is fused as [query | gate | key | value], with query and gate as contiguous halves rather than interleaved per head. The converter (section 8 item 8) owns producing that ordering; the block's two splits assume it.

The checkpoint does NOT store it that way, which is what makes this a contract rather than a description. Qwen/Qwen3.8-27B carries three separate tensors – q_proj [12288, 5120], k_proj and v_proj [1024, 5120] – and the query projection interleaves the two halves PER HEAD, [q_h0 | gate_h0 | q_h1 | gate_h1 | ...], because the reference views it as [..., heads, 2 * head_dim] and chunks on the last axis. So the converter concatenates the three and de-interleaves the query half; nothing here changes.

Function Documentation

makeQwenAttentionBlockWorkspace()

template<DeviceType TDeviceType, TensorDataType TPrecision>
requires PrecisionSupportedOnDevice<TPrecision, TDeviceType>

QwenAttentionBlockWorkspace< TDeviceType, TPrecision > Mila::Dnn::makeQwenAttentionBlockWorkspace ( const QwenConfig & config, DeviceId device, dim_t B, dim_t prefill_chunk, const std::string & name_prefix )

export

Allocate the shared attention-block workspace.

Lives here, beside the struct, rather than inside QwenTransformer, because the network is not the only thing that drives these blocks: a layer-streamed parity harness holds one block at a time and needs the identical slot geometry. A second copy of these eighteen widths would agree until the first time one of them changed, and then measure a different model than the one it is checking.