DeepSeek-V3 Administrator & Developer Guide

This guide is for ML researchers, infrastructure engineers, and developers operating or modifying the deepseek-ai/DeepSeek-V3 reference inference repository. It documents architecture, configuration, deployment, internals, and security at the depth necessary to run the code in production-adjacent settings without consulting the upstream paper.

The scope is strictly the source code at commit 9b4e978 (5 Python files, 4 JSON configs, one requirements file). Anything not present in that source is not documented here.

1. Architecture Overview

DeepSeek-V3's inference repo is a flat, layered architecture: a single-file Transformer at the core, two preparatory utilities (HF conversion and FP8 dequantization), one CLI driver, and a Triton kernel layer. There is no plugin system, no service layer, no dependency injection — module-level constants control behavior, and dist-aware Linears handle parallelism.

1.1 Inference pipeline diagram

  ┌──────────────────────────────────────────────────────────────────────────┐
  │                           CHECKPOINT PREPARATION                         │
  │                                                                          │
  │   HF release (safetensors)                                               │
  │        │                                                                 │
  │        │  fp8_cast_bf16.py  (optional, for non-FP8 hardware)             │
  │        │     - reads model.safetensors.index.json                        │
  │        │     - calls weight_dequant() (Triton) per FP8 weight            │
  │        │     - writes BF16 shards + new index                            │
  │        ▼                                                                 │
  │   HF (BF16 or original FP8)                                              │
  │        │                                                                 │
  │        │  convert.py                                                     │
  │        │     - renames params (self_attn->attn, mlp->ffn, ...)           │
  │        │     - skips model.layers.61 (MTP head, unused)                  │
  │        │     - shards by dim or by expert index                          │
  │        ▼                                                                 │
  │   /weights/.../model{rank}-mp{mp}.safetensors    + tokenizer files       │
  └──────────────────────────────────────────────────────────────────────────┘
                                      │
                                      ▼
  ┌──────────────────────────────────────────────────────────────────────────┐
  │                              RUNTIME (per rank)                          │
  │                                                                          │
  │   torchrun -> WORLD_SIZE / RANK / LOCAL_RANK env vars                    │
  │        │                                                                 │
  │        ▼                                                                 │
  │   generate.main()                                                        │
  │        │  dist.init_process_group("nccl") if WORLD_SIZE > 1              │
  │        │  torch.cuda.set_device(LOCAL_RANK)                              │
  │        │  set_default_dtype(bfloat16)                                    │
  │        │                                                                 │
  │        │  ModelArgs(**json.load(config))                                 │
  │        │                                                                 │
  │        ▼                                                                 │
  │   Transformer(args)        ┌────────────────────────────────────────┐    │
  │        │                   │  ParallelEmbedding (vocab // world_size)│    │
  │        │                   │  Block × n_layers                       │    │
  │        │                   │    ├── RMSNorm (attn_norm)              │    │
  │        │                   │    ├── MLA (Multi-head Latent Attn)     │    │
  │        │                   │    │     ColumnParallel(wq, wq_b,wkv_b) │    │
  │        │                   │    │     RowParallel(wo)                │    │
  │        │                   │    │     KV cache: kv_cache + pe_cache  │    │
  │        │                   │    │       (or k_cache + v_cache naive) │    │
  │        │                   │    ├── RMSNorm (ffn_norm)               │    │
  │        │                   │    └── MLP (dense layers <n_dense)      │    │
  │        │                   │        OR MoE (Gate + Experts + shared) │    │
  │        │                   │           Gate -> top-k experts         │    │
  │        │                   │           routed Experts (sharded)      │    │
  │        │                   │           shared MLP (replicated)       │    │
  │        │                   │  RMSNorm                                │    │
  │        │                   │  ColumnParallel head -> all_gather      │    │
  │        │                   └────────────────────────────────────────┘    │
  │        │                                                                 │
  │        │  AutoTokenizer.from_pretrained(ckpt_path)                       │
  │        │  warm-up generate("DeepSeek", 2 tokens)                         │
  │        │  load_model(model, model{rank}-mp{world_size}.safetensors)      │
  │        ▼                                                                 │
  │   {interactive REPL  ||  --input-file batch}                             │
  │        │                                                                 │
  │        │  per turn:                                                      │
  │        │     tokenizer.apply_chat_template(messages)                     │
  │        │     generate(model, tokens, max_new_tokens, eos, temperature)   │
  │        │        │                                                        │
  │        │        ▼                                                        │
  │        │     Transformer.forward(tokens[prev_pos:cur_pos], prev_pos)     │
  │        │        │                                                        │
  │        │        ▼                                                        │
  │        │     sample(logits, temperature)  OR  argmax                     │
  │        │                                                                 │
  │        ▼                                                                 │
  │   tokenizer.decode(...) -> printed completion                            │
  └──────────────────────────────────────────────────────────────────────────┘

1.2 Component inventory

1.3 Technology stack

Note

The repository has no test suite, no Dockerfile, no CI configuration, no setup script. It is shipped as runnable reference code; downstream users are expected to integrate it.

1.4 Data flow for a single token

For each step in generate():

1. Transformer.forward(tokens[:, prev_pos:cur_pos], prev_pos) is called on every rank with the same input.
2. ParallelEmbedding looks up rows in its local vocab slice, masks the rest, and all_reduces if world_size > 1.
3. freqs_cis is sliced to [start_pos:start_pos+seqlen].
4. Each Block runs attn(attn_norm(x), ...) then ffn(ffn_norm(x)). The first n_dense_layers blocks use MLP; the rest use MoE.
5. MLA computes Q (with optional LoRA), K/V latent + rope, fills KV cache, applies softmax-scaled attention, projects back via wo (RowParallel — has all_reduce).
6. MoE gates each token to n_activated_experts of n_routed_experts experts; only experts whose index falls in this rank's slice run; shared experts run on every rank; per-expert outputs are summed and all_reduced.
7. After all blocks, RMSNorm then ColumnParallelLinear head produces local logits; an all_gather concatenates the vocab slices across ranks.
8. The full vocab logits return to generate() which samples and writes the next token.

Note

Steps 6 and 7 are the only places per-token cross-rank communication is non-trivial: MoE's all_reduce over expert outputs and the head's all_gather over vocab slices.

2. System Requirements

2.1 Hardware

Note

FP8 (torch.float8_e4m3fn) is supported natively on Hopper (H100/H200), Ada (some L40-class), and Blackwell. On Ampere (A100, A40, RTX 3090) FP8 is emulated by Triton — use the BF16 cast for these.

2.2 Software

• Linux x86_64 (or any platform PyTorch 2.4.1 + Triton 3.0.0 supports). Windows is not officially supported by Triton.
• Python 3.10–3.12 (PyTorch 2.4 wheel coverage).
• CUDA toolkit compatible with PyTorch 2.4 (CUDA 12.1 wheel is the typical stack).
• NCCL bundled with PyTorch.

2.3 Capacity sizing

Working memory per rank for the 671B FP8 model on 8-way tensor parallel:

Important

Per-rank weight memory drops linearly with world_size for parallel-Linear weights, but dense layer weights and shared experts replicate across ranks. Going from 8 to 16 ranks does not halve total memory if your config has shared experts.

3. Configuration Reference

3.1 Configuration sources

3.2 ModelArgs — complete field reference

The ModelArgs dataclass in model.py (line 20) is the full architecture surface. Defaults are the 16B-shape values.

3.3 Per-config differences (shipping JSONs)

Important

Fields shown as (default ...) are not present in the JSON; they take their ModelArgs default. Keys that are present override the default.

3.4 Module-level constants in model.py

These are not configurable via JSON; they require a source edit.

Caution

Transformer.__init__ mutates module-level globals (world_size, rank) and class attributes (Linear.dtype, Linear.scale_fmt). This means constructing two Transformer instances with different args.dtype in the same process will silently corrupt the first. Don't mix configs in one Python process.

3.5 CLI flag reference (all scripts)

3.5.1 convert.py

3.5.2 fp8_cast_bf16.py

3.5.3 generate.py

4. Deployment & Scaling

4.1 Deployment topologies

The repo supports four launch shapes:

4.2 Pre-launch checklist

• [ ] --n-experts matches the JSON's n_routed_experts.
• [ ] --model-parallel (in convert.py) equals total ranks (--nnodes × --nproc-per-node).
• [ ] Sharded files exist: model0-mp{world_size}.safetensors … model{world_size-1}-mp{world_size}.safetensors.
• [ ] Tokenizer files (anything matching *token*) are inside --ckpt-path.
• [ ] vocab_size divides world_size evenly (102400 and 129280 accommodate 1, 2, 4, 5, 8, 10, 16, 20, 32 …).
• [ ] dim divides world_size (covers parallel Linear assertions).
• [ ] All ranks can reach the master address on the master port.
• [ ] nvidia-smi shows the expected GPUs as available.

4.3 Scaling considerations

Tip

Tensor-parallel degree controls memory footprint per rank but adds NCCL traffic per layer. For the 671B / V3.1 model:

• 8-way TP on a single node with NVLink/NVSwitch is the sweet spot.
• Going to 16-way across two nodes adds inter-node all_reduce per layer; expect a noticeable per-token latency hit if interconnect is sub-200 Gbps.
• There is no built-in pipeline parallelism, so you cannot get around inter-node bandwidth by partitioning layers across nodes.

Note

Throughput does not increase with --max-batch-size past the point where one rank's KV cache exceeds GPU memory. Pure-batch parallelism is constrained by KV memory, not compute. The repo offers no built-in continuous batching.

4.4 Determinism

generate.py calls torch.manual_seed(965) once at startup. Subsequent sample() calls draw from the global RNG. Determinism across runs requires:

• Identical config JSON.
• Identical seed (already fixed at 965).
• Identical hardware and PyTorch/Triton versions.
• Same --temperature and prompt sequence.

Warning

Triton autotune may select different kernel configs across runs (caches under ~/.triton/cache). Identical numerical output run-to-run is not guaranteed; bit-exactness on FP8 GEMM in particular is hardware-dependent.

5. Operations

5.1 Operating procedures

5.1.1 Cold start

1. Verify checkpoint integrity: confirm model{i}-mp{N}.safetensors exists for i in [0, N).
2. Confirm nvidia-smi reports all expected GPUs idle.
3. Launch: torchrun --nnodes ... --nproc-per-node ... generate.py --ckpt-path ... --config ... [--interactive | --input-file ...].
4. First user prompt may take longer than subsequent ones because of the warm-up generation, Triton autotune, and optional weight materialization paths. Subsequent forwards are steady-state.

5.1.2 Health checks

The repo does not ship liveness/readiness endpoints. To verify a running session:

• Rank 0 prints args (the resolved ModelArgs) immediately after construction.
• Rank 0 also accepts and responds at >>> once warm-up completes.
• nvidia-smi should show stable GPU utilization spikes during decode.

5.1.3 Graceful shutdown

• Interactive mode: type /exit at the prompt. The script breaks the loop and calls dist.destroy_process_group() if world_size > 1.
• Batch mode: the script terminates after processing all prompts.
• Forced: Ctrl-C on the launcher kills torchrun, which signals all workers. NCCL communicators may take a moment to tear down.

5.2 Logging

There is no structured logging. Diagnostic output:

• Rank 0 prints the resolved ModelArgs.
• Progress bars (tqdm / trange) appear in convert.py and fp8_cast_bf16.py.
• fp8_cast_bf16.py prints Warning: Missing scale_inv tensor for {name}, skipping conversion to stdout when an FP8 weight has no scale.
• generate.py prints completions in batch mode with Prompt: / Completion: labels.
• Non-zero ranks have print overridden to a no-op.

Tip

To capture rank-0 output for offline debugging, redirect with tee:

torchrun ... generate.py ... 2>&1 | tee run.log

Other ranks' stderr / stdout still go to the launcher's defaults — useful when chasing NCCL errors.

5.3 Known operational issues

Note

All four issues live in the released weights / tokenizer, not in this repo's Python. Mitigations are operational, not code fixes.

5.4 Backup & restore

Checkpoints are immutable artifacts; standard filesystem backup applies. The repo does not maintain mutable runtime state beyond the in-process messages list (cleared on /exit or /clear). There is nothing to back up between runs.

6. Performance

6.1 Where the time goes

A decode step is:

1. Transformer.forward(tokens[:, prev_pos:cur_pos], prev_pos) — typically cur_pos - prev_pos == 1 after prefill.
2. Per layer: one MLA.forward + one MLP.forward or MoE.forward.
3. MLA does Q/KV projections, applies rotary on q_pe / k_pe, computes attention scores via einsum, softmax, AV via einsum, and wo projection.
4. MoE.forward runs the gate (a single linear), bincount over selected indices, and n_activated_experts per-token feed-forward computations (loop over experts in this rank's slice).

6.2 Memory by model

Approximate weight footprint at the shipping dtypes:

Note

KV-cache figures are for the attn_impl = "absorb" path: kv_cache is (batch, max_seq_len, kv_lora_rank=512) plus pe_cache (batch, max_seq_len, qk_rope_head_dim=64). Switching to attn_impl = "naive" multiplies cache size by n_local_heads * (qk_head_dim + v_head_dim) / (kv_lora_rank + qk_rope_head_dim), often ~30-40× larger.

6.3 The FP8 path

When args.dtype == "fp8":

1. Transformer.__init__ sets Linear.dtype = torch.float8_e4m3fn (one byte per weight element). This applies to every Linear/ColumnParallelLinear/RowParallelLinear constructed afterward.
2. Each FP8 weight gets a companion scale parameter of shape (ceil(out / 128), ceil(in / 128)), dtype float32 — block-wise inverse scales.
3. linear() (line 131) inspects weight.element_size():
   • > 1 byte → standard F.linear.
   • == 1 and gemm_impl == "bf16" → dequantize via Triton weight_dequant and run F.linear in BF16.
   • == 1 and gemm_impl == "fp8" → quantize activations via act_quant and run fp8_gemm (autotuned Triton kernel).

Tip

The default gemm_impl = "bf16" is the safe path: it loads FP8 weights to save disk/memory but runs the actual matmul in BF16. To get the FP8 throughput win, set gemm_impl = "fp8" in model.py (it is a module-level constant).

6.4 MoE expert routing

The MoE.forward loop is sparse over experts:

counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist()
for i in range(self.experts_start_idx, self.experts_end_idx):
    if counts[i] == 0:
        continue
    expert = self.experts[i]
    idx, top = torch.where(indices == i)
    y[idx] += expert(x[idx]) * weights[idx, top, None]

Per rank, only experts with experts_start_idx <= i < experts_end_idx are constructed (others are None). Each rank does work proportional to the number of tokens whose top-k included an expert in its slice. The shared MLP runs on every rank.

Note

Per-rank load is uneven across ranks, especially with small batches: a rank may have zero tokens routed to its experts on a given step. The dist.all_reduce(y) at the end of MoE.forward waits for the slowest rank.

6.5 Triton autotune

fp8_gemm_kernel is decorated with:

@triton.autotune(configs=fp8_gemm_configs, key=['N', 'K'])

fp8_gemm_configs enumerates 36 (BLOCK_SIZE_M, BLOCK_SIZE_N, num_stages) combinations (3 × 3 × 4) at fixed BLOCK_SIZE_K=128, num_warps=8. Triton runs the search the first time a new (N, K) shape appears. Cached results live under ~/.triton/cache.

Tip

First-token latency includes autotuning. The 2-token warm-up generation in generate.main() triggers most paths before the user types — but odd context lengths or batch sizes can hit fresh (N, K) combinations and cause one-time pauses.

6.6 Performance tuning levers

7. Security

7.1 Licensing

The repository ships two LICENSE files (in the file tree, contents not surfaced in the briefing):

• LICENSE-CODE — the license for this Python code.
• LICENSE-MODEL — the license for the model weights (separate document).

Caution

Verify both licenses before commercial deployment. The model license is distinct from the code license, and DeepSeek's model licenses include use-case restrictions that do not apply to the code. Do not assume code-license terms govern the weights.

7.2 Weights provenance and integrity

load_model and safetensors.load_file perform structural validation but do not verify cryptographic signatures. Treat the checkpoint as code: it controls the model's behavior end-to-end.

Warning

safetensors files are safe to load (no pickle), but a malicious checkpoint with poisoned weights can produce arbitrary text output. If you take a checkpoint from an untrusted source, you are running an untrusted model. Hash-verify against the publisher's manifest.

7.3 Multi-tenant cautions

This repo is not designed for multi-tenant serving:

• messages is a single in-process list; there is no session isolation.
• The interactive REPL has no auth, no rate limiting, no input validation beyond what the tokenizer provides.
• KV cache is sized for max_batch_size at construction; mixing tenants in one process leaks attention context across them via the cache.

Important

If you need to expose this to multiple users, wrap each invocation in a separate Python process (or use a serving framework like vLLM that handles isolation properly). Do not multiplex tenants inside one generate.py.

7.4 What NOT to log

The repo itself logs nothing user-content-related by default. If you wrap it for production, do not log:

• Raw user prompts (PII, secrets, sensitive context).
• Generated completions verbatim (may leak training-data quotes or sensitive synthesis).
• The full messages list (combines prompts and replies).
• Tokenizer outputs (token IDs are reversible to text).

What's safe to log: latency, token counts, error types, rank/world identifiers, config name (config_*.json), --temperature, --max-new-tokens.

7.5 Network exposure

The repo opens no sockets directly. NCCL opens TCP connections among ranks (default master port 29500 or whatever you pass via --master-port). These should remain on a private interconnect.

Caution

Never expose NCCL master ports on a public network. NCCL has no authentication; an attacker reachable on the master port can disrupt or hijack communication. Bind NCCL to a private subnet or use Unix-domain sockets via PyTorch's launcher options.

7.6 Input validation

generate.py validates only:

• world_size divides vocab_size (assertion in ParallelEmbedding).
• World-size divisibility for parallel Linears (assertions).
• Prompt length ≤ max_seq_len (assertion in generate()).
• Prompt count ≤ max_batch_size in batch mode (assertion).
• At least one of --input-file/--interactive is set (assertion).
• --n-experts % --model-parallel == 0 (assertion in convert.py).

There is no content filter, prompt-injection guard, jailbreak detector, or output safety filter. All those are application concerns layered on top.

7.7 Dependency surface

inference/requirements.txt pins four direct dependencies. Their transitive surface is:

• torch==2.4.1 — PyTorch core.
• triton==3.0.0 — kernel compiler.
• transformers==4.46.3 — only AutoTokenizer is used.
• safetensors==0.4.5 — checkpoint I/O.
• tqdm (transitively, used in convert.py and fp8_cast_bf16.py).

Warning

Dependencies are pinned. Audit them against your security policy; in particular, transformers is a large package and you only use AutoTokenizer. Locking to 4.46.3 is a deliberate compatibility choice — newer versions can change tokenizer template behavior.

7.8 Threat model summary

8. Developer Internals

8.1 Repository layout

DeepSeek-V3/
├── LICENSE-CODE              # Code license (separate from model license).
├── LICENSE-MODEL             # Model-weights license.
├── README.md                 # Project README (not exposed in briefing).
├── README_WEIGHTS.md         # Weights-specific notes.
├── figures/                  # Plot images (benchmark.png, niah.png).
└── inference/                # The entire runnable codebase.
    ├── requirements.txt
    ├── convert.py            # HF -> sharded checkpoint.
    ├── fp8_cast_bf16.py      # FP8 -> BF16 cast.
    ├── generate.py           # CLI driver + decode loop.
    ├── kernel.py             # Triton kernels (act_quant, weight_dequant, fp8_gemm).
    ├── model.py              # Transformer, MLA, MoE, all parallel primitives.
    └── configs/
        ├── config_16B.json
        ├── config_236B.json
        ├── config_671B.json
        └── config_v3.1.json

8.2 Class hierarchy (model.py)

nn.Module
├── ParallelEmbedding
├── Linear
│   ├── ColumnParallelLinear
│   └── RowParallelLinear
├── RMSNorm
├── MLA
├── MLP
├── Gate
├── Expert
├── MoE
├── Block
└── Transformer

dataclass
└── ModelArgs

13 classes total. Linear is the base for both parallel variants, which only override __init__ (asserting divisibility and partitioning features) and forward (adding collectives where needed).

8.3 The MLA forward pass — annotated

MLA.forward(x, start_pos, freqs_cis, mask):

1. Compute Q. With q_lora_rank == 0, Q is a single wq @ x. Otherwise Q is wq_b @ q_norm(wq_a @ x) — a low-rank decomposition.
2. Reshape Q to (bsz, seqlen, n_local_heads, qk_head_dim) and split per-head into q_nope (no positional) and q_pe (rotary), sized qk_nope_head_dim and qk_rope_head_dim respectively.
3. Apply rotary embedding to q_pe.
4. Compute KV via wkv_a @ x, split into a kv_lora_rank-sized latent and a qk_rope_head_dim-sized k_pe. Apply rotary to k_pe.
5. Branch on attn_impl:
   • "naive" — project the latent through wkv_b to get full per-head K and V, write them to k_cache / v_cache, compute scores via einsum("bshd,bthd->bsht", q, k_cache).
   • "absorb" — fold wkv_b into Q's no-positional projection (q_nope = einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :qk_nope_head_dim])), write the latent to kv_cache and k_pe to pe_cache, compute scores as the sum of two einsums (one against kv_cache, one against pe_cache).
6. Add mask if present. Softmax in float32 (cast back to x.dtype).
7. Branch on attn_impl again for AV product:
   • "naive" — einsum("bsht,bthd->bshd", scores, v_cache).
   • "absorb" — einsum("bsht,btc->bshc", scores, kv_cache) then project to V via einsum("bshc,hdc->bshd", x, wkv_b[:, -v_head_dim:]).
8. wo(x.flatten(2)) — RowParallel projection back to dim (with all_reduce if world_size > 1).

Note

The "absorb" trick saves cache memory: instead of storing K and V at full head dimensionality, only the latent is stored. Reconstruction happens on-the-fly during the QK and AV einsums. This is the central design that distinguishes MLA from standard MHA.

8.4 The MoE forward pass — annotated

MoE.forward(x):

1. Flatten to (N, dim) where N = batch * seqlen.
2. Gate returns (weights, indices) of shapes (N, n_activated_experts).
3. Allocate y = zeros_like(x).
4. bincount(indices.flatten()) gives per-expert hit counts (a CPU-side tolist() is used to skip zero-count experts cheaply).
5. For each expert i in this rank's slice (experts_start_idx ≤ i < experts_end_idx): if its count is 0, skip. Otherwise, find the (idx, top) pairs where indices == i, run that batch of tokens through expert(x[idx]), weight by weights[idx, top], accumulate into y.
6. Run the shared MLP on the entire input (not gated): z = self.shared_experts(x).
7. dist.all_reduce(y) to sum routed-expert contributions across ranks.
8. Return (y + z).view(shape).

The Gate itself optionally implements group-limited routing:

• If n_groups > 1: scores are reshaped to (N, n_groups, n_routed_experts/n_groups), group scores are computed (max with no bias, top-2 sum with bias), top topk_groups groups are selected per token, the rest are masked to -inf, and finally top-k experts are picked across the (now-restricted) score tensor.
• This is what n_expert_groups: 8, n_limited_groups: 4 means in the 671B/V3.1 configs: tokens may only route to experts within 4 of the 8 groups.

8.5 The Triton kernels — annotated

8.5.1 act_quant_kernel

Per-program (i.e., per CUDA threadblock) work:

• Load a BLOCK_SIZE-sized chunk of activations as float32.
• Compute the absolute max, clamp to 1e-4, divide by 448.0 to get the scale.
• If scale_fmt == "ue8m0", round the scale up to the nearest power of two (2^ceil(log2(s))).
• Quantize: y = x / s, cast to float8_e4m3fn (y_ptr.dtype.element_ty).
• Store y and s.

448.0 is the maximum representable magnitude in e4m3fn.

8.5.2 weight_dequant_kernel

Per-program 2D tile:

• Load a BLOCK_SIZE × BLOCK_SIZE tile of FP8 weights as float32.
• Load the single corresponding scale from s_ptr.
• Multiply and store as the default dtype (BF16 in normal use).

The scale layout assumes one scale per (BLOCK_SIZE, BLOCK_SIZE) block of the weight matrix.

8.5.3 fp8_gemm_kernel

Standard tiled matmul, FP8 inputs, float32 accumulator, scaled per-block:

accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]

a_s is per-row, per-K-tile (one scale per BLOCK_SIZE_K chunk along K). b_s is per-block. The inner loop walks K, advancing all four pointer sets (a_ptrs, b_ptrs, a_s_ptrs, b_s_ptrs).

Note

The autotune search space is BLOCK_SIZE_M ∈ {16, 32, 64} × BLOCK_SIZE_N ∈ {32, 64, 128} × num_stages ∈ {3, 4, 5, 6} with fixed BLOCK_SIZE_K=128, num_warps=8. The autotuner re-runs only when (N, K) changes — M is treated as a runtime quantity and does not trigger a new search.

8.6 Code patterns and conventions

• Module-level globals control behavior. world_size, rank, gemm_impl, attn_impl, block_size are all module-level constants in model.py. world_size and rank are re-assigned at Transformer.__init__ time; the others are edit-source-only.
• Class attributes for FP8/BF16 selection. Linear.dtype and Linear.scale_fmt are class attributes set by Transformer.__init__ based on args.dtype / args.scale_fmt. Every Linear-derived layer constructed afterward inherits.
• Asserts as input validation. No raised exceptions, no error classes. The repo uses assert everywhere (assert vocab_size % world_size == 0, assert max(prompt_lens) <= model.max_seq_len, etc.). Strip-debugging (python -O) would silently disable these.
• Explicit per-rank shard naming. Files are model{rank}-mp{world_size}.safetensors; each rank loads its own. There is no central index that re-shards on demand.
• @torch.inference_mode() on the forward. Wraps both Transformer.forward and generate(). This disables autograd globally during decode — about a 1-2% memory saving over torch.no_grad().
• Hard-coded device="cuda". The decode loop allocates tokens = torch.full((..., total_len), -1, dtype=torch.long, device="cuda"). CPU-only execution is impossible without source edits.
• Rank-0-only printing. global print; if rank != 0: print = lambda *_, **__: None. Reassigns the builtin in the generate.main scope so all subsequent prints are silent on non-zero ranks.
• No type checker enforcement. Many functions have type hints, but there is no mypy config.
• No tests. No tests/, no pytest.ini, no test fixtures.

8.7 Module reference

8.7.1 inference/model.py

Top-level constants:

world_size = 1
rank = 0
block_size = 128
gemm_impl: Literal["bf16", "fp8"] = "bf16"
attn_impl: Literal["naive", "absorb"] = "absorb"

Public API (everything users typically import): ModelArgs, Transformer, Block, MLA, MLP, MoE, Gate, Expert, Linear, ColumnParallelLinear, RowParallelLinear, RMSNorm, ParallelEmbedding, precompute_freqs_cis, apply_rotary_emb, linear.

A __main__ block at the bottom builds a default-shaped Transformer with random tokens and prints the output shape — useful as a smoke test.

8.7.2 inference/generate.py

Public functions: sample, generate, main.

Global side effects in main(): re-bound print, torch.cuda.set_device, torch.set_default_dtype, torch.set_num_threads, torch.manual_seed. These persist for the lifetime of the process.

8.7.3 inference/convert.py

Public function: main(hf_ckpt_path, save_path, n_experts, mp).

The mapping dict at the top of the file is the source of truth for HF parameter renaming. Adding support for new parameter names means extending this dict.

8.7.4 inference/fp8_cast_bf16.py

Public function: main(fp8_path, bf16_path).

Imports weight_dequant from kernel. Caches loaded files in a dict, evicting in insertion order when more than 2 are resident.

8.7.5 inference/kernel.py

Public functions: act_quant, weight_dequant, fp8_gemm. Plus the underlying Triton kernels (act_quant_kernel, weight_dequant_kernel, fp8_gemm_kernel).

fp8_gemm_configs is a module-level autotune config list.

8.8 Testing strategy

There are no tests. To validate a change:

1. Run the __main__ block of model.py (random-weight smoke test): cd inference && python model.py. This builds a Transformer with default args, runs a forward, and prints the output shape torch.Size([2, 102400]).
2. Run generate.py with --input-file against a known-good checkpoint and compare completions byte-for-byte.
3. For kernel changes, write a small driver that compares Triton output to a PyTorch reference (torch.einsum for GEMM, manual quant for act_quant).

Tip

Before merging architecture changes, run the same prompt at --temperature 0 (greedy) on both old and new code. Any drift should be intentional and explained.

8.9 Potential improvements

The briefing's extracted/summary.tsv shows zero error classes and zero API endpoints — both expected for this style of repo. Concrete improvements that match what the code actually does:

• Configurable attn_impl / gemm_impl. Make these ModelArgs fields rather than module-level constants. Allows --attn-impl naive for debugging without source edits.
• Per-rank logging. A simple logging.getLogger() with per-rank prefix beats the current "silence all but rank 0" pattern when investigating multi-rank issues.
• Structured exceptions. Replace some asserts with named exceptions (ConfigMismatchError, PromptTooLongError) so callers can branch.
• Type-checked configs. pydantic over dataclass would catch JSON typos like "n_actived_experts" instead of silently using the default.
• Tests. Even a single end-to-end test that loads a tiny synthetic checkpoint and generates 5 tokens would catch regressions.
• Continuous batching / dynamic batching. Out of scope for a reference repo, but the obvious next step for serving.

8.10 Contributing

The repo has no CONTRIBUTING.md. Standard GitHub flow applies: fork, branch, PR. Active community issues focus on the model behavior, not the inference code (the most common labels are [BUG] filed against the running model, e.g. #1224, #1226).

9. Integration Points

9.1 Inputs

9.2 Outputs

9.3 Internal integrations

• generate.py → model.py. Imports Transformer, ModelArgs. Calls Transformer.forward, reads model.max_seq_len.
• generate.py → Hugging Face transformers. Imports AutoTokenizer. Uses from_pretrained, apply_chat_template, encode, decode, batch_decode, eos_token_id.
• generate.py → safetensors. Imports load_model (loads weights into a built model).
• model.py → kernel.py. Imports act_quant, weight_dequant, fp8_gemm. Used by the FP8 path of linear().
• fp8_cast_bf16.py → kernel.py. Imports weight_dequant.
• model.py → torch.distributed. Calls dist.is_initialized, dist.get_world_size, dist.get_rank, dist.all_reduce, dist.all_gather.
• generate.py → torch.distributed. Calls dist.init_process_group, dist.broadcast_object_list, dist.destroy_process_group.

9.4 External integrations

The repo does not directly integrate with vLLM, SGLang, TGI, or any other serving framework. Those frameworks integrate with the Hugging Face checkpoint (the source format), not with the post-convert.py sharded layout. To use this code as the inference engine for a serving framework, you would write a thin wrapper that:

1. Loads the model via Transformer(args) + load_model.
2. Exposes a generate(prompt_tokens, max_new_tokens, eos_id, temperature) function.
3. Bridges to the framework's expected request/response format (OpenAI-compatible JSON, gRPC, etc.).

There is no example wrapper in this repo.

9.5 Hooks for extension

Extension points the existing code naturally provides:

Tip

The smallest meaningful extension is a streaming generator: turn generate() into a Python generator with yield after each tokens[:, cur_pos] write. This unlocks token-by-token streaming in any wrapper, with about 5 lines of change.