Skip to content

vllm.models.deepseek_v41.nvidia.specialized_decode.runtime.model_state

The model state of the specialized decode.

DeepseekV41SpecializedDecodeModelState is DeepseekV41ModelState plus:

  • the decode executor, built and warmed up once the weights are loaded (the model runner creates the model state after loading the target and draft models, and before memory profiling);
  • CUDA graphs of the attention metadata builds of FULL decode steps (metadata_graphs), recorded with the inner CUDA graphs and replayed in place of the eager build;
  • no replay-start upload for decode steps whose starts are known to be zero.

Functions:

padded_decode_batch(key, input_buffers, max_model_len)

A dummy uniform decode batch padded like a runtime step of key. Like a capture, it takes worst-case sequence lengths, which bound host values only, so its recorded build serves every step of the key.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/model_state.py
def padded_decode_batch(
    key: GraphKey, input_buffers: InputBuffers, max_model_len: int
) -> InputBatch:
    """A dummy uniform decode batch padded like a runtime step of ``key``. Like
    a capture, it takes worst-case sequence lengths, which bound host values
    only, so its recorded build serves every step of the key."""
    num_reqs, num_reqs_padded, num_tokens, num_tokens_padded = key
    input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers)
    # The padded requests are empty.
    query_start_loc_np = np.full(num_reqs_padded + 1, num_tokens, dtype=np.int32)
    query_start_loc_np[: num_reqs + 1] = input_batch.query_start_loc_np
    return replace(
        input_batch,
        num_reqs_after_padding=num_reqs_padded,
        num_tokens_after_padding=num_tokens_padded,
        query_start_loc=input_buffers.query_start_loc[: num_reqs_padded + 1],
        query_start_loc_np=query_start_loc_np,
        seq_lens=input_buffers.seq_lens[:num_reqs_padded],
        seq_lens_cpu_upper_bound=torch.full(
            (num_reqs_padded,), max_model_len, dtype=torch.int32
        ),
        input_ids=input_buffers.input_ids[:num_tokens_padded],
        positions=input_buffers.positions[:num_tokens_padded],
        is_padding=input_buffers.is_padding[:num_tokens_padded],
    )