# How to Configure Needle 2 Model Parameters Using TransformerConfig

> Configure Needle 2 model parameters easily by instantiating TransformerConfig. Control dimensionality, attention, quantization, and engram memory settings.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: how-to-guide
- Published: 2026-08-14

---

**You configure Needle 2 model parameters by instantiating the `TransformerConfig` dataclass from [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) and passing it to `SimpleAttentionNetwork`, which encapsulates dimensionality, attention topology, quantization settings, and engram memory configuration.**

The Needle 2 architecture in the cactus-compute/needle repository provides a flexible transformer implementation designed for efficient inference and training on JAX. All tunable hyperparameters—from model width and attention heads to quantization bit-widths and memory cache settings—are centralized in the `TransformerConfig` dataclass. Understanding how to properly configure this single configuration object allows you to customize every aspect of the model without modifying the underlying source code.

## Core Architecture of TransformerConfig

The `TransformerConfig` dataclass acts as the single source of truth for the Needle 2 model architecture. When you instantiate `SimpleAttentionNetwork`, the constructor receives a config instance that supplies dimensionality, depth, attention topology, and precision settings to all submodules.

The configuration validates inputs during initialization (see `__init__` in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) lines 84-92) and automatically normalizes fields such as `attn_dim` to inherit from `d_model` when set to zero. The model subsequently accesses derived properties like `jax_dtype` (lines 92-95) to ensure consistent data typing across JAX operations.

## Key Model Parameters

TransformerConfig exposes granular control over the model's behavior through the following parameter categories.

### Dimensionality and Model Depth

- **d_model**: Base hidden dimension determining model width (default 768 for needle preset)
- **attn_dim**: Dimension of Q/K/V projections; inherits `d_model` when set to 0
- **num_layers**: Number of transformer blocks in the stack (typically 27)
- **vocab_size**: Token vocabulary size (default 8192)
- **max_seq_len**: Maximum sequence length supported (2048 default)
- **pad_token_id**: Token ID used for padding mask generation (0 default)
- **contrastive_dim**: Output dimensionality for the contrastive head (128 default)

### Attention Configuration

- **num_heads**: Number of parallel attention heads (12 in needle preset)
- **num_kv_heads**: Key/value heads for multi-query attention (6 in needle preset)
- **kv_window**: KV-cache window size; 0 enables automatic budget-based calculation via `effective_kv_window` (lines 14-17)
- **rope_theta**: Rotary position embedding scaling factor (1e5 default)
- **flash**: Boolean enabling Flash-Attention kernels when available

### Quantization and Precision

- **dtype**: JAX data type specification (`"float32"`, `"bfloat16"`, or `"float16"`)
- **kv_bits**: Bit-width for key/value quantization (e.g., 8 or 4)
- **act_bits**: Activation quantization precision
- **weight_bits**: Optional weight quantization override

### Engram Memory Settings

- **engram_layers**: Tuple of layer indices hosting learned token-level caches (e.g., `(2, 15)`)
- **engram_orders**: N-gram orders for Engram geometry (e.g., `(2, 3)`)
- **engram_slots**: Number of memory slots in the Engram cache (default 8192)
- **engram_heads**: Heads per Engram layer (0 for automatic selection)

### Performance Optimization

- **remat**: Enable rematerialization for memory-efficient training via `nn.scan`
- **scan_unroll**: Unroll factor for the layer scan loop
- **mhc_lanes**: Multi-head-correlation lane count

## Creating a Custom TransformerConfig

To configure Needle 2 model parameters, import the dataclass and instantiate it with your desired specifications. The constructor filters unknown keys automatically.

```python
from needle.model.architecture import TransformerConfig

custom_cfg = TransformerConfig(
    d_model=1024,          # Wider model than default

    num_heads=16,
    num_kv_heads=8,
    num_layers=32,
    dtype="float32",       # Use FP32 for numerical debugging

    flash=False,           # Disable Flash-Attention on CPU

    kv_bits=4,             # 4-bit KV quantization

    act_bits=4,
    weight_bits="4",       # 4-bit weight quantization

    engram_layers=(3, 12),
    engram_orders=(2, 3, 4),
)

```

You can also extend an existing preset by unpacking its dictionary and overriding specific fields.

```python
from needle.model.architecture import TransformerConfig, SimpleAttentionNetwork
import jax.numpy as jnp

# Start from needle preset and modify

base = TransformerConfig(
    **{"d_model": 768, "num_heads": 12, "num_kv_heads": 6, "num_layers": 27}
)
custom_cfg = TransformerConfig(
    **base.__dict__, 
    dtype="float32", 
    kv_bits=4, 
    act_bits=4
)

# Instantiate model

model = SimpleAttentionNetwork(config=custom_cfg)
tokens = jnp.array([[1, 5, 23, 0, 0]])
logits = model(tokens)  # Shape: (batch, seq_len, vocab_size)

```

Inspect derived properties to verify configuration.

```python
from needle.model.architecture import effective_kv_window

print("JAX dtype:", custom_cfg.jax_dtype)                  # → jnp.float32

print("Effective KV window:", custom_cfg.kv_window)        # → 0 (auto)

print("Budget-based window:", effective_kv_window(custom_cfg))

```

## Integrating Configuration with Model Components

The `TransformerConfig` object propagates through the entire Needle 2 architecture according to the cactus-compute/needle source code. The `Stack` module receives the config to allocate correct dtype and lane counts (lines 76-82 in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)). The `Engram` memory system queries `engram_layers`, `engram_orders`, and `engram_slots` to construct its cache tables (lines 90-95).

When `kv_window` is set to 0, the `effective_kv_window` budget function automatically calculates the optimal cache window based on available memory. The `MultiHeadAttention` layer respects `num_heads`, `num_kv_heads`, and `flash` settings to select appropriate kernel implementations.

## Summary

- **TransformerConfig** centralizes all Needle 2 hyperparameters in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)
- Configure dimensionality via `d_model`, `attn_dim`, and `num_layers`
- Control attention topology with `num_heads`, `num_kv_heads`, and `kv_window`
- Enable quantization using `kv_bits`, `act_bits`, and `weight_bits` for compressed inference
- Customize memory caching through `engram_layers`, `engram_orders`, and `engram_slots`
- Pass the config instance to `SimpleAttentionNetwork(config=...)` to apply all settings

## Frequently Asked Questions

### What is the default data type for Needle 2 models?

The default `dtype` is `"bfloat16"` for production deployments, though you can specify `"float32"` for debugging or `"float16"` for specific hardware. Access the processed JAX type via the `jax_dtype` property on your config instance, which converts the string to the appropriate `jnp.dtype`.

### How do I disable Flash-Attention for CPU compatibility?

Set `flash=False` in your `TransformerConfig` constructor. This forces the `MultiHeadAttention` module to use standard attention implementations instead of optimized Flash-Attention kernels, which is necessary when running on hardware without specialized GPU support.

### Can I use different quantization levels for keys, values, and weights?

Yes. The config supports independent quantization via `kv_bits` for key/value caches, `act_bits` for activations, and `weight_bits` for model parameters. Set these to integers such as 4 or 8 to enable quantization, or leave `weight_bits` as an empty string to disable weight quantization while quantizing KV caches.

### Where are the Engram memory layers configured?

Engram layers are specified using the `engram_layers` tuple parameter, which accepts layer indices (0-indexed) where the learned token-level cache should be inserted. The `engram_orders` tuple defines the n-gram window sizes, while `engram_slots` controls the cache capacity per layer according to the implementation in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py).