Skip to content

[Feat][Test] script: support check model support without loading weight - #1276

Open
flesher0813 wants to merge 4 commits into
ModelEngine-Group:developfrom
flesher0813:develop_sglang
Open

[Feat][Test] script: support check model support without loading weight#1276
flesher0813 wants to merge 4 commits into
ModelEngine-Group:developfrom
flesher0813:develop_sglang

Conversation

@flesher0813

@flesher0813 flesher0813 commented Aug 25, 2026

Copy link
Copy Markdown
Contributor

Purpose

Add no-weight model compatibility checkers for environments with limited accelerator memory.

The checkers validate whether a model's native vLLM KV-cache layout is compatible with UCM without loading weights. Unlike shape-only mock tests, they continue to exercise the production vLLM model structure, ModelRunner, Scheduler, KVCacheManager, and UCM dump/load paths.

This provides a faster way to identify model-specific KV-cache compatibility issues while avoiding the accelerator-memory cost of loading full model weights.

Modifications

Common compatibility-checking flow

Add examples/model_check/ucm_model_check_common.py to provide the platform-independent validation workflow:

  • Build the real VllmConfig and UCM connector configuration.
  • Construct the model structure on the meta device without invoking loading weights.
  • Obtain the model's native KVCacheSpec through the production ModelRunner.
  • Use vLLM's native grouping and configuration APIs to generate:
    • KVCacheGroupSpec
    • KVCacheConfig
    • layer-to-group mappings
    • tensor shapes, strides, and layouts
  • Allocate a bounded KV-cache pool through num_gpu_blocks_override.
  • Create a real vLLM Scheduler and submit synthetic source and target requests.
  • Use the Scheduler-owned KVCacheManager to allocate group-aware block tables.
  • Generate UCM dump/load metadata through the normal SchedulerOutput path.
  • Fill source KV blocks with deterministic values, dump them through UCM, load them into distinct target blocks, and compare the resulting tensors byte-for-byte.
  • Support both regular block metadata and FAWA segment metadata.
  • Provide version-tolerant handling for supported vLLM API families and report unsupported environments explicitly instead of falling back to handcrafted cache layouts.
  • Clean up distributed state and accelerator memory after the check.

CUDA checker

Add examples/model_check/ucm_model_check_cuda.py:

  • Use the production CUDA GPUModelRunner.
  • Support both the classic V1 runner and the V2 model-runner path.
  • Initialize distributed state with NCCL.
  • Construct model modules on meta while redirecting unintended device allocations.
  • Allocate and verify the real CUDA KV-cache tensors through UCM.

Ascend checker

Add examples/model_check/ucm_model_check_ascend.py:

  • Apply the required vLLM-Ascend platform initialization and patches.
  • Use the production NPUModelRunner.
  • Initialize distributed state with HCCL.
  • Construct the model on meta without loading checkpoint weights.
  • Handle Ascend UniformTypeKVCacheSpecs through the native unwrap and manager factory path.
  • Allocate and verify the real NPU KV-cache tensors through UCM.

User configuration

Each platform launcher exposes a small configuration section for:

  • model path
  • request token count
  • block size
  • layer-wise transfer mode
  • UCM store pipeline
  • storage backends
  • visible accelerator device

The compatibility checker intentionally rejects command-line arguments so that each test configuration remains explicit and reproducible.

Test

Edit the User configuration and UCM configuration sections in the corresponding launcher before running.
Tested on vllm 0.26.0 and vllm-ascend 0.23.0rc1/0.26.0 using models: Qwen2.5-14B、deepseek-v2-ite、GLM5.2、deepseek-v4-flash、Qwen-Next、Minimax-M2.5 etc.

Located the KV cache layout changed problem: #1259.

@qyh111

qyh111 commented Aug 31, 2026

Copy link
Copy Markdown
Contributor

I don’t think examples/ is the right location for this code. This is a model compatibility diagnostic tool with explicit PASS/FAIL semantics, substantial vLLM-version handling, device/distributed initialization, and real KV-cache allocation—not a minimal educational example.
Could we move it to toolkit/ucm_toolkit/tools/model_check/ and expose it through the toolkit CLI? For example:

toolkit/ucm_toolkit/tools/model_check/
├── common.py
├── cuda.py
├── ascend.py
├── cli.py
└── README.md

This would also provide a clearer place for unit tests and future extensions. I would keep it separate from the default environment precheck because it allocates real accelerator KV-cache tensors and is therefore relatively heavyweight.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants