IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python module

max.pipelines.architectures.deepseekV41

DeepSeek-V4.1-Flash mixture-of-experts architecture for text generation.

V4.1 changes the decoder block, hyper-connections, output head, config, attention modes and weight layout relative to V4, so it is a separate package that reuses V4’s unchanged layers.

DeepseekV41Config​

class max.pipelines.architectures.deepseekV41.DeepseekV41Config(*, dtype, kv_params, devices, max_seq_len, quantization_encoding=None, data_parallel_degree=1, ep_config=None, forward_max_rows=None, vocab_size=129280, hidden_size=5120, num_hidden_layers=40, rms_norm_eps=1e-20, sliding_window=128, n_routed_experts=384, n_shared_experts=1, moe_intermediate_size=2304, num_experts_per_tok=6, routed_scaling_factor=1.5, swiglu_limit=10.0, num_hash_layers=0, num_dspark_layers=3, dspark_noise_token_id=None, dspark_target_layer_ids=<factory>, dspark_block_size=5, dspark_markov_rank=256, dspark_n_routed_experts=128, dspark_num_experts_per_tok=3, expert_dtype=None, moe_quant_algo=None, compress_ratios=<factory>, kv_source_layer_ids=<factory>, index_source_layer_ids=<factory>, candidate_source_layer_id=-1, candidate_topk_blocks=0, candidate_block_size=0, engram_layer_ids=<factory>, engram_num_embeddings=<factory>, engram_max_ngram_size=0, engram_n_heads=0, engram_head_dim=0, engram_compressed_vocab_size=0, engram_pad_token_id=0, engram_vocab_size=0, index_topk=None, rope_theta=10000.0, rope_scaling=None, compress_rope_theta=160000.0, num_attention_heads=64, head_dim=512, qk_rope_head_dim=64, q_lora_rank=1024, o_groups=8, o_lora_rank=1024, index_n_heads=64, index_head_dim=128, hc_mult=4, hc_sinkhorn_iters=20, hc_eps=1e-06, return_logits=ReturnLogits.LAST_TOKEN, return_hidden_states=ReturnHiddenStates.NONE)

source

Bases: ArchConfigWithKVCache

Configuration for DeepSeek-V4.1-Flash models.

Core dimensions come from text_config in the checkpoint’s config.json.

Parameters:

  • dtype (DType)
  • kv_params (KVCacheParamInterface)
  • devices (list[DeviceRef])
  • max_seq_len (int)
  • quantization_encoding (SupportedEncoding | None)
  • data_parallel_degree (int)
  • ep_config (EPConfig | None)
  • forward_max_rows (int | None)
  • vocab_size (int)
  • hidden_size (int)
  • num_hidden_layers (int)
  • rms_norm_eps (float)
  • sliding_window (int)
  • n_routed_experts (int)
  • n_shared_experts (int)
  • moe_intermediate_size (int)
  • num_experts_per_tok (int)
  • routed_scaling_factor (float)
  • swiglu_limit (float)
  • num_hash_layers (int)
  • num_dspark_layers (int)
  • dspark_noise_token_id (int | None)
  • dspark_target_layer_ids (list[int])
  • dspark_block_size (int)
  • dspark_markov_rank (int)
  • dspark_n_routed_experts (int)
  • dspark_num_experts_per_tok (int)
  • expert_dtype (str | None)
  • moe_quant_algo (str | None)
  • compress_ratios (list[int])
  • kv_source_layer_ids (list[int])
  • index_source_layer_ids (list[int])
  • candidate_source_layer_id (int)
  • candidate_topk_blocks (int)
  • candidate_block_size (int)
  • engram_layer_ids (list[int])
  • engram_num_embeddings (list[int])
  • engram_max_ngram_size (int)
  • engram_n_heads (int)
  • engram_head_dim (int)
  • engram_compressed_vocab_size (int)
  • engram_pad_token_id (int)
  • engram_vocab_size (int)
  • index_topk (int | None)
  • rope_theta (float)
  • rope_scaling (dict[str, Any] | None)
  • compress_rope_theta (float)
  • num_attention_heads (int)
  • head_dim (int)
  • qk_rope_head_dim (int)
  • q_lora_rank (int)
  • o_groups (int)
  • o_lora_rank (int)
  • index_n_heads (int)
  • index_head_dim (int)
  • hc_mult (int)
  • hc_sinkhorn_iters (int)
  • hc_eps (float)
  • return_logits (ReturnLogits)
  • return_hidden_states (ReturnHiddenStates)

