DeepSeek V4.1 Flash Engineering, Part 2: From Weight Loading to an Eight-GPU Text Pipeline
ROCm integration in practice: interpreting FP8 weights and carrying shared KV, sparse selections, and mHC state across GPU boundaries
Part 1 covered Engram: keeping learned memory in host RAM and using lookup, AVX-512 BF16, and CPU/GPU overlap to make it part of inference. This article returns to the GPU backbone.
Splitting a model into eight stages assigns the work to devices. Correct text generation also requires preserving the model's state. DeepSeek V4.1 Flash shares compressed KV across layers, reuses sparse selections, and passes mHC mixing coefficients between sublayers. Device boundaries cut through these dependencies; they cannot disappear with a stage's local variables.
We integrated the official safetensors weights on eight AMD Radeon Pro W7900D GPUs, fixing weight interpretation, ROCm execution, session state, and the output head along the way. This account follows that implementation and its subsequent corrections. It does not fold other ongoing work into the conclusions of those integration tests.
Start with the complete inference round
The language backbone has 40 layers, a hidden dimension of 5,120, and four mHC residual streams. The eight-GPU configuration assigns consecutive complete layers to each device. Each layer_ends entry is an exclusive upper bound:
layer_ends: [5, 10, 15, 20, 25, 30, 35, 40]
| GPU | Backbone layers | State encountered in this stage |
|---|---|---|
| 0 | L0–L4 | Engram at L1; L2 publishes compressed KV and selections |
| 1 | L5–L9 | Continues with L2 state, then L8 publishes a new group |
| 2 | L10–L14 | Continues with L8 state; L14 applies Engram and publishes new state |
| 3 | L15–L19 | Uses L14's shared compressed KV and selections |
| 4 | L20–L24 | L20 publishes KV for the remaining layers; L24 updates selections |
| 5 | L25–L29 | Uses L20 KV; L28 updates selections |
| 6 | L30–L34 | Uses L20 KV; L32 updates selections |
| 7 | L35–L39 | Uses L20 KV; L36 updates selections, then the backbone finishes |
Layer numbers are zero-based. This is a pipeline of consecutive layers: the device responsible for a layer runs its attention, router, and expert computation. Five layers per GPU describes placement; it does not mean eight independent pieces of the same token execute simultaneously.
With speculative decoding disabled, the complete round is:
text → tokenizer / conversation encoding → embedding
→ GPU0: L0–L4 → GPU1: L5–L9 → … → GPU7: L35–L39
→ final mHC reduction → final norm → LM head → sampling
→ next token returns to the entry point
Prefill can place multiple input chunks into the pipeline. Ordinary single-request decode must wait for the sampled token before starting the next token's round. Eight GPUs provide capacity and a division of work; throughput still depends on the dependency chain, chunk size, and stage times.
Step 1: Assemble file tensors into usable model weights
Weight loading starts with three questions: what is the tensor called, what is its logical shape, and how should its bytes be interpreted?
Safetensors and GGUF solve different layers of the problem
Safetensors is primarily a safe, simple tensor container. It records tensor names, dtypes, shapes, and byte ranges, either in one file or across shards described by an index JSON. It does not execute deserialization code, works well for official training checkpoints, and supports tensor-level or row-level reads. The model configuration and implementation must still explain what a tensor such as layers.2.attn.wq_a.weight means in DeepSeek V4.1 inference.
GGUF combines a tensor directory with inference metadata that commonly describes the model architecture, context parameters, and tokenizer, while also defining quantization types used by the GGML ecosystem. It is closer to a distribution package for local inference: one file can tell an engine both what model it contains and how packed weight bytes are represented. Conversion has already chosen naming, quantization, and metadata conventions, however, so a third-party GGUF still needs validation that it fully represents a new architecture.
The distinction is not “Safetensors means high precision, GGUF means low precision.” Safetensors can store FP8, MXFP4, BF16, and other dtypes. GGUF can contain F32, F16, BF16, or a mixture of quantized tensors. Their fundamental differences concern container contracts, metadata conventions, and quantized representations.
| Dimension | Safetensors | GGUF |
|---|---|---|
| Primary role | General, safe tensor storage | Inference-oriented model and quantization container |
| Model structure | Usually interpreted with external configuration and model code | Usually described by file metadata |
| Tokenizer | Usually stored in external files | Commonly carried in GGUF metadata |
| Typical organization | Official checkpoint, often sharded | Single-file or size-sharded distribution package |
| zLLM access | Read tensors or selected rows by name, shape, and dtype | Parse metadata, tensor directory, and GGML quantization types, then access matrices on demand |
Format independence begins at zLLM's model-orchestration boundary
zLLM does not force Safetensors and GGUF to share a byte layout. The implementation has four layers:
File containers
SafetensorsStore / GgufReader
↓
Model weight assembly
tensor naming, shape validation, logical weight roles
↓
Unified prepared weights and backend capabilities
LinearWeight / PreparedLayer / ExpertSource
↓
DeepSeek V4.1 runtime
attention → shared KV / selections → MoE → mHC
The container layer handles metadata, tensor indexing, and raw byte access. The Safetensors path assembles BlockFP8, MXFP8, MXFP4, and dense tensors into model weights with explicit logical shapes. The GGUF path interprets its own quantization types and prepares linear weights and expert sources consumable by the same backend.
The paths may differ before weight preparation. The current DeepSeek implementation has separate Safetensors and GGUF layer-preparation functions because tensor naming, fusion, and quantized slicing differ. After preparation, both produce the same DeepSeekV4PreparedLayer and enter the same layer runtime. Attention, MoE, KV lifecycle, and mHC ordering are not reimplemented for each file extension.
Format independence therefore does not mean that any GGUF can automatically replace the official checkpoint. It means that model algorithms do not depend on a file container; container differences converge at the reading and weight-assembly boundary. A new format must still supply complete architecture metadata, tensor mappings, and backend support. If Engram, mHC, shared compressed KV, or a required quantized kernel is missing, the engine should reject it explicitly rather than generate under incorrect assumptions.
zLLM detects the top-level namespace. For example, the current loader uses a language_model. prefix when it finds language_model.embed.weight, and unprefixed names otherwise. This accommodates checkpoint organizations encountered during integration. The tensor directory determines the prefix, not the name of the download folder.
Likewise, knowing that weights are quantized does not determine their execution path:
| Representation | What loading must establish |
|---|---|
| BlockFP8 | Data matrix, two-dimensional scale grid, and block sizes along both axes |
| MXFP8 | Codes and scales grouped along each row's input dimension, rather than a two-dimensional block layout |
| MXFP4 experts | Packed data, grouped scales, expert index, and logical gate/up/down shapes |
| Unquantized tensors | dtype, shape, and their role in normalization, mixing, or projection |
Experts can also be stored separately or in merged tensors. The implementation reads individually named expert matrices or selects the appropriate row range from a merged rank-three expert tensor. For the FFN, gate/up have shape [intermediate, hidden], while down has shape [hidden, intermediate]. The relationship between packed columns and logical columns must remain intact.
The weight layer handles these differences. Model orchestration receives interpretable matrices and expert sources, and the backend selects the corresponding execution path.
Step 2: The FP8 bug was in the two-dimensional scale layout
A concrete integration fix concerned 32×32 BlockFP8 in the official V4.1 backbone linear weights. The official inference implementation explicitly uses this block layout. Official model implementation
For an [M, N] matrix, the scale grid has shape:
[ceil(M / 32), ceil(N / 32)]
Element (r, c) uses the scale at (r / 32, c / 32), with integer division. Both axes determine the scale. Row-independent MXFP8 groups values within each row; a group size of 32 does not make it the same representation.
We added a small test: all codes in a 64×64 matrix encode the same value, while its four 32×32 blocks receive scales corresponding to multipliers of 1, 2, 4, and 8. The decoded quadrants must contain those four values. This directly tests scale transitions along both rows and columns.
Small matrices introduce another subtlety: their scale shapes may not distinguish 32×32 from the older 128×128 layout. The loader must use model-version information as well as tensor shape.
The correction had two parts. 80b8d3e5 fixed weight assembly; 07cf53f1 completed the ROCm 32×32 BlockFP8 path. Correct CPU decoding does not automatically establish that a GPU kernel uses the same scale indices.
Step 3: Separate local KV, shared compressed KV, and selections
The familiar idea of one KV cache per layer needs to be expanded into several distinct objects.
Sliding-window KV belongs to each layer and serves its local history. The configured window is 128.
Compressed KV is produced by a designated source layer and reused by a group of layers. zLLM's V4.1 mapping is:
| Source layer | Consumers | Ratio |
|---|---|---|
| L2 | L2–L7 | 2 |
| L8 | L8–L13 | 2 |
| L14 | L14–L19 | 2 |
| L20 | L20–L39 | 1 |
L0 and L1 do not use this compressed branch. From L20 onward, ratio=1 still means consuming the 1:1 representation produced at L20. It does not disable the branch or make each layer produce an independent replacement. Ratio describes granularity; the source mapping describes provenance.
This can be understood as a further engineering step beyond DSA, or DeepSeek Sparse Attention. DSA first makes the history dimension sparse. As the context grows, the indexer builds a lighter representation of the history, scores it against the current query, and selects the Top-K positions. Attention then reads KV only from those positions instead of visiting the entire history. The longer the context, the more unnecessary history this mechanism can skip.
V4.1 extends sparsity from the time axis into network depth. Adjacent layers do not have to construct another compressed history, nor does every layer need to repeat the full selection process. A KV source layer publishes compressed KV and index keys for a group; index source layers recompute selections at a configured cadence; the other layers consume the latest published result. The design therefore has two independent dimensions of sparsity:
| Dimension | Problem it addresses | Runtime behavior |
|---|---|---|
| History sparsity | A query does not need to visit every previous token | The indexer selects Top-K history positions and attention reads only those positions |
| Cross-layer sparsity | Adjacent layers need not repeatedly build similar history representations and selections | Source layers publish compressed KV, index keys, or selections for later layers in the group |
This sharing preserves each layer's own attention computation. Every layer derives its own query from its hidden state and maintains its own sliding-window KV. Layers reuse the compressed representation of distant history and, within its valid range, the selected history positions. Layer-specific computation remains, while highly repetitive history preparation and retrieval move from once per layer to once per group with periodic refreshes.
The L20–L39 range makes this separation especially clear. Its compressed KV continues to come from L20, while selections are refreshed at L20, L24, L28, L32, and L36. In other words, which layer supplies history and when the model reassesses which history is worth reading follow separate schedules. The larger, more stable compressed history can span more layers, while the lighter selection result can be refreshed more often.
This follows a path similar to our GLM 5.3 DSA implementation. A full indexer layer in GLM 5.3 builds the history index and publishes a selection; subsequent IndexShare layers reuse the latest valid selection, and consecutive MTP steps can reuse it when their state is complete and continuous. The weights and layer maps differ between the models, but the engineering goal is the same: find the worthwhile positions in a long history, then let subsequent computation reuse that retrieval instead of rescanning the entire history at every layer.
Selections are the history positions chosen by the indexer for each query. The configuration updates them at L2, L8, L14, L20, L24, L28, L32, and L36; other layers reuse the latest published result. Selections and compressed KV are different state: L24 can publish new selections while continuing to read L20's compressed history.
Separating these objects lets the runtime express what a layer writes, whose history it reads, and whether it performs a new selection.
Step 4: Do not reset shared state at a GPU boundary
Consider the transition from GPU0 to GPU1.
L2 on GPU0 has already produced shared history and selections. GPU1 starts at L5, and L5–L7 need those results until L8 publishes the next group. Creating fresh local shared state at every stage entrance gives L5 its hidden tensor but loses the selection it should reuse.
Fix cc0be412 made selections part of the pipeline work item. A stage restores the preceding stage's published selections, updates them as needed, and passes them onward. Compressed-KV sources are resolved through the whole model's layer mapping instead of whichever source happened to appear inside the current stage.
| Data | How the pipeline manages it |
|---|---|
| Hidden state | Travels through stages with the current input chunk |
| Selections | Continue with the chunk and update at indexer layers |
| Compressed KV history | Lives in the session cache table; consumers resolve it by source layer |
| Layer-local sliding KV | Maintained by the session state responsible for that layer |
Continuing with a chunk is a logical dependency, not a requirement to download the data to the CPU. Subsequent optimization changed its physical placement.
Step 5: Share across stages within a session, isolate across sessions
Another lifecycle bug remained after fixing source resolution: each stage independently forked a full-layer cache table from its template.
Every table could be valid, but they were different instances. GPU0 wrote L2 history into one instance; GPU1 looked up L2 in another and found empty state.
Fix 307c5e85 introduced a coordinated session-opening path: fork one full-layer cache table for the session, then distribute references to that same table across all eight stages. Engram sequence state is also created at session scope and shared among that session's stages.
This defines the boundary of sharing. Model objects and read-only weights can be reused between sessions. Mutable sequence history belongs to a particular session. Resetting a session must handle shared history and Engram token state as well as layer-local caches.
Sharing only works when ownership, consumers, and write visibility are explicit.
Step 6: mHC pre-mix also crosses layers and GPUs
Once weights and KV are connected, checking the hidden tensor's shape is still insufficient.
V4.1 mHC carries four residual streams and reduces them with mixing coefficients for sublayer inputs. Our early path treated sublayers more independently; the correction passed pre-mix state continuously:
pre published by the preceding layer's FFN
→ input reduction for this layer's attention
pre published by this layer's attention
→ input reduction for this layer's FFN
pre published by this layer's FFN
→ the next layer's attention, or final output-head reduction
The first layer uses the defined initialization behavior when no upstream pre is available. After the final layer, the last FFN's pre reduces the expanded state before final norm and the LM head.
Fix 8919f797 made layer execution return both hidden state and the pre to propagate, and carried pre across stage boundaries. When GPU0 finishes L4, GPU1's L5 needs both L4's hidden tensor and the corresponding pre.
Shape checks alone are poor at catching this error: coefficients from different sublayers can have identical shapes while representing different points in execution. Numerical correctness requires checking who produced an input and who should consume it.
After correctness: keep shared data near its consumers
The first goal was ensuring that shared state existed and reached the right consumer. We then optimized physical access.
One change added local mirrors of compressed KV. Reading source-GPU history works, but long prefill repeatedly scans the remote data. Later code maintains a local copy at the consumer, transferring only appended rows while the layout and existing prefix match. Growth, layout changes, or a different prefix require rebuilding the necessary contents. This uses additional local VRAM in exchange for local reads during computation.
Another change keeps selections on GPUs. The early path downloaded indices into host vectors and uploaded them to later layers. The optimized ROCm selection holds a device-buffer reference, reuses it on the same GPU, and uses ordered P2P transfers across GPUs rather than repeatedly passing through the CPU.
Neither change is merely a pointer substitution. Asynchronous consumers must retain source data until completion. Production, copying, and consumption need correct ordering, and temporary allocations need a lifetime suitable for crossing stages. Separating logical sharing from physical placement creates room to move from working inference to faster inference.
How we verified the chain
Validation operated at several levels; producing text did not replace the other checks:
- Format: independent scale-quadrant tests for BlockFP8 interpretation.
- Orchestration: a small full-forward model test covering compressed sharing, candidate selection, and Engram, including prefill, decode, and the effect of enabling Engram.
- Real weights: loading the official checkpoint and executing all 40 layers, final reduction, head, and sampling.
- Optimization regression: fixed inputs and output budgets, complete output-text hashes, and separate first-token and generation timing.
One scope limit matters. The official model includes candidate-block filtering followed by a second selection stage. In the ROCm integration configuration inspected for this article, candidate_source_layer=None and candidate_topk_blocks=0 select a single-stage Top-K path. A CPU reference for the two-stage logic does not establish that the ROCm path reproduces every official selection step. Nor should the two strategies be called equivalent for all inputs without validation.
Part 1's approximately 26.06-second TTFT for 50K input and 20.45 tokens/s single-request generation came from later optimized fixed-input text tests with DSpark disabled. They are not the isolated performance benefit of an FP8 or state-management fix, and this writing session did not rerun those benchmarks.
The essential work from checkpoint to eight-GPU generation was preserving one model across locations: consistent byte interpretation, explicit history sources, correct session boundaries, and the right order of mixing state. Performance work can then change where data lives and when it moves while preserving those relationships.
The next article follows the image path: finding differences in preprocessing, vision-tower, and aligner intermediates, then connecting image-and-text inputs to this language backbone.
Engineering evidence: zLLM's V4.1 integration notes and changes in 80b8d3e5, 07cf53f1, cc0be412, 307c5e85, 8919f797, and e9b539e4. Model configuration reference: official inference configuration.