Add New Models

Adding a model family means registering the model implementation, adding sampler behavior when the update rule differs, wiring the config and CLI names, and verifying a small engine run before broad benchmark or serving work.

Choose the Model Name

Pick one stable model_name string. This value is used by:

  • Config.model_name;

  • AutoModelForDiffusionLM;

  • AutoSampler;

  • benchmark CLI choices;

  • benchmark YAML configs;

  • examples and docs.

Use lowercase names with underscores, matching existing names such as fast_dllm_v2 and sdar_moe.

Register the Model

Add the implementation under diffulex/model/ and register it with AutoModelForDiffusionLM.register.

Most model factories receive config.hf_config:

from diffulex.model.auto_model import AutoModelForDiffusionLM


@AutoModelForDiffusionLM.register("my_model")
class MyModel:
    def __init__(self, hf_config):
        ...

Use use_full_config=True only when construction needs Diffulex-specific runtime settings:

@AutoModelForDiffusionLM.register("my_model", use_full_config=True)
class MyModel:
    def __init__(self, config):
        ...

Register the Sampler

If the model needs custom token update logic, add a sampler under diffulex/sampler/ and register it with AutoSampler.register.

from diffulex.sampler.auto_sampler import AutoSampler


@AutoSampler.register("my_model")
class MySampler:
    ...

If the existing sampler path is enough, avoid adding a new sampler. A model family should only add custom sampler code when the update rule differs.

Update Config Validation

Add validation in diffulex.config.Config only for required constraints. Good validation catches invalid runtime states, such as unsupported strategy/mode combinations or model-specific block/page sizes.

Do not add model-specific defaults unless the generic defaults would produce an incorrect engine state.

Update CLI Choices and Configs

If the model should be used through benchmarks, add the model name to MODEL_NAME_CHOICES in diffulex_bench/arg_parser.py.

Add a benchmark YAML under diffulex_bench/configs/ only after a small offline generation works. Keep checkpoint paths as placeholders and keep strategy, sampling mode, block size, and thresholds explicit.

Verification

A practical verification sequence is:

  1. Import the model module and sampler module.

  2. Print available model and sampler registry keys.

  3. Construct a Config for the new model_name.

  4. Run one tiny Diffulex.generate call.

  5. Run python -m diffulex_bench.main --dataset-limit 1.

  6. Add tests for the model-specific behavior you introduced.

Do not start with full serving or full benchmark runs. They make registration and config mistakes slower to diagnose.