DEFAULT_ENCODING​

DEFAULT_ENCODING: ClassVar[SupportedEncoding] = 'float8_e4m3fn'

source

SUPPORTED_ENCODINGS​

SUPPORTED_ENCODINGS: ClassVar[set[SupportedEncoding]] = {'float32', 'float4_e2m1fnx2', 'float8_e4m3fn'}

source

calculate_max_seq_len()​

classmethod calculate_max_seq_len(huggingface_config, model_config)

source

Returns the resolved maximum sequence length.

Bounds or defaults the user’s model_config.max_length with the model’s own limits. Construction runs this once and stores the result on model_config.max_length; memory planning may lower it further, but only on the memory plan.

Parameters:

  • huggingface_config (AutoConfig) – The HuggingFace config to read model bounds from.
  • model_config (MAXModelConfig) – The model config whose max_length carries the user’s setting.

Return type:

int

candidate_block_size​

candidate_block_size: int = 0

source

candidate_source_layer_id​

candidate_source_layer_id: int = -1

source

candidate_topk_blocks​

candidate_topk_blocks: int = 0

source

compress_ratios​

compress_ratios: list[int]

source

compress_rope_theta​

compress_rope_theta: float = 160000.0

source

construct_kv_params()​

classmethod construct_kv_params(huggingface_config, pipeline_config, devices, kv_cache_config, cache_dtype)

source

Builds KV cache params from the nested text_config.

Callers pass the raw top-level huggingface_config.

Parameters:

Return type:

KVCacheParamInterface

data_parallel_degree​

data_parallel_degree: int = 1

source

devices​

devices: list[DeviceRef]

source

dspark_block_size​

dspark_block_size: int = 5

source

dspark_markov_rank​

dspark_markov_rank: int = 256

source

dspark_n_routed_experts​

dspark_n_routed_experts: int = 128

source

dspark_noise_token_id​

dspark_noise_token_id: int | None = None

source

dspark_num_experts_per_tok​

dspark_num_experts_per_tok: int = 3

source

dspark_target_layer_ids​

dspark_target_layer_ids: list[int]

source

dtype​

dtype: DType

source

engram_compressed_vocab_size​

engram_compressed_vocab_size: int = 0

source

engram_head_dim​

engram_head_dim: int = 0

source

engram_layer_ids​

engram_layer_ids: list[int]

source

engram_max_ngram_size​

engram_max_ngram_size: int = 0

source

engram_n_heads​

engram_n_heads: int = 0

source

engram_num_embeddings​

engram_num_embeddings: list[int]

source

engram_pad_token_id​

engram_pad_token_id: int = 0

source

engram_vocab_size​

engram_vocab_size: int = 0

source

ep_config​

ep_config: EPConfig | None = None

source

Expert-parallel dispatch of the routed experts; only under data parallelism (data_parallel_ep_config()).

estimate_vision_cache_entry_bytes()​

classmethod estimate_vision_cache_entry_bytes(huggingface_config)

source

One largest image span’s bf16 embeddings; 0 without a tower.

Parameters:

huggingface_config (AutoConfig)

Return type:

int

expert_dtype​

expert_dtype: str | None = None

source

forward_max_rows​

forward_max_rows: int | None = None

source

Most query rows one forward step carries (None: max_seq_len).

Bounds the indexer’s row-tile count.

get_kv_params()​

get_kv_params()

source

KV cache parameters to use when running the model.

Return type:

KVCacheParamInterface

get_max_seq_len()​

get_max_seq_len()

source

Returns the effective maximum sequence length for the model.

