Skip to content

[draft] support ATOM plugin for qwen3.5 DPxTPx/DPxEPx#1690

Draft
zovonoir wants to merge 1 commit into
mainfrom
qwen-dp
Draft

[draft] support ATOM plugin for qwen3.5 DPxTPx/DPxEPx#1690
zovonoir wants to merge 1 commit into
mainfrom
qwen-dp

Conversation

@zovonoir

Copy link
Copy Markdown
Contributor

Enable qwen3.5 series DPEP/DPTP feature, testing.

@github-actions

Copy link
Copy Markdown
Contributor

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

  • ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

  • ✅ Run after the PR is approved and Pre Checkin passes
  • ✅ Run immediately when an approval review is submitted
  • ✅ Can be requested before approval with labels
Label Tests
ci:full Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom Run native ATOM model accuracy tests
ci:vllm Run ATOM vLLM OOT model accuracy tests
ci:sglang Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 1690 --add-label <label>

Comment on lines +292 to +300
with plugin_runtime_scope(framework="sglang", atom_config=self.atom_config):
with SGLangPluginRuntime(
atom_config=self.atom_config,
forward_batch=forward_batch,
positions=positions,
input_ids=input_ids,
input_embeds=inputs_embeds,
set_forward_context=True,
) as runtime:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ [ruff] <SIM117> reported by reviewdog 🐶
Use a single with statement with multiple contexts instead of nested with statements

Comment on lines +315 to +316
with SGLangForwardBatchMetadata.bind(metadata):
with SGLangGDNForwardContext.bind(metadata):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ [ruff] <SIM117> reported by reviewdog 🐶
Use a single with statement with multiple contexts instead of nested with statements

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ [ruff] <I001> reported by reviewdog 🐶
Import block is un-sorted or un-formatted

from __future__ import annotations
from contextlib import contextmanager
from collections.abc import Iterable
from typing import Any, Optional
import torch
from torch import nn
from atom.plugin.sglang.runtime import (
SGLangForwardBatchMetadata,
SGLangPluginRuntime,
plugin_runtime_scope,
)
from aiter.dist.parallel_state import get_pp_group as _aiter_pp_group
from sglang.srt.layers.quantization.base_config import (
QuantizationConfig as SGLangQuantizationConfig,
)
from sglang.srt.model_executor.forward_batch_info import (
ForwardBatch,
PPProxyTensors,
)
from sglang.srt.models.qwen3_5 import (
Qwen3_5ForConditionalGeneration as _SglQwen35VL,
Qwen3_5MoeForConditionalGeneration as _SglQwen35MoeVL,
)
from atom.model_loader.loader import WeightsMapper
from atom.models.qwen3_5 import (
Qwen3_5ForCausalLM,
Qwen3_5ForCausalLMBase,
Qwen3_5Model,
Qwen3_5MoeForCausalLM,
detect_fused_expert_format,
get_fused_expert_mapping,
load_fused_expert_weights,
)
from atom.models.utils import IntermediateTensors
from atom.plugin.sglang.attention_backend.attention_gdn import (
SGLangGDNForwardContext,
)
from atom.plugin.sglang.runtime import (
SGLangForwardBatchMetadata,
)

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.

1 participant