For configs that store a deployment length, this is the value initialize received; for metadata-only configs it derives from the checkpoint.

Return type:

int

get_text_config()​

static get_text_config(huggingface_config)

source

DeepSeek-V4.1-Flash nests every text field under text_config.

Parameters:

huggingface_config (AutoConfig)

Return type:

AutoConfig

get_vision_cache_row_spec()​

classmethod get_vision_cache_row_spec(huggingface_config)

source

One bf16 text-width row per image span position.

Parameters:

huggingface_config (AutoConfig)

Return type:

tuple[int, DType] | None

hc_eps​

hc_eps: float = 1e-06

source

hc_mult​

hc_mult: int = 4

source

hc_sinkhorn_iters​

hc_sinkhorn_iters: int = 20

source

head_dim​

head_dim: int = 512

source

hidden_size​

hidden_size: int = 5120

source

index_head_dim​

index_head_dim: int = 128

source

index_n_heads​

index_n_heads: int = 64

source

index_source_layer_ids​

index_source_layer_ids: list[int]

source

index_topk​

index_topk: int | None = None

source

initialize()​

classmethod initialize(pipeline_config, model_config=None, *, max_seq_len)

source

Initialize the config from a PipelineConfig.

Parameters:

  • pipeline_config (PipelineConfig) – The pipeline configuration.
  • model_config (MAXModelConfig | None) – The model configuration to read from. When None (the default), pipeline_config.model is used. Pass an explicit config (such as pipeline_config.draft_model) to initialize the arch config for a different model.
  • max_seq_len (int) – The effective maximum sequence length to store on the config. The value is received, never derived here: the pipeline model passes the memory plan’s VRAM-clamped length, while memory planning (which runs before a plan exists) passes the construction-resolved model_config.max_length. Configs whose sequence length is pure model metadata (such as diffusion components) ignore it.

Return type:

Self

kv_params​

kv_params: KVCacheParamInterface

source

kv_source_layer_ids​

kv_source_layer_ids: list[int]

source

max_seq_len​

max_seq_len: int

source

moe_intermediate_size​

moe_intermediate_size: int = 2304

source

moe_quant_algo​

moe_quant_algo: str | None = None

source

n_routed_experts​

n_routed_experts: int = 384

source

n_shared_experts​

n_shared_experts: int = 1

source

num_attention_heads​

num_attention_heads: int = 64

source

num_dspark_layers​

num_dspark_layers: int = 3

source

num_experts_per_tok​

num_experts_per_tok: int = 6

source

num_hash_layers​

num_hash_layers: int = 0

source

num_hidden_layers​

num_hidden_layers: int = 40

source

o_groups​

o_groups: int = 8

source

o_lora_rank​

o_lora_rank: int = 1024

source

q_lora_rank​

q_lora_rank: int = 1024

source

qk_rope_head_dim​

qk_rope_head_dim: int = 64

source

quantization_encoding​

quantization_encoding: SupportedEncoding | None = None

source

resolve_quantization_fields()​

classmethod resolve_quantization_fields(huggingface_config)

source

Returns the config kwargs read from quantization_config.

quantization_config sits at the top level, beside text_config. Both fields are None for an unquantized checkpoint.

Parameters:

huggingface_config (AutoConfig)

Return type:

dict[str, Any]

resolve_text_fields()​

classmethod resolve_text_fields(text_config)

source

Returns the config kwargs read from text_config.

initialize() merges these with the fields it resolves itself.

Parameters:

text_config (AutoConfig)

Return type:

dict[str, Any]

return_hidden_states​

return_hidden_states: ReturnHiddenStates = 'none'

source

return_logits​

return_logits: ReturnLogits = 'last_token'

source

rms_norm_eps​

rms_norm_eps: float = 1e-20

source

rope_scaling​

rope_scaling: dict[str, Any] | None = None

source

rope_theta​

rope_theta: float = 10000.0

source

routed_scaling_factor​

routed_scaling_factor: float = 1.5

source

sliding_window​

sliding_window: int = 128

source

swiglu_limit​

swiglu_limit: float = 10.0

source

validate_num_speculative_tokens()​

validate_num_speculative_tokens(num_speculative_tokens)

source

Checks a speculative config against the DSpark drafter’s block.

The drafter proposes one token per block position, and the block is non-causal at its trained width, so the two must be equal.

Raises:

ValueError – If num_speculative_tokens != dspark_block_size.

Parameters:

num_speculative_tokens (int)

Return type:

None

vocab_size​

vocab_size: int = 129280

source

DeepseekV41Model​

class max.pipelines.architectures.deepseekV41.DeepseekV41Model(pipeline_config, session, devices, kv_cache_config, weights, *, memory_plan, adapter=None, return_logits=ReturnLogits.LAST_TOKEN, max_batch_size=1)

source

Bases: LogProbabilitiesMixin, PipelineModelWithKVCache[TextContext, DeepseekV41Config, Model]

A DeepSeek-V4.1-Flash pipeline model for text and image generation.

LogProbabilitiesMixin computes logprobs from the graph’s float32 logits. As a SupportsVisionEncoding model, the pipeline’s vision encoder cache drives vision_execute() and hands the language graph each step’s span rows and their token rows.

Parameters:

batch_processor_cls​

batch_processor_cls

source

alias of DeepseekV41BatchProcessor

empty_vision_embeddings()​

empty_vision_embeddings(devices)

source

Per-device zero-row span embeddings, cached across steps.

Parameters:

devices (list[Device])

Return type:

list[Buffer]

ep_comm_initializer​

ep_comm_initializer: EPCommInitializer | None = None

source

execute()​

execute(model_inputs)

source

Executes the graph with the given inputs.

Parameters:

model_inputs (ModelInputs) – The model inputs to execute, containing tensors and any other required data for model execution.

Returns:

ModelOutputs containing the pipeline’s output tensors.

Return type:

ModelOutputs

This is an abstract method that must be implemented by concrete PipelineModels to define their specific execution logic.

graph_specs()​

graph_specs(model_config)

source

Returns the decoder graph, tensor or data parallel across devices.

The model is built as the graph traces, since its Engram tables are graph inputs, and the trace ends by starting the deferred weight loads (_start_weight_loads()).

Parameters:

model_config (DeepseekV41Config)

Return type:

dict[str, GraphSpec]

load_graphs()​

load_graphs(session)

source

Compiles and initializes the graphs graph_specs() returns.

The decoder’s trace starts the deferred weight loads, so the compile waits for the routed-expert upload and overlaps the Engram pinning, which finishes before the host weights can be released.

Parameters:

session (InferenceSession) – The session that compiles and initializes the graphs.

Returns:

The initialized models, keyed by role.

Return type:

dict[str, Model]

load_model()​

load_model(session)

source

Loads the decoder graph_specs() returns.

A subclass that still overrides _build_graph_for_compile builds the graph through that hook, then compiles it while the deferred weight loads run on worker threads: the compile waits for the routed-expert upload (which holds the GIL) and overlaps the Engram pinning.

Parameters:

session (InferenceSession)

Return type:

Model

model​

model: Model

source

The compiled and initialized MAX Engine model.

model_config_cls​

model_config_cls

source

alias of DeepseekV41Config

pack_vision_inputs()​

pack_vision_inputs(selection, devices)

source

Packs the pipeline-selected uncached images to device 0.

Parameters:

Return type:

list[Buffer] | None

prepare_initial_token_inputs()​

prepare_initial_token_inputs(replica_batches, kv_cache_inputs=None, return_n_logits=1)

source

Returns the batch processor’s inputs plus the pinned Engram tables.

Keeps inputs.buffers the full input list for device-graph capture and replay.

Parameters:

Return type:

ModelInputs

vision_execute()​

vision_execute(selection, devices, packed)

source

Encodes the packed images into their span rows on every device.

Parameters:

Return type:

VisionEncodeResult

vision_model​

vision_model: CompiledCallable[..., Any] | None

source

The compiled vision model; None for a checkpoint without one.

Models

Cookbook

Puzzles

Blog

Community