diff --git a/pyproject.toml b/pyproject.toml
index a9d1197df..57b67da75 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -61,7 +61,7 @@ megatron = [
"nvidia-ml-py==13.580.82",
"nvidia-modelopt>=0.42.0a0 ; sys_platform != 'darwin'",
"nvidia-resiliency-ext<0.5 ; sys_platform == 'linux'",
- "transformers==5.6.2",
+ "transformers==5.12.1",
"ml-dtypes>=0.5.0 ; python_full_version < '3.13'",
]
@@ -194,10 +194,8 @@ name = "deep-ep"
version = "1.2.1+9af0e0d"
requires-dist = []
-# The Megatron Bridge source metadata currently requires Transformers 5.8.x,
-# but this branch is validated against Transformers 5.6.2 for Gemma 4.
# Keep Bridge's runtime deps explicit here and let ART's megatron extra own the
-# Transformers pin.
+# Transformers pin validated by model-support handlers in this branch.
[[tool.uv.dependency-metadata]]
name = "megatron-bridge"
version = "0.5.0+e1a207ac"
diff --git a/src/art/dev/__init__.py b/src/art/dev/__init__.py
index dcc62d3f6..87c0e7aad 100644
--- a/src/art/dev/__init__.py
+++ b/src/art/dev/__init__.py
@@ -9,10 +9,15 @@
TinkerNativeArgs,
TinkerTrainingClientArgs,
TrainerArgs,
+ VllmRuntimeArgs,
)
from .openai_server import OpenAIServerConfig, ServerArgs, get_openai_server_config
from .train import TrainConfig, TrainSFTConfig
-from .validate import is_dedicated_mode, validate_dedicated_config
+from .validate import (
+ is_dedicated_mode,
+ is_external_vllm_mode,
+ validate_dedicated_config,
+)
__all__ = [
"EngineArgs",
@@ -25,8 +30,10 @@
"TinkerNativeArgs",
"TinkerTrainingClientArgs",
"TrainerArgs",
+ "VllmRuntimeArgs",
"get_openai_server_config",
"is_dedicated_mode",
+ "is_external_vllm_mode",
"OpenAIServerConfig",
"ServerArgs",
"TrainSFTConfig",
diff --git a/src/art/dev/get_model_config.py b/src/art/dev/get_model_config.py
index 57eda12de..f478a3acc 100644
--- a/src/art/dev/get_model_config.py
+++ b/src/art/dev/get_model_config.py
@@ -1,4 +1,7 @@
-from ..megatron.model_support import default_target_modules_for_model
+from ..megatron.model_support import (
+ default_target_modules_for_model,
+ vllm_lora_config_for_model,
+)
from .engine import EngineArgs
from .model import (
PEFT_ARGS_MIGRATION_MESSAGE,
@@ -67,7 +70,11 @@ def get_model_config(
if rollout_weights_mode == "lora" and "lora_target_modules" not in config.get(
"engine_args", {}
):
- engine_args["lora_target_modules"] = merged_lora_config["target_modules"]
+ engine_args["lora_target_modules"] = vllm_lora_config_for_model(
+ base_model,
+ dict(merged_lora_config),
+ allow_unvalidated_arch=True,
+ )["target_modules"]
trainer_args = TrainerArgs(
adam_beta1=0.9,
adam_beta2=0.99,
@@ -101,4 +108,6 @@ def get_model_config(
result["trainer_gpu_ids"] = config["trainer_gpu_ids"]
if "inference_gpu_ids" in config:
result["inference_gpu_ids"] = config["inference_gpu_ids"]
+ if "vllm_runtime" in config:
+ result["vllm_runtime"] = config["vllm_runtime"]
return result
diff --git a/src/art/dev/model.py b/src/art/dev/model.py
index a042c2d47..881a2b3d7 100644
--- a/src/art/dev/model.py
+++ b/src/art/dev/model.py
@@ -6,6 +6,16 @@
from .engine import EngineArgs
RolloutWeightsMode = Literal["lora", "merged"]
+VllmRuntimeMode = Literal["managed", "external"]
+
+
+class VllmRuntimeArgs(TypedDict, total=False):
+ mode: Required[VllmRuntimeMode]
+ server_url: str
+ api_key: str | None
+ local_checkpoint_root: str | None
+ server_checkpoint_root: str | None
+ health_timeout_s: float
# Vendored from transformers.training_args.OptimizerNames
@@ -135,6 +145,8 @@ class InternalModelConfig(TypedDict, total=False):
chat_template_content_format: vLLM chat template content format.
chat_template_tool_schema_format: Tool schema rendering format used for
local training tokenization.
+ vllm_runtime: vLLM runtime location. Omit for ART-managed local runtime;
+ set mode="external" to attach to a pre-launched vLLM server.
allow_unvalidated_arch: Permit model-support validation workflows to run
architectures that are not yet in the supported-model registry.
"""
@@ -152,6 +164,7 @@ class InternalModelConfig(TypedDict, total=False):
chat_template_path: str
chat_template_content_format: str
chat_template_tool_schema_format: Literal["default", "vllm_openai"]
+ vllm_runtime: VllmRuntimeArgs
allow_unvalidated_arch: bool
diff --git a/src/art/dev/validate.py b/src/art/dev/validate.py
index 56e91c1df..ae68116fa 100644
--- a/src/art/dev/validate.py
+++ b/src/art/dev/validate.py
@@ -1,11 +1,30 @@
"""Validation functions for model configuration."""
-from .model import InternalModelConfig, RolloutWeightsMode
+from collections.abc import Mapping
+from typing import cast
+
+from .model import InternalModelConfig, RolloutWeightsMode, VllmRuntimeMode
+
+
+def _vllm_runtime_mode(config: InternalModelConfig) -> VllmRuntimeMode:
+ runtime_config = config.get("vllm_runtime", {})
+ if not isinstance(runtime_config, Mapping):
+ raise ValueError("vllm_runtime must be a mapping")
+ mode = runtime_config.get("mode", "managed")
+ if mode in {"managed", "external"}:
+ return cast(VllmRuntimeMode, mode)
+ raise ValueError("vllm_runtime.mode must be either 'managed' or 'external'")
+
+
+def is_external_vllm_mode(config: InternalModelConfig) -> bool:
+ return _vllm_runtime_mode(config) == "external"
def is_dedicated_mode(config: InternalModelConfig) -> bool:
"""Return True if the config specifies dedicated mode (separate training and inference GPUs)."""
- return "trainer_gpu_ids" in config and "inference_gpu_ids" in config
+ return is_external_vllm_mode(config) or (
+ "trainer_gpu_ids" in config and "inference_gpu_ids" in config
+ )
def _rollout_weights_mode(config: InternalModelConfig) -> RolloutWeightsMode:
@@ -24,6 +43,25 @@ def validate_dedicated_config(config: InternalModelConfig) -> None:
has_trainer = "trainer_gpu_ids" in config
has_inference = "inference_gpu_ids" in config
rollout_weights_mode = _rollout_weights_mode(config)
+ external = is_external_vllm_mode(config)
+
+ if external:
+ runtime_config = config.get("vllm_runtime", {})
+ assert isinstance(runtime_config, Mapping)
+ if not runtime_config.get("server_url"):
+ raise ValueError("vllm_runtime.server_url is required for external mode")
+ if rollout_weights_mode != "lora":
+ raise ValueError(
+ "vllm_runtime.mode='external' requires rollout_weights_mode='lora'"
+ )
+ if has_trainer and not config["trainer_gpu_ids"]:
+ raise ValueError("trainer_gpu_ids must be non-empty")
+ if "fast_inference" in config.get("init_args", {}):
+ raise ValueError(
+ "fast_inference is no longer supported; ART always uses an external "
+ "vLLM runtime"
+ )
+ return
if has_trainer != has_inference:
raise ValueError(
diff --git a/src/art/local/backend.py b/src/art/local/backend.py
index a874cf22d..a566e7ec7 100644
--- a/src/art/local/backend.py
+++ b/src/art/local/backend.py
@@ -43,6 +43,10 @@
pull_model_from_s3,
push_model_to_s3,
)
+from art.vllm_runtime import (
+ get_external_vllm_runtime_config,
+ openai_base_url_from_vllm_server_url,
+)
from mp_actors import close_proxy, move_to_child_process
from .. import dev
@@ -111,6 +115,16 @@ def _configured_chat_template_server_arg(
return chat_template_path or chat_template
+def _model_support_default_chat_template(
+ base_model: str,
+ internal_config: dev.InternalModelConfig,
+) -> str | None:
+ handler = _model_support_handler(base_model, internal_config)
+ if handler is None:
+ return None
+ return handler.default_chat_template()
+
+
def _apply_configured_chat_template(
tokenizer: PreTrainedTokenizerBase,
internal_config: dev.InternalModelConfig,
@@ -120,11 +134,37 @@ def _apply_configured_chat_template(
tokenizer.chat_template = chat_template
+def _model_support_handler(
+ base_model: str,
+ internal_config: dev.InternalModelConfig,
+) -> Any | None:
+ from ..megatron.model_support.registry import (
+ UnsupportedModelArchitectureError,
+ get_model_support_handler,
+ )
+
+ try:
+ return get_model_support_handler(
+ base_model,
+ allow_unvalidated_arch=bool(
+ internal_config.get("allow_unvalidated_arch", False)
+ ),
+ )
+ except UnsupportedModelArchitectureError:
+ return None
+
+
def _apply_configured_chat_template_server_args(
config_dict: dict,
internal_config: dev.InternalModelConfig,
+ *,
+ base_model: str | None = None,
) -> None:
chat_template = _configured_chat_template_server_arg(internal_config)
+ if chat_template is None and base_model is not None:
+ chat_template = _model_support_default_chat_template(
+ base_model, internal_config
+ )
if chat_template is None:
return
server_args = dict(config_dict.get("server_args", {}))
@@ -300,6 +340,19 @@ def _chat_template_tool_schema_format(
self._default_chat_template_tool_schema_format,
)
+ def _configure_training_tokenizer(
+ self,
+ tokenizer: PreTrainedTokenizerBase,
+ *,
+ model: AnyTrainableModel,
+ internal_config: dev.InternalModelConfig,
+ ) -> PreTrainedTokenizerBase:
+ _apply_configured_chat_template(tokenizer, internal_config)
+ handler = _model_support_handler(model.base_model, internal_config)
+ if handler is None:
+ return tokenizer
+ return handler.configure_tokenizer(tokenizer, internal_config=internal_config)
+
def __enter__(self) -> Self:
return self
@@ -524,8 +577,11 @@ def _get_packed_tensors(
internal_config = cast(dev.InternalModelConfig, model._internal_config or {})
tokenizer_key = _tokenizer_cache_key(model.base_model, internal_config)
if tokenizer_key not in self._tokenizers:
- tokenizer = AutoTokenizer.from_pretrained(model.base_model)
- _apply_configured_chat_template(tokenizer, internal_config)
+ tokenizer = self._configure_training_tokenizer(
+ AutoTokenizer.from_pretrained(model.base_model),
+ model=model,
+ internal_config=internal_config,
+ )
self._tokenizers[tokenizer_key] = tokenizer
if model.base_model not in self._image_processors:
try:
@@ -702,8 +758,15 @@ async def _prepare_backend_for_training(
service = await self._get_service(model)
host, port = await service.start_openai_server(config=resolved_config)
- base_url = f"http://{host}:{port}/v1"
- api_key = server_args.get("api_key") or "default"
+ external_runtime = get_external_vllm_runtime_config(internal_config)
+ if external_runtime is not None:
+ base_url = openai_base_url_from_vllm_server_url(external_runtime.server_url)
+ api_key = (
+ server_args.get("api_key") or external_runtime.api_key or "default"
+ )
+ else:
+ base_url = f"http://{host}:{port}/v1"
+ api_key = server_args.get("api_key") or "default"
return base_url, api_key
@@ -1130,8 +1193,11 @@ async def _train_sft(
internal_config = cast(dev.InternalModelConfig, model._internal_config or {})
tokenizer_key = _tokenizer_cache_key(model.base_model, internal_config)
if tokenizer_key not in self._tokenizers:
- tokenizer = AutoTokenizer.from_pretrained(model.base_model)
- _apply_configured_chat_template(tokenizer, internal_config)
+ tokenizer = self._configure_training_tokenizer(
+ AutoTokenizer.from_pretrained(model.base_model),
+ model=model,
+ internal_config=internal_config,
+ )
self._tokenizers[tokenizer_key] = tokenizer
tokenizer = self._tokenizers[tokenizer_key]
diff --git a/src/art/megatron/dsv4/__init__.py b/src/art/megatron/dsv4/__init__.py
new file mode 100644
index 000000000..8b1378917
--- /dev/null
+++ b/src/art/megatron/dsv4/__init__.py
@@ -0,0 +1 @@
+
diff --git a/src/art/megatron/dsv4/bridge.py b/src/art/megatron/dsv4/bridge.py
new file mode 100644
index 000000000..d0eec5609
--- /dev/null
+++ b/src/art/megatron/dsv4/bridge.py
@@ -0,0 +1,1273 @@
+from functools import lru_cache, partial
+import re
+from typing import Any, Iterable, Mapping, cast
+
+from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
+from megatron.bridge.models.conversion.model_bridge import WeightConversionTask
+from megatron.bridge.models.conversion.param_mapping import (
+ AutoMapping,
+ GatedMLPMapping,
+ ReplicatedMapping,
+ RowParallelMapping,
+ extract_expert_number_from_param,
+)
+from megatron.bridge.models.deepseek.deepseek_v3_bridge import DeepSeekV3Bridge
+from megatron.bridge.models.mla_provider import MLAModelProvider
+from megatron.core.models.gpt.gpt_model import GPTModel
+import torch
+
+from art.megatron.dsv4.spec import get_dsv4_decoder_block_spec
+
+_DSV4_FP4_TABLE = (
+ 0.0,
+ 0.5,
+ 1.0,
+ 1.5,
+ 2.0,
+ 3.0,
+ 4.0,
+ 6.0,
+ 0.0,
+ -0.5,
+ -1.0,
+ -1.5,
+ -2.0,
+ -3.0,
+ -4.0,
+ -6.0,
+)
+
+_DSV4_FUSED_EXPORT_SPECS = (
+ (
+ ".self_attn.q_a_proj.weight",
+ ".self_attn.kv_proj.weight",
+ ".attn.fused_wqa_wkv.weight",
+ ),
+ (
+ ".self_attn.compressor.kv_proj.weight",
+ ".self_attn.compressor.gate_proj.weight",
+ ".attn.compressor.fused_wkv_wgate.weight",
+ ),
+ (
+ ".self_attn.compressor.indexer.kv_proj.weight",
+ ".self_attn.compressor.indexer.gate_proj.weight",
+ ".attn.indexer.compressor.fused_wkv_wgate.weight",
+ ),
+ (
+ ".mlp.shared_experts.gate_proj.weight",
+ ".mlp.shared_experts.up_proj.weight",
+ ".ffn.shared_experts.gate_up_proj.weight",
+ ),
+)
+
+_DSV4_RENAMED_EXPORT_SUFFIXES = (
+ ("lm_head.weight", "head.weight"),
+ ("model.hc_head.hc_fn", "model.hc_head_fn"),
+ ("model.hc_head.hc_base", "model.hc_head_base"),
+ ("model.hc_head.hc_scale", "model.hc_head_scale"),
+ (".input_layernorm.weight", ".attn_norm.weight"),
+ (".post_attention_layernorm.weight", ".ffn_norm.weight"),
+ (".attn_hc.fn", ".hc_attn_fn"),
+ (".attn_hc.base", ".hc_attn_base"),
+ (".attn_hc.scale", ".hc_attn_scale"),
+ (".ffn_hc.fn", ".hc_ffn_fn"),
+ (".ffn_hc.base", ".hc_ffn_base"),
+ (".ffn_hc.scale", ".hc_ffn_scale"),
+ (".mlp.gate.weight", ".ffn.gate.weight"),
+ (".mlp.gate.tid2eid", ".ffn.gate.tid2eid"),
+ (".mlp.gate.e_score_correction_bias", ".ffn.gate.bias"),
+ (".mlp.shared_experts.down_proj.weight", ".ffn.shared_experts.w2.weight"),
+ (".self_attn.q_a_norm.weight", ".attn.q_norm.weight"),
+ (".self_attn.q_b_proj.weight", ".attn.wq_b.weight"),
+ (".self_attn.kv_norm.weight", ".attn.kv_norm.weight"),
+ (".self_attn.o_a_proj.weight", ".attn.wo_a.weight"),
+ (".self_attn.o_b_proj.weight", ".attn.wo_b.weight"),
+ (".self_attn.sinks", ".attn.attn_sink"),
+ (".self_attn.compressor.position_bias", ".attn.compressor.ape"),
+ (".self_attn.compressor.kv_norm.weight", ".attn.compressor.norm.weight"),
+ (
+ ".self_attn.compressor.indexer.q_b_proj.weight",
+ ".attn.indexer.wq_b.weight",
+ ),
+ (
+ ".self_attn.compressor.indexer.scorer.weights_proj.weight",
+ ".attn.indexer.weights_proj.weight",
+ ),
+ (
+ ".self_attn.compressor.indexer.position_bias",
+ ".attn.indexer.compressor.ape",
+ ),
+ (
+ ".self_attn.compressor.indexer.kv_norm.weight",
+ ".attn.indexer.compressor.norm.weight",
+ ),
+)
+_DSV4_LAYER_TYPE_TO_COMPRESS_RATIO = {
+ "sliding_attention": 0,
+ "compressed_sparse_attention": 4,
+ "heavily_compressed_attention": 128,
+}
+
+
+def _dsv4_compress_ratios_from_hf_config(hf_config: Any) -> list[int] | None:
+ ratios = getattr(hf_config, "compress_ratios", None)
+ if ratios is not None:
+ return [int(ratio) for ratio in ratios]
+ layer_types = getattr(hf_config, "layer_types", None)
+ if layer_types is None:
+ return None
+ compress_rates = getattr(hf_config, "compress_rates", None) or {}
+ return [
+ int(
+ compress_rates.get(
+ layer_type,
+ _DSV4_LAYER_TYPE_TO_COMPRESS_RATIO[layer_type],
+ )
+ )
+ for layer_type in layer_types
+ ]
+
+
+def _dequant_dsv4_mxfp4(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+ if weight.dtype not in (torch.int8, torch.uint8):
+ raise ValueError(
+ f"Expected MXFP4 packed int8/uint8 weight, got {weight.dtype}."
+ )
+ if scale.dtype != torch.float8_e8m0fnu:
+ raise ValueError(f"Expected MXFP4 E8M0 scale, got {scale.dtype}.")
+ if weight.ndim != 2 or scale.ndim != 2:
+ raise ValueError(
+ f"Expected 2-D MXFP4 weight and scale, got {weight.shape=} {scale.shape=}."
+ )
+
+ out_dim, in_bytes = weight.shape
+ in_dim = in_bytes * 2
+ if in_dim % 32 != 0 or tuple(scale.shape) != (out_dim, in_dim // 32):
+ raise ValueError(
+ "Unexpected MXFP4 scale shape: "
+ f"weight={tuple(weight.shape)} scale={tuple(scale.shape)}."
+ )
+
+ if torch.cuda.is_available():
+ from art.megatron.dsv4.dequant import dequant_mxfp4_cuda
+
+ return dequant_mxfp4_cuda(weight, scale)
+
+ table = torch.tensor(_DSV4_FP4_TABLE, dtype=torch.float32, device=weight.device)
+ packed = weight.contiguous().view(torch.uint8)
+ low = (packed & 0x0F).long()
+ high = ((packed >> 4) & 0x0F).long()
+ decoded = torch.stack((table[low], table[high]), dim=-1).reshape(out_dim, in_dim)
+ expanded_scale = scale.float().repeat_interleave(32, dim=1)
+ return (decoded * expanded_scale).to(torch.bfloat16)
+
+
+def _dequant_dsv4_block_fp8(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+ if weight.dtype != torch.float8_e4m3fn:
+ raise ValueError(f"Expected block-FP8 E4M3 weight, got {weight.dtype}.")
+ if scale.dtype != torch.float8_e8m0fnu:
+ raise ValueError(f"Expected block-FP8 E8M0 scale, got {scale.dtype}.")
+ if weight.ndim != 2 or scale.ndim != 2:
+ raise ValueError(
+ f"Expected 2-D block-FP8 weight and scale, got {weight.shape=} {scale.shape=}."
+ )
+
+ block = 128
+ out_dim, in_dim = weight.shape
+ expected_scale_shape = (
+ (out_dim + block - 1) // block,
+ (in_dim + block - 1) // block,
+ )
+ if tuple(scale.shape) != expected_scale_shape:
+ raise ValueError(
+ "Unexpected block-FP8 scale shape: "
+ f"weight={tuple(weight.shape)} scale={tuple(scale.shape)} "
+ f"expected={expected_scale_shape}."
+ )
+
+ if torch.cuda.is_available():
+ from art.megatron.dsv4.dequant import dequant_block_fp8_cuda
+
+ return dequant_block_fp8_cuda(weight, scale)
+
+ expanded_scale = (
+ scale.float().repeat_interleave(block, dim=0).repeat_interleave(block, dim=1)
+ )
+ expanded_scale = expanded_scale[:out_dim, :in_dim]
+ return (weight.float() * expanded_scale).to(torch.bfloat16)
+
+
+def _dequant_dsv4_weight(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+ if weight.dtype in (torch.int8, torch.uint8):
+ return _dequant_dsv4_mxfp4(weight, scale)
+ if weight.dtype == torch.float8_e4m3fn:
+ return _dequant_dsv4_block_fp8(weight, scale)
+ return weight
+
+
+def _quant_dsv4_mxfp4(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ if weight.ndim != 2 or weight.shape[1] % 32 != 0:
+ raise ValueError(
+ f"Expected 2-D MXFP4 weight with K % 32 == 0, got {weight.shape}."
+ )
+ if weight.device.type == "cuda":
+ from art.megatron.dsv4.dequant import quant_mxfp4_cuda
+
+ return quant_mxfp4_cuda(weight)
+
+ out_dim, in_dim = weight.shape
+ blocks = weight.float().contiguous().view(out_dim, in_dim // 32, 32)
+ amax = blocks.abs().amax(dim=-1).clamp_min(1e-4)
+ scale = torch.pow(2.0, torch.ceil(torch.log2(amax / 6.0)))
+ scale_e8m0 = scale.to(torch.float8_e8m0fnu)
+
+ scaled = blocks / scale_e8m0.float().unsqueeze(-1)
+ thresholds = torch.tensor(
+ (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0),
+ device=weight.device,
+ dtype=torch.float32,
+ )
+ codes = torch.bucketize(scaled.abs(), thresholds).to(torch.uint8)
+ codes = codes + ((scaled < 0) & (codes != 0)).to(torch.uint8) * 8
+ low = codes[..., 0::2]
+ high = codes[..., 1::2] << 4
+ return (low | high).reshape(out_dim, in_dim // 2), scale_e8m0
+
+
+def _is_dsv4_routed_expert_weight(name: str) -> bool:
+ return ".ffn.experts." in name and name.endswith(
+ (".w1.weight", ".w2.weight", ".w3.weight")
+ )
+
+
+def _dsv4_canonical_expert_source_name(name: str) -> str | None:
+ replacements = {
+ ".gate_proj.weight": ".w1.weight",
+ ".up_proj.weight": ".w3.weight",
+ ".down_proj.weight": ".w2.weight",
+ }
+ if ".mlp.experts." not in name:
+ return None
+ for canonical, source in replacements.items():
+ if name.endswith(canonical):
+ return (
+ name.replace(".mlp.experts.", ".ffn.experts.").removesuffix(canonical)
+ + source
+ )
+ return None
+
+
+def _is_dsv4_hash_router_table(name: str) -> bool:
+ return name.endswith(".ffn.gate.tid2eid")
+
+
+def _dsv4_expert_id(name: str) -> int:
+ match = re.search(r"\.experts\.(\d+)\.", name)
+ if match is None:
+ raise ValueError(f"Expected DSV4 expert id in weight name: {name}.")
+ return int(match.group(1))
+
+
+def _set_dsv4_expert_id(name: str, expert_id: int) -> str:
+ return re.sub(r"\.experts\.\d+\.", f".experts.{expert_id}.", name, count=1)
+
+
+def _export_dsv4_mxfp4_weight(
+ name: str, weight: torch.Tensor
+) -> dict[str, torch.Tensor]:
+ packed, scale = _quant_dsv4_mxfp4(weight)
+ return {
+ name: packed.contiguous(),
+ f"{name.removesuffix('.weight')}.scale": scale.contiguous(),
+ }
+
+
+def _dsv4_fused_export_key(name: str) -> tuple[str, int] | None:
+ for first, second, target in _DSV4_FUSED_EXPORT_SPECS:
+ if name.endswith(first):
+ return f"{name.removesuffix(first)}{target}", 0
+ if name.endswith(second):
+ return f"{name.removesuffix(second)}{target}", 1
+ return None
+
+
+def _concat_dsv4_fused_export_parts(
+ target: str, parts: dict[int, torch.Tensor]
+) -> torch.Tensor:
+ first, second = parts[0], parts[1]
+ if first.ndim != second.ndim or first.shape[1:] != second.shape[1:]:
+ raise ValueError(
+ f"Cannot fuse DSV4 export parts for {target}: "
+ f"{tuple(first.shape)} vs {tuple(second.shape)}."
+ )
+ return torch.cat((first, second), dim=0).contiguous()
+
+
+def _dsv4_source_export_name(name: str) -> str:
+ for canonical, source in _DSV4_RENAMED_EXPORT_SUFFIXES:
+ if name.endswith(canonical):
+ return f"{name.removesuffix(canonical)}{source}"
+ return name
+
+
+def _dsv4_full_parallel_shape(task: WeightConversionTask) -> list[int]:
+ param_weight = task.param_weight
+ if param_weight is None:
+ raise RuntimeError(f"Missing DSV4 export param for {task.global_param_name}")
+ shape = list(param_weight.shape)
+ if not bool(getattr(param_weight, "tensor_model_parallel", False)):
+ tp_size = int(getattr(task.mapping, "tp_size", 1) or 1)
+ if task.global_param_name in {
+ "embedding.word_embeddings.weight",
+ "output_layer.weight",
+ } or task.global_param_name.endswith(".self_attention.attn_sink"):
+ shape[0] *= tp_size
+ elif task.global_param_name.endswith(
+ (
+ ".self_attention.wq_b.weight",
+ ".self_attention.wo_a.weight",
+ ".self_attention.indexer.linear_wq_b.weight",
+ )
+ ):
+ shape[0] *= tp_size
+ elif task.global_param_name.endswith(
+ (
+ ".self_attention.wo_b.weight",
+ ".mlp.shared_experts.linear_fc2.weight",
+ )
+ ):
+ shape[1] *= tp_size
+ return shape
+ partition_dim = int(getattr(param_weight, "partition_dim", 0) or 0)
+ shape[partition_dim] *= int(getattr(task.mapping, "tp_size", 1) or 1)
+ return shape
+
+
+def _dsv4_gated_shape(task: WeightConversionTask) -> list[int]:
+ param_weight = task.param_weight
+ if param_weight is None:
+ raise RuntimeError(f"Missing DSV4 export param for {task.global_param_name}")
+ shape = list(param_weight.shape)
+ shape[0] *= int(getattr(task.mapping, "tp_size", 1) or 1)
+ if shape[0] % 2 != 0:
+ raise ValueError(
+ f"Expected even DSV4 gated export dim for {task.global_param_name}: {shape}"
+ )
+ shape[0] //= 2
+ return shape
+
+
+def _dsv4_expert_down_shape(task: WeightConversionTask) -> list[int]:
+ param_weight = task.param_weight
+ if param_weight is None:
+ raise RuntimeError(f"Missing DSV4 export param for {task.global_param_name}")
+ shape = list(param_weight.shape)
+ if len(shape) > 1:
+ shape[1] *= int(getattr(task.mapping, "tp_size", 1) or 1)
+ return shape
+
+
+def _dsv4_expert_names(
+ *,
+ task: WeightConversionTask,
+ export_name: str,
+) -> list[str]:
+ config = getattr(task.megatron_module, "config", None)
+ num_experts = int(getattr(config, "num_moe_experts", 0) or 0)
+ ep_size = int(getattr(task.mapping, "ep_size", 1) or 1)
+ if num_experts <= 0 or num_experts % ep_size != 0:
+ raise ValueError(
+ f"Cannot infer DSV4 expert metadata for {task.global_param_name}: "
+ f"num_experts={num_experts}, ep_size={ep_size}."
+ )
+ experts_per_rank = num_experts // ep_size
+ local_expert = (
+ extract_expert_number_from_param(task.mapping.megatron_param) % experts_per_rank
+ )
+ return [
+ _set_dsv4_expert_id(export_name, local_expert + experts_per_rank * ep_rank)
+ for ep_rank in range(ep_size)
+ ]
+
+
+def _dsv4_quantized_expert_metadata(
+ name: str,
+ shape: list[int],
+) -> list[tuple[str, torch.dtype, list[int]]]:
+ if len(shape) != 2 or shape[1] % 32 != 0:
+ raise ValueError(f"Expected 2-D K%32 DSV4 expert weight for {name}: {shape}")
+ return [
+ (name, torch.uint8, [shape[0], shape[1] // 2]),
+ (
+ f"{name.removesuffix('.weight')}.scale",
+ torch.float8_e8m0fnu,
+ [shape[0], shape[1] // 32],
+ ),
+ ]
+
+
+def _dsv4_modified_metadata(
+ pending_fused: dict[str, dict[int, tuple[torch.dtype, list[int]]]],
+ name: str,
+ dtype: torch.dtype,
+ shape: list[int],
+) -> list[tuple[str, torch.dtype, list[int]]]:
+ fused_key = _dsv4_fused_export_key(name)
+ if fused_key is not None:
+ target, part_index = fused_key
+ parts = pending_fused.setdefault(target, {})
+ if part_index in parts:
+ raise ValueError(
+ f"Duplicate DSV4 fused metadata part {part_index}: {name}."
+ )
+ parts[part_index] = (dtype, shape)
+ if len(parts) < 2:
+ return []
+ pending_fused.pop(target)
+ first_dtype, first_shape = parts[0]
+ second_dtype, second_shape = parts[1]
+ if first_dtype != second_dtype or first_shape[1:] != second_shape[1:]:
+ raise ValueError(
+ f"Cannot fuse DSV4 metadata parts for {target}: "
+ f"{first_dtype}/{first_shape} vs {second_dtype}/{second_shape}."
+ )
+ return [
+ (target, first_dtype, [first_shape[0] + second_shape[0], *first_shape[1:]])
+ ]
+
+ source_expert = _dsv4_canonical_expert_source_name(name)
+ source_name = (
+ source_expert if source_expert is not None else _dsv4_source_export_name(name)
+ )
+ if _is_dsv4_routed_expert_weight(source_name):
+ return _dsv4_quantized_expert_metadata(source_name, shape)
+ if _is_dsv4_hash_router_table(source_name):
+ return [(source_name, torch.int32, shape)]
+ return [(source_name, dtype, shape)]
+
+
+def _load_dsv4_hf_tensor(
+ hf_param: str, hf_state_dict: Mapping[str, torch.Tensor]
+) -> torch.Tensor:
+ weight = hf_state_dict[hf_param]
+ if not hf_param.endswith(".weight") or weight.dtype not in (
+ torch.int8,
+ torch.uint8,
+ torch.float8_e4m3fn,
+ ):
+ return weight
+
+ scale_param = f"{hf_param.removesuffix('.weight')}.scale"
+ try:
+ scale = hf_state_dict[scale_param]
+ except KeyError as exc:
+ raise ValueError(
+ f"Quantized DSV4 weight {hf_param} requires scale {scale_param}."
+ ) from exc
+ return _dequant_dsv4_weight(weight, scale)
+
+
+def _resolve_dsv4_hf_param(hf_param: Any, captures: tuple[str, ...]) -> Any:
+ def resolve_one(pattern: str) -> str:
+ resolved = pattern
+ capture_index = 0
+ while "**" in resolved and capture_index < len(captures):
+ resolved = resolved.replace("**", captures[capture_index], 1)
+ capture_index += 1
+ while "*" in resolved and capture_index < len(captures):
+ resolved = resolved.replace("*", captures[capture_index], 1)
+ capture_index += 1
+ return resolved
+
+ if isinstance(hf_param, str):
+ return resolve_one(hf_param)
+ return {key: resolve_one(value) for key, value in hf_param.items()}
+
+
+class _Dsv4AliasStateSource:
+ def __init__(self, source: Any, aliases: Mapping[str, str]):
+ self.source = source
+ self.aliases = dict(aliases)
+ self._all_keys_cache: list[str] | None = None
+
+ def __getattr__(self, name: str) -> Any:
+ return getattr(self.source, name)
+
+ def get_all_keys(self) -> list[str]:
+ if self._all_keys_cache is not None:
+ return self._all_keys_cache
+ keys = set(self.source.get_all_keys())
+ keys.update(alias for alias, target in self.aliases.items() if target in keys)
+ self._all_keys_cache = sorted(keys)
+ return self._all_keys_cache
+
+ def load_tensors(self, keys: list[str]) -> dict[str, torch.Tensor]:
+ source_keys = [self.aliases.get(key, key) for key in keys]
+ loaded = self.source.load_tensors(source_keys)
+ return {key: loaded[self.aliases.get(key, key)] for key in keys}
+
+ def has_glob(self, pattern: str) -> bool:
+ import fnmatch
+
+ return any(fnmatch.fnmatch(key, pattern) for key in self.get_all_keys())
+
+
+def _install_dsv4_source_aliases(hf_pretrained: Any) -> None:
+ state = hf_pretrained.state
+ source = getattr(state, "source", None)
+ if source is None or isinstance(source, _Dsv4AliasStateSource):
+ return
+ keys = set(source.get_all_keys())
+ aliases = {"model.norm.weight": "norm.weight"}
+ active_aliases = {
+ alias: target
+ for alias, target in aliases.items()
+ if alias not in keys and target in keys
+ }
+ if active_aliases:
+ state.source = _Dsv4AliasStateSource(source, active_aliases)
+
+
+class _Dsv4AutoMapping(AutoMapping):
+ def __init__(
+ self,
+ megatron_param: str,
+ hf_param: str,
+ export_hf_param: str | None = None,
+ permute_dims: tuple[int, ...] | None = None,
+ ):
+ super().__init__(megatron_param, hf_param, permute_dims)
+ self.export_hf_param = export_hf_param or hf_param
+
+ def megatron_to_hf(self, megatron_weights: Any, megatron_module: Any):
+ converted = super().megatron_to_hf(megatron_weights, megatron_module)
+ if not converted or self.export_hf_param == self.hf_param:
+ return converted
+ return {self.export_hf_param: next(iter(converted.values()))}
+
+ def resolve(self, captures: tuple[str, ...]):
+ resolved_megatron_param, resolved_hf_param = self._resolve_names(captures)
+ return type(self)(
+ resolved_megatron_param,
+ cast(str, resolved_hf_param),
+ cast(str, _resolve_dsv4_hf_param(self.export_hf_param, captures)),
+ self.permute_dims,
+ )
+
+
+class _Dsv4ReplicatedMapping(ReplicatedMapping):
+ def __init__(
+ self,
+ megatron_param: str,
+ hf_param: str,
+ export_hf_param: str | None = None,
+ ):
+ super().__init__(megatron_param, hf_param)
+ self.export_hf_param = export_hf_param or hf_param
+
+ def megatron_to_hf(self, megatron_weights: Any, megatron_module: Any):
+ converted = super().megatron_to_hf(megatron_weights, megatron_module)
+ if not converted or self.export_hf_param == self.hf_param:
+ return converted
+ return {self.export_hf_param: next(iter(converted.values()))}
+
+ def resolve(self, captures: tuple[str, ...]):
+ resolved_megatron_param, resolved_hf_param = self._resolve_names(captures)
+ return type(self)(
+ resolved_megatron_param,
+ cast(str, resolved_hf_param),
+ cast(str, _resolve_dsv4_hf_param(self.export_hf_param, captures)),
+ )
+
+
+class _Dsv4RowParallelMapping(RowParallelMapping):
+ def __init__(
+ self,
+ megatron_param: str,
+ hf_param: str,
+ export_hf_param: str | None = None,
+ ):
+ super().__init__(megatron_param, hf_param)
+ self.export_hf_param = export_hf_param or hf_param
+
+ def megatron_to_hf(self, megatron_weights: Any, megatron_module: Any):
+ converted = super().megatron_to_hf(megatron_weights, megatron_module)
+ if not converted or self.export_hf_param == self.hf_param:
+ return converted
+ return {self.export_hf_param: next(iter(converted.values()))}
+
+ def resolve(self, captures: tuple[str, ...]):
+ resolved_megatron_param, resolved_hf_param = self._resolve_names(captures)
+ return type(self)(
+ resolved_megatron_param,
+ cast(str, resolved_hf_param),
+ cast(str, _resolve_dsv4_hf_param(self.export_hf_param, captures)),
+ )
+
+
+class _Dsv4GatedMLPMapping(GatedMLPMapping):
+ def __init__(
+ self,
+ megatron_param: str,
+ gate: str,
+ up: str,
+ export_gate: str | None = None,
+ export_up: str | None = None,
+ ):
+ super().__init__(megatron_param, gate, up)
+ self.export_hf_param = {
+ "gate": export_gate or gate,
+ "up": export_up or up,
+ }
+
+ def megatron_to_hf(self, megatron_weights: Any, megatron_module: Any):
+ converted = super().megatron_to_hf(megatron_weights, megatron_module)
+ if not converted or self.export_hf_param == self.hf_param:
+ return converted
+ remapped: dict[str, torch.Tensor] = {}
+ source_hf_param = cast(dict[str, str], self.hf_param)
+ for source_key, target_key in zip(
+ source_hf_param.values(), self.export_hf_param.values(), strict=True
+ ):
+ if source_key in converted:
+ remapped[target_key] = converted[source_key]
+ return remapped
+
+ def resolve(self, captures: tuple[str, ...]):
+ resolved_megatron_param, resolved_hf_param = self._resolve_names(captures)
+ resolved_hf_param = cast(dict[str, str], resolved_hf_param)
+ resolved_export = cast(
+ dict[str, str], _resolve_dsv4_hf_param(self.export_hf_param, captures)
+ )
+ return type(self)(
+ resolved_megatron_param,
+ resolved_hf_param["gate"],
+ resolved_hf_param["up"],
+ resolved_export["gate"],
+ resolved_export["up"],
+ )
+
+
+@lru_cache(maxsize=1)
+def _art_dsv4_expert_mapping_types() -> tuple[type[Any], type[Any]]:
+ class _ArtDsv4ExpertGateUpMapping(GatedMLPMapping):
+ is_grouped_export = False
+
+ def __init__(
+ self,
+ megatron_param: str,
+ gate: str,
+ up: str,
+ export_gate: str,
+ export_up: str,
+ ):
+ super().__init__(megatron_param, gate, up)
+ self.export_hf_param = {
+ "gate": export_gate,
+ "up": export_up,
+ }
+
+ @property
+ def group_key(self) -> str:
+ return cast(dict[str, str], self.export_hf_param)["gate"]
+
+ def hf_to_megatron(
+ self,
+ hf_weights: Any,
+ megatron_module: Any,
+ ) -> torch.Tensor:
+ from megatron.bridge.models.conversion.param_mapping import (
+ _align_expert_weight_to_shape,
+ )
+ from megatron.bridge.models.conversion.utils import (
+ get_module_and_param_from_name,
+ )
+
+ normalized_param = self._normalize_expert_param_name(self.megatron_param)
+ target_param = get_module_and_param_from_name(
+ megatron_module, normalized_param
+ )[1]
+ full_target_shape = (
+ target_param.shape[0] * self.tp_size,
+ target_param.shape[1],
+ )
+ if full_target_shape[0] % 2 != 0:
+ raise ValueError(
+ "Expected even fused dim for "
+ f"{self.megatron_param}, got {full_target_shape}."
+ )
+ gate_target_shape = (full_target_shape[0] // 2, full_target_shape[1])
+ gate = _align_expert_weight_to_shape(
+ cast(torch.Tensor, hf_weights["gate"]),
+ torch.Size(gate_target_shape),
+ "gate",
+ transpose_hint=False,
+ )
+ up = _align_expert_weight_to_shape(
+ cast(torch.Tensor, hf_weights["up"]),
+ torch.Size(gate_target_shape),
+ "up",
+ transpose_hint=False,
+ )
+ return super().hf_to_megatron({"gate": gate, "up": up}, megatron_module)
+
+ def megatron_to_hf(self, megatron_weights: Any, megatron_module: Any):
+ converted = super().megatron_to_hf(megatron_weights, megatron_module)
+ if not converted:
+ return {}
+ hf_param = cast(dict[str, str], self.hf_param)
+ gate_suffix = hf_param["gate"].rpartition(".experts.")[2].split(".", 1)[1]
+ remapped: dict[str, torch.Tensor] = {}
+ export = cast(dict[str, str], self.export_hf_param)
+ for gate_key, gate in converted.items():
+ if not gate_key.endswith(gate_suffix):
+ continue
+ expert_id = _dsv4_expert_id(gate_key)
+ up_key = _set_dsv4_expert_id(hf_param["up"], expert_id)
+ up = converted.get(up_key)
+ if up is None:
+ raise ValueError(
+ f"Missing DSV4 gathered expert up weight {up_key} "
+ f"for gate weight {gate_key}."
+ )
+ remapped[_set_dsv4_expert_id(export["gate"], expert_id)] = gate
+ remapped[_set_dsv4_expert_id(export["up"], expert_id)] = up
+ return remapped
+
+ def resolve(self, captures: tuple[str, ...]):
+ resolved_megatron_param, resolved_hf_param = self._resolve_names(captures)
+ resolved_hf_param = cast(dict[str, str], resolved_hf_param)
+ resolved_export = cast(
+ dict[str, str], _resolve_dsv4_hf_param(self.export_hf_param, captures)
+ )
+ return type(self)(
+ resolved_megatron_param,
+ resolved_hf_param["gate"],
+ resolved_hf_param["up"],
+ resolved_export["gate"],
+ resolved_export["up"],
+ )
+
+ class _ArtDsv4ExpertDownMapping(AutoMapping):
+ is_grouped_export = False
+
+ def __init__(
+ self,
+ megatron_param: str,
+ hf_param: str,
+ export_hf_param: str,
+ ):
+ super().__init__(megatron_param, hf_param)
+ self.weight_hf_param = hf_param
+ self.hf_param = {"weight": hf_param}
+ self.export_hf_param = export_hf_param
+
+ @property
+ def group_key(self) -> str:
+ return self.export_hf_param
+
+ def hf_to_megatron(
+ self,
+ hf_weights: Any,
+ megatron_module: Any,
+ ) -> torch.Tensor:
+ from megatron.bridge.models.conversion.param_mapping import (
+ ColumnParallelMapping,
+ RowParallelMapping,
+ _align_expert_weight_to_shape,
+ )
+ from megatron.bridge.models.conversion.utils import (
+ get_module_and_param_from_name,
+ )
+
+ normalized_param = self._normalize_expert_param_name(self.megatron_param)
+ target_param = get_module_and_param_from_name(
+ megatron_module, normalized_param
+ )[1]
+ if self._mapping is None:
+ self._detected_type = self._detect_parallelism_type(megatron_module)
+ hf_param = self.hf_param
+ self.hf_param = self.weight_hf_param
+ try:
+ self._mapping = self._get_or_create_mapping(self._detected_type)
+ finally:
+ self.hf_param = hf_param
+ if isinstance(self._mapping, ColumnParallelMapping):
+ full_target_shape = (
+ target_param.shape[0] * self.tp_size,
+ target_param.shape[1],
+ )
+ elif isinstance(self._mapping, RowParallelMapping):
+ full_target_shape = (
+ target_param.shape[0],
+ target_param.shape[1] * self.tp_size,
+ )
+ else:
+ full_target_shape = tuple(target_param.shape)
+ hf_weight = (
+ cast(dict[str, torch.Tensor], hf_weights)["weight"]
+ if isinstance(hf_weights, dict)
+ else cast(torch.Tensor, hf_weights)
+ )
+ aligned = _align_expert_weight_to_shape(
+ hf_weight,
+ torch.Size(full_target_shape),
+ "down_proj",
+ transpose_hint=False,
+ )
+ return self._mapping.hf_to_megatron(aligned, megatron_module)
+
+ def megatron_to_hf(self, megatron_weights: Any, megatron_module: Any):
+ hf_param = self.hf_param
+ self.hf_param = self.weight_hf_param
+ try:
+ converted = super().megatron_to_hf(megatron_weights, megatron_module)
+ finally:
+ self.hf_param = hf_param
+ if not converted:
+ return {}
+ return {
+ _set_dsv4_expert_id(
+ self.export_hf_param,
+ _dsv4_expert_id(name),
+ ): weight
+ for name, weight in converted.items()
+ }
+
+ def resolve(self, captures: tuple[str, ...]):
+ resolved_megatron_param, resolved_hf_param = self._resolve_names(captures)
+ resolved_hf_param = cast(dict[str, str], resolved_hf_param)
+ return type(self)(
+ resolved_megatron_param,
+ resolved_hf_param["weight"],
+ cast(str, _resolve_dsv4_hf_param(self.export_hf_param, captures)),
+ )
+
+ return _ArtDsv4ExpertGateUpMapping, _ArtDsv4ExpertDownMapping
+
+
+def _register_dsv4_module_types() -> None:
+ AutoMapping.register_module_type("DeepSeekV4Attention", "column")
+ AutoMapping.register_module_type("DeepSeekV4Compressor", "replicated")
+ AutoMapping.register_module_type("Dsv4FinalNorm", "replicated")
+ AutoMapping.register_module_type("Dsv4Router", "replicated")
+ AutoMapping.register_module_type("Dsv4TransformerLayer", "replicated")
+ AutoMapping.register_module_type("HCHeadParams", "replicated")
+
+
+def _dsv4_mapping_registry() -> MegatronMappingRegistry:
+ _register_dsv4_module_types()
+ expert_gate_up_mapping, expert_down_mapping = _art_dsv4_expert_mapping_types()
+ mappings: list[Any] = [
+ _Dsv4AutoMapping(
+ "embedding.word_embeddings.weight",
+ "embed.weight",
+ "model.embed_tokens.weight",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.input_layernorm.weight",
+ "layers.*.attn_norm.weight",
+ "model.layers.*.input_layernorm.weight",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.pre_mlp_layernorm.weight",
+ "layers.*.ffn_norm.weight",
+ "model.layers.*.post_attention_layernorm.weight",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.mlp.router.weight",
+ "layers.*.ffn.gate.weight",
+ "model.layers.*.mlp.gate.weight",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.mlp.router.tid2eid",
+ "layers.*.ffn.gate.tid2eid",
+ "model.layers.*.mlp.gate.tid2eid",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.mlp.router.e_score_correction_bias",
+ "layers.*.ffn.gate.bias",
+ "model.layers.*.mlp.gate.e_score_correction_bias",
+ ),
+ expert_down_mapping(
+ "decoder.layers.*.mlp.experts.linear_fc2.weight*",
+ "layers.*.ffn.experts.*.w2.weight",
+ "model.layers.*.mlp.experts.*.down_proj.weight",
+ ),
+ _Dsv4RowParallelMapping(
+ "decoder.layers.*.mlp.shared_experts.linear_fc2.weight",
+ "layers.*.ffn.shared_experts.w2.weight",
+ "model.layers.*.mlp.shared_experts.down_proj.weight",
+ ),
+ _Dsv4AutoMapping("decoder.final_layernorm.weight", "model.norm.weight"),
+ _Dsv4AutoMapping(
+ "decoder.final_layernorm.hc_head_params.hc_head_fn",
+ "hc_head_fn",
+ "model.hc_head.hc_fn",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.final_layernorm.hc_head_params.hc_head_base",
+ "hc_head_base",
+ "model.hc_head.hc_base",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.final_layernorm.hc_head_params.hc_head_scale",
+ "hc_head_scale",
+ "model.hc_head.hc_scale",
+ ),
+ _Dsv4AutoMapping("output_layer.weight", "head.weight", "lm_head.weight"),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.hc_attn_fn",
+ "layers.*.hc_attn_fn",
+ "model.layers.*.attn_hc.fn",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.hc_attn_base",
+ "layers.*.hc_attn_base",
+ "model.layers.*.attn_hc.base",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.hc_attn_scale",
+ "layers.*.hc_attn_scale",
+ "model.layers.*.attn_hc.scale",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.hc_ffn_fn",
+ "layers.*.hc_ffn_fn",
+ "model.layers.*.ffn_hc.fn",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.hc_ffn_base",
+ "layers.*.hc_ffn_base",
+ "model.layers.*.ffn_hc.base",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.hc_ffn_scale",
+ "layers.*.hc_ffn_scale",
+ "model.layers.*.ffn_hc.scale",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.self_attention.wq_a.weight",
+ "layers.*.attn.wq_a.weight",
+ "model.layers.*.self_attn.q_a_proj.weight",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.self_attention.q_norm.weight",
+ "layers.*.attn.q_norm.weight",
+ "model.layers.*.self_attn.q_a_norm.weight",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.self_attention.wq_b.weight",
+ "layers.*.attn.wq_b.weight",
+ "model.layers.*.self_attn.q_b_proj.weight",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.self_attention.wkv.weight",
+ "layers.*.attn.wkv.weight",
+ "model.layers.*.self_attn.kv_proj.weight",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.self_attention.kv_norm.weight",
+ "layers.*.attn.kv_norm.weight",
+ "model.layers.*.self_attn.kv_norm.weight",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.self_attention.wo_a.weight",
+ "layers.*.attn.wo_a.weight",
+ "model.layers.*.self_attn.o_a_proj.weight",
+ ),
+ _Dsv4RowParallelMapping(
+ "decoder.layers.*.self_attention.wo_b.weight",
+ "layers.*.attn.wo_b.weight",
+ "model.layers.*.self_attn.o_b_proj.weight",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.self_attention.attn_sink",
+ "layers.*.attn.attn_sink",
+ "model.layers.*.self_attn.sinks",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.self_attention.compressor.ape",
+ "layers.*.attn.compressor.ape",
+ "model.layers.*.self_attn.compressor.position_bias",
+ ),
+ _Dsv4ReplicatedMapping(
+ "decoder.layers.*.self_attention.compressor.wkv.weight",
+ "layers.*.attn.compressor.wkv.weight",
+ "model.layers.*.self_attn.compressor.kv_proj.weight",
+ ),
+ _Dsv4ReplicatedMapping(
+ "decoder.layers.*.self_attention.compressor.wgate.weight",
+ "layers.*.attn.compressor.wgate.weight",
+ "model.layers.*.self_attn.compressor.gate_proj.weight",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.self_attention.compressor.norm.weight",
+ "layers.*.attn.compressor.norm.weight",
+ "model.layers.*.self_attn.compressor.kv_norm.weight",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.self_attention.indexer.linear_wq_b.weight",
+ "layers.*.attn.indexer.wq_b.weight",
+ "model.layers.*.self_attn.compressor.indexer.q_b_proj.weight",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.self_attention.indexer.linear_weights_proj.weight",
+ "layers.*.attn.indexer.weights_proj.weight",
+ "model.layers.*.self_attn.compressor.indexer.scorer.weights_proj.weight",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.self_attention.indexer.compressor.ape",
+ "layers.*.attn.indexer.compressor.ape",
+ "model.layers.*.self_attn.compressor.indexer.position_bias",
+ ),
+ _Dsv4ReplicatedMapping(
+ "decoder.layers.*.self_attention.indexer.compressor.wkv.weight",
+ "layers.*.attn.indexer.compressor.wkv.weight",
+ "model.layers.*.self_attn.compressor.indexer.kv_proj.weight",
+ ),
+ _Dsv4ReplicatedMapping(
+ "decoder.layers.*.self_attention.indexer.compressor.wgate.weight",
+ "layers.*.attn.indexer.compressor.wgate.weight",
+ "model.layers.*.self_attn.compressor.indexer.gate_proj.weight",
+ ),
+ _Dsv4AutoMapping(
+ "decoder.layers.*.self_attention.indexer.compressor.norm.weight",
+ "layers.*.attn.indexer.compressor.norm.weight",
+ "model.layers.*.self_attn.compressor.indexer.kv_norm.weight",
+ ),
+ expert_gate_up_mapping(
+ megatron_param="decoder.layers.*.mlp.experts.linear_fc1.weight*",
+ gate="layers.*.ffn.experts.*.w1.weight",
+ up="layers.*.ffn.experts.*.w3.weight",
+ export_gate="model.layers.*.mlp.experts.*.gate_proj.weight",
+ export_up="model.layers.*.mlp.experts.*.up_proj.weight",
+ ),
+ _Dsv4GatedMLPMapping(
+ megatron_param="decoder.layers.*.mlp.shared_experts.linear_fc1.weight",
+ gate="layers.*.ffn.shared_experts.w1.weight",
+ up="layers.*.ffn.shared_experts.w3.weight",
+ export_gate="model.layers.*.mlp.shared_experts.gate_proj.weight",
+ export_up="model.layers.*.mlp.shared_experts.up_proj.weight",
+ ),
+ ]
+ return MegatronMappingRegistry(*mappings)
+
+
+class ArtDeepSeekV4Bridge(DeepSeekV3Bridge):
+ def _maybe_collect_fused_export(
+ self, name: str, weight: torch.Tensor
+ ) -> dict[str, torch.Tensor] | None:
+ key = _dsv4_fused_export_key(name)
+ if key is None:
+ return None
+ target, part_index = key
+ pending = getattr(self, "_dsv4_fused_export_parts", None)
+ if pending is None:
+ pending = {}
+ self._dsv4_fused_export_parts = pending
+ parts = pending.setdefault(target, {})
+ if part_index in parts:
+ raise ValueError(f"Duplicate DSV4 fused export part {part_index}: {name}.")
+ parts[part_index] = weight
+ if len(parts) < 2:
+ return {}
+ pending.pop(target)
+ return {target: _concat_dsv4_fused_export_parts(target, parts)}
+
+ def provider_bridge(self, hf_pretrained: Any):
+ _install_dsv4_source_aliases(hf_pretrained)
+ hf_config = hf_pretrained.config
+ if not hasattr(hf_config, "first_k_dense_replace"):
+ hf_config.first_k_dense_replace = 0
+ provider = cast(Any, super().provider_bridge(hf_pretrained))
+ provider.transformer_layer_spec = partial(get_dsv4_decoder_block_spec)
+ provider.num_layers = hf_config.num_hidden_layers
+ provider.normalization = "RMSNorm"
+ provider.gated_linear_unit = True
+ provider.add_bias_linear = False
+ provider.share_embeddings_and_output_weights = False
+ provider.multi_latent_attention = False
+ provider.q_lora_rank = hf_config.q_lora_rank
+ provider.kv_lora_rank = hf_config.head_dim
+ provider.qk_pos_emb_head_dim = hf_config.qk_rope_head_dim
+ provider.num_attention_heads = hf_config.num_attention_heads
+ provider.num_query_groups = 1
+ provider.kv_channels = hf_config.head_dim
+ provider.num_moe_experts = hf_config.n_routed_experts
+ provider.moe_router_topk = hf_config.num_experts_per_tok
+ provider.moe_router_score_function = hf_config.scoring_func
+ provider.moe_router_topk_scaling_factor = hf_config.routed_scaling_factor
+ provider.moe_router_enable_expert_bias = False
+ provider.moe_router_fusion = False
+ provider.moe_layer_freq = [1] * hf_config.num_hidden_layers
+ provider.moe_ffn_hidden_size = hf_config.moe_intermediate_size
+ provider.ffn_hidden_size = hf_config.moe_intermediate_size
+ provider.moe_shared_expert_intermediate_size = (
+ hf_config.moe_intermediate_size * hf_config.n_shared_experts
+ )
+ provider.dsv4_hc_mult = getattr(hf_config, "hc_mult", 4)
+ provider.dsv4_hc_sinkhorn_iters = getattr(hf_config, "hc_sinkhorn_iters", 20)
+ provider.dsv4_hc_eps = getattr(hf_config, "hc_eps", 1e-6)
+ provider.dsv4_compress_ratios = _dsv4_compress_ratios_from_hf_config(hf_config)
+ provider.dsv4_compress_rope_theta = getattr(
+ hf_config, "compress_rope_theta", 160000
+ )
+ rope_scaling = getattr(hf_config, "rope_scaling", None) or {}
+ provider.rotary_scaling_factor = rope_scaling.get("factor", 16)
+ provider.original_max_position_embeddings = rope_scaling.get(
+ "original_max_position_embeddings", 65536
+ )
+ provider.beta_fast = rope_scaling.get("beta_fast", 32)
+ provider.beta_slow = rope_scaling.get("beta_slow", 1)
+ provider.dsv4_swiglu_limit = getattr(hf_config, "swiglu_limit", 0.0)
+ provider.dsv4_o_groups = getattr(hf_config, "o_groups", 16)
+ provider.dsv4_o_lora_rank = getattr(hf_config, "o_lora_rank", 1024)
+ provider.dsv4_n_hash_layers = getattr(hf_config, "n_hash_layers", 3)
+ provider.dsv4_window_size = getattr(hf_config, "sliding_window", 128)
+ provider.dsa_indexer_n_heads = getattr(hf_config, "index_n_heads", 64)
+ provider.dsa_indexer_head_dim = getattr(hf_config, "index_head_dim", 128)
+ provider.dsa_indexer_topk = getattr(hf_config, "index_topk", 1024)
+ if provider.dsv4_swiglu_limit > 0:
+ provider.bias_activation_fusion = False
+ provider.activation_func_clamp_value = provider.dsv4_swiglu_limit
+ provider.mtp_num_layers = None
+ return provider
+
+ def mapping_registry(self) -> MegatronMappingRegistry:
+ return _dsv4_mapping_registry()
+
+ def art_extra_hf_prefetch_keys(
+ self,
+ keys: Iterable[str],
+ hf_state_dict: Mapping[str, torch.Tensor],
+ ) -> list[str]:
+ all_keys = set(hf_state_dict.keys())
+ scale_keys: list[str] = []
+ for key in keys:
+ if not key.endswith(".weight"):
+ continue
+ scale_key = f"{key.removesuffix('.weight')}.scale"
+ if scale_key in all_keys:
+ scale_keys.append(scale_key)
+ return scale_keys
+
+ def maybe_modify_loaded_hf_weight(
+ self,
+ hf_param: str | dict[str, str],
+ hf_state_dict: Mapping[str, torch.Tensor],
+ ) -> torch.Tensor:
+ if isinstance(hf_param, str):
+ return _load_dsv4_hf_tensor(hf_param, hf_state_dict)
+ return cast(
+ torch.Tensor,
+ {
+ name: _load_dsv4_hf_tensor(param, hf_state_dict)
+ for name, param in hf_param.items()
+ },
+ )
+
+ def maybe_modify_converted_hf_weight(
+ self,
+ task: WeightConversionTask,
+ converted_weights_dict: dict[str, torch.Tensor],
+ hf_state_dict: Any,
+ ) -> dict[str, torch.Tensor]:
+ del task
+ if isinstance(hf_state_dict, dict):
+ return converted_weights_dict
+ remapped: dict[str, torch.Tensor] = {}
+ for name, weight in converted_weights_dict.items():
+ fused = self._maybe_collect_fused_export(name, weight)
+ if fused is not None:
+ remapped.update(fused)
+ else:
+ source_expert = _dsv4_canonical_expert_source_name(name)
+ source_name = (
+ source_expert
+ if source_expert is not None
+ else _dsv4_source_export_name(name)
+ )
+ if _is_dsv4_routed_expert_weight(source_name):
+ remapped.update(_export_dsv4_mxfp4_weight(source_name, weight))
+ elif _is_dsv4_hash_router_table(source_name):
+ remapped[source_name] = weight.to(torch.int32).contiguous()
+ else:
+ remapped[source_name] = weight
+ return remapped
+
+ def iter_merged_vllm_weight_metadata(
+ self,
+ weight_export: Any,
+ ) -> Iterable[tuple[str, torch.dtype, list[int]]]:
+ pending_fused: dict[str, dict[int, tuple[torch.dtype, list[int]]]] = {}
+ for task in weight_export.conversion_tasks:
+ mapping_name = type(task.mapping).__name__
+ dtype = task.param_weight.dtype
+ if mapping_name == "_ArtDsv4ExpertGateUpMapping":
+ shape = _dsv4_gated_shape(task)
+ export = cast(dict[str, str], task.mapping.export_hf_param)
+ for gate_name, up_name in zip(
+ _dsv4_expert_names(task=task, export_name=export["gate"]),
+ _dsv4_expert_names(task=task, export_name=export["up"]),
+ strict=True,
+ ):
+ yield from _dsv4_modified_metadata(
+ pending_fused, gate_name, dtype, shape
+ )
+ yield from _dsv4_modified_metadata(
+ pending_fused, up_name, dtype, shape
+ )
+ continue
+
+ if mapping_name == "_ArtDsv4ExpertDownMapping":
+ shape = _dsv4_expert_down_shape(task)
+ export_name = cast(str, task.mapping.export_hf_param)
+ for name in _dsv4_expert_names(task=task, export_name=export_name):
+ yield from _dsv4_modified_metadata(
+ pending_fused, name, dtype, shape
+ )
+ continue
+
+ if isinstance(task.mapping.export_hf_param, dict):
+ shape = _dsv4_gated_shape(task)
+ export = cast(dict[str, str], task.mapping.export_hf_param)
+ yield from _dsv4_modified_metadata(
+ pending_fused, export["gate"], dtype, shape
+ )
+ yield from _dsv4_modified_metadata(
+ pending_fused, export["up"], dtype, shape
+ )
+ continue
+
+ yield from _dsv4_modified_metadata(
+ pending_fused,
+ cast(str, task.mapping.export_hf_param),
+ dtype,
+ _dsv4_full_parallel_shape(task),
+ )
+ if pending_fused:
+ raise ValueError(
+ f"Incomplete DSV4 fused metadata parts: {sorted(pending_fused)}"
+ )
+
+
+_DSV4_BRIDGE_REGISTERED = False
+
+
+def ensure_dsv4_bridge_registered() -> None:
+ global _DSV4_BRIDGE_REGISTERED
+ if _DSV4_BRIDGE_REGISTERED:
+ return
+ from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge
+
+ from art.megatron.dsv4.hf_config import ensure_dsv4_hf_config_registered
+
+ ensure_dsv4_hf_config_registered()
+ MegatronModelBridge.register_bridge(
+ source="DeepseekV4ForCausalLM",
+ target=GPTModel,
+ provider=MLAModelProvider,
+ model_type="deepseek_v4",
+ )(ArtDeepSeekV4Bridge)
+ _DSV4_BRIDGE_REGISTERED = True
diff --git a/src/art/megatron/dsv4/compressor.py b/src/art/megatron/dsv4/compressor.py
new file mode 100644
index 000000000..b677f2521
--- /dev/null
+++ b/src/art/megatron/dsv4/compressor.py
@@ -0,0 +1,656 @@
+from typing import Any, NamedTuple, cast
+
+import einops
+from megatron.core.transformer.transformer_config import TransformerConfig
+from pydantic import BaseModel, ConfigDict, Field
+import torch
+import torch.nn as nn
+from torch.nn import Linear
+
+from art.megatron.dsv4.kernel.precision_aligned_ops import linear_bf16_fp32
+from art.megatron.dsv4.rope import (
+ apply_rotary_emb,
+ configure_rope_cache,
+ get_rope_cache,
+ get_rope_cache_at_positions,
+)
+from art.megatron.dsv4.utils import rotate_activation
+
+
+class Dsv4CompressionLayout(NamedTuple):
+ current_indices: torch.Tensor
+ previous_indices: torch.Tensor
+ entry_group_ids: torch.Tensor
+ entry_parent_visible: torch.Tensor
+ entry_start_positions: torch.Tensor
+ entry_end_positions: torch.Tensor
+ entry_valid: torch.Tensor
+
+
+class Dsv4SharedPrefixState(BaseModel):
+ model_config = ConfigDict(arbitrary_types_allowed=True)
+
+ compression_layouts: dict[int, Dsv4CompressionLayout]
+ topk_idx_cache: dict[Any, Any] = Field(default_factory=dict)
+
+
+def move_compression_layout_to_device(
+ layout: Dsv4CompressionLayout,
+ device: torch.device,
+) -> Dsv4CompressionLayout:
+ return Dsv4CompressionLayout(
+ *(tensor.to(device=device, non_blocking=True) for tensor in layout)
+ )
+
+
+def build_shared_prefix_compression_layouts(
+ *,
+ position_ids: torch.Tensor,
+ group_ids: torch.Tensor,
+ parent_ids: torch.Tensor,
+ device: torch.device,
+) -> dict[int, Dsv4CompressionLayout]:
+ return {
+ ratio: move_compression_layout_to_device(
+ build_shared_prefix_compression_layout(
+ position_ids=position_ids,
+ group_ids=group_ids,
+ parent_ids=parent_ids,
+ ratio=ratio,
+ duplicate_prompt_entries=ratio == 4,
+ ),
+ device,
+ )
+ for ratio in (4, 128)
+ }
+
+
+def _logical_window_indices(
+ *,
+ prompt_start: int,
+ prompt_len: int,
+ completion_start: int | None,
+ logical_start: int,
+ ratio: int,
+) -> list[int]:
+ indices: list[int] = []
+ for logical_pos in range(logical_start, logical_start + ratio):
+ if logical_pos < 0:
+ indices.append(-1)
+ elif logical_pos < prompt_len:
+ indices.append(prompt_start + logical_pos)
+ elif completion_start is not None:
+ indices.append(completion_start + logical_pos - prompt_len)
+ else:
+ indices.append(-1)
+ return indices
+
+
+def build_shared_prefix_compression_layout(
+ *,
+ position_ids: torch.Tensor,
+ group_ids: torch.Tensor,
+ parent_ids: torch.Tensor,
+ ratio: int,
+ duplicate_prompt_entries: bool = False,
+) -> Dsv4CompressionLayout:
+ device = position_ids.device
+ bsz, seqlen = position_ids.shape
+ position_cpu = position_ids.detach().cpu()
+ group_cpu = group_ids.detach().cpu()
+ parent_cpu = parent_ids.detach().cpu()
+ current_rows: list[list[list[int]]] = []
+ previous_rows: list[list[list[int]]] = []
+ group_rows: list[list[int]] = []
+ parent_visible_rows: list[list[bool]] = []
+ start_rows: list[list[int]] = []
+ end_rows: list[list[int]] = []
+
+ for b in range(bsz):
+ valid_mask = (group_cpu[b] != -1) & (parent_cpu[b] != -1)
+ padding = torch.nonzero(~valid_mask, as_tuple=False)
+ valid_len = int(padding[0].item()) if padding.numel() else seqlen
+ if bool(valid_mask[valid_len:].any().item()):
+ raise ValueError("DSV4 shared-prefix metadata must pad only at row end.")
+ segments: list[tuple[int, int, int, int, int, int]] = []
+ cursor = 0
+ while cursor < valid_len:
+ group = int(group_cpu[b, cursor].item())
+ parent = int(parent_cpu[b, cursor].item())
+ start = cursor
+ cursor += 1
+ while cursor < valid_len and int(group_cpu[b, cursor].item()) == group:
+ cursor += 1
+ end = cursor
+ start_pos = int(position_cpu[b, start].item())
+ end_pos = int(position_cpu[b, end - 1].item())
+ segments.append((group, parent, start, end, start_pos, end_pos))
+
+ row_current: list[list[int]] = []
+ row_previous: list[list[int]] = []
+ row_groups: list[int] = []
+ row_parent_visible: list[bool] = []
+ row_starts: list[int] = []
+ row_ends: list[int] = []
+
+ def append_entry(
+ *,
+ entry_group: int,
+ parent_visible: bool,
+ current_indices: list[int],
+ previous_indices: list[int],
+ logical_start: int,
+ ) -> None:
+ row_current.append(current_indices)
+ row_previous.append(previous_indices)
+ row_groups.append(entry_group)
+ row_parent_visible.append(parent_visible)
+ row_starts.append(logical_start)
+ row_ends.append(logical_start + ratio - 1)
+
+ seg_idx = 0
+ while seg_idx < len(segments):
+ group, parent, prompt_start, prompt_end, _, _ = segments[seg_idx]
+ if group != parent:
+ seg_idx += 1
+ continue
+ prompt_len = prompt_end - prompt_start
+ shared_usable = (prompt_len // ratio) * ratio
+ shared_windows: list[tuple[list[int], list[int], int]] = []
+ for logical_start in range(0, shared_usable, ratio):
+ current_indices = _logical_window_indices(
+ prompt_start=prompt_start,
+ prompt_len=prompt_len,
+ completion_start=None,
+ logical_start=logical_start,
+ ratio=ratio,
+ )
+ previous_indices = _logical_window_indices(
+ prompt_start=prompt_start,
+ prompt_len=prompt_len,
+ completion_start=None,
+ logical_start=logical_start - ratio,
+ ratio=ratio,
+ )
+ shared_windows.append(
+ (current_indices, previous_indices, logical_start)
+ )
+ append_entry(
+ entry_group=group,
+ parent_visible=not duplicate_prompt_entries,
+ current_indices=current_indices,
+ previous_indices=previous_indices,
+ logical_start=logical_start,
+ )
+
+ child_idx = seg_idx + 1
+ while child_idx < len(segments):
+ child_group, child_parent, child_start, _, _, child_end_pos = segments[
+ child_idx
+ ]
+ if child_group == child_parent or child_parent != group:
+ break
+ if duplicate_prompt_entries:
+ for (
+ current_indices,
+ previous_indices,
+ logical_start,
+ ) in shared_windows:
+ append_entry(
+ entry_group=child_group,
+ parent_visible=False,
+ current_indices=current_indices,
+ previous_indices=previous_indices,
+ logical_start=logical_start,
+ )
+ branch_usable = ((child_end_pos + 1) // ratio) * ratio
+ for logical_start in range(0, branch_usable, ratio):
+ if logical_start + ratio <= prompt_len:
+ continue
+ append_entry(
+ entry_group=child_group,
+ parent_visible=False,
+ current_indices=_logical_window_indices(
+ prompt_start=prompt_start,
+ prompt_len=prompt_len,
+ completion_start=child_start,
+ logical_start=logical_start,
+ ratio=ratio,
+ ),
+ previous_indices=_logical_window_indices(
+ prompt_start=prompt_start,
+ prompt_len=prompt_len,
+ completion_start=child_start,
+ logical_start=logical_start - ratio,
+ ratio=ratio,
+ ),
+ logical_start=logical_start,
+ )
+ child_idx += 1
+ seg_idx = child_idx
+
+ current_rows.append(row_current)
+ previous_rows.append(row_previous)
+ group_rows.append(row_groups)
+ parent_visible_rows.append(row_parent_visible)
+ start_rows.append(row_starts)
+ end_rows.append(row_ends)
+
+ max_entries = max((len(row) for row in current_rows), default=0)
+ current = torch.full((bsz, max_entries, ratio), -1, dtype=torch.long, device=device)
+ previous = torch.full_like(current, -1)
+ entry_groups = torch.full((bsz, max_entries), -1, dtype=torch.long, device=device)
+ parent_visible = torch.zeros((bsz, max_entries), dtype=torch.bool, device=device)
+ starts = torch.full_like(entry_groups, -1)
+ ends = torch.full_like(entry_groups, -1)
+ valid = torch.zeros((bsz, max_entries), dtype=torch.bool, device=device)
+ for b, row in enumerate(current_rows):
+ if not row:
+ continue
+ count = len(row)
+ current[b, :count] = torch.tensor(row, dtype=torch.long, device=device)
+ previous[b, :count] = torch.tensor(
+ previous_rows[b], dtype=torch.long, device=device
+ )
+ entry_groups[b, :count] = torch.tensor(
+ group_rows[b], dtype=torch.long, device=device
+ )
+ parent_visible[b, :count] = torch.tensor(
+ parent_visible_rows[b], dtype=torch.bool, device=device
+ )
+ starts[b, :count] = torch.tensor(start_rows[b], dtype=torch.long, device=device)
+ ends[b, :count] = torch.tensor(end_rows[b], dtype=torch.long, device=device)
+ valid[b, :count] = True
+ return Dsv4CompressionLayout(
+ current,
+ previous,
+ entry_groups,
+ parent_visible,
+ starts,
+ ends,
+ valid,
+ )
+
+
+def compressed_layout_visibility(
+ layout: Dsv4CompressionLayout,
+ *,
+ position_ids: torch.Tensor,
+ group_ids: torch.Tensor,
+ parent_ids: torch.Tensor,
+ q_start: int = 0,
+ q_end: int | None = None,
+) -> torch.Tensor:
+ if q_end is None:
+ q_end = position_ids.shape[1]
+ q_pos = position_ids[:, q_start:q_end].to(dtype=torch.long).unsqueeze(-1)
+ q_group = group_ids[:, q_start:q_end].to(dtype=torch.long).unsqueeze(-1)
+ q_parent = parent_ids[:, q_start:q_end].to(dtype=torch.long).unsqueeze(-1)
+ entry_group = layout.entry_group_ids.unsqueeze(1)
+ parent_visible = layout.entry_parent_visible.unsqueeze(1)
+ return (
+ layout.entry_valid.unsqueeze(1)
+ & (layout.entry_end_positions.unsqueeze(1) <= q_pos)
+ & ((entry_group == q_group) | (parent_visible & (entry_group == q_parent)))
+ )
+
+
+def compressed_layout_topk_idxs(
+ layout: Dsv4CompressionLayout,
+ *,
+ position_ids: torch.Tensor,
+ group_ids: torch.Tensor,
+ parent_ids: torch.Tensor,
+ offset: int,
+) -> torch.Tensor:
+ visibility = compressed_layout_visibility(
+ layout,
+ position_ids=position_ids,
+ group_ids=group_ids,
+ parent_ids=parent_ids,
+ )
+ entry_ids = torch.arange(
+ layout.entry_group_ids.shape[1],
+ device=layout.entry_group_ids.device,
+ dtype=torch.long,
+ ).view(1, 1, -1)
+ return torch.where(visibility, entry_ids + offset, torch.full_like(entry_ids, -1))
+
+
+class RMSNorm(nn.Module):
+ """
+ Kept in pure PyTorch with FP32 weights to match SGLang's compressor norm.
+
+ Args:
+ dim: Dimension of the input tensor.
+ eps: Epsilon for numerical stability. Defaults to ``1e-6``.
+ """
+
+ def __init__(self, dim: int, eps: float = 1e-6):
+ super().__init__()
+ self.dim = dim
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32))
+
+ def forward(self, x: torch.Tensor):
+ dtype = x.dtype
+ x = x.float()
+ var = x.square().mean(-1, keepdim=True)
+ x = x * torch.rsqrt(var + self.eps)
+ return (self.weight * x).to(dtype)
+
+
+def _overlap_transform(
+ tensor: torch.Tensor, *, compress_ratio: int, head_dim: int, value=0
+) -> torch.Tensor:
+ """Overlap-transform for compress_ratio=4: for each token group of size ``ratio``,
+ split into (first_half, second_half) halves along ``head_dim`` and re-arrange
+ them across a doubled ratio axis (`2 * ratio`), shifting the first half by one
+ group so that adjacent groups overlap by ``ratio`` positions.
+ """
+ b, s, _, _ = tensor.size()
+ new_tensor = tensor.new_full((b, s, 2 * compress_ratio, head_dim), value)
+ new_tensor[:, :, compress_ratio:] = tensor[:, :, :, head_dim:]
+ new_tensor[:, 1:, :compress_ratio] = tensor[:, :-1, :, :head_dim]
+ return new_tensor
+
+
+def _add_lora_if_present(
+ owner: nn.Module,
+ attr_name: str,
+ base: torch.Tensor,
+ x: torch.Tensor,
+) -> torch.Tensor:
+ lora = getattr(owner, attr_name, None)
+ if lora is None:
+ return base
+ return base + lora(x)
+
+
+def _gather_projected_tokens(
+ tensor: torch.Tensor, indices: torch.Tensor
+) -> torch.Tensor:
+ bsz, seqlen, channels = tensor.shape
+ if indices.numel() == 0:
+ return tensor.new_empty((*indices.shape, channels))
+ batch_offsets = (
+ torch.arange(bsz, device=tensor.device, dtype=torch.long).view(bsz, 1, 1)
+ * seqlen
+ )
+ safe_indices = indices.clamp(0, max(seqlen - 1, 0))
+ flat_indices = (safe_indices + batch_offsets).reshape(-1)
+ gathered = tensor.reshape(bsz * seqlen, channels).index_select(0, flat_indices)
+ return gathered.view(*indices.shape, channels)
+
+
+class DeepSeekV4Compressor(nn.Module):
+ def __init__(
+ self,
+ config: TransformerConfig,
+ head_dim: int,
+ compress_ratio: int,
+ rotate: bool,
+ cp_group: Any | None = None,
+ ):
+ super().__init__()
+
+ cfg = cast(Any, config)
+ dim = config.hidden_size
+ rope_head_dim = int(cfg.qk_pos_emb_head_dim)
+ norm_eps = config.layernorm_epsilon
+
+ assert head_dim in {128, 512}
+ assert rope_head_dim == 64
+ assert compress_ratio in {4, 128}
+ assert norm_eps == 1e-6
+
+ self.config = config
+ self.dim = dim
+ self.head_dim = head_dim
+ self.rope_head_dim = rope_head_dim
+ self.nope_head_dim = head_dim - rope_head_dim
+ self.compress_ratio = compress_ratio
+ self.overlap = compress_ratio == 4
+ self.rotate = rotate
+ coff = 1 + self.overlap
+
+ self.cp_group = cp_group
+ self.cp_size = cp_group.size() if cp_group is not None else 1
+ self.cp_rank = cp_group.rank() if cp_group is not None else 0
+
+ self.ape = nn.Parameter(
+ torch.empty(compress_ratio, coff * self.head_dim, dtype=torch.float32)
+ )
+ weight_dtype = config.params_dtype
+ if weight_dtype not in {torch.bfloat16, torch.float32}:
+ raise TypeError(
+ f"DeepSeek-V4 compressor requires bf16/fp32 params, got {weight_dtype}"
+ )
+ self.wkv = Linear(
+ self.dim, coff * self.head_dim, bias=False, dtype=weight_dtype
+ )
+ self.wgate = Linear(
+ self.dim, coff * self.head_dim, bias=False, dtype=weight_dtype
+ )
+ self.norm = RMSNorm(self.head_dim, norm_eps)
+
+ self._keep_fp32_parameters = ("ape",)
+ setattr(self.ape, "_keep_fp32", True)
+
+ base = cfg.dsv4_compress_rope_theta
+ assert rope_head_dim == 64
+ assert base == 160000
+ configure_rope_cache(self, config, rope_head_dim=rope_head_dim, base=base)
+
+ @property
+ def weight(self) -> torch.Tensor:
+ return self.ape
+
+ def overlap_transform_raw(self, tensor: torch.Tensor, value=0):
+ """Raw overlap transform without CP handling."""
+ return _overlap_transform(
+ tensor,
+ compress_ratio=self.compress_ratio,
+ head_dim=self.head_dim,
+ value=value,
+ )
+
+ def overlap_transform_with_cp(self, tensor: torch.Tensor, value=0) -> torch.Tensor:
+ """
+ Overlap transform with CP support.
+
+ Args:
+ tensor: [bsz, G_local, ratio, coff*d]
+ value: Fill value for overlap transform (0 for kv, -inf for score)
+
+ Returns:
+ [bsz, G_local, ratio, coff*d]
+ """
+ if self.cp_size != 1:
+ raise RuntimeError(
+ "DeepSeek-V4 non-CP compressor received context_parallel_size > 1."
+ )
+ return self.overlap_transform_raw(tensor, value)
+
+ def _project_raw(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ if self.ape.dtype != torch.float32:
+ raise TypeError(
+ f"DeepSeek-V4 compressor APE must stay fp32, got {self.ape.dtype}."
+ )
+ if self.wkv.weight.dtype not in {torch.bfloat16, torch.float32}:
+ raise TypeError(
+ "DeepSeek-V4 compressor KV projection requires bf16/fp32 weights, got "
+ f"{self.wkv.weight.dtype}."
+ )
+ if self.wgate.weight.dtype not in {torch.bfloat16, torch.float32}:
+ raise TypeError(
+ "DeepSeek-V4 compressor gate projection requires bf16/fp32 weights, "
+ f"got {self.wgate.weight.dtype}."
+ )
+
+ kv = _add_lora_if_present(
+ self, "kv_proj_lora", linear_bf16_fp32(x, self.wkv.weight), x
+ )
+ score = _add_lora_if_present(
+ self, "gate_proj_lora", linear_bf16_fp32(x, self.wgate.weight), x
+ )
+ return kv, score
+
+ def _compress_projected(
+ self,
+ kv: torch.Tensor,
+ score: torch.Tensor,
+ *,
+ freqs_cis: torch.Tensor,
+ ) -> torch.Tensor:
+ bsz, seqlen_local, _ = kv.size()
+ ratio, overlap, _ = self.compress_ratio, self.overlap, self.head_dim
+ dtype = kv.dtype
+
+ usable = (seqlen_local // ratio) * ratio
+ if usable == 0:
+ return kv.new_zeros((bsz, 0, self.head_dim))
+ kv = kv[:, :usable]
+ score = score[:, :usable]
+ if self.cp_size > 1:
+ assert usable % (ratio * 2) == 0
+
+ kv = kv.unflatten(1, (-1, ratio))
+ score = score.unflatten(1, (-1, ratio)) + self.ape
+
+ if overlap:
+ kv = self.overlap_transform_with_cp(kv, 0)
+ score = self.overlap_transform_with_cp(score, float("-inf"))
+
+ score_softmax = score.softmax(dim=2, dtype=torch.float32).to(kv.dtype)
+ kv = (kv * score_softmax).sum(dim=2)
+
+ kv = self.norm(kv.to(dtype))
+
+ apply_rotary_emb(kv[..., -self.rope_head_dim :], freqs_cis)
+
+ if self.rotate:
+ kv = rotate_activation(kv)
+
+ return kv
+
+ def _compress_shared_prefix_projected(
+ self,
+ kv: torch.Tensor,
+ score: torch.Tensor,
+ *,
+ layout: Dsv4CompressionLayout,
+ ) -> torch.Tensor:
+ dtype = kv.dtype
+ ratio = self.compress_ratio
+ current_valid = layout.current_indices >= 0
+ current_kv = _gather_projected_tokens(kv, layout.current_indices)
+ current_score = _gather_projected_tokens(score, layout.current_indices)
+ current_kv = torch.where(
+ current_valid.unsqueeze(-1), current_kv, torch.zeros_like(current_kv)
+ )
+ current_score = torch.where(
+ current_valid.unsqueeze(-1),
+ current_score,
+ torch.full_like(current_score, float("-inf")),
+ )
+ if self.overlap:
+ previous_valid = layout.previous_indices >= 0
+ previous_kv = _gather_projected_tokens(kv, layout.previous_indices)
+ previous_score = _gather_projected_tokens(score, layout.previous_indices)
+ previous_kv = torch.where(
+ previous_valid.unsqueeze(-1),
+ previous_kv,
+ torch.zeros_like(previous_kv),
+ )
+ previous_score = torch.where(
+ previous_valid.unsqueeze(-1),
+ previous_score,
+ torch.full_like(previous_score, float("-inf")),
+ )
+ current_score = current_score + self.ape.view(1, 1, ratio, -1)
+ previous_score = previous_score + self.ape.view(1, 1, ratio, -1)
+ slots_kv = torch.cat(
+ [previous_kv[..., : self.head_dim], current_kv[..., self.head_dim :]],
+ dim=2,
+ )
+ slots_score = torch.cat(
+ [
+ previous_score[..., : self.head_dim],
+ current_score[..., self.head_dim :],
+ ],
+ dim=2,
+ )
+ else:
+ slots_kv = current_kv
+ slots_score = current_score + self.ape.view(1, 1, ratio, -1)
+
+ slots_score = torch.where(
+ layout.entry_valid[:, :, None, None],
+ slots_score,
+ torch.zeros_like(slots_score),
+ )
+ score_softmax = slots_score.softmax(dim=2, dtype=torch.float32).to(
+ slots_kv.dtype
+ )
+ compressed = (slots_kv * score_softmax).sum(dim=2)
+ compressed = self.norm(compressed.to(dtype))
+ freqs_cis = get_rope_cache_at_positions(
+ self, position_ids=layout.entry_start_positions, device=kv.device
+ )
+ apply_rotary_emb(compressed[..., -self.rope_head_dim :], freqs_cis)
+ compressed = torch.where(
+ layout.entry_valid.unsqueeze(-1),
+ compressed,
+ torch.zeros_like(compressed),
+ )
+
+ if self.rotate:
+ compressed = rotate_activation(compressed)
+ return compressed
+
+ def forward_raw(
+ self,
+ x: torch.Tensor,
+ *,
+ position_ids: torch.Tensor | None = None,
+ shared_layout: Dsv4CompressionLayout | None = None,
+ ) -> torch.Tensor:
+ kv, score = self._project_raw(x)
+ if shared_layout is not None:
+ return self._compress_shared_prefix_projected(
+ kv, score, layout=shared_layout
+ )
+ usable = (x.shape[1] // self.compress_ratio) * self.compress_ratio
+ freqs_cis = get_rope_cache(self, seqlen=usable, device=x.device)[
+ : usable : self.compress_ratio
+ ]
+ if position_ids is not None:
+ freqs_cis = get_rope_cache_at_positions(
+ self,
+ position_ids=position_ids[:, : usable : self.compress_ratio],
+ device=x.device,
+ )
+ return self._compress_projected(kv, score, freqs_cis=freqs_cis)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ *,
+ position_ids: torch.Tensor | None = None,
+ shared_layout: Dsv4CompressionLayout | None = None,
+ ) -> torch.Tensor:
+ """
+ Args:
+ x: [seqlen, batch, dim] SBHD layout (Megatron standard)
+ Returns:
+ k: [floor(seqlen / compress_ratio), batch, head_dim] SBHD layout
+ """
+ x_bshd = einops.rearrange(x, "s b d -> b s d")
+ k_bshd = self.forward_raw(
+ x_bshd,
+ position_ids=position_ids,
+ shared_layout=shared_layout,
+ )
+ k = einops.rearrange(k_bshd, "b sc d -> sc b d")
+ return k
diff --git a/src/art/megatron/dsv4/deepseek_v4.py b/src/art/megatron/dsv4/deepseek_v4.py
new file mode 100644
index 000000000..85d9d313a
--- /dev/null
+++ b/src/art/megatron/dsv4/deepseek_v4.py
@@ -0,0 +1,679 @@
+import copy
+from typing import Any, cast
+
+import einops
+import torch
+import torch.nn as nn
+
+# Enable TF32 for fp32 matmul to match the precision of the TileKernels MHC
+# kernels (which use TF32 tensor-core GEMM for the HC fp32 mixer). Without
+# this, PyTorch's default ``allow_tf32=False`` keeps fp32 ``F.linear`` on the
+# SIMT path, which introduces a ~1e-4 mean-abs gap vs the TileKernels output;
+# matching TF32 brings the gap to <=1.5e-5 mean-abs (1 ULP bf16 max-abs).
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.extensions.transformer_engine import (
+ TEColumnParallelLinear,
+ TELinear,
+ TENorm,
+ TERowParallelLinear,
+)
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear
+from megatron.core.tensor_parallel.mappings import (
+ copy_to_tensor_model_parallel_region,
+ gather_from_sequence_parallel_region,
+ scatter_to_sequence_parallel_region,
+)
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
+
+from art.megatron.dsv4.compressor import (
+ DeepSeekV4Compressor,
+ Dsv4CompressionLayout,
+ Dsv4SharedPrefixState,
+ compressed_layout_topk_idxs,
+)
+from art.megatron.dsv4.kernel.tilelang_sparse_mla import sparse_attn_tilelang
+from art.megatron.dsv4.rope import (
+ apply_rotary_emb,
+ configure_rope_cache,
+ get_rope_cache,
+ get_rope_cache_at_positions,
+)
+from art.megatron.dsv4.v4_indexer import V4Indexer
+
+
+def _window_topk_idxs(
+ q_positions: torch.Tensor, *, window_size: int, bsz: int
+) -> torch.Tensor:
+ base = q_positions.unsqueeze(1)
+ offsets = torch.arange(
+ min(q_positions.numel(), window_size), device=q_positions.device
+ )
+ k_pos = (base - window_size + 1).clamp(0) + offsets
+ topk = torch.where(k_pos > base, -1, k_pos)
+ return topk.unsqueeze(0).expand(bsz, -1, -1).to(torch.int32)
+
+
+def _compress_topk_idxs(
+ q_positions: torch.Tensor, *, ratio: int, bsz: int
+) -> torch.Tensor:
+ seqlen = int(q_positions.numel())
+ offset = seqlen
+ k_group_idx = torch.arange(seqlen // ratio, device=q_positions.device).repeat(
+ seqlen, 1
+ )
+ q_first_invalid_group = (q_positions + 1).unsqueeze(1) // ratio
+ compress = torch.where(
+ k_group_idx >= q_first_invalid_group, -1, k_group_idx + offset
+ )
+ return compress.unsqueeze(0).expand(bsz, -1, -1).to(torch.int32)
+
+
+def _shared_prefix_tensors(
+ attention_bias: Any,
+ *,
+ bsz: int,
+ seqlen: int,
+ device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor] | None:
+ group_ids = getattr(attention_bias, "group_ids", None)
+ parent_ids = getattr(attention_bias, "parent_ids", None)
+ if group_ids is None or parent_ids is None:
+ return None
+ group_ids = group_ids.to(device=device, dtype=torch.long)
+ parent_ids = parent_ids.to(device=device, dtype=torch.long)
+ if group_ids.shape != (bsz, seqlen) or parent_ids.shape != (bsz, seqlen):
+ raise ValueError(
+ "DSV4 shared-prefix metadata must match attention input shape: "
+ f"group_ids={tuple(group_ids.shape)} parent_ids={tuple(parent_ids.shape)} "
+ f"expected={(bsz, seqlen)}."
+ )
+ return group_ids, parent_ids
+
+
+def _shared_prefix_window_topk_idxs(
+ position_ids: torch.Tensor,
+ group_ids: torch.Tensor,
+ parent_ids: torch.Tensor,
+ *,
+ window_size: int,
+) -> torch.Tensor:
+ bsz, seqlen = group_ids.shape
+ width = min(seqlen, window_size)
+ arange = torch.arange(seqlen, device=group_ids.device).expand(bsz, -1)
+ valid_token = (group_ids != -1) & (parent_ids != -1)
+ group_start = valid_token.clone()
+ group_start[:, 1:] &= group_ids[:, 1:] != group_ids[:, :-1]
+ start_values = torch.where(group_start, arange, torch.zeros_like(arange))
+ group_start_idx = torch.cummax(start_values, dim=1).values
+ prompt_start = group_start & (group_ids == parent_ids)
+ prompt_values = torch.where(prompt_start, arange, torch.zeros_like(arange))
+ prompt_start_idx = torch.cummax(prompt_values, dim=1).values
+
+ q_pos = position_ids.to(device=group_ids.device, dtype=torch.long)
+ group_start_pos = q_pos.gather(1, group_start_idx)
+ offsets = torch.arange(width, device=group_ids.device)
+ target_pos = q_pos.unsqueeze(-1) - width + 1 + offsets
+ from_group = target_pos >= group_start_pos.unsqueeze(-1)
+ group_k = group_start_idx.unsqueeze(-1) + target_pos - group_start_pos.unsqueeze(-1)
+ prompt_k = prompt_start_idx.unsqueeze(-1) + target_pos
+ k_idx = torch.where(from_group, group_k, prompt_k)
+ valid = (
+ valid_token.unsqueeze(-1)
+ & (target_pos >= 0)
+ & (target_pos <= q_pos.unsqueeze(-1))
+ & (k_idx >= 0)
+ & (k_idx < seqlen)
+ )
+ return torch.where(valid, k_idx, torch.full_like(k_idx, -1)).to(torch.int32)
+
+
+def _dsv4_shared_prefix_state(attention_bias: Any) -> Dsv4SharedPrefixState:
+ model_state = getattr(attention_bias, "model_state", None)
+ state = model_state.get("dsv4") if isinstance(model_state, dict) else None
+ if not isinstance(state, Dsv4SharedPrefixState):
+ raise RuntimeError(
+ "DSV4 shared-prefix state is missing. Build it once per packed "
+ "sequence through the model-support shared-prefix state hook."
+ )
+ return state
+
+
+def _dsv4_topk_cache(attention_bias: Any) -> dict[Any, Any]:
+ return _dsv4_shared_prefix_state(attention_bias).topk_idx_cache
+
+
+def _shared_prefix_window_topk_idxs_cached(
+ attention_bias: Any,
+ position_ids: torch.Tensor,
+ group_ids: torch.Tensor,
+ parent_ids: torch.Tensor,
+ *,
+ window_size: int,
+) -> torch.Tensor:
+ cache = _dsv4_topk_cache(attention_bias)
+ key = ("raw_swa", int(window_size))
+ cached = cache.get(key)
+ if cached is None:
+ cached = _shared_prefix_window_topk_idxs(
+ position_ids,
+ group_ids,
+ parent_ids,
+ window_size=window_size,
+ ).contiguous()
+ cache[key] = cached
+ return cached
+
+
+def _shared_prefix_compressed_topk_idxs_cached(
+ attention_bias: Any,
+ layout: Dsv4CompressionLayout,
+ *,
+ position_ids: torch.Tensor,
+ group_ids: torch.Tensor,
+ parent_ids: torch.Tensor,
+ ratio: int,
+ offset: int,
+) -> torch.Tensor:
+ cache = _dsv4_topk_cache(attention_bias)
+ key = ("compressed", int(ratio), int(offset))
+ cached = cache.get(key)
+ if cached is None:
+ cached = (
+ compressed_layout_topk_idxs(
+ layout,
+ position_ids=position_ids,
+ group_ids=group_ids,
+ parent_ids=parent_ids,
+ offset=offset,
+ )
+ .to(torch.int32)
+ .contiguous()
+ )
+ cache[key] = cached
+ return cached
+
+
+def _shared_prefix_i32_metadata_cached(
+ attention_bias: Any,
+ position_ids: torch.Tensor,
+ group_ids: torch.Tensor,
+ parent_ids: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ cache = _dsv4_topk_cache(attention_bias)
+ key = "indexer_q_metadata_i32"
+ cached = cache.get(key)
+ if cached is None:
+ cached = (
+ position_ids.to(torch.int32).contiguous(),
+ group_ids.to(torch.int32).contiguous(),
+ parent_ids.to(torch.int32).contiguous(),
+ )
+ cache[key] = cached
+ return cached
+
+
+def _shared_layout_indexer_metadata_cached(
+ attention_bias: Any,
+ layout: Dsv4CompressionLayout,
+ *,
+ ratio: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ cache = _dsv4_topk_cache(attention_bias)
+ key = ("indexer_layout_i32", int(ratio))
+ cached = cache.get(key)
+ if cached is None:
+ cached = (
+ layout.entry_group_ids.to(torch.int32).contiguous(),
+ layout.entry_parent_visible.to(torch.int32).contiguous(),
+ layout.entry_end_positions.to(torch.int32).contiguous(),
+ layout.entry_valid.to(torch.int32).contiguous(),
+ )
+ cache[key] = cached
+ return cached
+
+
+def _shared_prefix_compression_layout(
+ attention_bias: Any,
+ *,
+ ratio: int,
+) -> Dsv4CompressionLayout:
+ layouts = _dsv4_shared_prefix_state(attention_bias).compression_layouts
+ if ratio not in layouts:
+ raise RuntimeError(
+ "DSV4 shared-prefix compression layout was not prepared on the "
+ f"attention state for ratio={ratio}. Build it once per packed "
+ "sequence through the model-support shared-prefix state hook."
+ )
+ layout = layouts[ratio]
+ if not isinstance(layout, Dsv4CompressionLayout):
+ raise TypeError(f"Expected Dsv4CompressionLayout for ratio={ratio}.")
+ return layout
+
+
+def _add_lora_if_present(
+ owner: nn.Module,
+ attr_name: str,
+ base: torch.Tensor,
+ x: torch.Tensor,
+) -> torch.Tensor:
+ lora = getattr(owner, attr_name, None)
+ if lora is None:
+ return base
+ return base + lora(x)
+
+
+class DeepSeekV4Attention(MegatronModule):
+ def __init__(
+ self,
+ config: TransformerConfig,
+ submodules=None,
+ layer_number: int = 1,
+ attn_mask_type=None,
+ attention_type: str | None = None,
+ cp_comm_type: str | None = None,
+ pg_collection=None,
+ ):
+ super().__init__(config=config)
+ cfg = cast(Any, config)
+ init_method = config.init_method
+ if init_method is None:
+ raise RuntimeError("DeepSeek-V4 attention requires config.init_method.")
+
+ if pg_collection is None:
+ pg_collection = ProcessGroupCollection.use_mpu_process_groups(
+ required_pgs=["tp"]
+ )
+ else:
+ assert hasattr(pg_collection, "tp")
+ self.pg_collection = pg_collection
+ self.tp_group = self.pg_collection.tp
+ self.cp_group = pg_collection.cp if hasattr(pg_collection, "cp") else None
+ self.cp_size = self.cp_group.size() if self.cp_group else 1
+ if self.cp_size != 1:
+ raise RuntimeError(
+ "DeepSeek-V4 non-CP attention received context_parallel_size > 1."
+ )
+
+ layer_id = layer_number - 1
+ del layer_number
+
+ self.layer_id = layer_id
+ self.dim = config.hidden_size
+ self.n_heads = config.num_attention_heads
+ self.n_local_heads = self.n_heads // config.tensor_model_parallel_size
+ self.q_lora_rank = int(cfg.q_lora_rank)
+ self.o_lora_rank = int(cfg.dsv4_o_lora_rank)
+ self.head_dim = int(cfg.kv_lora_rank)
+ self.rope_head_dim = int(cfg.qk_pos_emb_head_dim)
+ self.nope_head_dim = self.head_dim - self.rope_head_dim
+ self.n_groups = int(cfg.dsv4_o_groups)
+ self.n_local_groups = self.n_groups // config.tensor_model_parallel_size
+ self.window_size = int(cfg.dsv4_window_size)
+ compress_ratios = cfg.dsv4_compress_ratios
+ self.compress_ratio = int(compress_ratios[layer_id]) if compress_ratios else 0
+ self.eps = config.layernorm_epsilon
+
+ assert self.o_lora_rank == 1024
+ assert self.head_dim == 512
+ assert self.rope_head_dim == 64
+ assert self.nope_head_dim == 448
+ assert self.window_size == 128
+
+ config_no_sp = copy.copy(config)
+ config_no_sp.sequence_parallel = False
+
+ attn_sink = torch.empty(self.n_local_heads, dtype=torch.float32)
+ self._keep_fp32_parameters = ("attn_sink",)
+ self._keep_fp32_buffers = ("attn_sink",)
+ self.attn_sink = nn.Parameter(attn_sink)
+ setattr(self.attn_sink, "_keep_fp32", True)
+
+ self.wq_a = TELinear(
+ self.dim,
+ self.q_lora_rank,
+ config=config,
+ init_method=init_method,
+ bias=False,
+ skip_bias_add=False,
+ skip_weight_param_allocation=False,
+ parallel_mode="duplicated",
+ )
+ self.q_norm = TENorm(config_no_sp, self.q_lora_rank, eps=self.eps)
+ self.wq_b = TEColumnParallelLinear(
+ self.q_lora_rank,
+ self.n_heads * self.head_dim,
+ config=config_no_sp,
+ init_method=init_method,
+ bias=False,
+ gather_output=False,
+ skip_bias_add=False,
+ is_expert=False,
+ tp_group=self.tp_group,
+ )
+ self.wkv = TELinear(
+ self.dim,
+ self.head_dim,
+ config=config,
+ init_method=init_method,
+ bias=False,
+ skip_bias_add=False,
+ skip_weight_param_allocation=False,
+ parallel_mode="duplicated",
+ )
+ self.kv_norm = TENorm(config_no_sp, self.head_dim, eps=self.eps)
+
+ for p in list(self.wq_a.parameters()) + list(self.wkv.parameters()):
+ setattr(p, "sequence_parallel", False)
+
+ self.wo_a = ColumnParallelLinear(
+ self.n_heads * self.head_dim // self.n_groups,
+ self.n_groups * self.o_lora_rank,
+ config=config_no_sp,
+ init_method=init_method,
+ bias=False,
+ gather_output=False,
+ )
+ self.wo_b = TERowParallelLinear(
+ self.n_groups * self.o_lora_rank,
+ self.dim,
+ config=config_no_sp,
+ init_method=init_method,
+ bias=False,
+ input_is_parallel=True,
+ skip_bias_add=False,
+ is_expert=False,
+ tp_group=self.tp_group,
+ )
+ self.softmax_scale = self.head_dim**-0.5
+ self.sequence_parallel = config.sequence_parallel
+
+ if self.compress_ratio:
+ self.compressor = DeepSeekV4Compressor(
+ config=config,
+ head_dim=self.head_dim,
+ compress_ratio=self.compress_ratio,
+ rotate=False,
+ cp_group=self.cp_group,
+ )
+ if self.compress_ratio == 4:
+ self.indexer = V4Indexer(config=config, pg_collection=pg_collection)
+ else:
+ self.indexer = None
+
+ rope_base = (
+ cfg.dsv4_compress_rope_theta if self.compress_ratio else cfg.rotary_base
+ )
+ configure_rope_cache(
+ self,
+ config,
+ rope_head_dim=self.rope_head_dim,
+ base=rope_base,
+ )
+ self._dsv4_position_ids: torch.Tensor | None = None
+
+ def set_position_ids(self, position_ids: torch.Tensor | None) -> None:
+ self._dsv4_position_ids = position_ids
+
+ def sharded_state_dict(
+ self,
+ prefix: str = "",
+ sharded_offsets: tuple = (),
+ metadata: dict | None = None,
+ ) -> ShardedStateDict:
+ ans = super().sharded_state_dict(prefix, sharded_offsets, metadata)
+ ans.update(
+ make_sharded_tensors_for_checkpoint(
+ state_dict={"attn_sink": self.attn_sink},
+ prefix=prefix,
+ tensor_parallel_layers_axis_map={"attn_sink": 0},
+ sharded_offsets=sharded_offsets,
+ tp_group=self.tp_group,
+ dp_cp_group=(metadata or {})["dp_cp_group"],
+ )
+ )
+ return ans
+
+ @torch.compiler.disable
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask=None,
+ inference_context=None,
+ rotary_pos_emb=None,
+ rotary_pos_cos=None,
+ rotary_pos_sin=None,
+ rotary_pos_cos_sin=None,
+ attention_bias=None,
+ packed_seq_params=None,
+ sequence_len_offset=None,
+ ) -> torch.Tensor:
+ """Run DSV4 attention eager inside compiled transformer layers.
+
+ Torch 2.11 currently miscompiles this TP+SP autograd graph: the
+ attention module output receives a nonzero gradient, but the
+ zero-initialized LoRA branches inside the DSV4 attention path get zero
+ tangents. Keeping only this model-specific attention forward eager
+ preserves compiled surrounding layer code and correct first-step LoRA
+ gradients.
+ """
+ if self.sequence_parallel:
+ hidden_states = gather_from_sequence_parallel_region(
+ hidden_states,
+ tensor_parallel_output_grad=False,
+ group=self.tp_group,
+ )
+
+ x = einops.rearrange(hidden_states, "s b d -> b s d")
+
+ bsz, seqlen_local, _ = x.size()
+ position_ids = self._dsv4_position_ids
+ if position_ids is not None:
+ if position_ids.shape != (bsz, seqlen_local):
+ raise ValueError(
+ "DSV4 position_ids must match attention input shape: "
+ f"position_ids={tuple(position_ids.shape)} expected={(bsz, seqlen_local)}."
+ )
+ freqs_cis = get_rope_cache_at_positions(
+ self, position_ids=position_ids, device=x.device
+ )
+ else:
+ freqs_cis = get_rope_cache(self, seqlen=seqlen_local, device=x.device)
+ win = self.window_size
+ ratio = self.compress_ratio
+ rd = self.rope_head_dim
+ shared_prefix = _shared_prefix_tensors(
+ attention_bias, bsz=bsz, seqlen=seqlen_local, device=x.device
+ )
+ shared_layout: Dsv4CompressionLayout | None = None
+ if (
+ self.compress_ratio
+ and shared_prefix is not None
+ and position_ids is not None
+ ):
+ shared_layout = _shared_prefix_compression_layout(
+ attention_bias,
+ ratio=ratio,
+ )
+
+ q_after_wq_a = _add_lora_if_present(self, "wq_a_lora", self.wq_a(x)[0], x)
+ qr = q = cast(Any, self.q_norm)(q_after_wq_a)
+ q_after_wq_b = _add_lora_if_present(self, "wq_b_lora", self.wq_b(q)[0], q)
+ q = q_after_wq_b.unflatten(-1, (self.n_local_heads, self.head_dim))
+ q_fp32 = q.float()
+ q = (
+ q_fp32 * torch.rsqrt(q_fp32.square().mean(-1, keepdim=True) + self.eps)
+ ).to(q.dtype)
+ q = q.clone()
+ apply_rotary_emb(q[..., -rd:], freqs_cis)
+
+ kv_after_wkv = _add_lora_if_present(self, "wkv_lora", self.wkv(x)[0], x)
+ kv_vanilla = cast(Any, self.kv_norm)(kv_after_wkv)
+ kv_vanilla = kv_vanilla.clone()
+ apply_rotary_emb(kv_vanilla[..., -rd:], freqs_cis)
+
+ seqlen_global = seqlen_local
+ q_positions = torch.arange(seqlen_local, device=x.device)
+
+ if shared_prefix is not None and position_ids is not None:
+ group_ids, parent_ids = shared_prefix
+ topk_idxs = _shared_prefix_window_topk_idxs_cached(
+ attention_bias,
+ position_ids,
+ group_ids,
+ parent_ids,
+ window_size=win,
+ )
+ else:
+ topk_idxs = _window_topk_idxs(q_positions, window_size=win, bsz=bsz)
+
+ if self.compress_ratio:
+ kv_compress_offset = seqlen_global
+ if self.indexer is not None:
+ x_sbd = einops.rearrange(x, "b s d -> s b d")
+ qr_sbd = einops.rearrange(qr, "b s d -> s b d")
+ if self.sequence_parallel:
+ x_sbd = scatter_to_sequence_parallel_region(
+ x_sbd, group=self.tp_group
+ )
+ qr_sbd = scatter_to_sequence_parallel_region(
+ qr_sbd, group=self.tp_group
+ )
+ if isinstance(self.indexer, V4Indexer):
+ group_ids = parent_ids = None
+ shared_layout_i32 = None
+ index_position_ids = position_ids
+ if shared_prefix is not None:
+ group_ids, parent_ids = shared_prefix
+ if position_ids is not None and shared_layout is not None:
+ index_position_ids, group_ids, parent_ids = (
+ _shared_prefix_i32_metadata_cached(
+ attention_bias,
+ position_ids,
+ group_ids,
+ parent_ids,
+ )
+ )
+ shared_layout_i32 = _shared_layout_indexer_metadata_cached(
+ attention_bias,
+ shared_layout,
+ ratio=ratio,
+ )
+ compress_topk_idxs = self.indexer(
+ x_sbd,
+ qr_sbd,
+ position_ids=index_position_ids,
+ shared_layout=shared_layout,
+ group_ids=group_ids,
+ parent_ids=parent_ids,
+ shared_layout_i32=shared_layout_i32,
+ )
+ else:
+ indexer_mask = self._compute_indexer_mask(
+ q_positions=q_positions, seqlen_global=seqlen_global
+ )
+ compress_topk_idxs = self.indexer(
+ x_sbd, qr_sbd, mask=indexer_mask, packed_seq_params=None
+ )
+ if shared_layout is None:
+ q_first_invalid_group = (q_positions + 1).unsqueeze(1) // ratio
+ topk_idx_mask = (compress_topk_idxs >= q_first_invalid_group) | (
+ compress_topk_idxs < 0
+ )
+ compress_topk_idxs = torch.where(
+ topk_idx_mask, -1, compress_topk_idxs + kv_compress_offset
+ )
+ else:
+ compress_topk_idxs = torch.where(
+ compress_topk_idxs < 0,
+ -1,
+ compress_topk_idxs + kv_compress_offset,
+ )
+ else:
+ if (
+ shared_layout is None
+ or shared_prefix is None
+ or position_ids is None
+ ):
+ compress_topk_idxs = _compress_topk_idxs(
+ q_positions, ratio=ratio, bsz=bsz
+ )
+ else:
+ group_ids, parent_ids = shared_prefix
+ compress_topk_idxs = _shared_prefix_compressed_topk_idxs_cached(
+ attention_bias,
+ shared_layout,
+ position_ids=position_ids,
+ group_ids=group_ids,
+ parent_ids=parent_ids,
+ ratio=ratio,
+ offset=kv_compress_offset,
+ )
+ topk_idxs = torch.cat([topk_idxs, compress_topk_idxs], dim=-1)
+ topk_idxs = topk_idxs.to(torch.int32)
+
+ kv_compress = None
+ if self.compress_ratio:
+ x_sbd = einops.rearrange(x, "b s d -> s b d")
+ kv_compress_sbd = self.compressor(
+ x_sbd,
+ position_ids=position_ids,
+ shared_layout=shared_layout,
+ )
+ if kv_compress_sbd is not None:
+ kv_compress = einops.rearrange(kv_compress_sbd, "s b d -> b s d")
+
+ if self.attn_sink.dtype != torch.float32:
+ raise TypeError(
+ "DeepSeek-V4 attention sink must stay fp32, got "
+ f"{self.attn_sink.dtype}."
+ )
+
+ if kv_compress is not None:
+ kv = torch.cat([kv_vanilla, kv_compress], dim=1)
+ if kv_compress_offset != kv_vanilla.size(1):
+ raise RuntimeError(
+ "DeepSeek-V4 compressed KV offset must equal raw KV length, got "
+ f"{kv_compress_offset} and {kv_vanilla.size(1)}."
+ )
+ else:
+ kv = kv_vanilla
+
+ kv = copy_to_tensor_model_parallel_region(kv, group=self.tp_group)
+
+ o = sparse_attn_tilelang(q, kv, self.attn_sink, topk_idxs, self.softmax_scale)
+
+ apply_rotary_emb(o[..., -rd:], freqs_cis, inverse=True)
+
+ o = o.view(bsz, seqlen_local, self.n_local_groups, -1)
+ wo_a_input = o
+ wo_a = cast(torch.Tensor, self.wo_a.weight).view(
+ self.n_local_groups, self.o_lora_rank, -1
+ )
+ o = torch.einsum("bsgd,grd->bsgr", o, wo_a)
+ o = _add_lora_if_present(self, "wo_a_lora", o, wo_a_input)
+ wo_b_input = o.flatten(2)
+ x, _ = self.wo_b(wo_b_input)
+ x = _add_lora_if_present(self, "wo_b_lora", x, wo_b_input)
+
+ output = einops.rearrange(x, "b s d -> s b d")
+
+ if self.sequence_parallel:
+ output = scatter_to_sequence_parallel_region(output, group=self.tp_group)
+
+ return output
+
+ def _compute_indexer_mask(
+ self, *, q_positions: torch.Tensor, seqlen_global: int
+ ) -> torch.Tensor:
+ """Dense causal mask for legacy DSAIndexer path."""
+ ratio = 4
+ device = q_positions.device
+ k_group_idx = torch.arange(seqlen_global // ratio, device=device).unsqueeze(0)
+ q_first_invalid_group = (q_positions.unsqueeze(1) + 1) // ratio
+ invalid_mask = k_group_idx >= q_first_invalid_group
+ return torch.where(invalid_mask, float("-inf"), 0.0)
diff --git a/src/art/megatron/dsv4/dequant.py b/src/art/megatron/dsv4/dequant.py
new file mode 100644
index 000000000..992928a4e
--- /dev/null
+++ b/src/art/megatron/dsv4/dequant.py
@@ -0,0 +1,222 @@
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+_DSV4_FP4_TABLE = (
+ 0.0,
+ 0.5,
+ 1.0,
+ 1.5,
+ 2.0,
+ 3.0,
+ 4.0,
+ 6.0,
+ 0.0,
+ -0.5,
+ -1.0,
+ -1.5,
+ -2.0,
+ -3.0,
+ -4.0,
+ -6.0,
+)
+_TABLE_CACHE: dict[int, torch.Tensor] = {}
+
+
+@triton.jit
+def _mxfp4_dequant_kernel(
+ weight_ptr,
+ scale_ptr,
+ table_ptr,
+ out_ptr,
+ total: tl.constexpr,
+ in_dim: tl.constexpr,
+ in_bytes: tl.constexpr,
+ scale_cols: tl.constexpr,
+ block: tl.constexpr,
+):
+ offsets = tl.program_id(0) * block + tl.arange(0, block)
+ mask = offsets < total
+ col = offsets % in_dim
+ row = offsets // in_dim
+ packed = tl.load(weight_ptr + row * in_bytes + col // 2, mask=mask, other=0)
+ nibble = tl.where((col & 1) == 0, packed & 0x0F, (packed >> 4) & 0x0F)
+ fp4 = tl.load(table_ptr + nibble)
+ raw_scale = tl.load(scale_ptr + row * scale_cols + col // 32, mask=mask, other=127)
+ scale = tl.where(
+ raw_scale == 255,
+ float("nan"),
+ tl.exp2(raw_scale.to(tl.float32) - 127.0),
+ )
+ tl.store(out_ptr + offsets, (fp4 * scale).to(tl.bfloat16), mask=mask)
+
+
+def _fp4_table(device: torch.device) -> torch.Tensor:
+ index = torch.device(device).index
+ if index is None:
+ index = torch.cuda.current_device()
+ table = _TABLE_CACHE.get(index)
+ if table is None:
+ table = torch.tensor(_DSV4_FP4_TABLE, dtype=torch.float32, device=device)
+ _TABLE_CACHE[index] = table
+ return table
+
+
+def dequant_mxfp4_cuda(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+ device = (
+ weight.device
+ if weight.device.type == "cuda"
+ else torch.device("cuda", torch.cuda.current_device())
+ )
+ weight = weight.contiguous().to(device=device, non_blocking=True).view(torch.uint8)
+ scale = scale.contiguous().to(device=device, non_blocking=True).view(torch.uint8)
+ out_dim, in_bytes = weight.shape
+ in_dim = in_bytes * 2
+ out = torch.empty((out_dim, in_dim), dtype=torch.bfloat16, device=device)
+ block = 256
+ _mxfp4_dequant_kernel[(triton.cdiv(out.numel(), block),)](
+ weight,
+ scale,
+ _fp4_table(device),
+ out,
+ out.numel(), # ty: ignore[invalid-argument-type]
+ in_dim, # ty: ignore[invalid-argument-type]
+ in_bytes, # ty: ignore[invalid-argument-type]
+ scale.shape[1], # ty: ignore[invalid-argument-type]
+ block, # ty: ignore[invalid-argument-type]
+ num_warps=4, # ty: ignore[unknown-argument]
+ )
+ return out
+
+
+@triton.jit
+def _block_fp8_dequant_kernel(
+ weight_ptr,
+ scale_ptr,
+ out_ptr,
+ total: tl.constexpr,
+ n_cols: tl.constexpr,
+ scale_cols: tl.constexpr,
+ block_elems: tl.constexpr,
+):
+ offsets = tl.program_id(0) * block_elems + tl.arange(0, block_elems)
+ mask = offsets < total
+ col = offsets % n_cols
+ row = offsets // n_cols
+ value = tl.load(weight_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+ raw_scale = tl.load(
+ scale_ptr + (row // 128) * scale_cols + col // 128,
+ mask=mask,
+ other=127,
+ )
+ scale = tl.where(
+ raw_scale == 255,
+ float("nan"),
+ tl.exp2(raw_scale.to(tl.float32) - 127.0),
+ )
+ tl.store(out_ptr + offsets, (value * scale).to(tl.bfloat16), mask=mask)
+
+
+def dequant_block_fp8_cuda(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+ device = (
+ weight.device
+ if weight.device.type == "cuda"
+ else torch.device("cuda", torch.cuda.current_device())
+ )
+ weight = weight.contiguous().to(device=device, non_blocking=True)
+ scale = scale.contiguous().to(device=device, non_blocking=True).view(torch.uint8)
+ out = torch.empty_like(weight, dtype=torch.bfloat16)
+ block_elems = 256
+ _block_fp8_dequant_kernel[(triton.cdiv(out.numel(), block_elems),)](
+ weight,
+ scale,
+ out,
+ out.numel(), # ty: ignore[invalid-argument-type]
+ weight.shape[1], # ty: ignore[invalid-argument-type]
+ scale.shape[1], # ty: ignore[invalid-argument-type]
+ block_elems, # ty: ignore[invalid-argument-type]
+ num_warps=4, # ty: ignore[unknown-argument]
+ )
+ return out
+
+
+@triton.jit
+def _mxfp4_quant_kernel(
+ weight_ptr,
+ packed_ptr,
+ scale_ptr,
+ total_blocks: tl.constexpr,
+ cols: tl.constexpr,
+ blocks_per_row: tl.constexpr,
+):
+ block_id = tl.program_id(0)
+ pair_offsets = tl.arange(0, 16)
+ row = block_id // blocks_per_row
+ block_col = block_id - row * blocks_per_row
+ even = tl.load(weight_ptr + row * cols + block_col * 32 + pair_offsets * 2).to(
+ tl.float32
+ )
+ odd = tl.load(weight_ptr + row * cols + block_col * 32 + pair_offsets * 2 + 1).to(
+ tl.float32
+ )
+ amax = tl.maximum(
+ tl.maximum(tl.max(tl.abs(even), axis=0), tl.max(tl.abs(odd), axis=0)),
+ 1.0e-4,
+ )
+ exponent = tl.ceil(tl.log2(amax / 6.0))
+ raw_scale = (exponent + 127.0).to(tl.uint8)
+ scale = tl.exp2(exponent)
+ scaled_even = even / scale
+ scaled_odd = odd / scale
+ abs_even = tl.abs(scaled_even)
+ abs_odd = tl.abs(scaled_odd)
+ code_even = tl.zeros((16,), dtype=tl.uint8)
+ code_odd = tl.zeros((16,), dtype=tl.uint8)
+ code_even += (abs_even > 0.25).to(tl.uint8)
+ code_even += (abs_even > 0.75).to(tl.uint8)
+ code_even += (abs_even > 1.25).to(tl.uint8)
+ code_even += (abs_even > 1.75).to(tl.uint8)
+ code_even += (abs_even > 2.5).to(tl.uint8)
+ code_even += (abs_even > 3.5).to(tl.uint8)
+ code_even += (abs_even > 5.0).to(tl.uint8)
+ code_even += ((scaled_even < 0.0) & (code_even != 0)).to(tl.uint8) * 8
+ code_odd += (abs_odd > 0.25).to(tl.uint8)
+ code_odd += (abs_odd > 0.75).to(tl.uint8)
+ code_odd += (abs_odd > 1.25).to(tl.uint8)
+ code_odd += (abs_odd > 1.75).to(tl.uint8)
+ code_odd += (abs_odd > 2.5).to(tl.uint8)
+ code_odd += (abs_odd > 3.5).to(tl.uint8)
+ code_odd += (abs_odd > 5.0).to(tl.uint8)
+ code_odd += ((scaled_odd < 0.0) & (code_odd != 0)).to(tl.uint8) * 8
+ tl.store(
+ packed_ptr + row * (cols // 2) + block_col * 16 + pair_offsets,
+ code_even | (code_odd << 4),
+ )
+ tl.store(scale_ptr + block_id, raw_scale, mask=block_id < total_blocks)
+
+
+def quant_mxfp4_cuda(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ if weight.ndim != 2 or weight.shape[1] % 32 != 0:
+ raise ValueError(
+ f"Expected 2-D MXFP4 weight with K % 32 == 0, got {weight.shape}."
+ )
+ if weight.device.type != "cuda":
+ raise ValueError("quant_mxfp4_cuda expects a CUDA tensor")
+ weight = weight.contiguous()
+ rows, cols = weight.shape
+ packed = torch.empty((rows, cols // 2), dtype=torch.uint8, device=weight.device)
+ scale_raw = torch.empty((rows, cols // 32), dtype=torch.uint8, device=weight.device)
+ blocks_per_row = cols // 32
+ total_blocks = rows * blocks_per_row
+ _mxfp4_quant_kernel[(total_blocks,)](
+ weight,
+ packed,
+ scale_raw,
+ total_blocks, # ty: ignore[invalid-argument-type]
+ cols, # ty: ignore[invalid-argument-type]
+ blocks_per_row, # ty: ignore[invalid-argument-type]
+ num_warps=1, # ty: ignore[unknown-argument]
+ )
+ return packed, scale_raw.view(torch.float8_e8m0fnu)
diff --git a/src/art/megatron/dsv4/encoding.py b/src/art/megatron/dsv4/encoding.py
new file mode 100644
index 000000000..6895771e2
--- /dev/null
+++ b/src/art/megatron/dsv4/encoding.py
@@ -0,0 +1,757 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa
+# fmt: off
+
+"""
+DeepSeek-V4 Encoding
+
+A self-contained implementation for encoding/decoding DeepSeek-V4 chat messages
+with tool calling, thinking mode, and quick instruction task support.
+"""
+
+from typing import Any, Dict, List, Union, Optional, Tuple
+import copy
+import json
+
+import regex as re
+
+# ============================================================
+# Special Tokens
+# ============================================================
+
+bos_token: str = "<｜begin▁of▁sentence｜>"
+eos_token: str = "<｜end▁of▁sentence｜>"
+thinking_start_token: str = ""
+thinking_end_token: str = ""
+dsml_token: str = "｜DSML｜"
+
+USER_SP_TOKEN = "<｜User｜>"
+ASSISTANT_SP_TOKEN = "<｜Assistant｜>"
+LATEST_REMINDER_SP_TOKEN = "<｜latest_reminder｜>"
+
+# Task special tokens for internal classification tasks
+DS_TASK_SP_TOKENS = {
+ "action": "<｜action｜>",
+ "query": "<｜query｜>",
+ "authority": "<｜authority｜>",
+ "domain": "<｜domain｜>",
+ "title": "<｜title｜>",
+ "read_url": "<｜read_url｜>",
+}
+VALID_TASKS = set(DS_TASK_SP_TOKENS.keys())
+
+# ============================================================
+# Templates
+# ============================================================
+
+system_msg_template: str = "{content}"
+user_msg_template: str = "{content}"
+latest_reminder_msg_template: str = "{content}"
+assistant_msg_template: str = "{reasoning}{content}{tool_calls}" + eos_token
+assistant_msg_wo_eos_template: str = "{reasoning}{content}{tool_calls}"
+thinking_template: str = "{reasoning}"
+
+response_format_template: str = (
+ "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}"
+)
+tool_call_template: str = (
+ "<{dsml_token}invoke name=\"{name}\">\n{arguments}\n"
+)
+tool_calls_template = (
+ "<{dsml_token}{tc_block_name}>\n{tool_calls}\n"
+)
+tool_calls_block_name: str = "tool_calls"
+
+tool_output_template: str = (
+ "{content}"
+)
+
+REASONING_EFFORT_MAX = (
+ "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n"
+ "You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n"
+ "Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n"
+)
+
+TOOLS_TEMPLATE = """## Tools
+
+You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following:
+
+<{dsml_token}tool_calls>
+<{dsml_token}invoke name="$TOOL_NAME">
+<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE
+...
+
+<{dsml_token}invoke name="$TOOL_NAME2">
+...
+
+
+
+String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
+
+If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response.
+
+Otherwise, output directly after {thinking_end_token} with tool calls or final response.
+
+### Available Tool Schemas
+
+{tool_schemas}
+
+You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
+"""
+
+# ============================================================
+# Utility Functions
+# ============================================================
+
+def to_json(value: Any) -> str:
+ """Serialize a value to JSON string."""
+ try:
+ return json.dumps(value, ensure_ascii=False)
+ except Exception:
+ return json.dumps(value, ensure_ascii=True)
+
+
+def tools_from_openai_format(tools):
+ """Extract function definitions from OpenAI-format tool list."""
+ return [tool["function"] for tool in tools]
+
+
+def tool_calls_from_openai_format(tool_calls):
+ """Convert OpenAI-format tool calls to internal format."""
+ return [
+ {
+ "name": tool_call["function"]["name"],
+ "arguments": tool_call["function"]["arguments"],
+ }
+ for tool_call in tool_calls
+ ]
+
+
+def tool_calls_to_openai_format(tool_calls):
+ """Convert internal tool calls to OpenAI format."""
+ return [
+ {
+ "type": "function",
+ "function": {
+ "name": tool_call["name"],
+ "arguments": tool_call["arguments"],
+ }
+ }
+ for tool_call in tool_calls
+ ]
+
+
+def encode_arguments_to_dsml(tool_call: Dict[str, Any]) -> str:
+ """
+ Encode tool call arguments into DSML parameter format.
+
+ Args:
+ tool_call: Dict with "name" and "arguments" keys.
+
+ Returns:
+ DSML-formatted parameter string.
+ """
+ p_dsml_template = '<{dsml_token}parameter name="{key}" string="{is_str}">{value}'
+ P_dsml_strs = []
+
+ if isinstance(tool_call["arguments"], str):
+ arguments = json.loads(tool_call["arguments"])
+ else:
+ arguments = tool_call["arguments"]
+
+ for k, v in arguments.items():
+ p_dsml_str = p_dsml_template.format(
+ dsml_token=dsml_token,
+ key=k,
+ is_str="true" if isinstance(v, str) else "false",
+ value=v if isinstance(v, str) else to_json(v),
+ )
+ P_dsml_strs.append(p_dsml_str)
+
+ return "\n".join(P_dsml_strs)
+
+
+def decode_dsml_to_arguments(tool_name: str, tool_args: Dict[str, Tuple[str, str]]) -> Dict[str, str]:
+ """
+ Decode DSML parameters back to a tool call dict.
+
+ Args:
+ tool_name: Name of the tool.
+ tool_args: Dict mapping param_name -> (value, is_string_flag).
+
+ Returns:
+ Dict with "name" and "arguments" (JSON string) keys.
+ """
+ def _decode_value(key: str, value: str, string: str):
+ if string == "true":
+ value = to_json(value)
+ return f"{to_json(key)}: {value}"
+
+ tool_args_json = "{" + ", ".join([_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]) + "}"
+ return dict(name=tool_name, arguments=tool_args_json)
+
+
+def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str:
+ """
+ Render tool schemas into the system prompt format.
+
+ Args:
+ tools: List of tool schema dicts (each with name, description, parameters).
+
+ Returns:
+ Formatted tools section string.
+ """
+ tools_json = [to_json(t) for t in tools]
+
+ return TOOLS_TEMPLATE.format(
+ tool_schemas="\n".join(tools_json),
+ dsml_token=dsml_token,
+ thinking_start_token=thinking_start_token,
+ thinking_end_token=thinking_end_token,
+ )
+
+
+def find_last_user_index(messages: List[Dict[str, Any]]) -> int:
+ """Find the index of the last user/developer message."""
+ last_user_index = -1
+ for idx in range(len(messages) - 1, -1, -1):
+ if messages[idx].get("role") in ["user", "developer"]:
+ last_user_index = idx
+ break
+ return last_user_index
+
+
+# ============================================================
+# Message Rendering
+# ============================================================
+
+def render_message(index: int, messages: List[Dict[str, Any]], thinking_mode: str, drop_thinking: bool = True, reasoning_effort: Optional[str] = None) -> str:
+ """
+ Render a single message at the given index into its encoded string form.
+
+ This is the core function that converts each message in the conversation
+ into the DeepSeek-V4 format.
+
+ Args:
+ index: Index of the message to render.
+ messages: Full list of messages in the conversation.
+ thinking_mode: Either "chat" or "thinking".
+ drop_thinking: Whether to drop reasoning content from earlier turns.
+ reasoning_effort: Optional reasoning effort level ("max", "high", or None).
+
+ Returns:
+ Encoded string for this message.
+ """
+ assert 0 <= index < len(messages)
+ assert thinking_mode in ["chat", "thinking"], f"Invalid thinking_mode `{thinking_mode}`"
+
+ prompt = ""
+ msg = messages[index]
+ last_user_idx = find_last_user_index(messages)
+
+ role = msg.get("role")
+ content = msg.get("content")
+ tools = msg.get("tools")
+ response_format = msg.get("response_format")
+ tool_calls = msg.get("tool_calls")
+ reasoning = msg.get("reasoning")
+ wo_eos = msg.get("wo_eos", False)
+
+ if tools:
+ tools = tools_from_openai_format(tools)
+ if tool_calls:
+ tool_calls = tool_calls_from_openai_format(tool_calls)
+
+ # Reasoning effort prefix (only at index 0 in thinking mode with max effort)
+ assert reasoning_effort in ['max', None, 'high'], f"Invalid reasoning effort: {reasoning_effort}"
+ if index == 0 and thinking_mode == "thinking" and reasoning_effort == 'max':
+ prompt += REASONING_EFFORT_MAX
+
+ if role == "system":
+ prompt += system_msg_template.format(content=content or "")
+ if tools:
+ prompt += "\n\n" + render_tools(tools)
+ if response_format:
+ prompt += "\n\n" + response_format_template.format(schema=to_json(response_format))
+
+ elif role == "developer":
+ assert content, f"Invalid message for role `{role}`: {msg}"
+
+ content_developer = USER_SP_TOKEN
+ content_developer += content
+
+ if tools:
+ content_developer += "\n\n" + render_tools(tools)
+ if response_format:
+ content_developer += "\n\n" + response_format_template.format(schema=to_json(response_format))
+
+ prompt += user_msg_template.format(content=content_developer)
+
+ elif role == "user":
+ prompt += USER_SP_TOKEN
+
+ # Handle content blocks (tool results mixed with text)
+ content_blocks = msg.get("content_blocks")
+ if content_blocks:
+ parts = []
+ for block in content_blocks:
+ block_type = block.get("type")
+ if block_type == "text":
+ parts.append(block.get("text", ""))
+ elif block_type == "tool_result":
+ tool_content = block.get("content", "")
+ if isinstance(tool_content, list):
+ text_parts = []
+ for b in tool_content:
+ if b.get("type") == "text":
+ text_parts.append(b.get("text", ""))
+ else:
+ text_parts.append(f"[Unsupported {b.get('type')}]")
+ tool_content = "\n\n".join(text_parts)
+ parts.append(tool_output_template.format(content=tool_content))
+ else:
+ parts.append(f"[Unsupported {block_type}]")
+ prompt += "\n\n".join(parts)
+ else:
+ prompt += content or ""
+
+ elif role == "latest_reminder":
+ prompt += LATEST_REMINDER_SP_TOKEN + latest_reminder_msg_template.format(content=content)
+
+ elif role == "tool":
+ raise NotImplementedError("deepseek_v4 merges tool messages into user; please preprocess with merge_tool_messages()")
+
+ elif role == "assistant":
+ thinking_part = ""
+ tc_content = ""
+
+ if tool_calls:
+ tc_list = [
+ tool_call_template.format(
+ dsml_token=dsml_token,
+ name=tc.get("name"),
+ arguments=encode_arguments_to_dsml(tc)
+ )
+ for tc in tool_calls
+ ]
+ tc_content += '\n\n' + tool_calls_template.format(
+ dsml_token=dsml_token,
+ tool_calls="\n".join(tc_list),
+ tc_block_name=tool_calls_block_name,
+ )
+
+ summary_content = content or ""
+ reasoning = reasoning or ""
+
+ # Check if previous message has a task - if so, this is a task output (no thinking)
+ prev_has_task = index - 1 >= 0 and messages[index - 1].get("task") is not None
+
+ if thinking_mode == "thinking" and not prev_has_task:
+ if not drop_thinking or index > last_user_idx:
+ thinking_part = thinking_template.format(reasoning=reasoning) + thinking_end_token
+ else:
+ thinking_part = ""
+
+ if wo_eos:
+ prompt += assistant_msg_wo_eos_template.format(
+ reasoning=thinking_part,
+ content=summary_content,
+ tool_calls=tc_content,
+ )
+ else:
+ prompt += assistant_msg_template.format(
+ reasoning=thinking_part,
+ content=summary_content,
+ tool_calls=tc_content,
+ )
+ else:
+ raise NotImplementedError(f"Unknown role: {role}")
+
+ # Append transition tokens based on what follows
+ if index + 1 < len(messages) and messages[index + 1].get("role") not in ["assistant", "latest_reminder"]:
+ return prompt
+
+ task = messages[index].get("task")
+ if task is not None:
+ # Task special token for internal classification tasks
+ assert task in VALID_TASKS, f"Invalid task: '{task}'. Valid tasks are: {list(VALID_TASKS)}"
+ task_sp_token = DS_TASK_SP_TOKENS[task]
+
+ if task != "action":
+ # Non-action tasks: append task sp token directly after the message
+ prompt += task_sp_token
+ else:
+ # Action task: append Assistant + thinking token + action sp token
+ prompt += ASSISTANT_SP_TOKEN
+ prompt += thinking_end_token if thinking_mode != "thinking" else thinking_start_token
+ prompt += task_sp_token
+
+ elif messages[index].get("role") in ["user", "developer"]:
+ # Normal generation: append Assistant + thinking token
+ prompt += ASSISTANT_SP_TOKEN
+ if not drop_thinking and thinking_mode == "thinking":
+ prompt += thinking_start_token
+ elif drop_thinking and thinking_mode == "thinking" and index >= last_user_idx:
+ prompt += thinking_start_token
+ else:
+ prompt += thinking_end_token
+
+ return prompt
+
+
+# ============================================================
+# Preprocessing
+# ============================================================
+
+def merge_tool_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+ """
+ Merge tool messages into the preceding user message using content_blocks format.
+
+ DeepSeek-V4 does not have a standalone "tool" role; instead, tool results
+ are encoded as blocks within user messages.
+
+ This function converts a standard OpenAI-format conversation (with separate
+ "tool" role messages) into V4 format where tool results are merged into
+ user messages.
+
+ Args:
+ messages: List of message dicts in OpenAI format.
+
+ Returns:
+ Processed message list with tool messages merged into user messages.
+ """
+ merged: List[Dict[str, Any]] = []
+
+ for msg in messages:
+ msg = copy.deepcopy(msg)
+ role = msg.get("role")
+
+ if role == "tool":
+ # Convert tool message to a user message with tool_result block
+ tool_block = {
+ "type": "tool_result",
+ "tool_use_id": msg.get("tool_call_id", ""),
+ "content": msg.get("content", ""),
+ }
+ # Merge into previous message if it's already a user (merged tool)
+ if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1]:
+ merged[-1]["content_blocks"].append(tool_block)
+ else:
+ merged.append({
+ "role": "user",
+ "content_blocks": [tool_block],
+ })
+ elif role == "user":
+ text_block = {"type": "text", "text": msg.get("content", "")}
+ if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1] and merged[-1].get("task") is None:
+ merged[-1]["content_blocks"].append(text_block)
+ else:
+ new_msg = {
+ "role": "user",
+ "content": msg.get("content", ""),
+ "content_blocks": [text_block],
+ }
+ # Preserve extra fields (task, wo_eos, mask, etc.)
+ for key in ("task", "wo_eos", "mask"):
+ if key in msg:
+ new_msg[key] = msg[key]
+ merged.append(new_msg)
+ else:
+ merged.append(msg)
+
+ return merged
+
+
+def sort_tool_results_by_call_order(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+ """
+ Sort tool_result blocks within user messages by the order of tool_calls
+ in the preceding assistant message.
+
+ Args:
+ messages: Preprocessed message list (after merge_tool_messages).
+
+ Returns:
+ Message list with sorted tool result blocks.
+ """
+ last_tool_call_order: Dict[str, int] = {}
+
+ for msg in messages:
+ role = msg.get("role")
+ if role == "assistant" and msg.get("tool_calls"):
+ last_tool_call_order = {}
+ for idx, tc in enumerate(msg["tool_calls"]):
+ tc_id = tc.get("id") or tc.get("function", {}).get("id", "")
+ if tc_id:
+ last_tool_call_order[tc_id] = idx
+
+ elif role == "user" and msg.get("content_blocks"):
+ tool_blocks = [b for b in msg["content_blocks"] if b.get("type") == "tool_result"]
+ if len(tool_blocks) > 1 and last_tool_call_order:
+ sorted_blocks = sorted(
+ tool_blocks,
+ key=lambda b: last_tool_call_order.get(b.get("tool_use_id", ""), 0)
+ )
+ sorted_idx = 0
+ new_blocks = []
+ for block in msg["content_blocks"]:
+ if block.get("type") == "tool_result":
+ new_blocks.append(sorted_blocks[sorted_idx])
+ sorted_idx += 1
+ else:
+ new_blocks.append(block)
+ msg["content_blocks"] = new_blocks
+
+ return messages
+
+
+# ============================================================
+# Main Encoding Function
+# ============================================================
+
+def encode_messages(
+ messages: List[Dict[str, Any]],
+ thinking_mode: str,
+ context: Optional[List[Dict[str, Any]]] = None,
+ drop_thinking: bool = True,
+ add_default_bos_token: bool = True,
+ reasoning_effort: Optional[str] = None,
+) -> str:
+ """
+ Encode a list of messages into the DeepSeek-V4 prompt format.
+
+ This is the main entry point for encoding conversations. It handles:
+ - BOS token insertion
+ - Thinking mode with optional reasoning content dropping
+ - Tool message merging into user messages
+ - Multi-turn conversation context
+
+ Args:
+ messages: List of message dicts to encode.
+ thinking_mode: Either "chat" or "thinking".
+ context: Optional preceding context messages (already encoded prefix).
+ drop_thinking: If True, drop reasoning from earlier assistant turns
+ (only keep reasoning for messages after the last user message).
+ add_default_bos_token: Whether to prepend BOS token at conversation start.
+ reasoning_effort: Optional reasoning effort level ("max", "high", or None).
+
+ Returns:
+ The encoded prompt string.
+ """
+ context = context if context else []
+
+ # Preprocess: merge tool messages and sort tool results
+ messages = merge_tool_messages(messages)
+ messages = sort_tool_results_by_call_order(context + messages)[len(context):]
+ if context:
+ context = merge_tool_messages(context)
+ context = sort_tool_results_by_call_order(context)
+
+ full_messages = context + messages
+
+ prompt = bos_token if add_default_bos_token and len(context) == 0 else ""
+
+ # Resolve drop_thinking: if any message has tools defined, don't drop thinking
+ effective_drop_thinking = drop_thinking
+ if any(m.get("tools") for m in full_messages):
+ effective_drop_thinking = False
+
+ if thinking_mode == "thinking" and effective_drop_thinking:
+ full_messages = _drop_thinking_messages(full_messages)
+ # After dropping, recalculate how many messages to render
+ # (context may have shrunk too)
+ num_to_render = len(full_messages) - len(_drop_thinking_messages(context))
+ context_len = len(full_messages) - num_to_render
+ else:
+ num_to_render = len(messages)
+ context_len = len(context)
+
+ for idx in range(num_to_render):
+ prompt += render_message(
+ idx + context_len,
+ full_messages,
+ thinking_mode=thinking_mode,
+ drop_thinking=effective_drop_thinking,
+ reasoning_effort=reasoning_effort,
+ )
+
+ return prompt
+
+
+def _drop_thinking_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+ """
+ Drop reasoning and non-essential messages before the last user message.
+
+ Behavior:
+ - Messages with role in ["user", "system", "tool", "latest_reminder"] are always kept.
+ - Messages at or after the last user index are always kept.
+ - Assistant messages before the last user get reasoning removed.
+ - Developer messages before the last user are dropped entirely.
+ """
+ last_user_idx = find_last_user_index(messages)
+ result = []
+ keep_roles = {"user", "system", "tool", "latest_reminder", "direct_search_results"}
+
+ for idx, msg in enumerate(messages):
+ role = msg.get("role")
+ if role in keep_roles or idx >= last_user_idx:
+ result.append(msg)
+ elif role == "assistant":
+ msg = copy.copy(msg)
+ msg.pop("reasoning", None)
+ result.append(msg)
+ # developer and other roles before last_user_idx are dropped
+
+ return result
+
+
+# ============================================================
+# Parsing (Decoding model output)
+# ============================================================
+
+def _read_until_stop(index: int, text: str, stop: List[str]) -> Tuple[int, str, Optional[str]]:
+ """
+ Read text from index until one of the stop strings is found.
+
+ Returns:
+ Tuple of (new_index, content_before_stop, matched_stop_string_or_None).
+ """
+ min_pos = len(text)
+ matched_stop = None
+
+ for s in stop:
+ pos = text.find(s, index)
+ if pos != -1 and pos < min_pos:
+ min_pos = pos
+ matched_stop = s
+
+ if matched_stop:
+ content = text[index:min_pos]
+ return min_pos + len(matched_stop), content, matched_stop
+ else:
+ content = text[index:]
+ return len(text), content, None
+
+
+def parse_tool_calls(index: int, text: str) -> Tuple[int, Optional[str], List[Dict[str, str]]]:
+ """
+ Parse DSML tool calls from text starting at the given index.
+
+ Args:
+ index: Starting position in text.
+ text: The full text to parse.
+
+ Returns:
+ Tuple of (new_index, last_stop_token, list_of_tool_call_dicts).
+ Each tool call dict has "name" and "arguments" keys.
+ """
+ tool_calls: List[Dict[str, Any]] = []
+ stop_token = None
+ tool_calls_end_token = f""
+
+ while index < len(text):
+ index, content_before, stop_token = _read_until_stop(index, text, [f"<{dsml_token}invoke", tool_calls_end_token])
+ if content_before != ">\n":
+ raise ValueError(f"Tool call format error: expected '>\\n' but got '{content_before}'")
+
+ if stop_token == tool_calls_end_token:
+ break
+
+ if stop_token is None:
+ raise ValueError("Missing special token in tool calls")
+
+ index, tool_name_content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"\n$', tool_name_content, flags=re.DOTALL)
+ if len(p_tool_name) != 1:
+ raise ValueError(f"Tool name format error: '{tool_name_content}'")
+ tool_name = p_tool_name[0]
+
+ tool_args: Dict[str, Tuple[str, str]] = {}
+ while stop_token == f"<{dsml_token}parameter":
+ index, param_content, stop_token = _read_until_stop(index, text, [f"/{dsml_token}parameter"])
+
+ param_kv = re.findall(r'^ name="(.*?)" string="(true|false)">(.*?)<$', param_content, flags=re.DOTALL)
+ if len(param_kv) != 1:
+ raise ValueError(f"Parameter format error: '{param_content}'")
+ param_name, string, param_value = param_kv[0]
+
+ if param_name in tool_args:
+ raise ValueError(f"Duplicate parameter name: '{param_name}'")
+ tool_args[param_name] = (param_value, string)
+
+ index, content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"\n":
+ raise ValueError(f"Parameter format error: expected '>\\n' but got '{content}'")
+
+ tool_call = decode_dsml_to_arguments(tool_name=tool_name, tool_args=tool_args)
+ tool_calls.append(tool_call)
+
+ return index, stop_token, tool_calls
+
+
+def parse_message_from_completion_text(text: str, thinking_mode: str) -> Dict[str, Any]:
+ """
+ Parse a model completion text into a structured assistant message.
+
+ This function takes the raw text output from the model (a single assistant turn)
+ and extracts:
+ - reasoning (thinking block)
+ - content (summary/response)
+ - tool_calls (if any)
+
+ NOTE: This function is designed to parse only correctly formatted strings and
+ will raise ValueError for malformed output.
+
+ Args:
+ text: The raw completion text (including EOS token).
+ thinking_mode: Either "chat" or "thinking".
+
+ Returns:
+ Dict with keys: "role", "content", "reasoning", "tool_calls".
+ tool_calls are in OpenAI format.
+ """
+ summary_content, reasoning = "", ""
+ tool_calls: List[Dict[str, str]] = []
+ index, stop_token = 0, None
+ tool_calls_start_token = f"\n\n<{dsml_token}{tool_calls_block_name}"
+
+ is_thinking = thinking_mode == "thinking"
+ is_tool_calling = False
+
+ if is_thinking:
+ index, content_delta, stop_token = _read_until_stop(index, text, [thinking_end_token, tool_calls_start_token])
+ reasoning = content_delta
+ if stop_token != thinking_end_token:
+ raise ValueError("Invalid thinking format: missing ")
+
+ index, content_delta, stop_token = _read_until_stop(index, text, [eos_token, tool_calls_start_token])
+ summary_content = content_delta
+ if stop_token == tool_calls_start_token:
+ is_tool_calling = True
+ else:
+ if stop_token != eos_token:
+ raise ValueError("Invalid format: missing EOS token")
+
+ if is_tool_calling:
+ index, stop_token, tool_calls = parse_tool_calls(index, text)
+
+ index, tool_ends_text, stop_token = _read_until_stop(index, text, [eos_token])
+ if tool_ends_text:
+ raise ValueError("Unexpected content after tool calls")
+
+ if len(text) != index or stop_token not in [eos_token, None]:
+ raise ValueError("Unexpected content at end")
+
+ for sp_token in [bos_token, eos_token, thinking_start_token, thinking_end_token, dsml_token]:
+ if sp_token in summary_content or sp_token in reasoning:
+ raise ValueError(f"Unexpected special token '{sp_token}' in content")
+
+ return {
+ "role": "assistant",
+ "content": summary_content,
+ "reasoning": reasoning,
+ "tool_calls": tool_calls_to_openai_format(tool_calls)
+ }
+
+# fmt: on
diff --git a/src/art/megatron/dsv4/hf_config.py b/src/art/megatron/dsv4/hf_config.py
new file mode 100644
index 000000000..75f6a0b63
--- /dev/null
+++ b/src/art/megatron/dsv4/hf_config.py
@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import sys
+from typing import Any
+
+import transformers
+
+_TORCHVISION_LIB: Any | None = None
+
+
+class DeepseekV4ForCausalLM:
+ """Bridge-dispatch marker used before native HF modeling is imported."""
+
+ def __init__(self, *args: Any, **kwargs: Any) -> None:
+ del args, kwargs
+ raise RuntimeError(
+ "DeepSeek-V4 requires native Transformers DeepseekV4ForCausalLM."
+ )
+
+
+def _add_marker_to_transformers_module(module: Any, model_class: type) -> None:
+ if module is None:
+ return
+ objects = getattr(module, "_objects", None)
+ if isinstance(objects, dict):
+ objects["DeepseekV4ForCausalLM"] = model_class
+ setattr(module, "DeepseekV4ForCausalLM", model_class)
+
+
+def _ensure_transformers_marker(model_class: type | None = None) -> None:
+ model_class = model_class or DeepseekV4ForCausalLM
+ _add_marker_to_transformers_module(transformers, model_class)
+ auto_bridge = sys.modules.get("megatron.bridge.models.conversion.auto_bridge")
+ if auto_bridge is not None:
+ _add_marker_to_transformers_module(
+ getattr(auto_bridge, "transformers", None), model_class
+ )
+
+
+def _ensure_torchvision_nms_schema() -> None:
+ global _TORCHVISION_LIB
+ if _TORCHVISION_LIB is not None:
+ return
+ import torch
+
+ try:
+ _TORCHVISION_LIB = torch.library.Library("torchvision", "DEF")
+ _TORCHVISION_LIB.define(
+ "nms(Tensor dets, Tensor scores, float iou_threshold) -> Tensor"
+ )
+ except RuntimeError as exc:
+ if "Only a single TORCH_LIBRARY" not in str(exc) and "already" not in str(exc):
+ raise
+ _TORCHVISION_LIB = torch.library.Library("torchvision", "FRAGMENT")
+ try:
+ _TORCHVISION_LIB.define(
+ "nms(Tensor dets, Tensor scores, float iou_threshold) -> Tensor"
+ )
+ except RuntimeError as define_exc:
+ if "already" not in str(define_exc):
+ raise
+
+
+def _native_dsv4_config_class() -> type:
+ try:
+ from transformers.models.deepseek_v4.configuration_deepseek_v4 import (
+ DeepseekV4Config,
+ )
+ except ModuleNotFoundError as exc:
+ raise RuntimeError(
+ "DeepSeek-V4 requires transformers with native deepseek_v4 support."
+ ) from exc
+ return DeepseekV4Config
+
+
+def _native_dsv4_model_class() -> type:
+ _ensure_torchvision_nms_schema()
+ try:
+ from transformers.models.deepseek_v4.modeling_deepseek_v4 import (
+ DeepseekV4ForCausalLM as NativeDeepseekV4ForCausalLM,
+ )
+ except ModuleNotFoundError as exc:
+ raise RuntimeError(
+ "DeepSeek-V4 requires transformers with native DeepseekV4ForCausalLM."
+ ) from exc
+ return NativeDeepseekV4ForCausalLM
+
+
+def ensure_dsv4_hf_config_registered() -> None:
+ _native_dsv4_config_class()
+ _ensure_transformers_marker()
+
+
+def ensure_dsv4_hf_model_registered() -> None:
+ _native_dsv4_config_class()
+ _ensure_transformers_marker(_native_dsv4_model_class())
+
+
+__all__ = [
+ "DeepseekV4ForCausalLM",
+ "ensure_dsv4_hf_config_registered",
+ "ensure_dsv4_hf_model_registered",
+]
diff --git a/src/art/megatron/dsv4/hyper_connection.py b/src/art/megatron/dsv4/hyper_connection.py
new file mode 100644
index 000000000..abe37397f
--- /dev/null
+++ b/src/art/megatron/dsv4/hyper_connection.py
@@ -0,0 +1,156 @@
+from typing import Any, cast
+
+import einops
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.transformer_config import TransformerConfig
+import torch
+from torch import Tensor
+import torch.nn.functional as F
+
+
+def _unweighted_rms_norm(x: Tensor, eps: float) -> Tensor:
+ return x * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + eps).to(x.dtype)
+
+
+class HCHeadParams(MegatronModule):
+ def __init__(self, config: TransformerConfig):
+ super().__init__(config)
+ cfg = cast(Any, config)
+ hc_mult = int(cfg.dsv4_hc_mult)
+ hc_dim = hc_mult * config.hidden_size
+ self.hc_head_fn = torch.nn.Parameter(
+ torch.empty(hc_mult, hc_dim, dtype=torch.float32)
+ )
+ self.hc_head_base = torch.nn.Parameter(
+ torch.empty(hc_mult, dtype=torch.float32)
+ )
+ self.hc_head_scale = torch.nn.Parameter(torch.empty(1, dtype=torch.float32))
+ self._keep_fp32_parameters = (
+ "hc_head_fn",
+ "hc_head_base",
+ "hc_head_scale",
+ )
+ for param in (self.hc_head_fn, self.hc_head_base, self.hc_head_scale):
+ setattr(param, "_keep_fp32", True)
+
+ def forward(self):
+ raise NotImplementedError
+
+
+class DeepSeekV4HyperConnectionUtil:
+ """DeepSeek-V4 manifold-constrained hyper-connection math.
+
+ This implements the HF reference equations directly in PyTorch. TileKernels
+ MHC currently requires a newer CUDA toolchain than this ART Megatron env
+ provides, so production training keeps the exact eager math here.
+ """
+
+ def __init__(self, config: TransformerConfig):
+ cfg = cast(Any, config)
+ self.norm_eps = config.layernorm_epsilon
+ self.hc_mult = int(cfg.dsv4_hc_mult)
+ self.hc_sinkhorn_iters = int(cfg.dsv4_hc_sinkhorn_iters)
+ self.hc_eps = float(cfg.dsv4_hc_eps)
+
+ def hc_pre_raw(
+ self,
+ x: Tensor,
+ hc_fn: Tensor,
+ hc_scale: Tensor,
+ hc_base: Tensor,
+ ) -> tuple[Tensor, Tensor, Tensor]:
+ dtype = x.dtype
+ hc = self.hc_mult
+ flat = _unweighted_rms_norm(x.flatten(start_dim=2).float(), self.norm_eps)
+ pre_w, post_w, comb_w = F.linear(flat, hc_fn.float()).split(
+ [hc, hc, hc * hc], dim=-1
+ )
+ pre_b, post_b, comb_b = hc_base.float().split([hc, hc, hc * hc])
+ pre_scale, post_scale, comb_scale = hc_scale.float().unbind(0)
+
+ pre = torch.sigmoid(pre_w * pre_scale + pre_b) + self.hc_eps
+ post = 2 * torch.sigmoid(post_w * post_scale + post_b)
+ comb_logits = comb_w.view(
+ *comb_w.shape[:-1], hc, hc
+ ) * comb_scale + comb_b.view(hc, hc)
+ comb = torch.softmax(comb_logits, dim=-1) + self.hc_eps
+ comb = comb / (comb.sum(dim=-2, keepdim=True) + self.hc_eps)
+ for _ in range(self.hc_sinkhorn_iters - 1):
+ comb = comb / (comb.sum(dim=-1, keepdim=True) + self.hc_eps)
+ comb = comb / (comb.sum(dim=-2, keepdim=True) + self.hc_eps)
+ layer_input = (pre.unsqueeze(-1) * x).sum(dim=2).to(dtype)
+ return layer_input, post, comb
+
+ def hc_post_raw(
+ self,
+ x: Tensor,
+ residual: Tensor,
+ post: Tensor,
+ comb: Tensor,
+ ) -> Tensor:
+ dtype = residual.dtype
+ return post.to(dtype).unsqueeze(-1) * x.unsqueeze(-2) + torch.matmul(
+ comb.to(dtype).transpose(-1, -2), residual
+ )
+
+ def hc_head_raw(
+ self,
+ x: Tensor,
+ hc_fn: Tensor,
+ hc_scale: Tensor,
+ hc_base: Tensor,
+ ) -> Tensor:
+ dtype = x.dtype
+ flat = _unweighted_rms_norm(x.flatten(start_dim=2).float(), self.norm_eps)
+ mixes = F.linear(flat, hc_fn.float())
+ pre = (
+ torch.sigmoid(mixes * hc_scale.float().reshape(1) + hc_base.float())
+ + self.hc_eps
+ )
+ return (pre.unsqueeze(-1) * x).sum(dim=2).to(dtype)
+
+ def layer_pre(
+ self,
+ hidden_states: Tensor,
+ hc_fn: Tensor,
+ hc_scale: Tensor,
+ hc_base: Tensor,
+ ) -> tuple[Tensor, Tensor, Tensor]:
+ x = einops.rearrange(hidden_states, "s b hc d -> b s hc d")
+ x, post, comb = self.hc_pre_raw(
+ x=x, hc_fn=hc_fn, hc_scale=hc_scale, hc_base=hc_base
+ )
+ return einops.rearrange(x, "b s d -> s b d"), post, comb
+
+ def layer_post(
+ self,
+ output_with_bias: Tensor | tuple[Tensor, Tensor | None],
+ residual: Tensor,
+ post: Tensor,
+ comb: Tensor,
+ ) -> Tensor:
+ if isinstance(output_with_bias, tuple):
+ out, bias = output_with_bias
+ assert bias is None
+ else:
+ out = output_with_bias
+ out = einops.rearrange(out, "s b d -> b s d")
+ residual_bshd = einops.rearrange(residual, "s b hc d -> b s hc d")
+ hidden_states = self.hc_post_raw(
+ x=out, residual=residual_bshd, post=post, comb=comb
+ )
+ return einops.rearrange(hidden_states, "b s hc d -> s b hc d")
+
+ def block_expand(self, hidden_states: Tensor) -> Tensor:
+ return einops.repeat(hidden_states, "s b d -> s b hc d", hc=self.hc_mult)
+
+ def block_head(
+ self,
+ hidden_states: Tensor,
+ hc_fn: Tensor,
+ hc_scale: Tensor,
+ hc_base: Tensor,
+ ) -> Tensor:
+ x = einops.rearrange(hidden_states, "s b hc d -> b s hc d")
+ x = self.hc_head_raw(x=x, hc_fn=hc_fn, hc_scale=hc_scale, hc_base=hc_base)
+ return einops.rearrange(x, "b s d -> s b d")
diff --git a/src/art/megatron/dsv4/kernel/__init__.py b/src/art/megatron/dsv4/kernel/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/src/art/megatron/dsv4/kernel/precision_aligned_ops.py b/src/art/megatron/dsv4/kernel/precision_aligned_ops.py
new file mode 100644
index 000000000..5cdb6d1c9
--- /dev/null
+++ b/src/art/megatron/dsv4/kernel/precision_aligned_ops.py
@@ -0,0 +1,46 @@
+from typing import Any
+
+import torch
+
+
+class _BFloat16LinearFP32Func(torch.autograd.Function):
+ # Forward matches SGLang's default DeepSeek-V4 compressor path
+ # (`sglang.jit_kernel.deepseek_v4.linear_bf16_fp32`, cublas backend):
+ # BF16 activation x BF16 weight -> FP32 output. This keeps Megatron's
+ # compressor log-prob computation aligned with SGLang rollout. Backward is
+ # only needed for training, so keep its gradient matmuls in FP32.
+ @staticmethod
+ def forward(ctx, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+ x_bf16 = x.to(torch.bfloat16)
+ weight_bf16 = weight.to(torch.bfloat16)
+ ctx.save_for_backward(x_bf16, weight_bf16)
+ ctx.input_shape = x.shape
+ ctx.input_dtype = x.dtype
+ ctx.weight_dtype = weight.dtype
+
+ x_2d = x_bf16.reshape(-1, x_bf16.shape[-1])
+ out = torch.mm(x_2d, weight_bf16.t(), out_dtype=torch.float32)
+ return out.view(*x.shape[:-1], weight_bf16.shape[0])
+
+ @staticmethod
+ def backward(ctx: Any, *grad_outputs: Any):
+ grad_output = grad_outputs[0]
+ x_bf16, weight_bf16 = ctx.saved_tensors
+ grad_output_2d = grad_output.reshape(-1, grad_output.shape[-1]).float()
+
+ grad_x = None
+ if ctx.needs_input_grad[0]:
+ grad_x = grad_output_2d.matmul(weight_bf16.float())
+ grad_x = grad_x.view(ctx.input_shape).to(ctx.input_dtype)
+
+ grad_weight = None
+ if ctx.needs_input_grad[1]:
+ x_2d = x_bf16.reshape(-1, x_bf16.shape[-1])
+ grad_weight = grad_output_2d.t().matmul(x_2d.float()).to(ctx.weight_dtype)
+
+ return grad_x, grad_weight
+
+
+@torch.compiler.disable
+def linear_bf16_fp32(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+ return _BFloat16LinearFP32Func.apply(x, weight)
diff --git a/src/art/megatron/dsv4/kernel/tilelang_import.py b/src/art/megatron/dsv4/kernel/tilelang_import.py
new file mode 100644
index 000000000..2f35faef1
--- /dev/null
+++ b/src/art/megatron/dsv4/kernel/tilelang_import.py
@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+import importlib
+import os
+from typing import Any
+
+_TILELANG_ENV_KEYS = (
+ "PYTHONPATH",
+ "TVM_IMPORT_PYTHON_PATH",
+ "TVM_LIBRARY_PATH",
+ "TL_CUTLASS_PATH",
+ "TL_TEMPLATE_PATH",
+ "TL_COMPOSABLE_KERNEL_PATH",
+)
+
+_TILELANG_PATH_MARKERS = ("/site-packages/tilelang/", "\\site-packages\\tilelang\\")
+
+
+def _drop_tilelang_paths(value: str | None) -> str | None:
+ if value is None:
+ return None
+ kept = [
+ part
+ for part in value.split(os.pathsep)
+ if not any(marker in part for marker in _TILELANG_PATH_MARKERS)
+ ]
+ return os.pathsep.join(kept) if kept else None
+
+
+def sanitize_tilelang_env() -> None:
+ """Remove TileLang's vendored TVM paths from env inherited by child processes."""
+ for key in _TILELANG_ENV_KEYS:
+ value = _drop_tilelang_paths(os.environ.get(key))
+ if value is None:
+ os.environ.pop(key, None)
+ else:
+ os.environ[key] = value
+
+
+def _restore_env(saved: dict[str, str | None]) -> None:
+ for key, value in saved.items():
+ if value is None:
+ os.environ.pop(key, None)
+ else:
+ os.environ[key] = value
+ sanitize_tilelang_env()
+
+
+@contextmanager
+def preserve_tilelang_env() -> Iterator[None]:
+ saved = {key: os.environ.get(key) for key in _TILELANG_ENV_KEYS}
+ try:
+ yield
+ finally:
+ _restore_env(saved)
+
+
+def import_tilelang() -> tuple[Any, Any]:
+ """Import TileLang without leaking its vendored TVM paths to child processes."""
+ with preserve_tilelang_env():
+ tilelang = importlib.import_module("tilelang")
+ language = importlib.import_module("tilelang.language")
+ return tilelang, language
diff --git a/src/art/megatron/dsv4/kernel/tilelang_indexer.py b/src/art/megatron/dsv4/kernel/tilelang_indexer.py
new file mode 100644
index 000000000..87018deba
--- /dev/null
+++ b/src/art/megatron/dsv4/kernel/tilelang_indexer.py
@@ -0,0 +1,116 @@
+# ruff: noqa
+"""TileLang-based DSA Indexer for DeepSeek-V4.
+
+Adapts GLM-5's lighting_indexer to V4's SBHD data layout and causal masking.
+Provides both a low-level per-sample interface and a batched autograd Function.
+"""
+
+from typing import Any
+
+import torch
+
+from art.megatron.dsv4.kernel.tilelang_import import preserve_tilelang_env
+
+
+def pytorch_extract_topk_scores(logits, topk_indices, dim=-1):
+ valid_mask = topk_indices != -1
+ safe_indices = topk_indices.clamp(min=0).to(torch.int64)
+ scores = torch.gather(logits, dim=dim, index=safe_indices)
+ scores = torch.where(valid_mask, scores, float("-inf"))
+ return scores
+
+
+class V4IndexerFunction(torch.autograd.Function):
+ """Autograd function for V4 tilelang indexer.
+
+ Inputs are in V4's native SBHD layout:
+ q: [seqlen, batch, heads, dim] bf16
+ k: [seqlen_kv, batch, dim] bf16
+ weights: [seqlen, batch, heads] fp32
+ """
+
+ @staticmethod
+ def forward(
+ ctx,
+ index_q: torch.Tensor,
+ index_k: torch.Tensor,
+ weights: torch.Tensor,
+ compress_ratio: int,
+ topk: int,
+ topk_indices: torch.Tensor | None = None,
+ ):
+ with preserve_tilelang_env():
+ from art.megatron.dsv4.kernel.tilelang_indexer_fwd import (
+ _make_causal_cu_seqlens,
+ batched_indexer_fwd,
+ )
+
+ seqlen_q = index_q.shape[0]
+ seq_len_kv = index_k.shape[0]
+
+ cu_seqlen_ks, cu_seqlen_ke = _make_causal_cu_seqlens(
+ seqlen_q, seq_len_kv, compress_ratio, index_q.device
+ )
+
+ # [batch, seqlen, seqlen_kv]
+ logits = batched_indexer_fwd(
+ index_q, index_k, weights, cu_seqlen_ks, cu_seqlen_ke
+ )
+
+ if topk_indices is None:
+ actual_topk = min(topk, seq_len_kv)
+ index_score, topk_indices = torch.topk(logits, actual_topk, dim=-1)
+ topk_indices = topk_indices.to(torch.int32)
+ topk_indices = topk_indices.masked_fill(index_score == -torch.inf, -1)
+
+ index_score = pytorch_extract_topk_scores(logits, topk_indices)
+
+ ctx.save_for_backward(
+ index_q, index_k, weights, cu_seqlen_ks, cu_seqlen_ke, topk_indices
+ )
+ ctx.compress_ratio = compress_ratio
+ ctx.topk = topk
+ return index_score, topk_indices
+
+ @staticmethod
+ def backward(ctx: Any, *grad_outputs: Any):
+ grad_scores = grad_outputs[0]
+ index_q, index_k, weights, cu_seqlen_ks, cu_seqlen_ke, topk_indices = (
+ ctx.saved_tensors
+ )
+ with preserve_tilelang_env():
+ from art.megatron.dsv4.kernel.tilelang_indexer_bwd import (
+ batched_indexer_bwd,
+ )
+
+ grad_q, grad_w, grad_k = batched_indexer_bwd(
+ index_q, weights, index_k, topk_indices, grad_scores
+ )
+ return grad_q, grad_k, grad_w, None, None, None
+
+
+def v4_lighting_indexer(
+ index_q: torch.Tensor,
+ index_k: torch.Tensor,
+ weights: torch.Tensor,
+ compress_ratio: int,
+ topk: int,
+ topk_indices: torch.Tensor | None = None,
+):
+ """Main entry point for V4 tilelang indexer.
+
+ Args:
+ index_q: [seqlen, batch, heads, dim] bf16
+ index_k: [seqlen_kv, batch, dim] bf16
+ weights: [seqlen, batch, heads] fp32
+ compress_ratio: compression ratio (4 for C4 layers)
+ topk: number of top-k indices to select
+ topk_indices: optional pre-computed topk indices [batch, seqlen, topk] int32
+
+ Returns:
+ index_score: [batch, seqlen, topk] fp32
+ topk_indices: [batch, seqlen, topk] int32
+ """
+ return V4IndexerFunction.apply(
+ index_q, index_k, weights, compress_ratio, topk, topk_indices
+ )
diff --git a/src/art/megatron/dsv4/kernel/tilelang_indexer_bwd.py b/src/art/megatron/dsv4/kernel/tilelang_indexer_bwd.py
new file mode 100644
index 000000000..d17f17065
--- /dev/null
+++ b/src/art/megatron/dsv4/kernel/tilelang_indexer_bwd.py
@@ -0,0 +1,270 @@
+# ruff: noqa
+# Adapted from miles_plugins/models/glm5/ops/tilelang_indexer_bwd.py for DeepSeek-V4.
+from typing import Any, cast
+
+import torch
+
+from art.megatron.dsv4.kernel.tilelang_import import (
+ import_tilelang,
+ preserve_tilelang_env,
+ sanitize_tilelang_env,
+)
+
+_tl, _T = import_tilelang()
+
+tl = cast(Any, _tl)
+T = cast(Any, _T)
+
+BF16 = T.bfloat16
+FP32 = T.float32
+INT32 = T.int32
+
+pass_configs = {
+ tl.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
+ tl.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+}
+
+
+@tl.jit(pass_configs=pass_configs)
+def tl_indexer_bwd_impl(
+ heads: int,
+ dim: int,
+ topk: int,
+ block_I: int = 32,
+ num_stages: int = 0,
+ num_threads: int = 128,
+):
+ assert num_stages == 0
+ assert topk == tl.math.next_power_of_2(topk)
+ assert topk % block_I == 0
+ assert heads <= 64 and heads % 8 == 0
+ seq_len = T.symbolic("seq_len")
+ q_seq_len = T.symbolic("q_seq_len")
+
+ dtype: str = BF16
+ accum_dtype: str = FP32
+ index_q_shape = [q_seq_len, heads, dim]
+ weights_shape = [q_seq_len, heads]
+ index_k_shape = [seq_len, dim]
+ shape_p = [q_seq_len, topk]
+ topk_indices_shape = [q_seq_len, topk]
+
+ pad_heads = heads
+ if heads < 16:
+ pad_heads = 16
+
+ @T.prim_func
+ def tl_indexer_bwd_kernel(
+ IndexQ: T.Tensor(index_q_shape, dtype), # type: ignore
+ IndexK: T.Tensor(index_k_shape, dtype), # type: ignore
+ Weights: T.Tensor(weights_shape, FP32), # type: ignore
+ TopkIndices: T.Tensor(topk_indices_shape, INT32), # type: ignore
+ OGrad: T.Tensor(shape_p, FP32), # type: ignore
+ dIndexQ: T.Tensor(index_q_shape, dtype), # type: ignore
+ dWeights: T.Tensor(weights_shape, FP32), # type: ignore
+ dIndexK: T.Tensor(index_k_shape, FP32), # type: ignore
+ ):
+
+ with T.Kernel(q_seq_len, threads=num_threads) as (bx):
+ index_q_shared = T.alloc_shared([pad_heads, dim], dtype=FP32)
+ weights_shared = T.alloc_shared([pad_heads], dtype=FP32)
+ index_k_shared = T.alloc_shared([block_I, dim], dtype=FP32)
+ indices_shared = T.alloc_shared([block_I], dtype=INT32)
+ d_index_q_frag = T.alloc_fragment([pad_heads, dim], dtype=accum_dtype)
+ d_weights_frag = T.alloc_fragment([pad_heads], dtype=accum_dtype)
+ d_index_k_frag = T.alloc_fragment([block_I, dim], dtype=accum_dtype)
+ logits = T.alloc_fragment((block_I, pad_heads), dtype=accum_dtype)
+ _logits = T.alloc_shared((block_I, pad_heads), dtype=accum_dtype)
+ grad = T.alloc_shared([block_I], dtype=FP32)
+
+ num_blocks = T.ceildiv(topk, block_I)
+ for i, j in T.Parallel(pad_heads, dim):
+ index_q_shared[i, j] = T.if_then_else(i < heads, IndexQ[bx, i, j], 0)
+ for i in T.Parallel(heads):
+ weights_shared[i] = Weights[bx, i]
+
+ T.fill(d_index_q_frag, 0)
+ T.fill(d_weights_frag, 0)
+
+ for bi_i in T.serial(num_blocks):
+ for i in T.Parallel(block_I):
+ if bi_i * block_I + i < topk:
+ indices_shared[i] = TopkIndices[bx, bi_i * block_I + i]
+ grad[i] = OGrad[bx, bi_i * block_I + i]
+
+ T.sync_threads()
+ for i, j in T.Parallel(block_I, dim):
+ index_k_shared[i, j] = T.if_then_else(
+ indices_shared[i] > -1 and indices_shared[i] < seq_len,
+ IndexK[indices_shared[i], j],
+ 0,
+ )
+
+ T.sync_threads()
+ T.gemm(
+ index_k_shared,
+ index_q_shared,
+ logits,
+ transpose_A=False,
+ transpose_B=True,
+ clear_accum=True,
+ )
+ for i, j in T.Parallel(block_I, heads):
+ logits[i, j] = T.max(logits[i, j], 0)
+
+ d_weights_i = T.alloc_fragment((block_I, pad_heads), accum_dtype)
+ for i, j in T.Parallel(block_I, heads):
+ d_weights_i[i, j] = grad[i] * logits[i, j]
+ T.reduce_sum(d_weights_i, d_weights_frag, dim=0, clear=False)
+
+ for i, j in T.Parallel(block_I, pad_heads):
+ _logits[i, j] = T.if_then_else(
+ logits[i, j] > 0 and j < heads, grad[i] * weights_shared[j], 0
+ )
+ T.sync_threads()
+ T.gemm(
+ _logits,
+ index_k_shared,
+ d_index_q_frag,
+ transpose_A=True,
+ transpose_B=False,
+ clear_accum=False,
+ )
+
+ T.gemm(
+ _logits,
+ index_q_shared,
+ d_index_k_frag,
+ transpose_A=False,
+ transpose_B=False,
+ clear_accum=True,
+ )
+
+ for i, j in T.Parallel(block_I, dim):
+ if indices_shared[i] > -1 and indices_shared[i] < seq_len:
+ T.atomic_add(
+ dIndexK[indices_shared[i], j], d_index_k_frag[i, j]
+ )
+
+ T.copy(d_index_q_frag[:heads, :], dIndexQ[bx, :, :])
+ T.copy(d_weights_frag[:heads], dWeights[bx, :])
+
+ return tl_indexer_bwd_kernel
+
+
+sanitize_tilelang_env()
+
+
+def indexer_bwd_interface(
+ index_q: torch.Tensor,
+ weights: torch.Tensor,
+ index_k: torch.Tensor,
+ topk_indices: torch.Tensor,
+ grad_scores: torch.Tensor,
+):
+ """Backward interface for a single batch element.
+
+ Args:
+ index_q: [seq_len, heads, dim] bf16
+ weights: [seq_len, heads] fp32
+ index_k: [seq_len_kv, dim] bf16
+ topk_indices: [seq_len, topk] int32
+ grad_scores: [seq_len, topk] fp32
+
+ Returns:
+ grad_q: [seq_len, heads, dim] bf16
+ grad_w: [seq_len, heads] fp32
+ grad_k: [seq_len_kv, dim] fp32
+ """
+ _, head_num, head_dim = index_q.shape
+ k_top = topk_indices.shape[1]
+
+ grad_scores = grad_scores.contiguous()
+ grad_q = torch.empty_like(index_q)
+ grad_w = torch.empty_like(weights, dtype=torch.float32)
+ grad_k = torch.zeros_like(index_k, dtype=torch.float32)
+
+ # Pad topk to block_I=32 boundary (kernel requires topk % block_I == 0 and topk >= 32)
+ padded_topk = max(k_top, 32)
+ padded_topk = ((padded_topk + 31) // 32) * 32
+ if padded_topk != k_top:
+ pad_size = padded_topk - k_top
+ topk_indices = torch.cat(
+ [
+ topk_indices,
+ torch.full(
+ (topk_indices.shape[0], pad_size),
+ -1,
+ device=topk_indices.device,
+ dtype=topk_indices.dtype,
+ ),
+ ],
+ dim=1,
+ ).contiguous()
+ grad_scores = torch.cat(
+ [
+ grad_scores,
+ torch.zeros(
+ (grad_scores.shape[0], pad_size),
+ device=grad_scores.device,
+ dtype=grad_scores.dtype,
+ ),
+ ],
+ dim=1,
+ ).contiguous()
+
+ with preserve_tilelang_env():
+ tl_indexer_bwd_impl(head_num, head_dim, padded_topk)(
+ index_q.contiguous(),
+ index_k.contiguous(),
+ weights.squeeze(-1).contiguous(),
+ topk_indices.contiguous(),
+ grad_scores,
+ grad_q,
+ grad_w.squeeze(-1),
+ grad_k,
+ )
+
+ return grad_q, grad_w, grad_k
+
+
+@torch.compiler.disable
+def batched_indexer_bwd(index_q, weights, index_k, topk_indices, grad_scores):
+ """Batched backward: loops over batch dim.
+
+ Args:
+ index_q: [seqlen, batch, heads, dim] bf16
+ weights: [seqlen, batch, heads] fp32
+ index_k: [seqlen_kv, batch, dim] bf16
+ topk_indices: [batch, seqlen, topk] int32
+ grad_scores: [batch, seqlen, topk] fp32
+
+ Returns:
+ grad_q: [seqlen, batch, heads, dim] bf16
+ grad_w: [seqlen, batch, heads] fp32
+ grad_k: [seqlen_kv, batch, dim] fp32
+ """
+ seqlen, batch, heads, dim = index_q.shape
+ seq_len_kv = index_k.shape[0]
+
+ all_grad_q = torch.empty_like(index_q)
+ all_grad_w = torch.empty(
+ seqlen, batch, heads, device=index_q.device, dtype=torch.float32
+ )
+ all_grad_k = torch.zeros(
+ seq_len_kv, batch, dim, device=index_q.device, dtype=torch.float32
+ )
+
+ for b in range(batch):
+ gq, gw, gk = indexer_bwd_interface(
+ index_q[:, b, :, :].contiguous(),
+ weights[:, b, :].contiguous(),
+ index_k[:, b, :].contiguous(),
+ topk_indices[b].contiguous(),
+ grad_scores[b].contiguous(),
+ )
+ all_grad_q[:, b, :, :] = gq
+ all_grad_w[:, b, :] = gw
+ all_grad_k[:, b, :] = gk
+
+ return all_grad_q, all_grad_w, all_grad_k
diff --git a/src/art/megatron/dsv4/kernel/tilelang_indexer_fwd.py b/src/art/megatron/dsv4/kernel/tilelang_indexer_fwd.py
new file mode 100644
index 000000000..6d2e4a51c
--- /dev/null
+++ b/src/art/megatron/dsv4/kernel/tilelang_indexer_fwd.py
@@ -0,0 +1,459 @@
+# ruff: noqa
+# Adapted from miles_plugins/models/glm5/ops/tilelang_indexer_fwd.py for DeepSeek-V4.
+# Key differences from GLM-5:
+# - Operates on [seqlen, batch, heads, dim] (SBHD) layout, batch handled externally
+# - Uses causal mask via cu_seqlens instead of variable-length packed sequences
+# - Supports compressed KV (seq_len_kv = seq_len_q / compress_ratio)
+from typing import Any, cast
+
+import torch
+
+from art.megatron.dsv4.kernel.tilelang_import import (
+ import_tilelang,
+ preserve_tilelang_env,
+ sanitize_tilelang_env,
+)
+
+_tilelang, _T = import_tilelang()
+
+tilelang = cast(Any, _tilelang)
+T = cast(Any, _T)
+
+
+@tilelang.jit(
+ pass_configs={
+ tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
+ },
+)
+def tl_indexer_fwd_impl(
+ heads,
+ index_dim,
+ block_N=256,
+ num_stages=3,
+ threads=512,
+ block_Q=None,
+):
+ if block_Q is None:
+ block_Q = 128 // heads
+ dtype = T.bfloat16
+ accum_dtype = T.float32
+ index_dtype = T.int32
+
+ seq_len = T.dynamic("seq_len")
+ seq_len_kv = T.dynamic("seq_len_kv")
+
+ index_q_shape = [seq_len * heads, index_dim]
+ index_k_shape = [seq_len_kv, index_dim]
+ logits_shape = [seq_len, seq_len_kv]
+
+ @T.prim_func
+ def tl_indexer_fwd_kernel(
+ IndexQ: T.Tensor(index_q_shape, dtype), # type: ignore
+ IndexK: T.Tensor(index_k_shape, dtype), # type: ignore
+ Logits: T.Tensor(logits_shape, accum_dtype), # type: ignore
+ Weights: T.Tensor([seq_len, heads], accum_dtype), # type: ignore
+ CuSeqLenKS: T.Tensor([seq_len], index_dtype), # type: ignore
+ CuSeqLenKE: T.Tensor([seq_len], index_dtype), # type: ignore
+ ):
+ with T.Kernel(T.ceildiv(seq_len, block_Q), threads=threads) as bx:
+ index_q_shared = T.alloc_shared([block_Q * heads, index_dim], dtype)
+ index_k_shared = T.alloc_shared([block_N, index_dim], dtype)
+ s = T.alloc_fragment([block_N, block_Q * heads], accum_dtype)
+ s_reshaped = T.reshape(s, (block_N, block_Q, heads))
+ logits = T.alloc_fragment([block_N, block_Q], accum_dtype)
+ weights = T.alloc_fragment([block_Q, heads], accum_dtype)
+
+ seq_len_i = bx * block_Q
+
+ cu_k_s_min = T.alloc_var(index_dtype)
+ cu_k_e_max = T.alloc_var(index_dtype)
+
+ cu_k_s_min = 2147483647
+ cu_k_e_max = -2147483648
+
+ for bq_i in T.serial(block_Q):
+ cu_k_s_min = T.min(
+ cu_k_s_min, T.min(CuSeqLenKS[seq_len_i + bq_i], seq_len_kv)
+ )
+ for bq_i in T.serial(block_Q):
+ cu_k_e_max = T.max(
+ cu_k_e_max, T.min(CuSeqLenKE[seq_len_i + bq_i], seq_len_kv)
+ )
+
+ T.copy(IndexQ[seq_len_i * heads, 0], index_q_shared)
+ T.copy(Weights[seq_len_i, 0], weights)
+
+ for nbn_i in T.Pipelined(
+ T.ceildiv(cu_k_e_max - cu_k_s_min, block_N), num_stages=num_stages
+ ):
+ T.copy(IndexK[cu_k_s_min + nbn_i * block_N, 0], index_k_shared)
+
+ T.gemm(
+ index_k_shared,
+ index_q_shared,
+ s,
+ transpose_B=True,
+ clear_accum=True,
+ policy=T.GemmWarpPolicy.FullCol,
+ )
+
+ for bn_i, bq_i, h_i in T.Parallel(block_N, block_Q, heads):
+ s_reshaped[bn_i, bq_i, h_i] = (
+ T.max(s_reshaped[bn_i, bq_i, h_i], 0) * weights[bq_i, h_i]
+ )
+
+ T.reduce_sum(s_reshaped, logits, dim=-1, clear=True)
+
+ for bq_i, bn_i in T.Parallel(block_Q, block_N):
+ Logits[seq_len_i + bq_i, cu_k_s_min + nbn_i * block_N + bn_i] = (
+ logits[bn_i, bq_i]
+ )
+
+ return tl_indexer_fwd_kernel
+
+
+@tilelang.jit
+def clean_logits_(
+ threads: int = 512,
+ block_K: int = 4096,
+):
+ seq_len = T.dynamic("seq_len")
+ seq_len_kv = T.dynamic("seq_len_kv")
+
+ dtype = T.float
+ indices_dtype = T.int32
+
+ @T.prim_func
+ def clean_logits_kernel(
+ Logits: T.Tensor([seq_len, seq_len_kv], dtype), # type: ignore
+ CuSeqLenKS: T.Tensor([seq_len], indices_dtype), # type: ignore
+ CuSeqLenKE: T.Tensor([seq_len], indices_dtype), # type: ignore
+ ):
+ with T.Kernel(seq_len, threads=threads) as bx:
+ tx = T.thread_binding(0, threads, thread="threadIdx.x")
+ cu_k_s = CuSeqLenKS[bx]
+ cu_k_e = CuSeqLenKE[bx]
+
+ for n_i in T.Pipelined(T.ceildiv(seq_len_kv, block_K)):
+ for k_i in T.serial(block_K // threads):
+ idx = n_i * block_K + k_i * threads + tx
+ if idx < cu_k_s or idx >= cu_k_e:
+ Logits[bx, idx] = -T.infinity(dtype)
+
+ return clean_logits_kernel
+
+
+sanitize_tilelang_env()
+
+
+def _make_causal_cu_seqlens(seq_len_q, seq_len_kv, compress_ratio, device):
+ """Generate cu_seqlens for causal masking on compressed KV positions.
+
+ For query at position p, valid compressed groups are [0, (p+1) // compress_ratio).
+ """
+ positions = torch.arange(seq_len_q, device=device, dtype=torch.int32)
+ cu_seqlen_ks = torch.zeros(seq_len_q, device=device, dtype=torch.int32)
+ cu_seqlen_ke = ((positions + 1) // compress_ratio).to(torch.int32)
+ return cu_seqlen_ks, cu_seqlen_ke
+
+
+def indexer_fwd_interface(
+ q, kv, weights, cu_seqlen_ks, cu_seqlen_ke, clean_logits=True
+):
+ """Forward interface matching GLM-5's API but for a single batch element.
+
+ Args:
+ q: [seq_len, heads, index_dim] bf16
+ kv: [seq_len_kv, index_dim] bf16
+ weights: [seq_len, heads] fp32
+ cu_seqlen_ks: [seq_len] int32 — start of valid KV range per query
+ cu_seqlen_ke: [seq_len] int32 — end of valid KV range per query
+
+ Returns:
+ logits: [seq_len, seq_len_kv] fp32
+ """
+ if q.dtype is torch.float32:
+ q = q.to(torch.bfloat16)
+ if kv.dtype is torch.float32:
+ kv = kv.to(torch.bfloat16)
+ if q.dtype != torch.bfloat16 or kv.dtype != torch.bfloat16:
+ raise TypeError(
+ f"DSV4 indexer TileLang launch requires bf16, got {q.dtype=}, {kv.dtype=}"
+ )
+ seq_len, heads, index_dim = q.shape
+ seq_len_kv = kv.shape[0]
+ block_q = max(1, 128 // heads)
+ if seq_len % block_q != 0:
+ raise ValueError(
+ f"DSV4 indexer TileLang query length must be divisible by {block_q}, "
+ f"got {seq_len}."
+ )
+
+ logits = torch.empty([seq_len, seq_len_kv], device=q.device, dtype=torch.float32)
+ with preserve_tilelang_env():
+ clean_logits_kernel = clean_logits_()
+ tl_indexer_fwd_kernel = tl_indexer_fwd_impl(heads=heads, index_dim=index_dim)
+ tl_indexer_fwd_kernel(
+ q.view(seq_len * heads, index_dim),
+ kv,
+ logits,
+ weights.float(),
+ cu_seqlen_ks,
+ cu_seqlen_ke,
+ )
+ if clean_logits:
+ clean_logits_kernel(logits, cu_seqlen_ks, cu_seqlen_ke)
+ return logits
+
+
+def _topk_indices_from_logits(logits, topk):
+ actual_topk = min(int(topk), int(logits.shape[-1]))
+ top_scores, top_indices = logits.topk(actual_topk, dim=-1)
+ return torch.where(
+ torch.isneginf(top_scores),
+ torch.full_like(top_indices, -1),
+ top_indices,
+ ).to(torch.int32)
+
+
+def indexer_topk_interface(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke, topk):
+ """Compute DSV4 indexer scores and return top-k compressed ids.
+
+ Keep topk owned by the indexer wrapper so callers do not materialize or
+ inspect dense logits.
+ """
+ return _topk_indices_from_logits(
+ indexer_fwd_interface(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke),
+ topk,
+ )
+
+
+@tilelang.jit(
+ pass_configs={
+ tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
+ tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
+ },
+)
+def tl_shared_prefix_indexer_fwd_impl(
+ heads,
+ index_dim,
+ block_N=256,
+ num_stages=3,
+ threads=512,
+ block_Q=None,
+):
+ if block_Q is None:
+ block_Q = 128 // heads
+ dtype = T.bfloat16
+ accum_dtype = T.float32
+ index_dtype = T.int32
+
+ seq_len = T.dynamic("seq_len")
+ seq_len_kv = T.dynamic("seq_len_kv")
+
+ index_q_shape = [seq_len * heads, index_dim]
+ index_k_shape = [seq_len_kv, index_dim]
+ logits_shape = [seq_len, seq_len_kv]
+
+ @T.prim_func
+ def tl_shared_prefix_indexer_fwd_kernel(
+ IndexQ: T.Tensor(index_q_shape, dtype), # type: ignore
+ IndexK: T.Tensor(index_k_shape, dtype), # type: ignore
+ Logits: T.Tensor(logits_shape, accum_dtype), # type: ignore
+ Weights: T.Tensor([seq_len, heads], accum_dtype), # type: ignore
+ QPosition: T.Tensor([seq_len], index_dtype), # type: ignore
+ QGroup: T.Tensor([seq_len], index_dtype), # type: ignore
+ QParent: T.Tensor([seq_len], index_dtype), # type: ignore
+ EntryGroup: T.Tensor([seq_len_kv], index_dtype), # type: ignore
+ EntryParentVisible: T.Tensor([seq_len_kv], index_dtype), # type: ignore
+ EntryEndPosition: T.Tensor([seq_len_kv], index_dtype), # type: ignore
+ EntryValid: T.Tensor([seq_len_kv], index_dtype), # type: ignore
+ ):
+ with T.Kernel(T.ceildiv(seq_len, block_Q), threads=threads) as bx:
+ index_q_shared = T.alloc_shared([block_Q * heads, index_dim], dtype)
+ index_k_shared = T.alloc_shared([block_N, index_dim], dtype)
+ s = T.alloc_fragment([block_N, block_Q * heads], accum_dtype)
+ s_reshaped = T.reshape(s, (block_N, block_Q, heads))
+ logits = T.alloc_fragment([block_N, block_Q], accum_dtype)
+ weights = T.alloc_fragment([block_Q, heads], accum_dtype)
+
+ seq_len_i = bx * block_Q
+
+ T.copy(IndexQ[seq_len_i * heads, 0], index_q_shared)
+ T.copy(Weights[seq_len_i, 0], weights)
+
+ for nbn_i in T.Pipelined(
+ T.ceildiv(seq_len_kv, block_N), num_stages=num_stages
+ ):
+ for bn_i, d_i in T.Parallel(block_N, index_dim):
+ k_i = nbn_i * block_N + bn_i
+ index_k_shared[bn_i, d_i] = T.if_then_else(
+ k_i < seq_len_kv,
+ IndexK[k_i, d_i],
+ 0,
+ )
+
+ T.gemm(
+ index_k_shared,
+ index_q_shared,
+ s,
+ transpose_B=True,
+ clear_accum=True,
+ policy=T.GemmWarpPolicy.FullCol,
+ )
+
+ for bn_i, bq_i, h_i in T.Parallel(block_N, block_Q, heads):
+ s_reshaped[bn_i, bq_i, h_i] = (
+ T.max(s_reshaped[bn_i, bq_i, h_i], 0) * weights[bq_i, h_i]
+ )
+
+ T.reduce_sum(s_reshaped, logits, dim=-1, clear=True)
+
+ for bq_i, bn_i in T.Parallel(block_Q, block_N):
+ q_i = seq_len_i + bq_i
+ k_i = nbn_i * block_N + bn_i
+ if k_i < seq_len_kv:
+ entry_group = EntryGroup[k_i]
+ visible = (
+ (EntryValid[k_i] != 0)
+ and (EntryEndPosition[k_i] <= QPosition[q_i])
+ and (
+ (entry_group == QGroup[q_i])
+ or (
+ (EntryParentVisible[k_i] != 0)
+ and (entry_group == QParent[q_i])
+ )
+ )
+ )
+ Logits[q_i, k_i] = T.if_then_else(
+ visible,
+ logits[bn_i, bq_i],
+ -T.infinity(accum_dtype),
+ )
+
+ return tl_shared_prefix_indexer_fwd_kernel
+
+
+def _i32_contiguous(tensor):
+ if tensor.dtype != torch.int32:
+ tensor = tensor.to(torch.int32)
+ return tensor.contiguous()
+
+
+def shared_prefix_indexer_fwd_interface(
+ q,
+ kv,
+ weights,
+ position_ids,
+ group_ids,
+ parent_ids,
+ entry_group_ids,
+ entry_parent_visible,
+ entry_end_positions,
+ entry_valid,
+):
+ """Compute shared-prefix-aware indexer logits for one batch element.
+
+ The output is intentionally block-local [query_block, compressed_kv] fp32.
+ Callers run topk on this block to avoid materializing full [S, compressed_kv]
+ logits for long shared-prefix packed sequences.
+ """
+ if q.dtype is torch.float32:
+ q = q.to(torch.bfloat16)
+ if kv.dtype is torch.float32:
+ kv = kv.to(torch.bfloat16)
+ if q.dtype != torch.bfloat16 or kv.dtype != torch.bfloat16:
+ raise TypeError(
+ f"DSV4 indexer TileLang launch requires bf16, got {q.dtype=}, {kv.dtype=}"
+ )
+ seq_len, heads, index_dim = q.shape
+ seq_len_kv = kv.shape[0]
+ block_q = max(1, 128 // heads)
+ if seq_len % block_q != 0:
+ raise ValueError(
+ "DSV4 shared-prefix indexer TileLang query length must be divisible "
+ f"by {block_q}, got {seq_len}."
+ )
+
+ logits = torch.empty([seq_len, seq_len_kv], device=q.device, dtype=torch.float32)
+ with preserve_tilelang_env():
+ tl_indexer_fwd_kernel = tl_shared_prefix_indexer_fwd_impl(
+ heads=heads,
+ index_dim=index_dim,
+ )
+ tl_indexer_fwd_kernel(
+ q.view(seq_len * heads, index_dim),
+ kv,
+ logits,
+ weights.float(),
+ _i32_contiguous(position_ids),
+ _i32_contiguous(group_ids),
+ _i32_contiguous(parent_ids),
+ _i32_contiguous(entry_group_ids),
+ _i32_contiguous(entry_parent_visible),
+ _i32_contiguous(entry_end_positions),
+ _i32_contiguous(entry_valid),
+ )
+ return logits
+
+
+def shared_prefix_indexer_topk_interface(
+ q,
+ kv,
+ weights,
+ position_ids,
+ group_ids,
+ parent_ids,
+ entry_group_ids,
+ entry_parent_visible,
+ entry_end_positions,
+ entry_valid,
+ topk,
+):
+ """Compute shared-prefix-aware DSV4 indexer top-k compressed ids."""
+ return _topk_indices_from_logits(
+ shared_prefix_indexer_fwd_interface(
+ q,
+ kv,
+ weights,
+ position_ids,
+ group_ids,
+ parent_ids,
+ entry_group_ids,
+ entry_parent_visible,
+ entry_end_positions,
+ entry_valid,
+ ),
+ topk,
+ )
+
+
+@torch.compiler.disable
+def batched_indexer_fwd(q, k, weights, cu_seqlen_ks, cu_seqlen_ke):
+ """Batched forward: loops over batch dim.
+
+ Args:
+ q: [seqlen, batch, heads, dim] bf16
+ k: [seqlen_kv, batch, dim] bf16
+ weights: [seqlen, batch, heads] fp32
+ cu_seqlen_ks: [seqlen] int32
+ cu_seqlen_ke: [seqlen] int32
+
+ Returns:
+ logits: [batch, seqlen, seqlen_kv] fp32
+ """
+ seqlen, batch, heads, dim = q.shape
+ seq_len_kv = k.shape[0]
+
+ all_logits = torch.empty(
+ [batch, seqlen, seq_len_kv], device=q.device, dtype=torch.float32
+ )
+ for b in range(batch):
+ all_logits[b] = indexer_fwd_interface(
+ q[:, b, :, :].contiguous(),
+ k[:, b, :].contiguous(),
+ weights[:, b, :].contiguous(),
+ cu_seqlen_ks,
+ cu_seqlen_ke,
+ )
+ return all_logits
diff --git a/src/art/megatron/dsv4/kernel/tilelang_sparse_mla.py b/src/art/megatron/dsv4/kernel/tilelang_sparse_mla.py
new file mode 100644
index 000000000..20fefe44f
--- /dev/null
+++ b/src/art/megatron/dsv4/kernel/tilelang_sparse_mla.py
@@ -0,0 +1,117 @@
+from typing import Any
+
+import torch
+
+from art.megatron.dsv4.kernel.tilelang_import import preserve_tilelang_env
+
+
+def _sparse_attn_torch(q, kv, attn_sink, topk_idxs, sm_scale):
+ if sm_scale is None:
+ sm_scale = q.shape[-1] ** -0.5
+ bsz, seqlen, _, dim = q.shape
+ safe_idxs = topk_idxs.clamp_min(0)
+ selected_kv = torch.gather(
+ kv[:, None].expand(-1, seqlen, -1, -1),
+ 2,
+ safe_idxs[..., None].expand(-1, -1, -1, dim),
+ )
+ scores = torch.einsum("bshd,bskd->bshk", q.float(), selected_kv.float())
+ scores = scores * float(sm_scale)
+ scores = scores.masked_fill(topk_idxs[:, :, None, :] < 0, float("-inf"))
+ sinks = attn_sink.view(1, 1, -1, 1).expand(bsz, seqlen, -1, -1)
+ probs = torch.softmax(torch.cat([scores, sinks], dim=-1), dim=-1)
+ attn_probs = probs[..., :-1]
+ return torch.einsum("bshk,bskd->bshd", attn_probs, selected_kv.float())
+
+
+def _pad_topk_idxs(topk_idxs: torch.Tensor, block_size: int = 64) -> torch.Tensor:
+ topk = int(topk_idxs.shape[-1])
+ padded_topk = (topk + block_size - 1) // block_size * block_size
+ if padded_topk == topk:
+ return topk_idxs
+ pad = torch.full(
+ (*topk_idxs.shape[:-1], padded_topk - topk),
+ -1,
+ device=topk_idxs.device,
+ dtype=topk_idxs.dtype,
+ )
+ return torch.cat([topk_idxs, pad], dim=-1).contiguous()
+
+
+class DeepSeekV4SparseAttention(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, q, kv, attn_sink, topk_idxs, sm_scale=None, output_dtype=None):
+ with preserve_tilelang_env():
+ from art.megatron.dsv4.kernel import (
+ tilelang_sparse_mla_fwd as sparse_mla_fwd,
+ )
+
+ o, lse = sparse_mla_fwd.sparse_mqa_fwd_interface(
+ q, kv, attn_sink, topk_idxs, sm_scale=sm_scale
+ )
+
+ output = o if output_dtype is None else o.to(output_dtype)
+ ctx.save_for_backward(q, kv, attn_sink, topk_idxs, output.clone(), lse)
+ ctx.sm_scale = sm_scale
+
+ return output
+
+ @staticmethod
+ def backward(ctx: Any, *grad_outputs: Any):
+ do = grad_outputs[0]
+ q, kv, attn_sink, topk_idxs, output, lse = ctx.saved_tensors
+ sm_scale = ctx.sm_scale
+
+ with preserve_tilelang_env():
+ from art.megatron.dsv4.kernel import (
+ tilelang_sparse_mla_bwd as sparse_mla_bwd,
+ )
+
+ dq, dkv, d_attn_sink = sparse_mla_bwd.sparse_mqa_bwd_interface(
+ q,
+ kv,
+ attn_sink,
+ output.to(q.dtype),
+ do.to(q.dtype),
+ topk_idxs,
+ lse,
+ sm_scale=sm_scale,
+ )
+
+ return dq, dkv, d_attn_sink, None, None, None
+
+
+@torch.compiler.disable
+def sparse_attn_tilelang(q, kv, attn_sink, topk_idxs, sm_scale=None):
+ """Run TileLang sparse MLA outside TorchDynamo tracing.
+
+ TileLang's TVM FFI adapter uses non-literal string objects internally, which
+ Dynamo cannot represent as constants. Keep only this kernel boundary eager
+ while allowing the surrounding DSV4 transformer layer to compile.
+ """
+ output_dtype = q.dtype
+ if q.dtype is torch.float32:
+ return _sparse_attn_torch(q, kv, attn_sink, topk_idxs, sm_scale)
+ if kv.dtype != q.dtype:
+ kv = kv.to(q.dtype)
+ q = q.contiguous()
+ kv = kv.contiguous()
+ topk_idxs = _pad_topk_idxs(topk_idxs.contiguous())
+ head_count = int(q.shape[2])
+ if head_count < 16:
+ pad_heads = 16 - head_count
+ q = torch.cat(
+ [
+ q,
+ q.new_zeros((*q.shape[:2], pad_heads, q.shape[3])),
+ ],
+ dim=2,
+ ).contiguous()
+ attn_sink = torch.cat(
+ [attn_sink, attn_sink.new_zeros(pad_heads)],
+ dim=0,
+ ).contiguous()
+ out = DeepSeekV4SparseAttention.apply(
+ q, kv, attn_sink, topk_idxs, sm_scale, output_dtype
+ )
+ return out[:, :, :head_count, :]
diff --git a/src/art/megatron/dsv4/kernel/tilelang_sparse_mla_bwd.py b/src/art/megatron/dsv4/kernel/tilelang_sparse_mla_bwd.py
new file mode 100644
index 000000000..0d14e5e4a
--- /dev/null
+++ b/src/art/megatron/dsv4/kernel/tilelang_sparse_mla_bwd.py
@@ -0,0 +1,406 @@
+# ruff: noqa
+# Adapted from miles_plugins/models/glm5/ops/tilelang_sparse_mla_bwd.py for DeepSeek-V4.
+# Key differences from GLM-5:
+# - attn_sink: gradient computation for learnable per-head scalar
+# - Single-head KV: kv shape [B, S_kv, D] (no kv_group, no D/D_tail split)
+# - Index shape: [B, S, topk] (no kv_group dim)
+# - Outputs: dQ [B, S, H, D], dKV [B, S_kv, D], dAttnSink [H]
+from typing import Any, cast
+
+import torch
+
+from art.megatron.dsv4.kernel.tilelang_import import (
+ import_tilelang,
+ preserve_tilelang_env,
+ sanitize_tilelang_env,
+)
+
+_tilelang, _T = import_tilelang()
+
+tilelang = cast(Any, _tilelang)
+T = cast(Any, _T)
+
+
+@tilelang.jit(out_idx=[-1])
+def preprocess(
+ B,
+ S,
+ H,
+ D,
+ block_ND=32,
+ num_stages=5,
+ dtype=T.bfloat16,
+ accum_dtype=T.float32,
+):
+ assert dtype == T.bfloat16
+ assert accum_dtype == T.float32
+ shape = [B, S, H, D]
+
+ @T.prim_func
+ def preprocess_kernel(
+ O: T.Tensor(shape, dtype), # type: ignore
+ dO: T.Tensor(shape, dtype), # type: ignore
+ Delta: T.Tensor([B, S, H], accum_dtype), # type: ignore
+ ):
+ with T.Kernel(H, T.ceildiv(S, block_ND), B) as (bx, by, bz):
+ o = T.alloc_fragment([block_ND, block_ND], accum_dtype)
+ do = T.alloc_fragment([block_ND, block_ND], accum_dtype)
+ delta = T.alloc_fragment([block_ND], accum_dtype)
+ acc = T.alloc_fragment([block_ND, block_ND], accum_dtype)
+ T.clear(acc)
+ for k in T.Pipelined(T.ceildiv(D, block_ND), num_stages=num_stages):
+ T.copy(
+ O[
+ bz,
+ by * block_ND : (by + 1) * block_ND,
+ bx,
+ k * block_ND : (k + 1) * block_ND,
+ ],
+ o,
+ )
+ T.copy(
+ dO[
+ bz,
+ by * block_ND : (by + 1) * block_ND,
+ bx,
+ k * block_ND : (k + 1) * block_ND,
+ ],
+ do,
+ )
+ for i, j in T.Parallel(block_ND, block_ND):
+ acc[i, j] += o[i, j] * do[i, j]
+ T.reduce_sum(acc, delta, 1)
+ T.copy(delta, Delta[bz, by * block_ND : (by + 1) * block_ND, bx])
+
+ return preprocess_kernel
+
+
+@tilelang.jit(out_idx=[-1])
+def postprocess(
+ B,
+ S_kv,
+ D,
+ block_N=64,
+ threads=128,
+ dtype=T.bfloat16,
+ accum_dtype=T.float32,
+):
+ assert dtype == T.bfloat16
+ assert accum_dtype == T.float32
+ dkv_shape = [B, S_kv, D]
+
+ @T.prim_func
+ def postprocess_kernel(
+ dKV: T.Tensor(dkv_shape, accum_dtype), # type: ignore
+ dKV_out: T.Tensor(dkv_shape, dtype), # type: ignore
+ ):
+ with T.Kernel(T.ceildiv(S_kv, block_N), B, threads=threads) as (bx, by):
+ T.copy(
+ dKV[by, bx * block_N : (bx + 1) * block_N, :],
+ dKV_out[by, bx * block_N : (bx + 1) * block_N, :],
+ )
+
+ return postprocess_kernel
+
+
+@tilelang.jit(
+ out_idx=[-3],
+ pass_configs={
+ tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
+ tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+ tilelang.PassConfigKey.TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE: True,
+ },
+)
+def bwd(
+ B,
+ S,
+ S_kv,
+ H,
+ D,
+ topk,
+ sm_scale=None,
+ block_size=32,
+ num_stages=0,
+ threads=128,
+ indices_dtype=T.int32,
+ dtype=T.bfloat16,
+ accum_dtype=T.float32,
+):
+ assert topk % block_size == 0, (
+ f"topk ({topk}) must be divisible by block_size ({block_size})"
+ )
+ assert dtype == T.bfloat16
+ assert accum_dtype == T.float32
+
+ if sm_scale is None:
+ sm_scale = D ** (-0.5)
+ sm_scale_mul_reciprocal_log2 = sm_scale * 1.44269504 # log2(e)
+
+ q_shape = [B, S, H, D]
+ kv_shape = [B, S_kv, D]
+ o_shape = [B, S, H, D]
+ indices_shape = [B, S, topk]
+ delta_shape = [B, S, H]
+ lse_shape = [B, S, H]
+ attn_sink_shape = [H]
+
+ padded_H = max(tilelang.math.next_power_of_2(H), 16)
+ block_H = min(64, padded_H)
+ assert padded_H % block_H == 0
+ NH = padded_H // block_H
+ BS = block_size
+ NS = tilelang.cdiv(topk, block_size)
+
+ split_store = 2
+
+ @T.prim_func
+ def sparse_mqa_bwd_kernel(
+ Q: T.Tensor(q_shape, dtype), # type: ignore
+ KV: T.Tensor(kv_shape, dtype), # type: ignore
+ dO: T.Tensor(o_shape, dtype), # type: ignore
+ AttnSink: T.Tensor(attn_sink_shape, accum_dtype), # type: ignore
+ Indices: T.Tensor(indices_shape, indices_dtype), # type: ignore
+ Lse: T.Tensor(lse_shape, accum_dtype), # type: ignore
+ Delta: T.Tensor(delta_shape, accum_dtype), # type: ignore
+ dQ: T.Tensor(q_shape, dtype), # type: ignore
+ dKV: T.Tensor(kv_shape, accum_dtype), # type: ignore
+ dAttnSink: T.Tensor(attn_sink_shape, accum_dtype), # type: ignore
+ ):
+ with T.Kernel(S, B, NH, threads=threads) as (s_i, by, bz):
+ Q_shared = T.alloc_shared([block_H, D], dtype)
+ KV_shared = T.alloc_shared([BS, D], dtype)
+ dO_shared = T.alloc_shared([block_H, D], dtype)
+ mask = T.alloc_fragment([BS], "bool")
+
+ P_shared_cast = T.alloc_shared([block_H, BS], dtype)
+ dP_shared_cast = T.alloc_shared([block_H, BS], dtype)
+ dQ_shared = T.alloc_shared([block_H, D], dtype)
+
+ acc_p = T.alloc_fragment([block_H, BS], accum_dtype)
+ acc_dp = T.alloc_fragment([block_H, BS], accum_dtype)
+ acc_dq = T.alloc_fragment([block_H, D], accum_dtype)
+ acc_dq_i = T.alloc_fragment([block_H, D], accum_dtype)
+ acc_dkv = T.alloc_fragment([BS, D], accum_dtype)
+ acc_dkv_shared = T.alloc_shared([BS // split_store, D], accum_dtype)
+ safe_indices = T.alloc_fragment([BS], indices_dtype)
+
+ T.copy(Q[by, s_i, bz * block_H : (bz + 1) * block_H, :D], Q_shared)
+ T.copy(dO[by, s_i, bz * block_H : (bz + 1) * block_H, :D], dO_shared)
+
+ T.clear(acc_dq)
+
+ for i_i in T.Pipelined(NS, num_stages=num_stages):
+ for bi_i in T.Parallel(BS):
+ mask[bi_i] = Indices[by, s_i, i_i * BS + bi_i] != -1
+ safe_indices[bi_i] = T.if_then_else(
+ mask[bi_i], Indices[by, s_i, i_i * BS + bi_i], 0
+ )
+
+ T.clear(acc_p)
+
+ for bi_i, d_i in T.Parallel(BS, D):
+ KV_shared[bi_i, d_i] = KV[by, safe_indices[bi_i], d_i]
+
+ T.gemm(
+ Q_shared,
+ KV_shared,
+ acc_p,
+ transpose_B=True,
+ policy=T.GemmWarpPolicy.FullCol,
+ )
+ for h_i, bi_i in T.Parallel(block_H, BS):
+ acc_p[h_i, bi_i] = T.if_then_else(
+ mask[bi_i], acc_p[h_i, bi_i], -T.infinity(acc_p.dtype)
+ )
+
+ # P = exp2(scores * sm_scale_log2e - LSE)
+ for h_i, bi_i in T.Parallel(block_H, BS):
+ acc_p[h_i, bi_i] = T.exp2(
+ acc_p[h_i, bi_i] * sm_scale_mul_reciprocal_log2
+ - Lse[by, s_i, bz * block_H + h_i]
+ )
+
+ T.copy(acc_p, P_shared_cast)
+
+ # dP = P * (dO @ KV^T - Delta)
+ T.gemm(
+ dO_shared,
+ KV_shared,
+ acc_dp,
+ transpose_B=True,
+ policy=T.GemmWarpPolicy.FullCol,
+ clear_accum=True,
+ )
+
+ for h_i, bi_i in T.Parallel(block_H, BS):
+ acc_dp[h_i, bi_i] = (
+ acc_p[h_i, bi_i]
+ * (acc_dp[h_i, bi_i] - Delta[by, s_i, bz * block_H + h_i])
+ * sm_scale
+ )
+
+ T.copy(acc_dp, dP_shared_cast)
+
+ # dQ += dP @ KV
+ T.gemm(
+ dP_shared_cast,
+ KV_shared,
+ acc_dq_i,
+ policy=T.GemmWarpPolicy.FullCol,
+ clear_accum=True,
+ )
+ for h_i, d_i in T.Parallel(block_H, D):
+ acc_dq[h_i, d_i] += acc_dq_i[h_i, d_i]
+
+ # dKV += dP^T @ Q + P^T @ dO
+ T.gemm(
+ dP_shared_cast,
+ Q_shared,
+ acc_dkv,
+ transpose_A=True,
+ policy=T.GemmWarpPolicy.FullCol,
+ clear_accum=True,
+ )
+ T.gemm(
+ P_shared_cast,
+ dO_shared,
+ acc_dkv,
+ transpose_A=True,
+ policy=T.GemmWarpPolicy.FullCol,
+ )
+
+ # Atomic store dKV with split to reduce register pressure
+ for s in range(split_store):
+ for bi_i, d_i in T.Parallel(BS, D):
+ if bi_i < BS // split_store:
+ src_i = bi_i + s * (BS // split_store)
+ acc_dkv_shared[bi_i, d_i] = acc_dkv[src_i, d_i]
+
+ for bi_i, d_i in T.Parallel(BS // split_store, D // 4):
+ src_i = bi_i + s * (BS // split_store)
+ T.atomic_addx4(
+ dKV[
+ by,
+ safe_indices[src_i],
+ d_i * 4,
+ ],
+ acc_dkv_shared[bi_i, d_i * 4],
+ )
+
+ # Store dQ
+ T.copy(acc_dq, dQ_shared)
+ T.copy(dQ_shared, dQ[by, s_i, bz * block_H : (bz + 1) * block_H, :D])
+
+ # dAttnSink[h] = -sum_{b,s}( Delta[b,s,h] * p_sink[b,s,h] )
+ # where p_sink = exp(attn_sink[h]) / Z = exp2(attn_sink[h]*log2e - LSE)
+ # attn_sink is a pre-scaled logit, so only convert to log2 base (no sm_scale)
+ for h_i in T.Parallel(block_H):
+ T.atomic_add(
+ dAttnSink[bz * block_H + h_i],
+ -Delta[by, s_i, bz * block_H + h_i]
+ * T.exp2(
+ AttnSink[bz * block_H + h_i] * 1.44269504
+ - Lse[by, s_i, bz * block_H + h_i]
+ ),
+ )
+
+ return sparse_mqa_bwd_kernel
+
+
+sanitize_tilelang_env()
+
+
+def _tilelang_input_dtype(torch_dtype):
+ if torch_dtype is torch.bfloat16:
+ return T.bfloat16
+ raise TypeError(f"DSV4 sparse MLA TileLang launch requires bf16, got {torch_dtype}")
+
+
+def sparse_mqa_bwd_interface(q, kv, attn_sink, o, do, topk_idxs, lse, sm_scale=None):
+ """Backward interface for V4 sparse MQA attention.
+
+ Args:
+ q: [B, S, H, D] bf16
+ kv: [B, S_kv, D] bf16
+ attn_sink: [H] fp32
+ o: [B, S, H, D] bf16 (forward output)
+ do: [B, S, H, D] bf16 (grad of output)
+ topk_idxs: [B, S, topk] int32
+ lse: [B, S, H] fp32 (log-sum-exp from forward)
+ sm_scale: float or None
+
+ Returns:
+ dq: [B, S, H, D] bf16
+ dkv: [B, S_kv, D] bf16
+ d_attn_sink: [H] fp32
+ """
+ assert q.is_contiguous() and kv.is_contiguous()
+ assert topk_idxs.is_contiguous() and lse.is_contiguous()
+ B, S, H, D = q.shape
+ _, S_kv, _ = kv.shape
+ topk = topk_idxs.shape[-1]
+ dtype = _tilelang_input_dtype(q.dtype)
+ assert kv.dtype == q.dtype and o.dtype == q.dtype and do.dtype == q.dtype
+
+ # Pad topk to next multiple of block_size (kernel requires divisibility)
+ block_size = 64
+ padded_topk = (topk + block_size - 1) // block_size * block_size
+ if padded_topk != topk:
+ pad = torch.full(
+ (B, S, padded_topk - topk),
+ -1,
+ device=topk_idxs.device,
+ dtype=topk_idxs.dtype,
+ )
+ topk_idxs = torch.cat([topk_idxs, pad], dim=-1).contiguous()
+ topk = padded_topk
+
+ with preserve_tilelang_env():
+ preprocess_kernel = preprocess(B, S, H, D, dtype=dtype)
+ postprocess_kernel = postprocess(B, S_kv, D, dtype=dtype)
+ delta = preprocess_kernel(o, do)
+ dkv = torch.zeros_like(kv, dtype=torch.float32)
+ d_attn_sink = torch.zeros_like(attn_sink)
+ if topk <= block_size:
+ with preserve_tilelang_env():
+ bwd_kernel = bwd(
+ B,
+ S,
+ S_kv,
+ H,
+ D,
+ topk,
+ sm_scale,
+ block_size=block_size,
+ dtype=dtype,
+ )
+ dq = bwd_kernel(
+ q, kv, do, attn_sink, topk_idxs, lse, delta, dkv, d_attn_sink
+ )
+ else:
+ dq_accum = torch.zeros_like(q, dtype=torch.float32)
+ chunk_count = topk // block_size
+ with preserve_tilelang_env():
+ bwd_kernel = bwd(
+ B,
+ S,
+ S_kv,
+ H,
+ D,
+ block_size,
+ sm_scale,
+ block_size=block_size,
+ dtype=dtype,
+ )
+ for start in range(0, topk, block_size):
+ chunk = topk_idxs[:, :, start : start + block_size].contiguous()
+ dq_i = bwd_kernel(
+ q, kv, do, attn_sink, chunk, lse, delta, dkv, d_attn_sink
+ )
+ dq_accum.add_(dq_i.float())
+ dq = dq_accum.to(q.dtype)
+ d_attn_sink.div_(chunk_count)
+ with preserve_tilelang_env():
+ dkv = postprocess_kernel(dkv)
+
+ return dq, dkv, d_attn_sink
diff --git a/src/art/megatron/dsv4/kernel/tilelang_sparse_mla_fwd.py b/src/art/megatron/dsv4/kernel/tilelang_sparse_mla_fwd.py
new file mode 100644
index 000000000..34cf531db
--- /dev/null
+++ b/src/art/megatron/dsv4/kernel/tilelang_sparse_mla_fwd.py
@@ -0,0 +1,240 @@
+# ruff: noqa
+# Adapted from miles_plugins/models/glm5/ops/tilelang_sparse_mla_fwd.py for DeepSeek-V4.
+# Key differences from GLM-5:
+# - attn_sink: learnable per-head scalar added to softmax denominator
+# - Single-head KV: kv shape [B, S_kv, D] (no kv_group, no D/D_tail split)
+# - Index shape: [B, S, topk] (no kv_group dim)
+# - Output: [B, S, H, D] + LSE [B, S, H]
+from typing import Any, cast
+
+import torch
+
+from art.megatron.dsv4.kernel.tilelang_import import (
+ import_tilelang,
+ preserve_tilelang_env,
+ sanitize_tilelang_env,
+)
+
+_tilelang, _T = import_tilelang()
+
+tilelang = cast(Any, _tilelang)
+T = cast(Any, _T)
+
+
+@tilelang.jit(
+ out_idx=[-2, -1],
+ pass_configs={
+ tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
+ tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+ },
+)
+def sparse_mqa_fwd(
+ heads,
+ dim,
+ topk,
+ sm_scale=None,
+ block_I=64,
+ num_stages=2,
+ threads=256,
+ dtype=T.bfloat16,
+):
+ assert dim == tilelang.math.next_power_of_2(dim), (
+ f"dim must be power of 2, got {dim}"
+ )
+ assert topk % block_I == 0, (
+ f"topk ({topk}) must be divisible by block_I ({block_I})"
+ )
+ if sm_scale is None:
+ sm_scale = (1.0 / dim) ** 0.5 * 1.44269504 # log2(e)
+ else:
+ sm_scale = sm_scale * 1.44269504 # log2(e)
+
+ batch = T.dynamic("batch")
+ seq_len = T.dynamic("seq_len")
+ seq_len_kv = T.dynamic("seq_len_kv")
+
+ q_shape = [batch, seq_len, heads, dim]
+ kv_shape = [batch, seq_len_kv, dim]
+ o_shape = [batch, seq_len, heads, dim]
+ indices_shape = [batch, seq_len, topk]
+ lse_shape = [batch, seq_len, heads]
+ attn_sink_shape = [heads]
+ indices_dtype = T.int32
+ assert dtype == T.bfloat16
+ accum_dtype = T.float32
+
+ H = heads
+ padded_H = max(tilelang.math.next_power_of_2(heads), 16)
+ BI = block_I
+ NI = tilelang.cdiv(topk, block_I)
+ D = dim
+
+ if heads > 64:
+ assert heads % 64 == 0, "heads should be a multiple of 64"
+ REPLICATE_H = heads // 64
+ else:
+ REPLICATE_H = 1
+
+ H_per_block = padded_H if REPLICATE_H == 1 else 64
+
+ @T.prim_func
+ def main(
+ Q: T.Tensor(q_shape, dtype), # type: ignore
+ KV: T.Tensor(kv_shape, dtype), # type: ignore
+ AttnSink: T.Tensor(attn_sink_shape, accum_dtype), # type: ignore
+ Indices: T.Tensor(indices_shape, indices_dtype), # type: ignore
+ Output: T.Tensor(o_shape, dtype), # type: ignore
+ Lse: T.Tensor(lse_shape, accum_dtype), # type: ignore
+ ):
+ with T.Kernel(seq_len * REPLICATE_H, batch, threads=threads) as (bx, by):
+ Q_shared = T.alloc_shared([H_per_block, D], dtype)
+ KV_shared = T.alloc_shared([BI, D], dtype)
+ O_shared = T.alloc_shared([H_per_block, D], dtype)
+ Lse_shared = T.alloc_shared([H_per_block], accum_dtype)
+ mask = T.alloc_fragment([BI], "bool")
+
+ acc_o = T.alloc_fragment([H_per_block, D], accum_dtype)
+ acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype)
+ S_shared = T.alloc_shared([H_per_block, BI], dtype)
+ sumexp = T.alloc_fragment([H_per_block], accum_dtype)
+ sumexp_i = T.alloc_fragment([H_per_block], accum_dtype)
+ alpha = T.alloc_fragment([H_per_block], accum_dtype)
+ m_i = T.alloc_fragment([H_per_block], accum_dtype)
+ m_i_prev = T.alloc_fragment([H_per_block], accum_dtype)
+
+ T.fill(acc_o, 0)
+ T.fill(sumexp, 0)
+ T.fill(m_i, -(2**30))
+
+ b_i = by
+ s_i = bx if REPLICATE_H == 1 else (bx // REPLICATE_H)
+
+ H0 = 0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * 64
+ H1 = H0 + H_per_block
+
+ T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared)
+
+ for i_i in T.Pipelined(NI, num_stages=num_stages):
+ for bi_i in T.Parallel(BI):
+ mask[bi_i] = Indices[b_i, s_i, i_i * BI + bi_i] != -1
+
+ for bi_i, d_i in T.Parallel(BI, D):
+ KV_shared[bi_i, d_i] = KV[
+ b_i,
+ T.if_then_else(
+ Indices[b_i, s_i, i_i * BI + bi_i] != -1,
+ Indices[b_i, s_i, i_i * BI + bi_i],
+ 0,
+ ),
+ d_i,
+ ]
+
+ T.clear(acc_s)
+ T.gemm(
+ Q_shared,
+ KV_shared,
+ acc_s,
+ transpose_B=True,
+ policy=T.GemmWarpPolicy.FullRow,
+ )
+ for h_i, bi_i in T.Parallel(H_per_block, BI):
+ acc_s[h_i, bi_i] = T.if_then_else(
+ mask[bi_i], acc_s[h_i, bi_i], -T.infinity(acc_s.dtype)
+ )
+ T.copy(m_i, m_i_prev)
+ T.reduce_max(acc_s, m_i, dim=1, clear=False)
+ for h_i in T.Parallel(H_per_block):
+ m_i[h_i] = T.max(m_i[h_i], m_i_prev[h_i])
+ for h_i in T.Parallel(H_per_block):
+ alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale)
+ for h_i, bi_i in T.Parallel(H_per_block, BI):
+ acc_s[h_i, bi_i] = T.exp2(
+ acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale
+ )
+ T.reduce_sum(acc_s, sumexp_i, dim=1)
+ for h_i in T.Parallel(H_per_block):
+ sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i]
+ for h_i, d_i in T.Parallel(H_per_block, D):
+ acc_o[h_i, d_i] = acc_o[h_i, d_i] * alpha[h_i]
+
+ T.copy(acc_s, S_shared)
+ T.gemm(S_shared, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)
+
+ # attn_sink: add exp(attn_sink[h] - max_scaled) to softmax denominator
+ # attn_sink is a pre-scaled logit (same space as scores*sm_scale), so only convert to log2 base
+ for h_i in T.Parallel(H_per_block):
+ sumexp[h_i] += T.exp2(
+ AttnSink[H0 + h_i] * 1.44269504 - m_i[h_i] * sm_scale
+ )
+
+ # Rescale output
+ for h_i, d_i in T.Parallel(H_per_block, D):
+ acc_o[h_i, d_i] /= sumexp[h_i]
+ # LSE = log2(sumexp) + m_i * sm_scale (in log2 space)
+ for h_i in T.Parallel(H_per_block):
+ sumexp[h_i] = T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale
+
+ T.copy(acc_o, Output[b_i, s_i, H0:H1, :])
+ T.copy(sumexp, Lse[b_i, s_i, H0:H1])
+
+ return main
+
+
+sanitize_tilelang_env()
+
+
+def _tilelang_input_dtype(torch_dtype):
+ if torch_dtype is torch.bfloat16:
+ return T.bfloat16
+ raise TypeError(f"DSV4 sparse MLA TileLang launch requires bf16, got {torch_dtype}")
+
+
+def sparse_mqa_fwd_interface(
+ q, kv, attn_sink, topk_idxs, sm_scale=None, block_I=64, num_stages=2, threads=256
+):
+ """Forward interface for V4 sparse MQA attention.
+
+ Args:
+ q: [B, S, H, D] bf16
+ kv: [B, S_kv, D] bf16
+ attn_sink: [H] fp32
+ topk_idxs: [B, S, topk] int32
+ sm_scale: float or None (defaults to 1/sqrt(D))
+
+ Returns:
+ out: [B, S, H, D] bf16
+ lse: [B, S, H] fp32
+ """
+ assert q.is_contiguous() and kv.is_contiguous() and topk_idxs.is_contiguous()
+ batch, seq_len, heads, dim = q.shape
+ _, seq_len_kv, kv_dim = kv.shape
+ assert kv_dim == dim
+ assert kv.dtype == q.dtype
+ _, _, topk = topk_idxs.shape
+ dtype = _tilelang_input_dtype(q.dtype)
+
+ # Pad topk to next multiple of block_I (kernel requires divisibility)
+ padded_topk = (topk + block_I - 1) // block_I * block_I
+ if padded_topk != topk:
+ pad = torch.full(
+ (batch, seq_len, padded_topk - topk),
+ -1,
+ device=topk_idxs.device,
+ dtype=topk_idxs.dtype,
+ )
+ topk_idxs = torch.cat([topk_idxs, pad], dim=-1).contiguous()
+ topk = padded_topk
+
+ with preserve_tilelang_env():
+ kernel = sparse_mqa_fwd(
+ heads,
+ dim,
+ topk,
+ sm_scale,
+ block_I=block_I,
+ num_stages=num_stages,
+ threads=threads,
+ dtype=dtype,
+ )
+ out, lse = kernel(q, kv, attn_sink, topk_idxs)
+ return out, lse
diff --git a/src/art/megatron/dsv4/layer.py b/src/art/megatron/dsv4/layer.py
new file mode 100644
index 000000000..f7a42c0a7
--- /dev/null
+++ b/src/art/megatron/dsv4/layer.py
@@ -0,0 +1,346 @@
+import types
+from typing import Any, cast
+
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.moe.moe_layer import MoELayer
+from megatron.core.transformer.moe.router import TopKRouter
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.transformer_layer import TransformerLayer
+from megatron.core.typed_torch import apply_module
+from megatron.core.utils import make_viewless_tensor
+import torch
+from torch import Tensor
+import torch.nn.functional as F
+
+from art.megatron.dsv4.hyper_connection import (
+ DeepSeekV4HyperConnectionUtil,
+ HCHeadParams,
+)
+from art.megatron.dsv4.utils import freeze_parameters_as_buffers
+
+
+def _first_tensor(value: Any) -> Tensor:
+ return value[0] if isinstance(value, tuple) else value
+
+
+def _input_ids_sbd(
+ input_ids: Tensor, hidden_states: Tensor, tp_group: Any | None
+) -> Tensor:
+ if input_ids.ndim != 2:
+ raise ValueError(
+ f"DSV4 hash routing expects 2D input_ids, got {tuple(input_ids.shape)}."
+ )
+ if input_ids.shape == hidden_states.shape[:2]:
+ return input_ids
+ if input_ids.t().shape == hidden_states.shape[:2]:
+ return input_ids.t().contiguous()
+ input_ids_sbd = input_ids.t().contiguous()
+ if (
+ tp_group is not None
+ and tp_group.size() > 1
+ and input_ids_sbd.shape[1] == hidden_states.shape[1]
+ and input_ids_sbd.shape[0] % tp_group.size() == 0
+ ):
+ local_s = input_ids_sbd.shape[0] // tp_group.size()
+ start = tp_group.rank() * local_s
+ local_input_ids = input_ids_sbd[start : start + local_s]
+ if local_input_ids.shape == hidden_states.shape[:2]:
+ return local_input_ids
+ raise ValueError(
+ "DSV4 hash routing input_ids do not match hidden states: "
+ f"input_ids={tuple(input_ids.shape)} hidden={tuple(hidden_states.shape)}."
+ )
+
+
+class Dsv4Router(TopKRouter):
+ def __init__(self, config: TransformerConfig, *args: Any, **kwargs: Any) -> None:
+ super().__init__(config, *args, **kwargs)
+ self._dsv4_input_ids: Tensor | None = None
+ self._dsv4_is_hash_layer = False
+ if self.topk == 1:
+ freeze_parameters_as_buffers(self)
+
+ def set_layer_number(self, layer_number: int):
+ super().set_layer_number(layer_number)
+ self._dsv4_is_hash_layer = self._compute_is_hash_layer(layer_number)
+ cfg = cast(Any, self.config)
+ if self._is_hash_layer():
+ if "tid2eid" not in self._buffers:
+ vocab_size = getattr(cfg, "vocab_size", None)
+ if vocab_size is None:
+ vocab_size = cfg.padded_vocab_size
+ self.register_buffer(
+ "tid2eid",
+ torch.zeros(int(vocab_size), self.topk, dtype=torch.long),
+ persistent=True,
+ )
+ expert_pattern = (
+ torch.arange(self.topk, dtype=torch.long) % int(cfg.num_moe_experts)
+ ).expand(int(vocab_size), -1)
+ cast(Tensor, self.tid2eid).copy_(expert_pattern)
+ routing_fn = Dsv4Router._hash_routing
+ else:
+ if "e_score_correction_bias" not in self._buffers:
+ self.register_buffer(
+ "e_score_correction_bias",
+ torch.zeros(int(cfg.num_moe_experts), dtype=torch.float32),
+ persistent=True,
+ )
+ routing_fn = Dsv4Router._moe_routing
+ object.__setattr__(self, "routing", types.MethodType(routing_fn, self))
+
+ def set_input_ids(self, input_ids: Tensor | None) -> None:
+ self._dsv4_input_ids = input_ids
+
+ def _compute_is_hash_layer(self, layer_number: int | None) -> bool:
+ cfg = cast(Any, self.config)
+ return layer_number is not None and layer_number <= int(cfg.dsv4_n_hash_layers)
+
+ def _is_hash_layer(self) -> bool:
+ return self._dsv4_is_hash_layer
+
+ def _scores(self, logits: Tensor) -> Tensor:
+ if self.score_function == "sigmoid":
+ return torch.sigmoid(logits)
+ if self.score_function == "softmax":
+ return F.softmax(logits, dim=-1, dtype=torch.float32).to(logits.dtype)
+ if self.score_function == "sqrtsoftplus":
+ return F.softplus(logits.float()).sqrt().to(logits.dtype)
+ raise ValueError(
+ f"Unsupported DSV4 router score function {self.score_function!r}."
+ )
+
+ def _select_indices(
+ self, selection_scores: Tensor, default_indices: Tensor
+ ) -> Tensor:
+ router_replay = getattr(self, "router_replay", None)
+ if router_replay is None:
+ return default_indices.long()
+
+ def default_compute_topk(
+ local_scores: Tensor,
+ topk: int,
+ num_groups: int | None = None,
+ group_topk: int | None = None,
+ ) -> tuple[Tensor, Tensor]:
+ del num_groups, group_topk
+ del topk
+ return local_scores.gather(1, default_indices), default_indices
+
+ _selected_scores, indices = router_replay.get_replay_topk(
+ selection_scores,
+ self.topk,
+ None,
+ None,
+ default_compute_topk,
+ )
+ return indices.long()
+
+ def _finish_routing(
+ self,
+ scores: Tensor,
+ indices: Tensor,
+ padding_mask: Tensor | None,
+ num_moe_experts: int,
+ ) -> tuple[Tensor, Tensor]:
+ cfg = cast(Any, self.config)
+ if indices.shape[-1] != self.topk:
+ raise RuntimeError(
+ "DSV4 router selected an invalid number of experts: "
+ f"selected={indices.shape[-1]} expected={self.topk}."
+ )
+ selected_probs = scores.gather(1, indices)
+ if self.score_function != "softmax":
+ selected_probs = selected_probs / (
+ selected_probs.sum(dim=-1, keepdim=True) + 1e-20
+ )
+ selected_probs = selected_probs * float(cfg.moe_router_topk_scaling_factor)
+ probs = torch.zeros_like(scores).scatter(1, indices, selected_probs)
+ routing_map = F.one_hot(indices, num_classes=num_moe_experts).sum(dim=1).bool()
+ if padding_mask is not None:
+ valid = padding_mask.reshape(-1).bool()
+ probs = torch.where(valid.unsqueeze(-1), probs, torch.zeros_like(probs))
+ routing_map = routing_map & valid.unsqueeze(-1)
+ return probs, routing_map
+
+ def _hash_routing(self, logits: Tensor, padding_mask: Tensor | None = None):
+ cfg = cast(Any, self.config)
+ num_moe_experts = int(cfg.num_moe_experts)
+ scores = self._scores(logits.view(-1, num_moe_experts))
+ if self._dsv4_input_ids is None:
+ raise RuntimeError(
+ "DSV4 hash router requires input_ids for hash-moe layers."
+ )
+ tid2eid = cast(Tensor, self.tid2eid)
+ default_indices = tid2eid[self._dsv4_input_ids.reshape(-1)].long()
+ indices = self._select_indices(scores, default_indices)
+ return self._finish_routing(scores, indices, padding_mask, num_moe_experts)
+
+ def _moe_routing(self, logits: Tensor, padding_mask: Tensor | None = None):
+ cfg = cast(Any, self.config)
+ num_moe_experts = int(cfg.num_moe_experts)
+ scores = self._scores(logits.view(-1, num_moe_experts))
+ selection_scores = scores
+ e_score_correction_bias = getattr(self, "e_score_correction_bias", None)
+ if e_score_correction_bias is not None:
+ selection_scores = selection_scores + e_score_correction_bias
+ default_indices = selection_scores.topk(self.topk, dim=-1, sorted=False).indices
+ indices = self._select_indices(selection_scores, default_indices)
+ return self._finish_routing(scores, indices, padding_mask, num_moe_experts)
+
+ def routing(self, logits: Tensor, padding_mask: Tensor | None = None):
+ if self._is_hash_layer():
+ return self._hash_routing(logits, padding_mask)
+ return self._moe_routing(logits, padding_mask)
+
+
+class Dsv4MoELayer(MoELayer):
+ _dsv4_input_ids: Tensor | None = None
+
+ def set_input_ids(self, input_ids: Tensor | None) -> None:
+ self._dsv4_input_ids = input_ids
+
+ def forward(
+ self,
+ hidden_states: Tensor,
+ intermediate_tensors=None,
+ padding_mask: Tensor | None = None,
+ input_ids: Tensor | None = None,
+ ):
+ if isinstance(self.router, Dsv4Router):
+ input_ids = input_ids if input_ids is not None else self._dsv4_input_ids
+ router_input_ids = None
+ if input_ids is not None:
+ router_input_ids = _input_ids_sbd(
+ input_ids, hidden_states, getattr(self, "attn_tp_group", None)
+ )
+ self.router.set_input_ids(router_input_ids)
+ try:
+ return super().forward(
+ hidden_states,
+ intermediate_tensors=intermediate_tensors,
+ padding_mask=padding_mask,
+ )
+ finally:
+ if isinstance(self.router, Dsv4Router):
+ self.router.set_input_ids(None)
+
+
+class Dsv4TransformerLayer(TransformerLayer):
+ def __init__(self, *args: Any, **kwargs: Any) -> None:
+ super().__init__(*args, **kwargs)
+ cfg = cast(Any, self.config)
+ hc = int(cfg.dsv4_hc_mult)
+ hc_dim = hc * self.config.hidden_size
+ mix = (2 + hc) * hc
+ self.hc_attn_fn = torch.nn.Parameter(
+ torch.empty(mix, hc_dim, dtype=torch.float32)
+ )
+ self.hc_attn_base = torch.nn.Parameter(torch.empty(mix, dtype=torch.float32))
+ self.hc_attn_scale = torch.nn.Parameter(torch.empty(3, dtype=torch.float32))
+ self.hc_ffn_fn = torch.nn.Parameter(
+ torch.empty(mix, hc_dim, dtype=torch.float32)
+ )
+ self.hc_ffn_base = torch.nn.Parameter(torch.empty(mix, dtype=torch.float32))
+ self.hc_ffn_scale = torch.nn.Parameter(torch.empty(3, dtype=torch.float32))
+ self._keep_fp32_parameters = (
+ "hc_attn_fn",
+ "hc_attn_base",
+ "hc_attn_scale",
+ "hc_ffn_fn",
+ "hc_ffn_base",
+ "hc_ffn_scale",
+ )
+ for param in (
+ self.hc_attn_fn,
+ self.hc_attn_base,
+ self.hc_attn_scale,
+ self.hc_ffn_fn,
+ self.hc_ffn_base,
+ self.hc_ffn_scale,
+ ):
+ setattr(param, "_keep_fp32", True)
+ self.hc_util = DeepSeekV4HyperConnectionUtil(self.config)
+
+ def forward(
+ self,
+ hidden_states: Tensor,
+ attention_mask: Tensor | None = None,
+ context: Tensor | None = None,
+ context_mask: Tensor | None = None,
+ rotary_pos_emb: Tensor | None = None,
+ rotary_pos_cos: Tensor | None = None,
+ rotary_pos_sin: Tensor | None = None,
+ rotary_pos_cos_sin: Tensor | None = None,
+ attention_bias: Tensor | None = None,
+ inference_context: Any = None,
+ packed_seq_params: Any = None,
+ sequence_len_offset: Tensor | None = None,
+ padding_mask: Tensor | None = None,
+ input_ids: Tensor | None = None,
+ **_: Any,
+ ):
+ if hidden_states.ndim == 3:
+ hidden_states = self.hc_util.block_expand(hidden_states)
+
+ attn_input, post, comb = self.hc_util.layer_pre(
+ hidden_states, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
+ )
+ attn_input = _first_tensor(apply_module(self.input_layernorm)(attn_input))
+ attn_output = self.self_attention(
+ attn_input,
+ attention_mask=attention_mask,
+ inference_context=inference_context,
+ rotary_pos_emb=rotary_pos_emb,
+ rotary_pos_cos=rotary_pos_cos,
+ rotary_pos_sin=rotary_pos_sin,
+ rotary_pos_cos_sin=rotary_pos_cos_sin,
+ attention_bias=attention_bias,
+ packed_seq_params=packed_seq_params,
+ sequence_len_offset=sequence_len_offset,
+ )
+ hidden_states = self.hc_util.layer_post(attn_output, hidden_states, post, comb)
+
+ mlp_input, post, comb = self.hc_util.layer_pre(
+ hidden_states, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base
+ )
+ mlp_input = _first_tensor(apply_module(self.pre_mlp_layernorm)(mlp_input))
+ if isinstance(self.mlp, Dsv4MoELayer):
+ mlp_output = self.mlp(
+ mlp_input, padding_mask=padding_mask, input_ids=input_ids
+ )
+ else:
+ mlp_output = self.mlp(mlp_input, padding_mask=padding_mask)
+ hidden_states = self.hc_util.layer_post(mlp_output, hidden_states, post, comb)
+ return make_viewless_tensor(
+ inp=hidden_states,
+ requires_grad=hidden_states.requires_grad,
+ keep_graph=True,
+ ), context
+
+
+class Dsv4FinalNorm(MegatronModule):
+ def __init__(
+ self,
+ config: TransformerConfig,
+ hidden_size: int,
+ eps: float,
+ ) -> None:
+ super().__init__(config=config)
+ self.weight = torch.nn.Parameter(torch.ones(hidden_size))
+ self.eps = eps
+ self.hc_head_params = HCHeadParams(config)
+ self.hc_util = DeepSeekV4HyperConnectionUtil(config)
+
+ def forward(self, hidden_states: Tensor) -> Tensor:
+ if hidden_states.ndim == 4:
+ hidden_states = self.hc_util.block_head(
+ hidden_states,
+ self.hc_head_params.hc_head_fn,
+ self.hc_head_params.hc_head_scale,
+ self.hc_head_params.hc_head_base,
+ )
+ dtype = hidden_states.dtype
+ normed = hidden_states.float()
+ normed = normed * torch.rsqrt(normed.square().mean(-1, keepdim=True) + self.eps)
+ return (normed * self.weight.float()).to(dtype)
diff --git a/src/art/megatron/dsv4/lora.py b/src/art/megatron/dsv4/lora.py
new file mode 100644
index 000000000..658077081
--- /dev/null
+++ b/src/art/megatron/dsv4/lora.py
@@ -0,0 +1,665 @@
+from __future__ import annotations
+
+import math
+from typing import Any
+
+from megatron.core import parallel_state as ps
+from megatron.core.tensor_parallel.mappings import (
+ copy_to_tensor_model_parallel_region,
+ reduce_from_tensor_model_parallel_region,
+)
+import torch
+
+from art.megatron.lora import (
+ GRAD_SYNC_OP_NONE,
+ GRAD_SYNC_OP_SUM,
+ LORA_ALPHA,
+ TP_DEFAULT_GRAD_SYNC_DOMAIN,
+ LoRA,
+ LoRAParallelSpec,
+)
+
+
+def _weight_device_dtype(weight: torch.Tensor) -> tuple[torch.device, torch.dtype]:
+ return weight.device, weight.dtype
+
+
+def replicated_lora(
+ *,
+ adapter_model_prefix: str,
+ weight: torch.Tensor,
+ in_features: int,
+ out_features: int,
+ rank: int,
+ alpha: int = LORA_ALPHA,
+) -> LoRA:
+ device, dtype = _weight_device_dtype(weight)
+ sync = LoRAParallelSpec(
+ grad_sync_domain=TP_DEFAULT_GRAD_SYNC_DOMAIN,
+ grad_sync_op=GRAD_SYNC_OP_NONE,
+ )
+ return LoRA(
+ adapter_model_prefix=adapter_model_prefix,
+ in_features=in_features,
+ out_features=out_features,
+ rank=rank,
+ alpha=alpha,
+ dtype=dtype,
+ device=device,
+ a_parallel_spec=sync,
+ b_parallel_spec=sync,
+ allreduce=True,
+ )
+
+
+def column_parallel_lora(
+ *,
+ adapter_model_prefix: str,
+ weight: torch.Tensor,
+ in_features: int,
+ out_features: int,
+ rank: int,
+ alpha: int = LORA_ALPHA,
+) -> LoRA:
+ device, dtype = _weight_device_dtype(weight)
+ a_spec = LoRAParallelSpec(
+ grad_sync_domain=TP_DEFAULT_GRAD_SYNC_DOMAIN,
+ grad_sync_op=GRAD_SYNC_OP_SUM,
+ )
+ b_spec = LoRAParallelSpec(
+ sharded=True,
+ shard_dim=-1,
+ grad_sync_domain=TP_DEFAULT_GRAD_SYNC_DOMAIN,
+ grad_sync_op=GRAD_SYNC_OP_NONE,
+ )
+ return LoRA(
+ adapter_model_prefix=adapter_model_prefix,
+ in_features=in_features,
+ out_features=out_features,
+ rank=rank,
+ alpha=alpha,
+ dtype=dtype,
+ device=device,
+ a_parallel_spec=a_spec,
+ b_parallel_spec=b_spec,
+ allreduce=True,
+ )
+
+
+def row_parallel_lora(
+ *,
+ adapter_model_prefix: str,
+ weight: torch.Tensor,
+ in_features: int,
+ out_features: int,
+ rank: int,
+ alpha: int = LORA_ALPHA,
+) -> LoRA:
+ device, dtype = _weight_device_dtype(weight)
+ a_spec = LoRAParallelSpec(
+ sharded=True,
+ shard_dim=-2,
+ grad_sync_domain=TP_DEFAULT_GRAD_SYNC_DOMAIN,
+ grad_sync_op=GRAD_SYNC_OP_NONE,
+ )
+ b_spec = LoRAParallelSpec(
+ grad_sync_domain=TP_DEFAULT_GRAD_SYNC_DOMAIN,
+ grad_sync_op=GRAD_SYNC_OP_SUM,
+ )
+ return LoRA(
+ adapter_model_prefix=adapter_model_prefix,
+ in_features=in_features,
+ out_features=out_features,
+ rank=rank,
+ alpha=alpha,
+ dtype=dtype,
+ device=device,
+ a_parallel_spec=a_spec,
+ b_parallel_spec=b_spec,
+ allreduce=True,
+ )
+
+
+class Dsv4GroupedOutputLoRA(torch.nn.Module):
+ def __init__(
+ self,
+ *,
+ adapter_model_prefix: str,
+ weight: torch.Tensor,
+ in_features: int,
+ out_features: int,
+ local_groups: int,
+ rank: int,
+ alpha: int = LORA_ALPHA,
+ ) -> None:
+ super().__init__()
+ if out_features % local_groups != 0:
+ raise ValueError(
+ f"{adapter_model_prefix}: out_features={out_features} is not "
+ f"divisible by local_groups={local_groups}"
+ )
+ self.local_groups = local_groups
+ self.out_per_group = out_features // local_groups
+ self.lora = column_parallel_lora(
+ adapter_model_prefix=adapter_model_prefix,
+ weight=weight,
+ in_features=in_features,
+ out_features=out_features,
+ rank=rank,
+ alpha=alpha,
+ )
+
+ @property
+ def adapter_model_prefix(self) -> str:
+ return self.lora.adapter_model_prefix
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ rank = self.lora.A_T.shape[-1]
+ hidden = x @ self.lora.A_T
+ weight = self.lora.B_T.view(rank, self.local_groups, self.out_per_group)
+ out = torch.einsum("bsgr,rgo->bsgo", hidden, weight)
+ return out if self.lora.scale == 1.0 else out * self.lora.scale
+
+
+class Dsv4RowOutputLoRA(torch.nn.Module):
+ def __init__(
+ self,
+ *,
+ adapter_model_prefix: str,
+ weight: torch.Tensor,
+ in_features: int,
+ out_features: int,
+ tp_group: Any,
+ rank: int,
+ alpha: int = LORA_ALPHA,
+ ) -> None:
+ super().__init__()
+ self.tp_group = tp_group
+ self.lora = row_parallel_lora(
+ adapter_model_prefix=adapter_model_prefix,
+ weight=weight,
+ in_features=in_features,
+ out_features=out_features,
+ rank=rank,
+ alpha=alpha,
+ )
+
+ @property
+ def adapter_model_prefix(self) -> str:
+ return self.lora.adapter_model_prefix
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ out = self.lora(x)
+ if self.tp_group is not None and self.tp_group.size() > 1:
+ out = reduce_from_tensor_model_parallel_region(out, group=self.tp_group)
+ return out
+
+
+class Dsv4PlainLoRA(torch.nn.Module):
+ def __init__(self, lora: LoRA, *, sync_input_grad_group: Any | None = None) -> None:
+ super().__init__()
+ self.lora = lora
+ self.sync_input_grad_group = sync_input_grad_group
+
+ @property
+ def adapter_model_prefix(self) -> str:
+ return self.lora.adapter_model_prefix
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if (
+ self.sync_input_grad_group is not None
+ and self.sync_input_grad_group.size() > 1
+ ):
+ x = copy_to_tensor_model_parallel_region(
+ x, group=self.sync_input_grad_group
+ )
+ return self.lora(x)
+
+
+def _targets_include(target_modules: set[str], *names: str) -> bool:
+ return not target_modules or any(name in target_modules for name in names)
+
+
+def _attach_lora(
+ module: torch.nn.Module, attr_name: str, lora: torch.nn.Module
+) -> None:
+ if hasattr(module, attr_name):
+ raise RuntimeError(f"{module.__class__.__name__}.{attr_name} is already set")
+ setattr(module, attr_name, lora)
+
+
+def apply_dsv4_attention_lora(
+ attention: Any,
+ *,
+ adapter_model_prefix: str,
+ target_modules: set[str],
+ rank: int,
+ alpha: int,
+) -> None:
+ if _targets_include(target_modules, "q_a_proj"):
+ _attach_lora(
+ attention,
+ "wq_a_lora",
+ Dsv4PlainLoRA(
+ replicated_lora(
+ adapter_model_prefix=f"{adapter_model_prefix}.self_attn.q_a_proj",
+ weight=attention.wq_a.weight,
+ in_features=attention.wq_a.weight.shape[1],
+ out_features=attention.wq_a.weight.shape[0],
+ rank=rank,
+ alpha=alpha,
+ )
+ ),
+ )
+ if _targets_include(target_modules, "q_b_proj"):
+ _attach_lora(
+ attention,
+ "wq_b_lora",
+ Dsv4PlainLoRA(
+ column_parallel_lora(
+ adapter_model_prefix=f"{adapter_model_prefix}.self_attn.q_b_proj",
+ weight=attention.wq_b.weight,
+ in_features=attention.wq_b.weight.shape[1],
+ out_features=attention.wq_b.weight.shape[0],
+ rank=rank,
+ alpha=alpha,
+ ),
+ sync_input_grad_group=attention.tp_group,
+ ),
+ )
+ if _targets_include(target_modules, "kv_proj"):
+ _attach_lora(
+ attention,
+ "wkv_lora",
+ Dsv4PlainLoRA(
+ replicated_lora(
+ adapter_model_prefix=f"{adapter_model_prefix}.self_attn.kv_proj",
+ weight=attention.wkv.weight,
+ in_features=attention.wkv.weight.shape[1],
+ out_features=attention.wkv.weight.shape[0],
+ rank=rank,
+ alpha=alpha,
+ )
+ ),
+ )
+ if _targets_include(target_modules, "o_a_proj"):
+ _attach_lora(
+ attention,
+ "wo_a_lora",
+ Dsv4GroupedOutputLoRA(
+ adapter_model_prefix=f"{adapter_model_prefix}.self_attn.o_a_proj",
+ weight=attention.wo_a.weight,
+ in_features=attention.wo_a.weight.shape[1],
+ out_features=attention.wo_a.weight.shape[0],
+ local_groups=attention.n_local_groups,
+ rank=rank,
+ alpha=alpha,
+ ),
+ )
+ if _targets_include(target_modules, "o_b_proj"):
+ _attach_lora(
+ attention,
+ "wo_b_lora",
+ Dsv4RowOutputLoRA(
+ adapter_model_prefix=f"{adapter_model_prefix}.self_attn.o_b_proj",
+ weight=attention.wo_b.weight,
+ in_features=attention.wo_b.weight.shape[1],
+ out_features=attention.wo_b.weight.shape[0],
+ tp_group=attention.tp_group,
+ rank=rank,
+ alpha=alpha,
+ ),
+ )
+ compressor = getattr(attention, "compressor", None)
+ if compressor is None:
+ return
+ if _targets_include(target_modules, "compressor.kv_proj"):
+ _attach_lora(
+ compressor,
+ "kv_proj_lora",
+ Dsv4PlainLoRA(
+ replicated_lora(
+ adapter_model_prefix=(
+ f"{adapter_model_prefix}.self_attn.compressor.kv_proj"
+ ),
+ weight=compressor.wkv.weight,
+ in_features=compressor.wkv.weight.shape[1],
+ out_features=compressor.wkv.weight.shape[0],
+ rank=rank,
+ alpha=alpha,
+ )
+ ),
+ )
+ if _targets_include(target_modules, "compressor.gate_proj"):
+ _attach_lora(
+ compressor,
+ "gate_proj_lora",
+ Dsv4PlainLoRA(
+ replicated_lora(
+ adapter_model_prefix=(
+ f"{adapter_model_prefix}.self_attn.compressor.gate_proj"
+ ),
+ weight=compressor.wgate.weight,
+ in_features=compressor.wgate.weight.shape[1],
+ out_features=compressor.wgate.weight.shape[0],
+ rank=rank,
+ alpha=alpha,
+ )
+ ),
+ )
+
+
+def _unwrap_lora(module: torch.nn.Module | None) -> LoRA | None:
+ if isinstance(module, LoRA):
+ return module
+ nested = getattr(module, "lora", None) if module is not None else None
+ return nested if isinstance(nested, LoRA) else None
+
+
+def _adapter_alpha_dim(lora: LoRA) -> tuple[int, int]:
+ dim = int(lora.A_T.shape[-1])
+ alpha = float(lora.scale) * dim
+ rounded = round(alpha)
+ if not math.isclose(alpha, rounded):
+ raise RuntimeError(f"{lora.adapter_model_prefix}: non-integral alpha={alpha}")
+ return rounded, dim
+
+
+def _adapter_param_prefix(base_prefix: str, adapter_key: str | None) -> str:
+ if adapter_key is None:
+ return f"{base_prefix}.adapter"
+ return f"{base_prefix}.adapter.{adapter_key}"
+
+
+def _adapter_weight(
+ base_prefix: str,
+ lora: LoRA,
+ *,
+ adapter_key: str | None = None,
+) -> Any:
+ from megatron.bridge.models.conversion.model_bridge import MegatronWeightTuple
+ from megatron.bridge.models.conversion.peft_bridge import AdapterWeight
+
+ alpha, dim = _adapter_alpha_dim(lora)
+ linear_in = lora.A_T.transpose(-1, -2).contiguous()
+ linear_out = lora.B_T.transpose(-1, -2).contiguous()
+ param_prefix = _adapter_param_prefix(base_prefix, adapter_key)
+ return AdapterWeight(
+ global_base_prefix=base_prefix,
+ adapter_key=adapter_key,
+ alpha=alpha,
+ dim=dim,
+ linear_in_weight=MegatronWeightTuple(
+ param_name=f"{param_prefix}.linear_in.weight",
+ weight=linear_in,
+ vp_stage=0,
+ ),
+ linear_out_weight=MegatronWeightTuple(
+ param_name=f"{param_prefix}.linear_out.weight",
+ weight=linear_out,
+ vp_stage=0,
+ ),
+ )
+
+
+def _tensor_parallel_group() -> Any | None:
+ try:
+ return ps.get_tensor_model_parallel_group()
+ except Exception:
+ return None
+
+
+def _gather_tp_input_shards(
+ linear_in: torch.Tensor, tp_group: Any | None
+) -> torch.Tensor:
+ dist = torch.distributed
+ if not (dist.is_available() and dist.is_initialized()):
+ return linear_in
+ tp_group = tp_group or _tensor_parallel_group()
+ world_size = tp_group.size() if tp_group is not None else dist.get_world_size()
+ if world_size <= 1:
+ return linear_in
+ gathered = [torch.empty_like(linear_in) for _ in range(world_size)]
+ dist.all_gather(gathered, linear_in.contiguous(), group=tp_group)
+ return torch.cat(gathered, dim=-1)
+
+
+def _row_parallel_adapter_weight(
+ base_prefix: str,
+ lora: LoRA,
+ *,
+ tp_group: Any | None,
+) -> Any:
+ from megatron.bridge.models.conversion.model_bridge import MegatronWeightTuple
+ from megatron.bridge.models.conversion.peft_bridge import AdapterWeight
+
+ alpha, dim = _adapter_alpha_dim(lora)
+ linear_in = _gather_tp_input_shards(
+ lora.A_T.transpose(-1, -2).contiguous(), tp_group
+ )
+ linear_out = lora.B_T.transpose(-1, -2).contiguous()
+ param_prefix = _adapter_param_prefix(base_prefix, None)
+ return AdapterWeight(
+ global_base_prefix=base_prefix,
+ adapter_key=None,
+ alpha=alpha,
+ dim=dim,
+ linear_in_weight=MegatronWeightTuple(
+ param_name=f"{param_prefix}.linear_in.weight",
+ weight=linear_in,
+ vp_stage=0,
+ ),
+ linear_out_weight=MegatronWeightTuple(
+ param_name=f"{param_prefix}.linear_out.weight",
+ weight=linear_out,
+ vp_stage=0,
+ ),
+ )
+
+
+def _add_adapter(
+ adapter_weights_by_base: dict[str, list[Any]],
+ *,
+ base_prefix: str,
+ lora_module: torch.nn.Module | None,
+) -> None:
+ lora = _unwrap_lora(lora_module)
+ if lora is None:
+ return
+ adapter_weights_by_base[f"{base_prefix}.weight"] = [
+ _adapter_weight(base_prefix, lora)
+ ]
+
+
+def _add_row_parallel_adapter(
+ adapter_weights_by_base: dict[str, list[Any]],
+ *,
+ base_prefix: str,
+ lora_module: torch.nn.Module | None,
+) -> None:
+ lora = _unwrap_lora(lora_module)
+ if lora is None:
+ return
+ adapter_weights_by_base[f"{base_prefix}.weight"] = [
+ _row_parallel_adapter_weight(
+ base_prefix,
+ lora,
+ tp_group=getattr(lora_module, "tp_group", None),
+ )
+ ]
+
+
+def add_dsv4_attention_adapter_weights(
+ adapter_weights_by_base: dict[str, list[Any]],
+ *,
+ layer_prefix: str,
+ attention: Any,
+) -> None:
+ attn_prefix = f"{layer_prefix}.self_attention"
+ _add_adapter(
+ adapter_weights_by_base,
+ base_prefix=f"{attn_prefix}.wq_a",
+ lora_module=getattr(attention, "wq_a_lora", None),
+ )
+ _add_adapter(
+ adapter_weights_by_base,
+ base_prefix=f"{attn_prefix}.wq_b",
+ lora_module=getattr(attention, "wq_b_lora", None),
+ )
+ _add_adapter(
+ adapter_weights_by_base,
+ base_prefix=f"{attn_prefix}.wkv",
+ lora_module=getattr(attention, "wkv_lora", None),
+ )
+ _add_adapter(
+ adapter_weights_by_base,
+ base_prefix=f"{attn_prefix}.wo_a",
+ lora_module=getattr(attention, "wo_a_lora", None),
+ )
+ _add_row_parallel_adapter(
+ adapter_weights_by_base,
+ base_prefix=f"{attn_prefix}.wo_b",
+ lora_module=getattr(attention, "wo_b_lora", None),
+ )
+ compressor = getattr(attention, "compressor", None)
+ if compressor is None:
+ return
+ _add_adapter(
+ adapter_weights_by_base,
+ base_prefix=f"{attn_prefix}.compressor.wkv",
+ lora_module=getattr(compressor, "kv_proj_lora", None),
+ )
+ _add_adapter(
+ adapter_weights_by_base,
+ base_prefix=f"{attn_prefix}.compressor.wgate",
+ lora_module=getattr(compressor, "gate_proj_lora", None),
+ )
+
+
+def add_dsv4_shared_experts_adapter_weights(
+ adapter_weights_by_base: dict[str, list[Any]],
+ *,
+ layer_prefix: str,
+ shared_experts: Any,
+) -> None:
+ from art.megatron.lora import SharedExpertsLinearFC1LoRA, SharedExpertsLinearFC2LoRA
+
+ linear_fc1 = getattr(shared_experts, "linear_fc1", None)
+ if isinstance(linear_fc1, SharedExpertsLinearFC1LoRA):
+ base_prefix = f"{layer_prefix}.mlp.shared_experts.linear_fc1"
+ adapter_weights_by_base[f"{base_prefix}.weight"] = [
+ _adapter_weight(
+ base_prefix,
+ linear_fc1.gate_lora,
+ adapter_key="adapter_gate",
+ ),
+ _adapter_weight(
+ base_prefix,
+ linear_fc1.up_lora,
+ adapter_key="adapter_up",
+ ),
+ ]
+
+ linear_fc2 = getattr(shared_experts, "linear_fc2", None)
+ if isinstance(linear_fc2, SharedExpertsLinearFC2LoRA):
+ base_prefix = f"{layer_prefix}.mlp.shared_experts.linear_fc2"
+ adapter_weights_by_base[f"{base_prefix}.weight"] = [
+ _row_parallel_adapter_weight(
+ base_prefix,
+ linear_fc2.row_parallel_lora.lora,
+ tp_group=_tensor_parallel_group(),
+ )
+ ]
+
+
+def _triton_autotune_key(
+ autotuner: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
+) -> tuple[Any, ...]:
+ nargs = dict(zip(autotuner.arg_names, args))
+ all_args = {**nargs, **kwargs}
+ named_args = {
+ key: value for key, value in all_args.items() if key in autotuner.arg_names
+ }
+ key = [named_args[name] for name in autotuner.keys if name in named_args]
+ for arg in named_args.values():
+ if hasattr(arg, "dtype"):
+ key.append(str(arg.dtype))
+ return tuple(key)
+
+
+def _dsv4_te_permutation_config(hidden_size: int) -> Any:
+ import triton
+
+ block_size = 1 << max(6, min(12, (max(1, hidden_size) - 1).bit_length()))
+ return triton.Config({"BLOCK_SIZE": block_size})
+
+
+def _install_dsv4_triton_static_config(autotuner: Any) -> None:
+ """Bypass TE Triton permutation autotune during DSV4 ETP backward.
+
+ DSV4 ETP validation can make every rank hit the same first-use TE MoE
+ permutation backward path at once. Serializing autotune inside autograd is
+ unsafe because other ranks can be waiting at distributed edges of the same
+ backward graph. Instead, cache a deterministic BLOCK_SIZE before Triton's
+ Autotuner benchmarks. The failure risk is performance from a non-autotuned
+ copy/unpermute tile, not a change in model math.
+ """
+
+ if bool(getattr(autotuner, "_art_dsv4_static_config_wrapped", False)):
+ return
+ original_run = autotuner.run
+
+ def static_config_run(*args: Any, **kwargs: Any) -> Any:
+ configs = getattr(autotuner, "configs", ())
+ if len(configs) <= 1:
+ return original_run(*args, **kwargs)
+ key = _triton_autotune_key(autotuner, args, kwargs)
+ cache = getattr(autotuner, "cache", None)
+ if isinstance(cache, dict) and key not in cache and key:
+ cache[key] = _dsv4_te_permutation_config(int(key[0]))
+ return original_run(*args, **kwargs)
+
+ static_config_run._art_dsv4_static_config_wrapped = True # type: ignore[attr-defined]
+ autotuner.run = static_config_run
+ autotuner._art_dsv4_static_config_wrapped = True
+
+
+def install_dsv4_te_permutation_static_configs() -> None:
+ """Install DSV4-only static configs for TE MoE permutation Triton autotuners.
+
+ This is called from the DSV4 handler only when expert tensor parallelism is
+ enabled. It is intentionally not a generic compile workaround: the current
+ correctness blocker is DSV4 ETP first-backward autotune inside distributed
+ autograd, and the performance goal is to bypass only cold autotune while
+ leaving normal launches on Triton's regular JIT path.
+ """
+
+ from transformer_engine.common.triton import permutation as te_permutation
+
+ for name in (
+ "_permute_kernel",
+ "_unpermute_kernel",
+ "_unpermute_bwd_with_merging_probs_kernel",
+ "_sort_chunks_by_map_kernel",
+ ):
+ autotuner = getattr(te_permutation, name, None)
+ if hasattr(autotuner, "run"):
+ _install_dsv4_triton_static_config(autotuner)
+
+
+def disable_dsv4_etp_shared_expert_lora_compile(shared_experts: Any) -> None:
+ """Keep DSV4 ETP shared-expert down LoRA outside compiled layer graphs.
+
+ Torch 2.11 can spend unbounded time compiling the row-parallel shared-expert
+ LoRA path when DSV4 runs with expert tensor parallelism. The barrier is
+ instance-local and leaves non-ETP and generic MoE shared experts unchanged.
+ """
+ from art.megatron.lora import SharedExpertsLinearFC2LoRA
+
+ linear_fc2 = getattr(shared_experts, "linear_fc2", None)
+ if not isinstance(linear_fc2, SharedExpertsLinearFC2LoRA):
+ return
+ if bool(getattr(linear_fc2, "_art_dsv4_compile_disabled", False)):
+ return
+ linear_fc2.forward = torch.compiler.disable(linear_fc2.forward)
+ linear_fc2._art_dsv4_compile_disabled = True
diff --git a/src/art/megatron/dsv4/rope.py b/src/art/megatron/dsv4/rope.py
new file mode 100644
index 000000000..5f1a07d71
--- /dev/null
+++ b/src/art/megatron/dsv4/rope.py
@@ -0,0 +1,253 @@
+from functools import lru_cache
+import math
+from typing import Any, cast
+
+from megatron.core.transformer import TransformerConfig
+import torch
+from torch import nn
+
+_DEVICE_ROPE_CACHES: dict[tuple[object, ...], torch.Tensor] = {}
+
+
+@lru_cache(2)
+def precompute_freqs_cis(
+ dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow
+) -> torch.Tensor:
+ """Precompute the complex rotary frequencies for RoPE, with optional YaRN smoothing.
+
+ When ``original_seq_len > 0``, applies YaRN factor rescaling interpolated
+ by a linear ramp between ``beta_fast`` and ``beta_slow``. Otherwise the
+ base frequencies are used verbatim.
+ """
+
+ def find_correction_dim(num_rotations, dim, base, max_seq_len):
+ return (
+ dim
+ * math.log(max_seq_len / (num_rotations * 2 * math.pi))
+ / (2 * math.log(base))
+ )
+
+ def find_correction_range(low_rot, high_rot, dim, base, max_seq_len):
+ low = math.floor(find_correction_dim(low_rot, dim, base, max_seq_len))
+ high = math.ceil(find_correction_dim(high_rot, dim, base, max_seq_len))
+ return max(low, 0), min(high, dim - 1)
+
+ def linear_ramp_factor(min, max, dim):
+ if min == max:
+ max += 0.001
+ linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)
+ ramp_func = torch.clamp(linear_func, 0, 1)
+ return ramp_func
+
+ freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+ if original_seq_len > 0:
+ low, high = find_correction_range(
+ beta_fast, beta_slow, dim, base, original_seq_len
+ )
+ smooth = 1 - linear_ramp_factor(low, high, dim // 2)
+ freqs = freqs / factor * (1 - smooth) + freqs * smooth
+
+ t = torch.arange(seqlen)
+ freqs = torch.outer(t, freqs)
+ freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+ return freqs_cis
+
+
+def apply_rotary_emb(
+ x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False
+) -> torch.Tensor:
+ """Apply RoPE in-place to the last dim of ``x``.
+
+ ``x`` has shape ``[..., dim]`` where ``dim`` is even; the last-dim pairs are
+ treated as complex numbers multiplied by ``freqs_cis``. When ``inverse=True``
+ the conjugate rotation is applied (used for the indexer's inverse rope).
+ """
+ y = x
+ x = torch.view_as_complex(x.float().unflatten(-1, (-1, 2)))
+ if inverse:
+ freqs_cis = freqs_cis.conj()
+ if x.ndim == 3:
+ if freqs_cis.ndim == 2:
+ freqs_cis = freqs_cis.view(1, x.size(1), x.size(-1))
+ else:
+ freqs_cis = freqs_cis.view(x.size(0), x.size(1), x.size(-1))
+ else:
+ if freqs_cis.ndim == 2:
+ freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1))
+ else:
+ freqs_cis = freqs_cis.view(x.size(0), x.size(1), 1, x.size(-1))
+ x = torch.view_as_real(x * freqs_cis).flatten(-2)
+ y.copy_(x)
+ return y
+
+
+def wrapped_precompute_freqs_cis(
+ config: TransformerConfig,
+ rope_head_dim: int,
+ base: float,
+ yarn_disabled: bool = False,
+ seqlen: int = 65536,
+):
+ cfg = cast(Any, config)
+
+ # yarn_disabled=True makes precompute_freqs_cis skip YaRN interpolation.
+ # DSV4 Flash keeps YaRN enabled for sliding and compressed attention.
+ original_seq_len = 0 if yarn_disabled else cfg.original_max_position_embeddings
+
+ inputs = dict(
+ dim=rope_head_dim,
+ seqlen=seqlen,
+ original_seq_len=original_seq_len,
+ base=base,
+ factor=cfg.rotary_scaling_factor,
+ beta_fast=cfg.beta_fast,
+ beta_slow=cfg.beta_slow,
+ )
+
+ assert cfg.rotary_scaling_factor in (4, 16), (
+ f"Unexpected rotary_scaling_factor: {cfg.rotary_scaling_factor}"
+ )
+ expected_original = 0 if yarn_disabled else 65536
+ assert inputs == dict(
+ dim=rope_head_dim,
+ seqlen=seqlen,
+ original_seq_len=expected_original,
+ base=base,
+ factor=cfg.rotary_scaling_factor,
+ beta_fast=32,
+ beta_slow=1,
+ )
+
+ return precompute_freqs_cis(**inputs)
+
+
+def _rope_cache_key(
+ config: TransformerConfig,
+ *,
+ rope_head_dim: int,
+ base: float,
+ yarn_disabled: bool,
+ device: torch.device,
+) -> tuple[object, ...]:
+ cfg = cast(Any, config)
+ device_index = (
+ torch.cuda.current_device()
+ if device.type == "cuda" and device.index is None
+ else device.index
+ )
+ original_seq_len = 0 if yarn_disabled else cfg.original_max_position_embeddings
+ return (
+ device.type,
+ device_index,
+ int(rope_head_dim),
+ int(original_seq_len),
+ float(base),
+ float(cfg.rotary_scaling_factor),
+ int(cfg.beta_fast),
+ int(cfg.beta_slow),
+ )
+
+
+def _get_device_rope_cache(
+ config: TransformerConfig,
+ *,
+ rope_head_dim: int,
+ base: float,
+ yarn_disabled: bool,
+ device: torch.device,
+ seqlen: int,
+) -> torch.Tensor:
+ seqlen = max(1, int(seqlen))
+ key = _rope_cache_key(
+ config,
+ rope_head_dim=rope_head_dim,
+ base=base,
+ yarn_disabled=yarn_disabled,
+ device=device,
+ )
+ cached = _DEVICE_ROPE_CACHES.get(key)
+ if cached is None or cached.shape[0] < seqlen:
+ cached = wrapped_precompute_freqs_cis(
+ config,
+ rope_head_dim=rope_head_dim,
+ base=base,
+ yarn_disabled=yarn_disabled,
+ seqlen=seqlen,
+ ).to(device)
+ _DEVICE_ROPE_CACHES[key] = cached
+ return cached
+
+
+def configure_rope_cache(
+ module: nn.Module,
+ config: TransformerConfig,
+ *,
+ rope_head_dim: int,
+ base: float,
+ yarn_disabled: bool = False,
+) -> None:
+ setattr(module, "_dsv4_rope_config", config)
+ setattr(module, "_dsv4_rope_head_dim", int(rope_head_dim))
+ setattr(module, "_dsv4_rope_base", float(base))
+ setattr(module, "_dsv4_rope_yarn_disabled", bool(yarn_disabled))
+ module.register_buffer(
+ "freqs_cis", torch.empty(0, dtype=torch.complex64), persistent=False
+ )
+
+
+def materialize_rope_cache(
+ module: nn.Module, device: torch.device | None = None, seqlen: int = 65536
+) -> bool:
+ config = getattr(module, "_dsv4_rope_config", None)
+ if config is None:
+ return False
+ if device is None:
+ try:
+ device = next(module.parameters()).device
+ except StopIteration:
+ device = torch.device("cpu")
+ if device.type == "meta":
+ return False
+ freqs_cis = _get_device_rope_cache(
+ config,
+ rope_head_dim=int(getattr(module, "_dsv4_rope_head_dim")),
+ base=float(getattr(module, "_dsv4_rope_base")),
+ yarn_disabled=bool(getattr(module, "_dsv4_rope_yarn_disabled")),
+ device=device,
+ seqlen=seqlen,
+ )
+ module.freqs_cis = freqs_cis
+ return True
+
+
+def get_rope_cache(
+ module: nn.Module, *, seqlen: int, device: torch.device
+) -> torch.Tensor:
+ freqs_cis = cast(torch.Tensor, module.freqs_cis)
+ if (
+ freqs_cis.numel() == 0
+ or freqs_cis.device != device
+ or freqs_cis.shape[0] < seqlen
+ ):
+ materialize_rope_cache(module, device, seqlen=seqlen)
+ freqs_cis = cast(torch.Tensor, module.freqs_cis)
+ return freqs_cis[:seqlen]
+
+
+def get_rope_cache_at_positions(
+ module: nn.Module, *, position_ids: torch.Tensor, device: torch.device
+) -> torch.Tensor:
+ freqs_cis = cast(torch.Tensor, module.freqs_cis)
+ safe_positions = position_ids.to(device=device, dtype=torch.long).clamp_min(0)
+ seqlen = max(1, safe_positions.numel())
+ if (
+ freqs_cis.numel() == 0
+ or freqs_cis.device != device
+ or freqs_cis.shape[0] < seqlen
+ ):
+ materialize_rope_cache(module, device, seqlen=seqlen)
+ freqs_cis = cast(torch.Tensor, module.freqs_cis)
+ return freqs_cis.index_select(0, safe_positions.reshape(-1)).view(
+ *safe_positions.shape,
+ freqs_cis.shape[-1],
+ )
diff --git a/src/art/megatron/dsv4/spec.py b/src/art/megatron/dsv4/spec.py
new file mode 100644
index 000000000..2c6f6204b
--- /dev/null
+++ b/src/art/megatron/dsv4/spec.py
@@ -0,0 +1,48 @@
+from copy import deepcopy
+from typing import Any, cast
+
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
+from megatron.core.transformer.moe.moe_layer import MoELayer
+from megatron.core.transformer.spec_utils import ModuleSpec
+
+from art.megatron.dsv4.deepseek_v4 import DeepSeekV4Attention
+from art.megatron.dsv4.layer import (
+ Dsv4FinalNorm,
+ Dsv4MoELayer,
+ Dsv4Router,
+ Dsv4TransformerLayer,
+)
+
+try:
+ import transformer_engine # noqa: F401
+
+ HAVE_TE = True
+except (ImportError, ModuleNotFoundError):
+ HAVE_TE = False
+
+
+def get_dsv4_decoder_block_spec(config: Any, vp_stage: int | None = None) -> Any:
+ config.moe_layer_freq = [1] * int(config.num_layers)
+ block_spec = deepcopy(
+ get_gpt_decoder_block_spec(
+ config,
+ use_transformer_engine=HAVE_TE,
+ normalization="RMSNorm",
+ vp_stage=vp_stage,
+ )
+ )
+ block_spec.layer_norm = Dsv4FinalNorm
+ for layer_spec in block_spec.layer_specs or ():
+ layer_spec.module = Dsv4TransformerLayer
+ submodules = cast(Any, layer_spec.submodules)
+ submodules.input_layernorm = submodules.pre_mlp_layernorm
+ submodules.self_attention = ModuleSpec(
+ module=DeepSeekV4Attention,
+ submodules=None,
+ metainfo={"fuse_input_layernorm": False},
+ )
+ mlp = submodules.mlp
+ if isinstance(mlp, ModuleSpec) and mlp.module == MoELayer:
+ mlp.module = Dsv4MoELayer
+ mlp.submodules.router = Dsv4Router
+ return block_spec
diff --git a/src/art/megatron/dsv4/tokenizer.py b/src/art/megatron/dsv4/tokenizer.py
new file mode 100644
index 000000000..1fb33adf2
--- /dev/null
+++ b/src/art/megatron/dsv4/tokenizer.py
@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import copy
+from typing import Any
+
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+
+from art.megatron.dsv4.encoding import encode_messages
+
+DSV4_CHAT_TEMPLATE_MARKER = "deepseek_v4_python_encoder enable_thinking"
+
+
+def has_configured_chat_template(internal_config: Any) -> bool:
+ return (
+ internal_config.get("chat_template") is not None
+ or internal_config.get("chat_template_path") is not None
+ )
+
+
+def get_dsv4_tokenizer(
+ tokenizer: PreTrainedTokenizerBase,
+) -> PreTrainedTokenizerBase:
+ if getattr(tokenizer, "_art_dsv4_chat_template_wrapped", False):
+ return tokenizer
+
+ wrapped = copy.copy(tokenizer)
+ added_vocab = tokenizer.get_added_vocab()
+ added_vocab_size = len(added_vocab)
+ tokenizer_vocab_size = tokenizer.vocab_size
+
+ class _ArtDsv4Tokenizer(tokenizer.__class__): # type: ignore[misc, valid-type]
+ def apply_chat_template(
+ self,
+ messages: list[dict[str, Any]],
+ tools: list[dict[str, Any]] | None = None,
+ **kwargs: Any,
+ ) -> str | list[int]:
+ thinking = bool(kwargs.get("thinking", False)) or bool(
+ kwargs.get("enable_thinking", False)
+ )
+ thinking_mode = "thinking" if thinking else "chat"
+ conversation = kwargs.get("conversation", messages)
+ rendered_messages = list(conversation)
+ if tools:
+ rendered_messages.insert(0, {"role": "system", "tools": tools})
+
+ reasoning_effort = kwargs.get("reasoning_effort")
+ if not isinstance(reasoning_effort, str):
+ reasoning_effort = None
+ elif reasoning_effort == "none":
+ thinking_mode = "chat"
+ reasoning_effort = None
+ elif reasoning_effort in ("max", "xhigh"):
+ reasoning_effort = "max"
+ else:
+ reasoning_effort = "high"
+
+ prompt = encode_messages(
+ rendered_messages,
+ thinking_mode=thinking_mode,
+ drop_thinking=kwargs.get("drop_thinking", True),
+ reasoning_effort=reasoning_effort,
+ )
+ if not kwargs.get("tokenize", True):
+ return prompt
+ tokenizer_kwargs = {
+ key: kwargs[key]
+ for key in ("truncation", "max_length")
+ if key in kwargs
+ }
+ return self.encode(
+ prompt,
+ add_special_tokens=False,
+ **tokenizer_kwargs,
+ )
+
+ def num_special_tokens_to_add(self) -> int:
+ return len(self.encode(""))
+
+ def __len__(self) -> int:
+ return tokenizer_vocab_size + added_vocab_size
+
+ def get_added_vocab(self) -> dict[str, int]:
+ return added_vocab.copy()
+
+ def __reduce__(self) -> Any:
+ return get_dsv4_tokenizer, (tokenizer,)
+
+ _ArtDsv4Tokenizer.__name__ = f"ArtDsv4{tokenizer.__class__.__name__}"
+ wrapped.__class__ = _ArtDsv4Tokenizer
+ wrapped.chat_template = DSV4_CHAT_TEMPLATE_MARKER
+ setattr(wrapped, "_art_dsv4_chat_template_wrapped", True)
+ return wrapped
diff --git a/src/art/megatron/dsv4/utils.py b/src/art/megatron/dsv4/utils.py
new file mode 100644
index 000000000..3be094e67
--- /dev/null
+++ b/src/art/megatron/dsv4/utils.py
@@ -0,0 +1,32 @@
+import torch
+from torch import nn
+
+
+def freeze_parameters_as_buffers(module: nn.Module) -> None:
+ for child in module.modules():
+ for name, param in list(child.named_parameters(recurse=False)):
+ del child._parameters[name]
+ child.register_buffer(name, param.detach(), persistent=True)
+
+
+def rotate_activation(x: torch.Tensor) -> torch.Tensor:
+ """Scaled Hadamard transform over the last dimension.
+
+ DeepSeek-V4 uses this before activation FP8 simulation in the indexer and
+ compressor. The supported ART dimensions are powers of two, specifically
+ 128 and 512 in the DSV4 Flash config.
+ """
+ if x.dtype not in (torch.bfloat16, torch.float32):
+ raise TypeError(f"rotate_activation supports bf16/fp32, got {x.dtype}")
+ width = int(x.size(-1))
+ if width <= 0 or width & (width - 1):
+ raise ValueError(f"Hadamard width must be a power of two, got {width}.")
+ y = x.float()
+ h = 1
+ while h < width:
+ y = y.reshape(*x.shape[:-1], width // (2 * h), 2, h)
+ left, right = y.unbind(dim=-2)
+ y = torch.stack((left + right, left - right), dim=-2)
+ y = y.reshape(*x.shape[:-1], width)
+ h *= 2
+ return (y * (width**-0.5)).to(x.dtype)
diff --git a/src/art/megatron/dsv4/v4_indexer.py b/src/art/megatron/dsv4/v4_indexer.py
new file mode 100644
index 000000000..6a90a055e
--- /dev/null
+++ b/src/art/megatron/dsv4/v4_indexer.py
@@ -0,0 +1,378 @@
+from typing import Any, cast
+
+import einops
+from megatron.core.extensions.transformer_engine import TELinear
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.transformer_config import TransformerConfig
+import torch
+
+from art.megatron.dsv4.compressor import (
+ DeepSeekV4Compressor,
+ Dsv4CompressionLayout,
+ compressed_layout_visibility,
+)
+from art.megatron.dsv4.kernel.tilelang_indexer_fwd import (
+ indexer_topk_interface,
+ shared_prefix_indexer_topk_interface,
+)
+from art.megatron.dsv4.rope import (
+ apply_rotary_emb,
+ configure_rope_cache,
+ get_rope_cache,
+ get_rope_cache_at_positions,
+)
+from art.megatron.dsv4.utils import freeze_parameters_as_buffers, rotate_activation
+
+_INDEXER_QUERY_BLOCK = 512
+_INDEXER_HEAD_BLOCK = 8
+
+
+def _make_causal_cu_seqlens(seq_len_q, seq_len_kv, compress_ratio, device):
+ del seq_len_kv
+ positions = torch.arange(seq_len_q, device=device, dtype=torch.int32)
+ cu_seqlen_ks = torch.zeros(seq_len_q, device=device, dtype=torch.int32)
+ cu_seqlen_ke = ((positions + 1) // compress_ratio).to(torch.int32)
+ return cu_seqlen_ks, cu_seqlen_ke
+
+
+@torch.compiler.disable
+def _exact_indexer_topk(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ weights: torch.Tensor,
+ cu_seqlen_ks: torch.Tensor,
+ cu_seqlen_ke: torch.Tensor,
+ topk: int,
+ *,
+ shared_layout: Dsv4CompressionLayout | None = None,
+ position_ids: torch.Tensor | None = None,
+ group_ids: torch.Tensor | None = None,
+ parent_ids: torch.Tensor | None = None,
+) -> torch.Tensor:
+ """Return frozen DSV4 indexer topk using the reference score equation.
+
+ The indexer routing is discrete, so small BF16 fused-kernel score drift can
+ choose different compressed KV rows and send CSA gradients to different
+ windows. Keep the production path exact against the HF/Megatron reference:
+ fp32 dot, ReLU, head weighting, compressed-causal mask, then topk. Query
+ and head chunking avoid materializing the full [S, H, S / 4] score tensor.
+ """
+
+ seqlen, batch, heads, _ = q.shape
+ seqlen_kv = k.shape[0]
+ actual_topk = min(topk, seqlen_kv)
+ out = torch.empty(batch, seqlen, actual_topk, device=q.device, dtype=torch.int32)
+ if actual_topk == 0:
+ return out
+
+ kv_ids = torch.arange(seqlen_kv, device=q.device)
+ q_block = _INDEXER_QUERY_BLOCK
+ h_block = _INDEXER_HEAD_BLOCK
+ with torch.no_grad():
+ for b in range(batch):
+ k_b = k[:, b].float()
+ for q_start in range(0, seqlen, q_block):
+ q_end = min(q_start + q_block, seqlen)
+ scores = torch.zeros(
+ q_end - q_start, seqlen_kv, device=q.device, dtype=torch.float32
+ )
+ for h_start in range(0, heads, h_block):
+ h_end = min(h_start + h_block, heads)
+ dot = torch.einsum(
+ "qhd,kd->qhk",
+ q[q_start:q_end, b, h_start:h_end].float(),
+ k_b,
+ ).relu_()
+ scores += (
+ dot
+ * weights[q_start:q_end, b, h_start:h_end].float().unsqueeze(-1)
+ ).sum(dim=1)
+ if shared_layout is None:
+ visible = (
+ kv_ids.unsqueeze(0) >= cu_seqlen_ks[q_start:q_end, None]
+ ) & (kv_ids.unsqueeze(0) < cu_seqlen_ke[q_start:q_end, None])
+ else:
+ if position_ids is None or group_ids is None or parent_ids is None:
+ raise ValueError(
+ "DSV4 shared-prefix indexer requires position/group metadata."
+ )
+ visible = compressed_layout_visibility(
+ shared_layout,
+ position_ids=position_ids[b : b + 1],
+ group_ids=group_ids[b : b + 1],
+ parent_ids=parent_ids[b : b + 1],
+ q_start=q_start,
+ q_end=q_end,
+ )[0]
+ scores.masked_fill_(~visible, float("-inf"))
+ top_scores, top_indices = scores.topk(actual_topk, dim=-1)
+ out[b, q_start:q_end] = torch.where(
+ torch.isneginf(top_scores),
+ torch.full_like(top_indices, -1),
+ top_indices,
+ )
+ return out
+
+
+@torch.compiler.disable
+def _tilelang_indexer_topk(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ weights: torch.Tensor,
+ cu_seqlen_ks: torch.Tensor,
+ cu_seqlen_ke: torch.Tensor,
+ topk: int,
+ *,
+ shared_layout: Dsv4CompressionLayout | None = None,
+ position_ids: torch.Tensor | None = None,
+ group_ids: torch.Tensor | None = None,
+ parent_ids: torch.Tensor | None = None,
+ shared_layout_i32: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
+ | None = None,
+) -> torch.Tensor:
+ seqlen, batch, heads, _ = q.shape
+ seqlen_kv = k.shape[0]
+ actual_topk = min(topk, seqlen_kv)
+ out = torch.empty(batch, seqlen, actual_topk, device=q.device, dtype=torch.int32)
+ if actual_topk == 0:
+ return out
+
+ internal_q_block = max(1, 128 // heads)
+ q_block = (_INDEXER_QUERY_BLOCK // internal_q_block) * internal_q_block
+ q_block = max(internal_q_block, q_block)
+ fast_seqlen = (seqlen // internal_q_block) * internal_q_block
+
+ with torch.no_grad():
+ for b in range(batch):
+ if shared_layout is not None:
+ if position_ids is None or group_ids is None or parent_ids is None:
+ raise ValueError(
+ "DSV4 shared-prefix indexer requires position/group metadata."
+ )
+ if shared_layout_i32 is None:
+ entry_group_ids = shared_layout.entry_group_ids.to(torch.int32)
+ entry_parent_visible = shared_layout.entry_parent_visible.to(
+ torch.int32
+ )
+ entry_end_positions = shared_layout.entry_end_positions.to(
+ torch.int32
+ )
+ entry_valid = shared_layout.entry_valid.to(torch.int32)
+ else:
+ (
+ entry_group_ids,
+ entry_parent_visible,
+ entry_end_positions,
+ entry_valid,
+ ) = shared_layout_i32
+ position_b = position_ids[b].to(torch.int32).contiguous()
+ group_b = group_ids[b].to(torch.int32).contiguous()
+ parent_b = parent_ids[b].to(torch.int32).contiguous()
+ entry_group_b = entry_group_ids[b].contiguous()
+ entry_parent_visible_b = entry_parent_visible[b].contiguous()
+ entry_end_b = entry_end_positions[b].contiguous()
+ entry_valid_b = entry_valid[b].contiguous()
+ for q_start in range(0, fast_seqlen, q_block):
+ q_end = min(q_start + q_block, fast_seqlen)
+ q_slice = q[q_start:q_end, b].contiguous()
+ k_slice = k[:, b].contiguous()
+ weights_slice = weights[q_start:q_end, b].contiguous()
+ if shared_layout is None:
+ out[b, q_start:q_end] = indexer_topk_interface(
+ q_slice,
+ k_slice,
+ weights_slice,
+ cu_seqlen_ks[q_start:q_end].contiguous(),
+ cu_seqlen_ke[q_start:q_end].contiguous(),
+ topk,
+ )
+ else:
+ out[b, q_start:q_end] = shared_prefix_indexer_topk_interface(
+ q_slice,
+ k_slice,
+ weights_slice,
+ position_b[q_start:q_end],
+ group_b[q_start:q_end],
+ parent_b[q_start:q_end],
+ entry_group_b,
+ entry_parent_visible_b,
+ entry_end_b,
+ entry_valid_b,
+ topk,
+ )
+ if fast_seqlen != seqlen:
+ out[:, fast_seqlen:] = _exact_indexer_topk(
+ q[fast_seqlen:],
+ k,
+ weights[fast_seqlen:],
+ cu_seqlen_ks[fast_seqlen:],
+ cu_seqlen_ke[fast_seqlen:],
+ topk,
+ shared_layout=shared_layout,
+ position_ids=None
+ if position_ids is None
+ else position_ids[:, fast_seqlen:],
+ group_ids=None if group_ids is None else group_ids[:, fast_seqlen:],
+ parent_ids=None if parent_ids is None else parent_ids[:, fast_seqlen:],
+ )
+ return out
+
+
+class V4Indexer(MegatronModule):
+ """DSA Indexer for DeepSeek-V4 C4 layers."""
+
+ def __init__(self, config: TransformerConfig, pg_collection=None):
+ super().__init__(config=config)
+ cfg = cast(Any, config)
+ init_method = config.init_method
+ if init_method is None:
+ raise RuntimeError("DeepSeek-V4 indexer requires config.init_method.")
+
+ self.hidden_size = config.hidden_size
+ self.q_lora_rank = (
+ int(cfg.q_lora_rank) if cfg.q_lora_rank is not None else config.hidden_size
+ )
+ self.index_n_heads = int(cfg.dsa_indexer_n_heads)
+ self.index_head_dim = int(cfg.dsa_indexer_head_dim)
+ self.index_topk = int(cfg.dsa_indexer_topk)
+ self.rope_head_dim = int(cfg.qk_pos_emb_head_dim)
+ self.compress_ratio = 4
+
+ if pg_collection is None:
+ pg_collection = ProcessGroupCollection.use_mpu_process_groups(
+ required_pgs=["tp"]
+ )
+ self.pg_collection = pg_collection
+
+ self.linear_wq_b = TELinear(
+ self.q_lora_rank,
+ self.index_n_heads * self.index_head_dim,
+ config=config,
+ init_method=init_method,
+ bias=False,
+ skip_bias_add=False,
+ skip_weight_param_allocation=False,
+ parallel_mode="duplicated",
+ )
+
+ self.linear_weights_proj = TELinear(
+ self.hidden_size,
+ self.index_n_heads,
+ config=config,
+ init_method=init_method,
+ bias=False,
+ skip_bias_add=False,
+ skip_weight_param_allocation=False,
+ parallel_mode="duplicated",
+ )
+
+ self.compressor = DeepSeekV4Compressor(
+ config=config,
+ head_dim=self.index_head_dim,
+ compress_ratio=self.compress_ratio,
+ rotate=True,
+ cp_group=None,
+ )
+
+ rope_base = (
+ cfg.dsv4_compress_rope_theta if self.compress_ratio else cfg.rotary_base
+ )
+ configure_rope_cache(
+ self, config, rope_head_dim=self.rope_head_dim, base=rope_base
+ )
+ freeze_parameters_as_buffers(self)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ qr: torch.Tensor,
+ mask=None,
+ packed_seq_params=None,
+ position_ids: torch.Tensor | None = None,
+ shared_layout: Dsv4CompressionLayout | None = None,
+ group_ids: torch.Tensor | None = None,
+ parent_ids: torch.Tensor | None = None,
+ shared_layout_i32: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
+ | None = None,
+ ):
+ """Forward pass.
+
+ Args:
+ x: hidden states [seqlen, batch, hidden_size]
+ qr: low-rank query [seqlen, batch, q_lora_rank]
+ mask: unused (causal mask generated internally via cu_seqlens)
+ packed_seq_params: unused
+
+ Returns:
+ topk_indices: [batch, seqlen, index_topk] int64
+ """
+
+ # =========================================
+ # Gather inputs if SP is enabled
+ # =========================================
+ if self.config.sequence_parallel and self.pg_collection.tp.size() > 1:
+ x = gather_from_sequence_parallel_region(x, group=self.pg_collection.tp)
+ qr = gather_from_sequence_parallel_region(qr, group=self.pg_collection.tp)
+
+ seqlen, bsz, _ = x.size()
+
+ q, _ = self.linear_wq_b(qr)
+ q = q.reshape(seqlen, bsz, self.index_n_heads, self.index_head_dim)
+
+ rd = self.rope_head_dim
+ cp_group = getattr(self.pg_collection, "cp", None)
+ if cp_group is not None and cp_group.size() != 1:
+ raise RuntimeError(
+ "DeepSeek-V4 non-CP indexer received context_parallel_size > 1."
+ )
+ if position_ids is None:
+ freqs_cis = get_rope_cache(self, seqlen=seqlen, device=x.device)
+ else:
+ freqs_cis = get_rope_cache_at_positions(
+ self, position_ids=position_ids, device=x.device
+ )
+ q = q.clone()
+ q = einops.rearrange(q, "s b ... -> b s ...")
+ apply_rotary_emb(q[..., -rd:], freqs_cis)
+ q = einops.rearrange(q, "b s ... -> s b ...")
+
+ q = rotate_activation(q)
+
+ k = self.compressor(
+ x,
+ position_ids=position_ids,
+ shared_layout=shared_layout,
+ )
+ if k.shape[0] == 0:
+ return torch.empty(
+ bsz,
+ seqlen,
+ 0,
+ device=x.device,
+ dtype=torch.long,
+ )
+
+ weights, _ = self.linear_weights_proj(x)
+ softmax_scale = self.index_head_dim**-0.5
+ weights = weights * (self.index_n_heads**-0.5) * softmax_scale
+
+ seqlen_global = seqlen
+ seqlen_kv = k.shape[0]
+ cu_ks, cu_ke = _make_causal_cu_seqlens(
+ seqlen_global, seqlen_kv, self.compress_ratio, q.device
+ )
+ return _tilelang_indexer_topk(
+ q,
+ k,
+ weights.float(),
+ cu_ks,
+ cu_ke,
+ self.index_topk,
+ shared_layout=shared_layout,
+ position_ids=position_ids,
+ group_ids=group_ids,
+ parent_ids=parent_ids,
+ shared_layout_i32=shared_layout_i32,
+ )
diff --git a/src/art/megatron/kernels/cute_grouped_lora_quack.py b/src/art/megatron/kernels/cute_grouped_lora_quack.py
index f93bdb663..c0c9a70d6 100644
--- a/src/art/megatron/kernels/cute_grouped_lora_quack.py
+++ b/src/art/megatron/kernels/cute_grouped_lora_quack.py
@@ -1,6 +1,12 @@
from __future__ import annotations
+from getpass import getuser
+import hashlib
+import importlib
import os
+from pathlib import Path
+import pickle
+import tempfile
from typing import Any, cast
from quack.gemm import gemm as quack_gemm
@@ -29,6 +35,65 @@ def _env_positive_int(name: str) -> int | None:
return value
+def _quack_compile_lock_dir() -> Path:
+ path = (
+ Path(os.environ["ART_QUACK_COMPILE_LOCK_DIR"])
+ if "ART_QUACK_COMPILE_LOCK_DIR" in os.environ
+ else Path(tempfile.gettempdir()) / getuser() / "art_quack_compile_locks"
+ )
+ path.mkdir(parents=True, exist_ok=True)
+ return path
+
+
+def _cache_key(args: tuple[Any, ...], kwargs: dict[str, Any]) -> tuple[Any, ...]:
+ return args + tuple(sorted(kwargs.items())) if kwargs else args
+
+
+def _key_hash(key: tuple[Any, ...]) -> str:
+ try:
+ data = pickle.dumps(key)
+ except Exception:
+ data = repr(key).encode()
+ return hashlib.sha256(data).hexdigest()
+
+
+def _install_quack_compile_lock() -> None:
+ """Serialize QuACK CuTe compilation misses across distributed ranks.
+
+ QuACK protects disk-cache load/store with file locks, but it compiles cache
+ misses outside that lock. ETP runs can make every rank compile the same
+ grouped-LoRA backward GEMM at once, which has crashed inside CuTe DSL codegen.
+ Hot in-process cache hits bypass this guard, so kernel launches stay unlocked.
+ """
+ import quack.cache_utils as quack_cache_utils
+
+ gemm_module = importlib.import_module("quack.gemm")
+ compile_gemm = getattr(gemm_module, "_compile_gemm")
+ if bool(getattr(compile_gemm, "_art_compile_lock_wrapped", False)):
+ return
+ cache = getattr(compile_gemm, "cache", None)
+
+ def locked_compile_gemm(*args: Any, **kwargs: Any) -> Any:
+ key = _cache_key(args, kwargs)
+ if isinstance(cache, dict) and key in cache:
+ return compile_gemm(*args, **kwargs)
+ lock_path = _quack_compile_lock_dir() / f"{_key_hash(key)}.lock"
+ with quack_cache_utils.FileLock(lock_path, exclusive=True, timeout=600):
+ return compile_gemm(*args, **kwargs)
+
+ locked_compile_gemm_any = cast(Any, locked_compile_gemm)
+ locked_compile_gemm_any.cache = getattr(compile_gemm, "cache", None)
+ locked_compile_gemm_any.cache_clear = getattr(compile_gemm, "cache_clear", None)
+ locked_compile_gemm_any.cache_info = getattr(compile_gemm, "cache_info", None)
+ locked_compile_gemm_any._art_compile_lock_wrapped = True
+ setattr(gemm_module, "_compile_gemm", locked_compile_gemm)
+
+
+def _quack_gemm(*args: Any, **kwargs: Any) -> Any:
+ _install_quack_compile_lock()
+ return quack_gemm(*args, **kwargs)
+
+
def _tokens_per_expert_to_tensor(
tokens_per_expert: list[int] | torch.Tensor,
) -> torch.Tensor:
@@ -216,7 +281,7 @@ def _varlen_quack_gemm(
raise ValueError(
f"Output tensor must match input device/dtype, got {out.device}/{out.dtype}"
)
- quack_gemm(
+ _quack_gemm(
a,
b,
out,
@@ -252,7 +317,7 @@ def _varlen_quack_gemm_k(
device=a.device,
dtype=a.dtype,
)
- quack_gemm(
+ _quack_gemm(
a,
b,
out,
diff --git a/src/art/megatron/model_support/__init__.py b/src/art/megatron/model_support/__init__.py
index de569af7d..7e1a8b374 100644
--- a/src/art/megatron/model_support/__init__.py
+++ b/src/art/megatron/model_support/__init__.py
@@ -24,8 +24,10 @@
is_model_support_registered,
list_model_support_specs,
model_requires_merged_rollout,
+ model_supports_context_parallel,
model_uses_expert_parallel,
native_vllm_lora_status_for_model,
+ vllm_lora_config_for_model,
)
from art.megatron.model_support.spec import (
ArchitectureReport,
@@ -89,7 +91,9 @@ def __getattr__(name: str):
"is_model_support_registered",
"list_model_support_specs",
"model_uses_expert_parallel",
+ "model_supports_context_parallel",
"model_requires_merged_rollout",
"native_vllm_lora_status_for_model",
+ "vllm_lora_config_for_model",
"summarize_layer_families",
]
diff --git a/src/art/megatron/model_support/handlers/__init__.py b/src/art/megatron/model_support/handlers/__init__.py
index e45e74761..a8da0158d 100644
--- a/src/art/megatron/model_support/handlers/__init__.py
+++ b/src/art/megatron/model_support/handlers/__init__.py
@@ -48,6 +48,14 @@
"art.megatron.model_support.handlers.qwen3_5",
"Qwen35MoeHandler",
),
+ "DSV4_HANDLER": (
+ "art.megatron.model_support.handlers.dsv4",
+ "DSV4_HANDLER",
+ ),
+ "Dsv4Handler": (
+ "art.megatron.model_support.handlers.dsv4",
+ "Dsv4Handler",
+ ),
"GPT_OSS_MOE_HANDLER": (
"art.megatron.model_support.handlers.gpt_oss",
"GPT_OSS_MOE_HANDLER",
diff --git a/src/art/megatron/model_support/handlers/default_dense.py b/src/art/megatron/model_support/handlers/default_dense.py
index bf77d50e7..cc7128215 100644
--- a/src/art/megatron/model_support/handlers/default_dense.py
+++ b/src/art/megatron/model_support/handlers/default_dense.py
@@ -1,4 +1,4 @@
-from typing import Any, Sequence
+from typing import Any, Literal, Sequence
import torch
@@ -10,6 +10,7 @@
LayerFamilyInstance,
RolloutWeightsMode,
SharedExpertCompileState,
+ SharedPrefixModelStateContext,
)
_CONTEXT_PARALLEL_ATTENTION_WORKAROUND_FLAG = "context_parallel_attention"
@@ -37,6 +38,7 @@ class DefaultDenseHandler:
key = "default_dense"
build_gdn_execution_spec = False
is_moe = False
+ cp_supported = True
native_vllm_lora_status = "disabled"
def identity_lora_model_config(self, base_config: Any) -> Any:
@@ -96,6 +98,18 @@ def configure_provider_for_runtime(self, provider: Any) -> None:
del provider
return None
+ def default_chat_template(self) -> str | None:
+ return None
+
+ def configure_tokenizer(
+ self,
+ tokenizer: Any,
+ *,
+ internal_config: Any,
+ ) -> Any:
+ del internal_config
+ return tokenizer
+
def vllm_engine_args(
self,
*,
@@ -111,6 +125,23 @@ def install_preprocess_patch(self, model_chunks: Sequence[Any]) -> None:
del model_chunks
return None
+ def build_shared_prefix_model_state(
+ self,
+ context: SharedPrefixModelStateContext,
+ ) -> dict[str, Any]:
+ del context
+ return {}
+
+ def correctness_precision(self) -> Literal["bf16", "fp32"]:
+ return "fp32"
+
+ def correctness_use_fp32_lora_reference(self) -> bool:
+ return True
+
+ def correctness_phase_pass_fns(self, oracle_harness: Any) -> dict[str, Any] | None:
+ del oracle_harness
+ return None
+
def to_vllm_lora_tensors(
self,
tensors: dict[str, torch.Tensor],
@@ -119,6 +150,9 @@ def to_vllm_lora_tensors(
) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
return tensors, adapter_config
+ def to_vllm_lora_config(self, adapter_config: dict[str, Any]) -> dict[str, Any]:
+ return adapter_config
+
def from_vllm_lora_tensors(
self,
tensors: dict[str, torch.Tensor],
diff --git a/src/art/megatron/model_support/handlers/dsv4.py b/src/art/megatron/model_support/handlers/dsv4.py
new file mode 100644
index 000000000..500df992e
--- /dev/null
+++ b/src/art/megatron/model_support/handlers/dsv4.py
@@ -0,0 +1,1198 @@
+from __future__ import annotations
+
+import hashlib
+import os
+import re
+from typing import Any, Literal, Sequence, cast
+
+import torch
+
+from art.megatron.model_support.handlers.default_dense import (
+ DefaultMoeHandler,
+ _compile_workaround_flags_for_provider,
+ _require_moe_experts,
+)
+from art.megatron.model_support.spec import (
+ CompileWorkaroundConfig,
+ LayerFamilyInstance,
+ SharedPrefixModelStateContext,
+)
+
+_ORACLE_HIDDEN_SIZE = 512
+_ORACLE_Q_LORA_RANK = 128
+_ORACLE_NUM_ATTENTION_HEADS = 16
+_ORACLE_NUM_EXPERTS = 4
+_ORACLE_NUM_EXPERTS_PER_TOK = 1
+_ORACLE_FFN_HIDDEN_SIZE = 128
+_ORACLE_INDEX_HEADS = 1
+_ORACLE_INDEX_TOPK = 1024
+_VALIDATION_NUM_LAYERS_ENV = "ART_DSV4_VALIDATION_NUM_LAYERS"
+_ORACLE_EXPERT_WEIGHT_RE = re.compile(r"\.mlp\.experts\..*\.weight(?P\d+)$")
+_DSV4_ART_MOE_EXPERT_KEY_RE = re.compile(
+ r"^(?P.*\.mlp\.experts)\.(?P\d+)\."
+ r"(?Pgate_proj|up_proj|down_proj)\.(?Plora_[AB])\.weight$"
+)
+_DSV4_VLLM_MOE_KEY_RE = re.compile(
+ r"^(?P.*\.mlp\.experts)\."
+ r"(?:(?Pbase_layer)\.)?(?Plora_[AB])\.weight$"
+)
+_DSV4_VLLM_MOE_EXPERT_KEY_RE = re.compile(
+ r"^(?P.*\.ffn\.experts)\.(?P\d+)\."
+ r"(?Pw1|w2|w3)\.(?Plora_[AB])\.weight$"
+)
+_DSV4_MOE_COMPILE_WORKAROUND_FLAGS = (
+ "alltoall_dtoh",
+ "alltoall_dispatch_preprocess",
+ "deepep_dispatch_combine",
+ "deepep_permute_restore",
+ "te_triton_permute_with_mask_map",
+)
+
+
+class Dsv4Handler(DefaultMoeHandler):
+ key = "dsv4"
+ is_moe = True
+ cp_supported = False
+ native_vllm_lora_status = "validated"
+
+ def identity_lora_model_config(self, base_config: Any) -> Any:
+ self.ensure_hf_reference_registered()
+ return base_config
+
+ def patch_provider(self, provider: Any, bridge: Any) -> None:
+ del bridge
+ from art.megatron.dsv4.spec import get_dsv4_decoder_block_spec
+
+ provider.transformer_layer_spec = get_dsv4_decoder_block_spec
+ if int(getattr(provider, "context_parallel_size", 1) or 1) != 1:
+ raise RuntimeError(
+ "DSV4 model support in this worktree does not implement context parallelism."
+ )
+
+ def configure_provider_for_runtime(self, provider: Any) -> None:
+ provider.mtp_num_layers = None
+ provider.moe_shared_expert_overlap = False
+ raw_num_layers = os.environ.get(_VALIDATION_NUM_LAYERS_ENV)
+ if raw_num_layers is None:
+ return
+ num_layers = int(raw_num_layers)
+ if num_layers < 1:
+ raise ValueError(f"{_VALIDATION_NUM_LAYERS_ENV} must be positive")
+ provider.num_layers = num_layers
+ provider.moe_layer_freq = [1] * num_layers
+ ratios = list(getattr(provider, "dsv4_compress_ratios", ()) or ())
+ if ratios:
+ if num_layers > len(ratios):
+ raise ValueError(
+ f"{_VALIDATION_NUM_LAYERS_ENV}={num_layers} exceeds "
+ f"dsv4_compress_ratios length {len(ratios)}"
+ )
+ provider.dsv4_compress_ratios = ratios[:num_layers]
+
+ def default_chat_template(self) -> str | None:
+ return None
+
+ def configure_tokenizer(
+ self,
+ tokenizer: Any,
+ *,
+ internal_config: Any,
+ ) -> Any:
+ from art.megatron.dsv4.tokenizer import (
+ get_dsv4_tokenizer,
+ has_configured_chat_template,
+ )
+
+ if has_configured_chat_template(internal_config):
+ return tokenizer
+ return get_dsv4_tokenizer(tokenizer)
+
+ def build_shared_prefix_model_state(
+ self,
+ context: SharedPrefixModelStateContext,
+ ) -> dict[str, Any]:
+ if context.input_pos is None:
+ raise RuntimeError(
+ "DSV4 shared-prefix compression layouts require input_pos."
+ )
+ from art.megatron.dsv4.compressor import (
+ Dsv4SharedPrefixState,
+ build_shared_prefix_compression_layouts,
+ )
+
+ return {
+ "dsv4": Dsv4SharedPrefixState(
+ compression_layouts=build_shared_prefix_compression_layouts(
+ position_ids=context.input_pos,
+ group_ids=context.group_ids,
+ parent_ids=context.parent_ids,
+ device=context.device,
+ )
+ )
+ }
+
+ def correctness_precision(self) -> Literal["bf16", "fp32"]:
+ return "bf16"
+
+ def correctness_use_fp32_lora_reference(self) -> bool:
+ return False
+
+ def correctness_phase_pass_fns(self, oracle_harness: Any) -> dict[str, Any]:
+ non_zero_scales = {"typical_abs_scale": 0.0, "candidate_abs_scale": 0.0}
+ fwd = oracle_harness.MetricThresholdRule(
+ limits={"mean_abs_pct": 3.0},
+ minimums=non_zero_scales,
+ )
+ loss = oracle_harness.MetricThresholdRule(limits={"mean_abs_pct": 3.0})
+ grad = oracle_harness.MetricThresholdRule(
+ limits={"mean_abs_pct": 5.0},
+ minimums=non_zero_scales,
+ )
+ router_topk = oracle_harness.MetricThresholdRule(
+ limits={"topk_mismatch_fraction": 0.0, "top1_mismatch_fraction": 0.0}
+ )
+ return {
+ "forward": fwd,
+ "outputs": fwd,
+ "losses": loss,
+ "grads": grad,
+ "deltas": grad,
+ "router_scores": fwd,
+ "router_topk_ids": router_topk,
+ }
+
+ def identity_lora_target_parameters(
+ self,
+ model: Any,
+ *,
+ target_modules: list[str],
+ ) -> list[str]:
+ target_set = set(target_modules)
+
+ def include(name: str) -> bool:
+ if ".self_attn.compressor.indexer." in name:
+ return False
+ if "q_a_proj" in target_set and name.endswith(".self_attn.q_a_proj.weight"):
+ return True
+ if "q_b_proj" in target_set and name.endswith(".self_attn.q_b_proj.weight"):
+ return True
+ if "kv_proj" in target_set and name.endswith(".self_attn.kv_proj.weight"):
+ return True
+ if "o_a_proj" in target_set and name.endswith(".self_attn.o_a_proj.weight"):
+ return True
+ if "o_b_proj" in target_set and name.endswith(".self_attn.o_b_proj.weight"):
+ return True
+ if "compressor.kv_proj" in target_set and name.endswith(
+ ".self_attn.compressor.kv_proj.weight"
+ ):
+ return True
+ if "compressor.gate_proj" in target_set and name.endswith(
+ ".self_attn.compressor.gate_proj.weight"
+ ):
+ return True
+ if (
+ "gate_proj" in target_set
+ and ".mlp." in name
+ and name.endswith(".gate_proj.weight")
+ ):
+ return True
+ if (
+ "up_proj" in target_set
+ and ".mlp." in name
+ and name.endswith(".up_proj.weight")
+ ):
+ return True
+ if (
+ "down_proj" in target_set
+ and ".mlp." in name
+ and name.endswith(".down_proj.weight")
+ ):
+ return True
+ if "experts" in target_set and name.endswith(
+ (".mlp.experts.gate_up_proj", ".mlp.experts.down_proj")
+ ):
+ return True
+ return False
+
+ return [name for name, _ in model.named_parameters() if include(name)]
+
+ def install_preprocess_patch(self, model_chunks: Sequence[Any]) -> None:
+ from megatron.core.models.gpt.gpt_model import GPTModel
+
+ from art.megatron.dsv4.deepseek_v4 import DeepSeekV4Attention
+ from art.megatron.dsv4.layer import Dsv4MoELayer
+
+ for chunk in list(model_chunks):
+ module: Any = chunk
+ while hasattr(module, "module"):
+ module = module.module
+ gpt_module = (
+ module
+ if isinstance(module, GPTModel)
+ else cast(GPTModel, getattr(module, "language_model"))
+ )
+ preprocess = gpt_module._preprocess
+
+ def preprocess_hook(
+ *args: Any, _preprocess=preprocess, _gpt=gpt_module, **kwargs: Any
+ ):
+ input_ids = kwargs.get("input_ids")
+ position_ids = kwargs.get("position_ids")
+ for child in _gpt.decoder.modules():
+ if isinstance(child, Dsv4MoELayer):
+ child.set_input_ids(
+ input_ids if isinstance(input_ids, torch.Tensor) else None
+ )
+ if isinstance(child, DeepSeekV4Attention):
+ child.set_position_ids(
+ position_ids
+ if isinstance(position_ids, torch.Tensor)
+ else None
+ )
+ preproc_output = list(_preprocess(*args, **kwargs))
+ decoder_input = cast(torch.Tensor, preproc_output[0])
+ if not decoder_input.requires_grad and decoder_input.is_leaf:
+ decoder_input.requires_grad_(True)
+ table = preproc_output[1]
+ if isinstance(position_ids, torch.Tensor) and torch.is_tensor(table):
+ embedding_dim = int(table.shape[-1])
+ batch_size, sequence_length = position_ids.shape
+ gathered = table.view(table.shape[0], embedding_dim).index_select(
+ 0, position_ids.reshape(-1)
+ )
+ preproc_output[1] = (
+ gathered.view(batch_size, sequence_length, embedding_dim)
+ .permute(1, 0, 2)
+ .contiguous()
+ .unsqueeze(2)
+ )
+ return tuple(preproc_output)
+
+ gpt_module._preprocess = preprocess_hook # type: ignore[attr-defined]
+
+ def collect_layer_families(self, provider: Any) -> list[LayerFamilyInstance]:
+ ratios = list(getattr(provider, "dsv4_compress_ratios", ()) or ())
+
+ def first_layer_index(ratio: int) -> int | None:
+ try:
+ return ratios.index(ratio)
+ except ValueError:
+ return None
+
+ return [
+ LayerFamilyInstance(
+ key="dsv4_sliding_attention", layer_index=first_layer_index(0)
+ ),
+ LayerFamilyInstance(
+ key="dsv4_csa_attention", layer_index=first_layer_index(4)
+ ),
+ LayerFamilyInstance(
+ key="dsv4_hca_attention", layer_index=first_layer_index(128)
+ ),
+ LayerFamilyInstance(key="grouped_moe_mlp", layer_index=0),
+ LayerFamilyInstance(key="shared_experts_mlp", layer_index=0),
+ ]
+
+ def apply_lora_adapters(
+ self,
+ model_chunks: Sequence[Any],
+ provider: Any,
+ *,
+ target_modules: list[str],
+ rank: int,
+ alpha: int,
+ ) -> None:
+ from art.megatron.dsv4.layer import Dsv4TransformerLayer
+ from art.megatron.dsv4.lora import (
+ apply_dsv4_attention_lora,
+ disable_dsv4_etp_shared_expert_lora_compile,
+ install_dsv4_te_permutation_static_configs,
+ )
+ from art.megatron.lora import (
+ _adapter_model_prefix,
+ wrap_grouped_moe_experts,
+ wrap_shared_experts_mlp,
+ )
+
+ target_set = set(target_modules)
+ etp_enabled = int(getattr(provider, "expert_tensor_parallel_size", 1) or 1) > 1
+ if etp_enabled:
+ install_dsv4_te_permutation_static_configs()
+ for chunk in model_chunks:
+ for module in chunk.modules():
+ if not isinstance(module, Dsv4TransformerLayer):
+ continue
+ adapter_model_prefix = _adapter_model_prefix(module)
+ apply_dsv4_attention_lora(
+ module.self_attention,
+ adapter_model_prefix=adapter_model_prefix,
+ target_modules=target_set,
+ rank=rank,
+ alpha=alpha,
+ )
+ wrap_grouped_moe_experts(
+ _require_moe_experts(module),
+ adapter_model_prefix=adapter_model_prefix,
+ target_modules=target_set,
+ rank=rank,
+ alpha=alpha,
+ )
+ if getattr(module.mlp, "shared_experts", None) is not None:
+ wrap_shared_experts_mlp(
+ module.mlp.shared_experts,
+ adapter_model_prefix=adapter_model_prefix,
+ provider=provider,
+ target_modules=target_set,
+ rank=rank,
+ alpha=alpha,
+ )
+ if etp_enabled:
+ disable_dsv4_etp_shared_expert_lora_compile(
+ module.mlp.shared_experts
+ )
+
+ def build_adapter_weights_by_base(
+ self, model_chunks: Sequence[Any]
+ ) -> dict[str, list[Any]]:
+ from art.megatron.dsv4.layer import Dsv4TransformerLayer
+ from art.megatron.dsv4.lora import (
+ add_dsv4_attention_adapter_weights,
+ add_dsv4_shared_experts_adapter_weights,
+ )
+ from art.megatron.weights.adapter_export import (
+ add_grouped_moe_adapter_weights,
+ layer_base_prefix,
+ )
+
+ adapter_weights_by_base: dict[str, list[Any]] = {}
+ for chunk in model_chunks:
+ for module_name, module in chunk.named_modules():
+ if not isinstance(module, Dsv4TransformerLayer):
+ continue
+ layer_prefix = layer_base_prefix(module, module_name=module_name)
+ add_dsv4_attention_adapter_weights(
+ adapter_weights_by_base,
+ layer_prefix=layer_prefix,
+ attention=module.self_attention,
+ )
+ add_grouped_moe_adapter_weights(
+ adapter_weights_by_base,
+ layer_prefix=layer_prefix,
+ experts=_require_moe_experts(module),
+ )
+ if getattr(module.mlp, "shared_experts", None) is not None:
+ add_dsv4_shared_experts_adapter_weights(
+ adapter_weights_by_base,
+ layer_prefix=layer_prefix,
+ shared_experts=module.mlp.shared_experts,
+ )
+ return adapter_weights_by_base
+
+ def iter_merged_vllm_weight_metadata(
+ self,
+ weight_export: Any,
+ ) -> Any:
+ bridge = getattr(weight_export.bridge, "_model_bridge", None)
+ metadata_iter = getattr(bridge, "iter_merged_vllm_weight_metadata", None)
+ if metadata_iter is None:
+ return None
+ return metadata_iter(weight_export)
+
+ def from_vllm_lora_tensors(
+ self,
+ tensors: dict[str, torch.Tensor],
+ *,
+ adapter_config: dict[str, Any],
+ ) -> dict[str, torch.Tensor]:
+ return _dsv4_from_vllm_lora_tensors(
+ tensors,
+ adapter_config=adapter_config,
+ )
+
+ def to_vllm_lora_tensors(
+ self,
+ tensors: dict[str, torch.Tensor],
+ *,
+ adapter_config: dict[str, Any],
+ ) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
+ return _dsv4_to_vllm_lora_tensors(tensors, adapter_config=adapter_config)
+
+ def to_vllm_lora_config(self, adapter_config: dict[str, Any]) -> dict[str, Any]:
+ """Translate ART training targets only for restrictive vLLM launches.
+
+ A vLLM-format DSV4 adapter can be loaded by a vLLM server whose
+ ``target_modules`` filter is unset. ART-managed vLLM launches set that
+ filter for performance/memory control, so the filter must use
+ vLLM/Miles module names rather than ART/Megatron training target names.
+ """
+ return _dsv4_vllm_lora_config(adapter_config)
+
+ def compile_workaround_config(self, provider: Any) -> CompileWorkaroundConfig:
+ return CompileWorkaroundConfig(
+ flags=_compile_workaround_flags_for_provider(
+ provider,
+ _DSV4_MOE_COMPILE_WORKAROUND_FLAGS,
+ ),
+ shared_expert_state=self._shared_expert_compile_state(provider),
+ )
+
+ def ensure_hf_reference_registered(self) -> None:
+ from art.megatron.dsv4.hf_config import ensure_dsv4_hf_model_registered
+
+ ensure_dsv4_hf_model_registered()
+
+ def prepare_hf_reference_config(self, config: Any) -> None:
+ """Puts native HF parity in eager training mode with reduced fit-only axes."""
+ if hasattr(config, "quantization_config"):
+ delattr(config, "quantization_config")
+ config._experts_implementation = "eager"
+ self._apply_oracle_shape_overrides(config)
+
+ def hf_reference_from_pretrained_kwargs(
+ self, *, config: Any, dtype: torch.dtype
+ ) -> dict[str, Any]:
+ del config, dtype
+ return {
+ "experts_implementation": "eager",
+ "ignore_mismatched_sizes": True,
+ "key_mapping": _dsv4_source_to_hf_key_mapping(),
+ }
+
+ def use_hf_reference_state_for_hf_parity(self) -> bool:
+ """DSV4 parity seeds Megatron from the reduced canonical HF oracle state.
+
+ The public checkpoint uses Miles/RadixArk source names and full model
+ shapes, while the validation oracle uses canonical HF names and reduced
+ fit-only axes. This hook is validation-only; production loading remains
+ tied to the normal Bridge checkpoint source.
+ """
+ return True
+
+ def normalize_hf_reference_state_for_hf_parity(
+ self,
+ state: dict[str, torch.Tensor],
+ *,
+ config: Any,
+ ) -> None:
+ _add_dsv4_hf_reference_source_aliases(state, config)
+
+ def configure_oracle_provider(self, provider: Any, *, case_config: Any) -> None:
+ """Mirrors HF oracle reductions while keeping DSV4 hard kernel invariants."""
+ hooks = list(getattr(provider, "_pre_wrap_hooks", []))
+ kept = [hook for hook in hooks if not self._is_bridge_hf_load_hook(hook)]
+ if len(kept) != len(hooks):
+ provider._pre_wrap_hooks = kept
+ provider.register_pre_wrap_hook(
+ lambda chunks: self._initialize_oracle_base_weights(
+ chunks,
+ seed=int(case_config.seed),
+ )
+ )
+ self._apply_oracle_shape_overrides(provider)
+ provider.kv_lora_rank = 512
+ provider.kv_channels = 512
+ provider.qk_pos_emb_head_dim = 64
+ provider.num_query_groups = _ORACLE_NUM_ATTENTION_HEADS
+ provider.num_moe_experts = _ORACLE_NUM_EXPERTS
+ provider.moe_ffn_hidden_size = _ORACLE_FFN_HIDDEN_SIZE
+ provider.ffn_hidden_size = _ORACLE_FFN_HIDDEN_SIZE
+ provider.moe_shared_expert_intermediate_size = _ORACLE_FFN_HIDDEN_SIZE
+ provider.moe_router_topk = _ORACLE_NUM_EXPERTS_PER_TOK
+ provider.dsv4_o_groups = _ORACLE_NUM_ATTENTION_HEADS
+ provider.dsv4_o_lora_rank = 1024
+ provider.dsa_indexer_n_heads = _ORACLE_INDEX_HEADS
+ provider.dsa_indexer_head_dim = 128
+ provider.dsa_indexer_topk = _ORACLE_INDEX_TOPK
+
+ @staticmethod
+ def _is_bridge_hf_load_hook(hook: Any) -> bool:
+ fn = getattr(hook, "func", hook)
+ name = getattr(fn, "__name__", "")
+ qualname = getattr(fn, "__qualname__", "")
+ return name in {
+ "load_weights_hf_to_megatron",
+ "_optimized_load_weights_hf_to_megatron",
+ } or qualname.endswith(".load_weights_hf_to_megatron")
+
+ def _apply_oracle_shape_overrides(self, config: Any) -> None:
+ """Reduces memory-heavy axes only; head_dim/window/o-rank stay production-sized."""
+ config.hidden_size = _ORACLE_HIDDEN_SIZE
+ config.q_lora_rank = _ORACLE_Q_LORA_RANK
+ config.num_attention_heads = _ORACLE_NUM_ATTENTION_HEADS
+ config.n_routed_experts = _ORACLE_NUM_EXPERTS
+ config.num_experts_per_tok = _ORACLE_NUM_EXPERTS_PER_TOK
+ config.moe_intermediate_size = _ORACLE_FFN_HIDDEN_SIZE
+ config.o_groups = _ORACLE_NUM_ATTENTION_HEADS
+ config.index_n_heads = _ORACLE_INDEX_HEADS
+ config.index_head_dim = 128
+ config.index_topk = _ORACLE_INDEX_TOPK
+
+ def _initialize_oracle_base_weights(
+ self,
+ model_chunks: Sequence[Any],
+ *,
+ seed: int,
+ ) -> Sequence[Any]:
+ """Seeds reduced DSV4 oracle base tensors after meta-device materialization.
+
+ DSV4 correctness runs strip the full-checkpoint Bridge load because the
+ public tensors are production-sized and quantized. Since ART materializes
+ meta models with empty storage, every base tensor must be explicitly
+ initialized before freeze/LoRA hooks run; otherwise the oracle exercises a
+ zero model and cannot produce meaningful adapter gradients.
+ """
+ from megatron.core import parallel_state as ps
+
+ ep_rank = ps.get_expert_model_parallel_rank()
+ ep_size = ps.get_expert_model_parallel_world_size()
+ with torch.no_grad():
+ for chunk in model_chunks:
+ for name, param in chunk.named_parameters():
+ if self._is_oracle_lora_tensor(name):
+ continue
+ init_name = self._oracle_base_tensor_name(
+ name,
+ ep_rank=ep_rank,
+ ep_size=ep_size,
+ )
+ self._initialize_oracle_tensor(
+ init_name,
+ param,
+ seed=seed,
+ )
+ for name, buffer in chunk.named_buffers():
+ self._initialize_oracle_buffer(name, buffer, seed=seed)
+ return model_chunks
+
+ @staticmethod
+ def _is_oracle_lora_tensor(name: str) -> bool:
+ return "_lora." in name or ".lora." in name
+
+ @staticmethod
+ def _oracle_base_tensor_name(name: str, *, ep_rank: int, ep_size: int) -> str:
+ if ep_size <= 1:
+ return name
+ match = _ORACLE_EXPERT_WEIGHT_RE.search(name)
+ if match is None:
+ return name
+ local_expert = int(match.group("expert"))
+ local_expert_count = max(1, _ORACLE_NUM_EXPERTS // ep_size)
+ global_expert = ep_rank * local_expert_count + local_expert
+ return f"{name[: match.start('expert')]}{global_expert}"
+
+ def _initialize_oracle_buffer(
+ self,
+ name: str,
+ tensor: torch.Tensor,
+ *,
+ seed: int,
+ ) -> None:
+ if name.endswith("freqs_cis"):
+ return
+ if name.endswith("tid2eid"):
+ self._initialize_oracle_tid2eid(tensor)
+ return
+ if not torch.is_floating_point(tensor):
+ return
+ self._initialize_oracle_tensor(name, tensor, seed=seed)
+
+ @staticmethod
+ def _initialize_oracle_tid2eid(tensor: torch.Tensor) -> None:
+ if tensor.ndim != 2:
+ raise RuntimeError(
+ f"Expected DSV4 tid2eid to be 2D, got {tuple(tensor.shape)}"
+ )
+ token_ids = torch.arange(tensor.shape[0], device=tensor.device).unsqueeze(1)
+ offsets = torch.arange(tensor.shape[1], device=tensor.device).unsqueeze(0)
+ tensor.copy_(
+ (token_ids + offsets).remainder(_ORACLE_NUM_EXPERTS).to(tensor.dtype)
+ )
+
+ def _initialize_oracle_tensor(
+ self,
+ name: str,
+ tensor: torch.Tensor,
+ *,
+ seed: int,
+ ) -> None:
+ if tensor.is_meta:
+ raise RuntimeError(f"DSV4 oracle tensor was not materialized: {name}")
+ if not torch.is_floating_point(tensor):
+ return
+ if self._is_oracle_identity_weight(name):
+ tensor.fill_(1)
+ return
+ if name.endswith(
+ ("bias", "attn_sink", "_base", "_scale", "e_score_correction_bias")
+ ):
+ tensor.zero_()
+ return
+ logical_shape, partition_dim, partition_rank = (
+ self._oracle_logical_tensor_for_rank(name, tensor)
+ )
+ digest = hashlib.sha256(f"{seed}:{name}".encode("utf-8")).digest()
+ key_seed = int.from_bytes(digest[:8], "little") % (2**31)
+ generator = torch.Generator(device=tensor.device).manual_seed(key_seed)
+ values = torch.randn(
+ logical_shape,
+ generator=generator,
+ device=tensor.device,
+ dtype=torch.float32,
+ )
+ if partition_dim is not None:
+ values = self._oracle_slice_logical_tensor(
+ name,
+ values,
+ tensor,
+ partition_dim=partition_dim,
+ partition_rank=partition_rank,
+ )
+ tensor.copy_((0.02 * values).to(dtype=tensor.dtype))
+
+ @staticmethod
+ def _oracle_slice_logical_tensor(
+ name: str,
+ values: torch.Tensor,
+ tensor: torch.Tensor,
+ *,
+ partition_dim: int,
+ partition_rank: int,
+ ) -> torch.Tensor:
+ if partition_dim == 0 and Dsv4Handler._oracle_is_fused_fc1(name):
+ if values.shape[0] % 2 != 0 or tensor.shape[0] % 2 != 0:
+ raise RuntimeError(
+ "DSV4 fused FC1 oracle tensor must have an even "
+ f"gate/up dimension, got logical={tuple(values.shape)} "
+ f"local={tuple(tensor.shape)} for {name}."
+ )
+ local_component = tensor.shape[0] // 2
+ gate, up = values.chunk(2, dim=0)
+ start = partition_rank * local_component
+ return torch.cat(
+ (
+ gate.narrow(0, start, local_component),
+ up.narrow(0, start, local_component),
+ ),
+ dim=0,
+ )
+ return values.narrow(
+ partition_dim,
+ partition_rank * tensor.shape[partition_dim],
+ tensor.shape[partition_dim],
+ )
+
+ @staticmethod
+ def _oracle_logical_tensor_for_rank(
+ name: str,
+ tensor: torch.Tensor,
+ ) -> tuple[tuple[int, ...], int | None, int]:
+ """Returns the logical full tensor shape and rank-local TP slice metadata.
+
+ DSV4 oracle base weights are validation-only random tensors. TP ranks
+ must receive slices of the same logical tensor as the TP1 oracle; if each
+ rank independently initializes its local shard, TP2 is not mathematically
+ comparable to TP1 even when the layer implementation is correct.
+ """
+ partition_dim = getattr(tensor, "partition_dim", None)
+ from megatron.core import parallel_state as ps
+
+ if _ORACLE_EXPERT_WEIGHT_RE.search(name) is not None:
+ etp_group = ps.get_expert_tensor_parallel_group(check_initialized=False)
+ etp_size = etp_group.size() if etp_group is not None else 1
+ if etp_size <= 1:
+ return tuple(tensor.shape), None, 0
+ dim = (
+ int(partition_dim)
+ if partition_dim is not None and int(partition_dim) >= 0
+ else Dsv4Handler._oracle_dsv4_expert_missing_etp_partition_dim(name)
+ )
+ if dim is None:
+ return tuple(tensor.shape), None, 0
+ logical_shape = list(tensor.shape)
+ logical_shape[dim] *= etp_size
+ return tuple(logical_shape), dim, etp_group.rank()
+
+ tp_size = ps.get_tensor_model_parallel_world_size()
+ tp_rank = ps.get_tensor_model_parallel_rank()
+ if (
+ bool(getattr(tensor, "tensor_model_parallel", False))
+ and partition_dim is not None
+ and int(partition_dim) >= 0
+ ):
+ dim = int(partition_dim)
+ else:
+ dim = Dsv4Handler._oracle_dsv4_missing_tp_partition_dim(name)
+ if dim is None or tp_size <= 1:
+ return tuple(tensor.shape), None, 0
+ logical_shape = list(tensor.shape)
+ logical_shape[dim] *= tp_size
+ return tuple(logical_shape), dim, tp_rank
+
+ @staticmethod
+ def _oracle_dsv4_missing_tp_partition_dim(name: str) -> int | None:
+ """Mirrors DSV4 Bridge TP shape rules for custom modules without TP attrs."""
+ if name.endswith(("embedding.word_embeddings.weight", "output_layer.weight")):
+ return 0
+ if name.endswith(
+ (
+ ".self_attention.wq_b.weight",
+ ".self_attention.wo_a.weight",
+ ".mlp.shared_experts.linear_fc1.weight",
+ ".mlp.shared_experts.linear_fc1.linear_fc1.weight",
+ )
+ ):
+ return 0
+ if name.endswith(
+ (
+ ".self_attention.wo_b.weight",
+ ".mlp.shared_experts.linear_fc2.weight",
+ ".mlp.shared_experts.linear_fc2.row_parallel_lora.linear_proj.weight",
+ )
+ ):
+ return 1
+ return None
+
+ @staticmethod
+ def _oracle_dsv4_expert_missing_etp_partition_dim(name: str) -> int | None:
+ """Mirrors TE grouped-expert ETP shards when TE omits TP metadata."""
+ if ".mlp.experts.linear_fc1." in name:
+ return 0
+ if ".mlp.experts.linear_fc2." in name:
+ return 1
+ return None
+
+ @staticmethod
+ def _oracle_is_shared_expert_fused_fc1(name: str) -> bool:
+ return name.endswith(
+ (
+ ".mlp.shared_experts.linear_fc1.weight",
+ ".mlp.shared_experts.linear_fc1.linear_fc1.weight",
+ )
+ )
+
+ @staticmethod
+ def _oracle_is_grouped_expert_fused_fc1(name: str) -> bool:
+ return ".mlp.experts.linear_fc1." in name and bool(
+ _ORACLE_EXPERT_WEIGHT_RE.search(name)
+ )
+
+ @staticmethod
+ def _oracle_is_fused_fc1(name: str) -> bool:
+ return Dsv4Handler._oracle_is_shared_expert_fused_fc1(
+ name
+ ) or Dsv4Handler._oracle_is_grouped_expert_fused_fc1(name)
+
+ @staticmethod
+ def _is_oracle_identity_weight(name: str) -> bool:
+ return name.endswith(".weight") and any(
+ token in name for token in ("layernorm", "_norm", ".norm.")
+ )
+
+
+def ensure_dsv4_bridge_registered() -> None:
+ from art.megatron.dsv4.bridge import ensure_dsv4_bridge_registered as _ensure
+
+ _ensure()
+
+
+def _ensure_dsv4_hf_config_registered() -> None:
+ from art.megatron.dsv4.hf_config import ensure_dsv4_hf_config_registered
+
+ ensure_dsv4_hf_config_registered()
+
+
+def _sanitize_dsv4_child_process_env() -> None:
+ from art.megatron.dsv4.kernel.tilelang_import import sanitize_tilelang_env
+
+ sanitize_tilelang_env()
+
+
+_sanitize_dsv4_child_process_env()
+_ensure_dsv4_hf_config_registered()
+DSV4_HANDLER = Dsv4Handler()
+
+
+def _dsv4_source_to_hf_key_mapping() -> dict[str, str]:
+ layer = r"layers\.(\d+)"
+ target = r"model.layers.\1"
+ return {
+ r"^embed\.weight$": "model.embed_tokens.weight",
+ r"^head\.weight$": "lm_head.weight",
+ r"^norm\.weight$": "model.norm.weight",
+ r"^hc_head_fn$": "model.hc_head.hc_fn",
+ r"^hc_head_base$": "model.hc_head.hc_base",
+ r"^hc_head_scale$": "model.hc_head.hc_scale",
+ rf"^{layer}\.attn_norm\.weight$": rf"{target}.input_layernorm.weight",
+ rf"^{layer}\.ffn_norm\.weight$": rf"{target}.post_attention_layernorm.weight",
+ rf"^{layer}\.hc_attn_fn$": rf"{target}.attn_hc.fn",
+ rf"^{layer}\.hc_attn_base$": rf"{target}.attn_hc.base",
+ rf"^{layer}\.hc_attn_scale$": rf"{target}.attn_hc.scale",
+ rf"^{layer}\.hc_ffn_fn$": rf"{target}.ffn_hc.fn",
+ rf"^{layer}\.hc_ffn_base$": rf"{target}.ffn_hc.base",
+ rf"^{layer}\.hc_ffn_scale$": rf"{target}.ffn_hc.scale",
+ rf"^{layer}\.attn\.wq_a\.weight$": rf"{target}.self_attn.q_a_proj.weight",
+ rf"^{layer}\.attn\.q_norm\.weight$": rf"{target}.self_attn.q_a_norm.weight",
+ rf"^{layer}\.attn\.wq_b\.weight$": rf"{target}.self_attn.q_b_proj.weight",
+ rf"^{layer}\.attn\.wkv\.weight$": rf"{target}.self_attn.kv_proj.weight",
+ rf"^{layer}\.attn\.kv_norm\.weight$": rf"{target}.self_attn.kv_norm.weight",
+ rf"^{layer}\.attn\.wo_a\.weight$": rf"{target}.self_attn.o_a_proj.weight",
+ rf"^{layer}\.attn\.wo_b\.weight$": rf"{target}.self_attn.o_b_proj.weight",
+ rf"^{layer}\.attn\.attn_sink$": rf"{target}.self_attn.sinks",
+ rf"^{layer}\.ffn\.gate\.weight$": rf"{target}.mlp.gate.weight",
+ rf"^{layer}\.ffn\.gate\.tid2eid$": rf"{target}.mlp.gate.tid2eid",
+ rf"^{layer}\.ffn\.gate\.bias$": (rf"{target}.mlp.gate.e_score_correction_bias"),
+ rf"^{layer}\.ffn\.shared_experts\.w1\.weight$": (
+ rf"{target}.mlp.shared_experts.gate_proj.weight"
+ ),
+ rf"^{layer}\.ffn\.shared_experts\.w3\.weight$": (
+ rf"{target}.mlp.shared_experts.up_proj.weight"
+ ),
+ rf"^{layer}\.ffn\.shared_experts\.w2\.weight$": (
+ rf"{target}.mlp.shared_experts.down_proj.weight"
+ ),
+ rf"^{layer}\.attn\.compressor\.ape$": (
+ rf"{target}.self_attn.compressor.position_bias"
+ ),
+ rf"^{layer}\.attn\.compressor\.wkv\.weight$": (
+ rf"{target}.self_attn.compressor.kv_proj.weight"
+ ),
+ rf"^{layer}\.attn\.compressor\.wgate\.weight$": (
+ rf"{target}.self_attn.compressor.gate_proj.weight"
+ ),
+ rf"^{layer}\.attn\.compressor\.norm\.weight$": (
+ rf"{target}.self_attn.compressor.kv_norm.weight"
+ ),
+ rf"^{layer}\.attn\.indexer\.wq_b\.weight$": (
+ rf"{target}.self_attn.compressor.indexer.q_b_proj.weight"
+ ),
+ rf"^{layer}\.attn\.indexer\.weights_proj\.weight$": (
+ rf"{target}.self_attn.compressor.indexer.scorer.weights_proj.weight"
+ ),
+ rf"^{layer}\.attn\.indexer\.compressor\.ape$": (
+ rf"{target}.self_attn.compressor.indexer.position_bias"
+ ),
+ rf"^{layer}\.attn\.indexer\.compressor\.wkv\.weight$": (
+ rf"{target}.self_attn.compressor.indexer.kv_proj.weight"
+ ),
+ rf"^{layer}\.attn\.indexer\.compressor\.wgate\.weight$": (
+ rf"{target}.self_attn.compressor.indexer.gate_proj.weight"
+ ),
+ rf"^{layer}\.attn\.indexer\.compressor\.norm\.weight$": (
+ rf"{target}.self_attn.compressor.indexer.kv_norm.weight"
+ ),
+ }
+
+
+def _add_dsv4_hf_reference_source_aliases(
+ state: dict[str, torch.Tensor],
+ config: Any,
+) -> None:
+ def add(source: str, canonical: str) -> None:
+ if canonical in state and source not in state:
+ state[source] = state[canonical]
+
+ add("embed.weight", "model.embed_tokens.weight")
+ add("head.weight", "lm_head.weight")
+ add("hc_head_fn", "model.hc_head.hc_fn")
+ add("hc_head_base", "model.hc_head.hc_base")
+ add("hc_head_scale", "model.hc_head.hc_scale")
+ for layer_idx in range(int(config.num_hidden_layers)):
+ source = f"layers.{layer_idx}"
+ canonical = f"model.layers.{layer_idx}"
+ add(f"{source}.attn_norm.weight", f"{canonical}.input_layernorm.weight")
+ add(f"{source}.ffn_norm.weight", f"{canonical}.post_attention_layernorm.weight")
+ add(f"{source}.hc_attn_fn", f"{canonical}.attn_hc.fn")
+ add(f"{source}.hc_attn_base", f"{canonical}.attn_hc.base")
+ add(f"{source}.hc_attn_scale", f"{canonical}.attn_hc.scale")
+ add(f"{source}.hc_ffn_fn", f"{canonical}.ffn_hc.fn")
+ add(f"{source}.hc_ffn_base", f"{canonical}.ffn_hc.base")
+ add(f"{source}.hc_ffn_scale", f"{canonical}.ffn_hc.scale")
+ add(f"{source}.attn.wq_a.weight", f"{canonical}.self_attn.q_a_proj.weight")
+ add(f"{source}.attn.q_norm.weight", f"{canonical}.self_attn.q_a_norm.weight")
+ add(f"{source}.attn.wq_b.weight", f"{canonical}.self_attn.q_b_proj.weight")
+ add(f"{source}.attn.wkv.weight", f"{canonical}.self_attn.kv_proj.weight")
+ add(f"{source}.attn.kv_norm.weight", f"{canonical}.self_attn.kv_norm.weight")
+ add(f"{source}.attn.wo_a.weight", f"{canonical}.self_attn.o_a_proj.weight")
+ add(f"{source}.attn.wo_b.weight", f"{canonical}.self_attn.o_b_proj.weight")
+ add(f"{source}.attn.attn_sink", f"{canonical}.self_attn.sinks")
+ add(f"{source}.ffn.gate.weight", f"{canonical}.mlp.gate.weight")
+ add(f"{source}.ffn.gate.tid2eid", f"{canonical}.mlp.gate.tid2eid")
+ add(
+ f"{source}.ffn.gate.bias",
+ f"{canonical}.mlp.gate.e_score_correction_bias",
+ )
+ add(
+ f"{source}.ffn.shared_experts.w1.weight",
+ f"{canonical}.mlp.shared_experts.gate_proj.weight",
+ )
+ add(
+ f"{source}.ffn.shared_experts.w3.weight",
+ f"{canonical}.mlp.shared_experts.up_proj.weight",
+ )
+ add(
+ f"{source}.ffn.shared_experts.w2.weight",
+ f"{canonical}.mlp.shared_experts.down_proj.weight",
+ )
+ add(
+ f"{source}.attn.compressor.ape",
+ f"{canonical}.self_attn.compressor.position_bias",
+ )
+ add(
+ f"{source}.attn.compressor.wkv.weight",
+ f"{canonical}.self_attn.compressor.kv_proj.weight",
+ )
+ add(
+ f"{source}.attn.compressor.wgate.weight",
+ f"{canonical}.self_attn.compressor.gate_proj.weight",
+ )
+ add(
+ f"{source}.attn.compressor.norm.weight",
+ f"{canonical}.self_attn.compressor.kv_norm.weight",
+ )
+ add(
+ f"{source}.attn.indexer.wq_b.weight",
+ f"{canonical}.self_attn.compressor.indexer.q_b_proj.weight",
+ )
+ add(
+ f"{source}.attn.indexer.weights_proj.weight",
+ f"{canonical}.self_attn.compressor.indexer.scorer.weights_proj.weight",
+ )
+ add(
+ f"{source}.attn.indexer.compressor.ape",
+ f"{canonical}.self_attn.compressor.indexer.position_bias",
+ )
+ add(
+ f"{source}.attn.indexer.compressor.wkv.weight",
+ f"{canonical}.self_attn.compressor.indexer.kv_proj.weight",
+ )
+ add(
+ f"{source}.attn.indexer.compressor.wgate.weight",
+ f"{canonical}.self_attn.compressor.indexer.gate_proj.weight",
+ )
+ add(
+ f"{source}.attn.indexer.compressor.norm.weight",
+ f"{canonical}.self_attn.compressor.indexer.kv_norm.weight",
+ )
+ gate_up = state.get(f"{canonical}.mlp.experts.gate_up_proj")
+ if gate_up is not None:
+ gate, up = gate_up.chunk(2, dim=1)
+ for expert_idx in range(int(gate.shape[0])):
+ state.setdefault(
+ f"{source}.ffn.experts.{expert_idx}.w1.weight",
+ gate[expert_idx].contiguous(),
+ )
+ state.setdefault(
+ f"{source}.ffn.experts.{expert_idx}.w3.weight",
+ up[expert_idx].contiguous(),
+ )
+ down = state.get(f"{canonical}.mlp.experts.down_proj")
+ if down is not None:
+ for expert_idx in range(int(down.shape[0])):
+ state.setdefault(
+ f"{source}.ffn.experts.{expert_idx}.w2.weight",
+ down[expert_idx].contiguous(),
+ )
+
+
+def _dsv4_unpack_vllm_3d_lora_b(
+ tensor: torch.Tensor,
+ *,
+ num_experts: int,
+ rank: int,
+) -> torch.Tensor:
+ return tensor.reshape(tensor.shape[0], rank, num_experts).permute(2, 0, 1)
+
+
+def _dsv4_clone(tensor: torch.Tensor) -> torch.Tensor:
+ return tensor.clone().contiguous()
+
+
+def _dsv4_to_vllm_lora_key(key: str) -> str:
+ match = _DSV4_ART_MOE_EXPERT_KEY_RE.match(key)
+ if match is not None:
+ module = {
+ "gate_proj": "w1",
+ "down_proj": "w2",
+ "up_proj": "w3",
+ }[match.group("module")]
+ prefix = match.group("prefix").replace(".mlp.experts", ".ffn.experts", 1)
+ return f"{prefix}.{match.group('expert')}.{module}.{match.group('lora')}.weight"
+
+ replacements = (
+ (".self_attn.compressor.kv_proj.", ".attn.mla_attn.compressor.wkv."),
+ (".self_attn.compressor.gate_proj.", ".attn.mla_attn.compressor.wgate."),
+ (".self_attn.q_a_proj.", ".attn.wq_a."),
+ (".self_attn.q_b_proj.", ".attn.wq_b."),
+ (".self_attn.kv_proj.", ".attn.wkv."),
+ (".self_attn.o_a_proj.", ".attn.wo_a."),
+ (".self_attn.o_b_proj.", ".attn.wo_b."),
+ (".mlp.shared_expert.", ".ffn.shared_experts."),
+ (".mlp.shared_experts.", ".ffn.shared_experts."),
+ )
+ for old, new in replacements:
+ if old in key:
+ return key.replace(old, new, 1)
+ return key
+
+
+def _dsv4_from_vllm_lora_key(key: str) -> str:
+ match = _DSV4_VLLM_MOE_EXPERT_KEY_RE.match(key)
+ if match is not None:
+ module = {
+ "w1": "gate_proj",
+ "w2": "down_proj",
+ "w3": "up_proj",
+ }[match.group("module")]
+ prefix = match.group("prefix").replace(".ffn.experts", ".mlp.experts", 1)
+ return f"{prefix}.{match.group('expert')}.{module}.{match.group('lora')}.weight"
+
+ replacements = (
+ (".attn.mla_attn.compressor.wkv.", ".self_attn.compressor.kv_proj."),
+ (".attn.mla_attn.compressor.wgate.", ".self_attn.compressor.gate_proj."),
+ (".attn.wq_a.", ".self_attn.q_a_proj."),
+ (".attn.wq_b.", ".self_attn.q_b_proj."),
+ (".attn.wkv.", ".self_attn.kv_proj."),
+ (".attn.wo_a.", ".self_attn.o_a_proj."),
+ (".attn.wo_b.", ".self_attn.o_b_proj."),
+ (".ffn.shared_experts.", ".mlp.shared_expert."),
+ (".mlp.shared_experts.", ".mlp.shared_expert."),
+ )
+ for old, new in replacements:
+ if old in key:
+ return key.replace(old, new, 1)
+ return key
+
+
+def _dsv4_vllm_lora_config(adapter_config: dict[str, Any]) -> dict[str, Any]:
+ target_modules = adapter_config.get("target_modules")
+ if not isinstance(target_modules, (list, tuple, set)):
+ return adapter_config
+ transformed: list[str] = []
+ ordered_target_modules = (
+ sorted(target_modules) if isinstance(target_modules, set) else target_modules
+ )
+ for module in ordered_target_modules:
+ if module in {"q_a_proj", "kv_proj"}:
+ transformed.append("fused_wqa_wkv")
+ elif module == "q_b_proj":
+ transformed.append("wq_b")
+ elif module == "o_a_proj":
+ transformed.append("wo_a")
+ elif module == "o_b_proj":
+ transformed.append("wo_b")
+ elif module in {"compressor.kv_proj", "compressor.gate_proj"}:
+ transformed.append("fused_wkv_wgate")
+ elif module in {"gate_proj", "up_proj"}:
+ transformed.extend(("gate_up_proj", "experts"))
+ elif module == "down_proj":
+ transformed.extend(("down_proj", "experts"))
+ elif module == "experts":
+ transformed.append("experts")
+ else:
+ transformed.append(module)
+ config = dict(adapter_config)
+ config["target_modules"] = list(dict.fromkeys(transformed))
+ return config
+
+
+def _dsv4_to_vllm_lora_tensors(
+ tensors: dict[str, torch.Tensor],
+ *,
+ adapter_config: dict[str, Any],
+) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
+ transformed: dict[str, torch.Tensor] = {}
+ for key, tensor in tensors.items():
+ vllm_key = _dsv4_to_vllm_lora_key(key)
+ if vllm_key in transformed:
+ raise RuntimeError(
+ f"Duplicate DSV4 LoRA tensor after conversion: {vllm_key}"
+ )
+ transformed[vllm_key] = tensor
+ return transformed, _dsv4_vllm_lora_config(adapter_config)
+
+
+def _dsv4_from_vllm_lora_tensors(
+ tensors: dict[str, torch.Tensor],
+ *,
+ adapter_config: dict[str, Any],
+) -> dict[str, torch.Tensor]:
+ grouped: dict[str, dict[str, torch.Tensor]] = {}
+ for key, tensor in tensors.items():
+ match = _DSV4_VLLM_MOE_KEY_RE.match(key)
+ if match is None:
+ continue
+ slot = (
+ f"{'base_layer.' if match.group('base_layer') else ''}{match.group('lora')}"
+ )
+ grouped.setdefault(match.group("prefix"), {})[slot] = tensor
+ if not grouped:
+ return {
+ _dsv4_from_vllm_lora_key(key): tensor for key, tensor in tensors.items()
+ }
+
+ rank = int(adapter_config["r"])
+ transformed: dict[str, torch.Tensor] = {}
+ used_keys: set[str] = set()
+ for prefix, slots in grouped.items():
+ try:
+ gate_up_a = slots["base_layer.lora_A"]
+ gate_up_b = slots["base_layer.lora_B"]
+ down_a = slots["lora_A"]
+ down_b = slots["lora_B"]
+ except KeyError as exc:
+ raise RuntimeError(
+ f"Incomplete DSV4 vLLM MoE LoRA block for {prefix}"
+ ) from exc
+ if gate_up_a.shape[0] % rank != 0:
+ raise RuntimeError(
+ f"{prefix}: gate/up lora_A rows {gate_up_a.shape[0]} are not "
+ f"divisible by rank {rank}"
+ )
+ if gate_up_b.shape[0] % 2 != 0:
+ raise RuntimeError(
+ f"{prefix}: gate/up lora_B rows {gate_up_b.shape[0]} are not even"
+ )
+ num_experts = gate_up_a.shape[0] // rank
+ gate_up_b_by_expert = _dsv4_unpack_vllm_3d_lora_b(
+ gate_up_b,
+ num_experts=num_experts,
+ rank=rank,
+ )
+ down_b_by_expert = _dsv4_unpack_vllm_3d_lora_b(
+ down_b,
+ num_experts=num_experts,
+ rank=rank,
+ )
+ for expert in range(num_experts):
+ row = expert * rank
+ gate_up_a_block = gate_up_a[row : row + rank]
+ down_a_block = down_a[row : row + rank]
+ gate_b, up_b = gate_up_b_by_expert[expert].chunk(2, dim=0)
+ transformed[f"{prefix}.{expert}.gate_proj.lora_A.weight"] = _dsv4_clone(
+ gate_up_a_block
+ )
+ transformed[f"{prefix}.{expert}.gate_proj.lora_B.weight"] = _dsv4_clone(
+ gate_b
+ )
+ transformed[f"{prefix}.{expert}.up_proj.lora_A.weight"] = _dsv4_clone(
+ gate_up_a_block
+ )
+ transformed[f"{prefix}.{expert}.up_proj.lora_B.weight"] = _dsv4_clone(up_b)
+ transformed[f"{prefix}.{expert}.down_proj.lora_A.weight"] = _dsv4_clone(
+ down_a_block
+ )
+ transformed[f"{prefix}.{expert}.down_proj.lora_B.weight"] = _dsv4_clone(
+ down_b_by_expert[expert]
+ )
+ used_keys.update(
+ {
+ f"{prefix}.base_layer.lora_A.weight",
+ f"{prefix}.base_layer.lora_B.weight",
+ f"{prefix}.lora_A.weight",
+ f"{prefix}.lora_B.weight",
+ }
+ )
+ for key, tensor in tensors.items():
+ if key not in used_keys:
+ transformed[_dsv4_from_vllm_lora_key(key)] = tensor
+ return transformed
diff --git a/src/art/megatron/model_support/handlers/qwen3_moe.py b/src/art/megatron/model_support/handlers/qwen3_moe.py
index 548419d94..41245bc16 100644
--- a/src/art/megatron/model_support/handlers/qwen3_moe.py
+++ b/src/art/megatron/model_support/handlers/qwen3_moe.py
@@ -35,6 +35,9 @@ def to_vllm_lora_tensors(
) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
return _to_vllm_lora_tensors(tensors, adapter_config=adapter_config)
+ def to_vllm_lora_config(self, adapter_config: dict[str, Any]) -> dict[str, Any]:
+ return _qwen3_moe_config(adapter_config)
+
def install_preprocess_patch(self, model_chunks: Sequence[Any]) -> None:
install_qwen3_text_preprocess_patch(model_chunks)
diff --git a/src/art/megatron/model_support/registry.py b/src/art/megatron/model_support/registry.py
index 716573e64..0c9926ab6 100644
--- a/src/art/megatron/model_support/registry.py
+++ b/src/art/megatron/model_support/registry.py
@@ -14,6 +14,7 @@
_QWEN3_5_MOE_HANDLER_KEY = "qwen3_5_moe"
_GEMMA4_DENSE_HANDLER_KEY = "gemma4_dense"
_GEMMA4_MOE_HANDLER_KEY = "gemma4_moe"
+_DSV4_HANDLER_KEY = "dsv4"
_GPT_OSS_MOE_HANDLER_KEY = "gpt_oss_moe"
_VALIDATED_NATIVE_VLLM_LORA_STATUS: NativeVllmLoraStatus = "validated"
_WIP_NATIVE_VLLM_LORA_STATUS: NativeVllmLoraStatus = "wip"
@@ -59,6 +60,19 @@
"up_proj",
"down_proj",
)
+_DSV4_TARGET_MODULES = (
+ "q_a_proj",
+ "q_b_proj",
+ "kv_proj",
+ "o_a_proj",
+ "o_b_proj",
+ "compressor.kv_proj",
+ "compressor.gate_proj",
+ "gate_proj",
+ "up_proj",
+ "down_proj",
+ "experts",
+)
DEFAULT_DENSE_SPEC = ModelSupportSpec(
key="default_dense",
@@ -166,6 +180,21 @@
),
)
+DSV4_SPEC = ModelSupportSpec(
+ key="dsv4",
+ handler_key=_DSV4_HANDLER_KEY,
+ is_moe=True,
+ model_names=(
+ "deepseek-ai/DeepSeek-V4-Flash",
+ "deepseek-ai/DeepSeek-V4-Flash-Base",
+ "deepseek-ai/DeepSeek-V4-Pro",
+ "deepseek-ai/DeepSeek-V4-Pro-Base",
+ ),
+ default_target_modules=_DSV4_TARGET_MODULES,
+ native_vllm_lora_status=_VALIDATED_NATIVE_VLLM_LORA_STATUS,
+ dependency_floor=DependencyFloor(transformers="5.12.1"),
+)
+
GPT_OSS_MOE_SPEC = ModelSupportSpec(
key="gpt_oss_moe",
handler_key=_GPT_OSS_MOE_HANDLER_KEY,
@@ -189,6 +218,7 @@
QWEN3_5_DENSE_SPEC,
GEMMA4_MOE_SPEC,
GEMMA4_DENSE_SPEC,
+ DSV4_SPEC,
GPT_OSS_MOE_SPEC,
)
PROBE_ONLY_MODEL_SUPPORT_SPECS = ()
@@ -237,6 +267,10 @@
"art.megatron.model_support.handlers.gemma4",
"GEMMA4_DENSE_HANDLER",
),
+ _DSV4_HANDLER_KEY: (
+ "art.megatron.model_support.handlers.dsv4",
+ "DSV4_HANDLER",
+ ),
_GPT_OSS_MOE_HANDLER_KEY: (
"art.megatron.model_support.handlers.gpt_oss",
"GPT_OSS_MOE_HANDLER",
@@ -259,6 +293,10 @@
"art.megatron.model_support.handlers.gemma4",
"ensure_gemma4_text_only_bridge_registered",
),
+ "dsv4": (
+ "art.megatron.model_support.handlers.dsv4",
+ "ensure_dsv4_bridge_registered",
+ ),
}
_HANDLERS_BY_KEY: dict[str, ModelSupportHandler] = {}
_REGISTERED_BRIDGE_KEYS: set[str] = set()
@@ -270,6 +308,7 @@
QWEN3_5_MODELS = QWEN3_5_DENSE_MODELS | QWEN3_5_MOE_MODELS
GEMMA4_MOE_MODELS = frozenset(GEMMA4_MOE_SPEC.model_names)
GEMMA4_DENSE_MODELS = frozenset(GEMMA4_DENSE_SPEC.model_names)
+DSV4_MODELS = frozenset(DSV4_SPEC.model_names)
GPT_OSS_MOE_MODELS = frozenset(GPT_OSS_MOE_SPEC.model_names)
@@ -354,6 +393,18 @@ def default_target_modules_for_model(
)
+def vllm_lora_config_for_model(
+ base_model: str,
+ adapter_config: dict,
+ *,
+ allow_unvalidated_arch: bool = False,
+) -> dict:
+ return get_model_support_handler(
+ base_model,
+ allow_unvalidated_arch=allow_unvalidated_arch,
+ ).to_vllm_lora_config(adapter_config)
+
+
def native_vllm_lora_status_for_model(
base_model: str,
*,
@@ -390,6 +441,18 @@ def model_uses_expert_parallel(
).is_moe
+def model_supports_context_parallel(
+ base_model: str,
+ *,
+ allow_unvalidated_arch: bool = False,
+) -> bool:
+ spec = get_model_support_spec(
+ base_model,
+ allow_unvalidated_arch=allow_unvalidated_arch,
+ )
+ return bool(get_model_support_handler_for_spec(spec).cp_supported)
+
+
def is_model_support_registered(base_model: str) -> bool:
return base_model in _SPECS_BY_MODEL
diff --git a/src/art/megatron/model_support/spec.py b/src/art/megatron/model_support/spec.py
index eb10dfe79..e9d0e8b4b 100644
--- a/src/art/megatron/model_support/spec.py
+++ b/src/art/megatron/model_support/spec.py
@@ -1,6 +1,6 @@
from typing import TYPE_CHECKING, Any, Literal, Protocol, Sequence, runtime_checkable
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
if TYPE_CHECKING:
from megatron.bridge import AutoBridge
@@ -46,6 +46,18 @@ class ArchitectureReport(BaseModel):
unresolved_risks: list[str] = Field(default_factory=list)
+class SharedPrefixModelStateContext(BaseModel):
+ model_config = ConfigDict(arbitrary_types_allowed=True)
+
+ group_ids: Any
+ parent_ids: Any
+ input_pos: Any | None = None
+ device: Any
+ attention_token_layout_index: Any | None = None
+ attention_head_dim: int | None = None
+ attention_value_head_dim: int | None = None
+
+
class CompileWorkaroundConfig(BaseModel):
flags: tuple[str, ...] = ()
unconditional_flags: tuple[str, ...] = ()
@@ -92,6 +104,7 @@ class ModelSupportSpec(BaseModel):
class ModelSupportHandler(Protocol):
key: str
is_moe: bool
+ cp_supported: bool
native_vllm_lora_status: NativeVllmLoraStatus
def identity_lora_model_config(self, base_config: Any) -> Any: ...
@@ -121,6 +134,15 @@ def patch_provider(
def configure_provider_for_runtime(self, provider: "GPTModelProvider") -> None: ...
+ def default_chat_template(self) -> str | None: ...
+
+ def configure_tokenizer(
+ self,
+ tokenizer: Any,
+ *,
+ internal_config: Any,
+ ) -> Any: ...
+
def vllm_engine_args(
self,
*,
@@ -131,6 +153,19 @@ def vllm_server_args(self) -> dict[str, object]: ...
def install_preprocess_patch(self, model_chunks: Sequence[Any]) -> None: ...
+ def build_shared_prefix_model_state(
+ self,
+ context: SharedPrefixModelStateContext,
+ ) -> dict[str, Any]: ...
+
+ def correctness_precision(self) -> Literal["bf16", "fp32"]: ...
+
+ def correctness_use_fp32_lora_reference(self) -> bool: ...
+
+ def correctness_phase_pass_fns(
+ self, oracle_harness: Any
+ ) -> dict[str, Any] | None: ...
+
def collect_layer_families(
self,
provider: "GPTModelProvider",
@@ -158,6 +193,11 @@ def to_vllm_lora_tensors(
adapter_config: dict[str, Any],
) -> tuple[dict[str, Any], dict[str, Any]]: ...
+ def to_vllm_lora_config(
+ self,
+ adapter_config: dict[str, Any],
+ ) -> dict[str, Any]: ...
+
def expert_packed_lora_groups(self) -> tuple[ExpertPackedLoraGroup, ...]: ...
def from_vllm_lora_tensors(
diff --git a/src/art/megatron/model_support/tokenizer.py b/src/art/megatron/model_support/tokenizer.py
new file mode 100644
index 000000000..9e6da103c
--- /dev/null
+++ b/src/art/megatron/model_support/tokenizer.py
@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any, cast
+
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+
+from .registry import UnsupportedModelArchitectureError, get_model_support_handler
+
+
+def _has_configured_chat_template(internal_config: Mapping[str, Any]) -> bool:
+ return (
+ internal_config.get("chat_template") is not None
+ or internal_config.get("chat_template_path") is not None
+ )
+
+
+def configure_tokenizer_for_model_support(
+ tokenizer: PreTrainedTokenizerBase,
+ *,
+ base_model: str,
+ internal_config: Mapping[str, Any],
+) -> PreTrainedTokenizerBase:
+ try:
+ handler = get_model_support_handler(
+ base_model,
+ allow_unvalidated_arch=bool(
+ internal_config.get("allow_unvalidated_arch", False)
+ ),
+ )
+ except UnsupportedModelArchitectureError:
+ return tokenizer
+
+ if not _has_configured_chat_template(internal_config) and not isinstance(
+ getattr(tokenizer, "chat_template", None), str
+ ):
+ default = handler.default_chat_template()
+ if default is not None:
+ tokenizer.chat_template = default
+
+ return cast(
+ PreTrainedTokenizerBase,
+ handler.configure_tokenizer(tokenizer, internal_config=internal_config),
+ )
diff --git a/src/art/megatron/provider.py b/src/art/megatron/provider.py
index 822b34d6e..a3cf5a531 100644
--- a/src/art/megatron/provider.py
+++ b/src/art/megatron/provider.py
@@ -364,10 +364,18 @@ def _resolve_default_deepep_num_sms(provider: GPTModelProvider) -> int:
return sm_count if sm_count >= 2 else 20
-def _apply_default_parallel_topology(provider: GPTModelProvider) -> None:
+def _handler_cp_supported(handler: Any) -> bool:
+ return bool(getattr(handler, "cp_supported", True))
+
+
+def _apply_default_parallel_topology(
+ provider: GPTModelProvider,
+ handler: Any,
+) -> None:
visible_gpu_count = max(torch.cuda.device_count(), 1)
- provider.tensor_model_parallel_size = 1
- provider.context_parallel_size = visible_gpu_count
+ cp_supported = _handler_cp_supported(handler)
+ provider.tensor_model_parallel_size = 1 if cp_supported else visible_gpu_count
+ provider.context_parallel_size = visible_gpu_count if cp_supported else 1
provider.pipeline_model_parallel_size = 1
provider.expert_model_parallel_size = (
visible_gpu_count
@@ -377,12 +385,32 @@ def _apply_default_parallel_topology(provider: GPTModelProvider) -> None:
provider.expert_tensor_parallel_size = 1
-def _apply_art_training_runtime_prepare_defaults(provider: GPTModelProvider) -> None:
+def _apply_art_training_runtime_prepare_defaults(
+ provider: GPTModelProvider,
+ handler: Any,
+) -> None:
provider.recompute_granularity = "full"
provider.recompute_method = "uniform"
provider.recompute_num_layers = 1
provider.moe_shared_expert_overlap = True
- _apply_default_parallel_topology(provider)
+ _apply_default_parallel_topology(provider, handler)
+
+
+def _validate_context_parallel_support(
+ handler: Any,
+ runtime_env: _ProviderRuntimeEnv,
+) -> None:
+ if _handler_cp_supported(handler):
+ return
+ if (
+ runtime_env.is_set("context_parallel_size")
+ and runtime_env.context_parallel_size is not None
+ and runtime_env.context_parallel_size > 1
+ ):
+ raise RuntimeError(
+ f"{handler.key} model support does not implement context parallelism; "
+ "set ART_MEGATRON_CONTEXT_PARALLEL_SIZE=1."
+ )
def _apply_art_training_runtime_finalize_defaults(
@@ -582,8 +610,9 @@ def prepare_provider_bundle(
provider.calculate_per_token_loss = True
provider.cross_entropy_loss_fusion = True
provider.cross_entropy_fusion_impl = "te"
- _apply_art_training_runtime_prepare_defaults(provider)
+ _apply_art_training_runtime_prepare_defaults(provider, bundle.handler)
bundle.handler.configure_provider_for_runtime(provider)
+ _validate_context_parallel_support(bundle.handler, runtime_env)
_apply_runtime_env_overrides(provider, runtime_env)
provider.art_flex_compile_crash_config = (
bundle.handler.flex_attention_compile_crash_config(provider)
diff --git a/src/art/megatron/routing_replay.py b/src/art/megatron/routing_replay.py
index acc8ff2aa..69ce534bc 100644
--- a/src/art/megatron/routing_replay.py
+++ b/src/art/megatron/routing_replay.py
@@ -34,6 +34,39 @@ def _active_routing_replay_controller() -> Any | None:
return _ACTIVE_ROUTING_REPLAY_CONTROLLER
+def _branch_rows_without_future_scored_tokens(
+ *,
+ group_ids: torch.Tensor,
+ parent_ids: torch.Tensor,
+ logprobs: torch.Tensor,
+) -> torch.Tensor:
+ """Rows whose MoE route cannot affect any later trainable token.
+
+ vLLM returns routes for forwards it actually ran. During generation it does
+ not run a forward for the final generated token, because no later token is
+ sampled from that row. ART shared-prefix packing may still leave a masked
+ suffix token in the same branch, so a physical group boundary is not enough
+ to identify rows that can safely use synthetic replay routes.
+ """
+ non_padding = group_ids != -1
+ branch = non_padding & (group_ids != parent_ids)
+ scored = non_padding & torch.isfinite(logprobs)
+ allowed = torch.zeros_like(branch)
+ for sample_index in range(int(group_ids.shape[0])):
+ future_scored_by_group: dict[int, bool] = {}
+ for position in range(int(group_ids.shape[1]) - 1, -1, -1):
+ group_id = int(group_ids[sample_index, position].item())
+ if group_id == -1:
+ continue
+ has_future_scored = future_scored_by_group.get(group_id, False)
+ allowed[sample_index, position] = (
+ bool(branch[sample_index, position].item()) and not has_future_scored
+ )
+ if bool(scored[sample_index, position].item()):
+ future_scored_by_group[group_id] = True
+ return allowed
+
+
def _to_tensor_cpu_contiguous(
tensor: torch.Tensor, *, dtype: torch.dtype
) -> torch.Tensor:
@@ -318,7 +351,15 @@ def from_dir(cls, bundle_dir: str | Path) -> "MoeRoutingReplayBundle":
steps: dict[int, StepRoutes] = {}
for step_index_str, step_info in manifest["steps"].items():
step_index = int(step_index_str)
- step_tensors = load_file(str(base_dir / step_info["file"]))
+ loaded_tensors = load_file(str(base_dir / step_info["file"]))
+ # Own CPU storage immediately. Safetensors CPU loads can keep
+ # file-backed storage alive, which makes shared-filesystem cleanup
+ # fail while long-lived Megatron ranks still hold replay bundles.
+ step_tensors = {
+ key: tensor.detach().clone().contiguous()
+ for key, tensor in loaded_tensors.items()
+ }
+ del loaded_tensors
if GLOBAL_TOKEN_UIDS_KEY not in step_tensors:
raise RuntimeError(
f"Missing tensor key '{GLOBAL_TOKEN_UIDS_KEY}' for step={step_index}"
@@ -450,7 +491,13 @@ def build_moe_routing_replay_bundle_from_packed_tensors(
terminal_completion = (
non_padding & (group_ids != parent_ids) & (group_ids != next_group_ids)
)
- unexpected_missing = non_padding & ~token_mask & ~terminal_completion
+ no_future_scored = _branch_rows_without_future_scored_tokens(
+ group_ids=group_ids,
+ parent_ids=parent_ids,
+ logprobs=packed_tensors["logprobs"],
+ )
+ allowed_missing = terminal_completion | no_future_scored
+ unexpected_missing = non_padding & ~token_mask & ~allowed_missing
if bool(unexpected_missing.any().item()):
raise RuntimeError(
"Packed tensors are missing MoE routes outside terminal completion "
@@ -478,8 +525,8 @@ def build_moe_routing_replay_bundle_from_packed_tensors(
if bool(missing_rows.any().item()):
# Megatron Core RouterReplay replays only top-k ids and does
# not consume an expert mask. Rows without vLLM routes are
- # allowed only for padding or terminal completion query
- # positions, whose next-token logits are not scored.
+ # allowed only for padding or branch-tail query positions
+ # that cannot affect any later scored token.
missing_positions = torch.nonzero(
missing_rows, as_tuple=False
).flatten()
@@ -688,6 +735,26 @@ def __init__(
self._target_copy_waited: bool = True
self._active_token_uid_key: str | None = None
+ def update_bundle(self, *, bundle: MoeRoutingReplayBundle, strict: bool) -> None:
+ self.bundle = bundle
+ self.strict = strict
+ self.clear_replay_state()
+ if self.strict:
+ missing = sorted(
+ router_key
+ for router_key in self._local_router_keys
+ if router_key not in self.bundle.router_keys
+ )
+ if missing:
+ raise RuntimeError(
+ "Router keys from model are missing in replay bundle: "
+ f"router_keys={missing}"
+ )
+
+ def clear_replay_state(self) -> None:
+ self._clear_native_router_replay_state()
+ self._reset_step_state()
+
def _target_device(self) -> torch.device:
if self._device is not None:
return self._device
@@ -754,6 +821,28 @@ def _prepare_native_target_for_bound_router(
_prepare_native_target_for_bound_router
)
+ def _hash_routing_with_replay_target(
+ router_module: Any,
+ *args: Any,
+ _original_routing: Any = original_routing,
+ _prepare_native_target: Any = prepare_native_target,
+ **kwargs: Any,
+ ) -> Any:
+ del router_module
+ _prepare_native_target()
+ return _original_routing(*args, **kwargs)
+
+ def _moe_routing_with_replay_target(
+ router_module: Any,
+ *args: Any,
+ _original_routing: Any = original_routing,
+ _prepare_native_target: Any = prepare_native_target,
+ **kwargs: Any,
+ ) -> Any:
+ del router_module
+ _prepare_native_target()
+ return _original_routing(*args, **kwargs)
+
def _routing_with_replay_target(
router_module: Any,
*args: Any,
@@ -768,7 +857,16 @@ def _routing_with_replay_target(
_prepare_native_target()
return _original_routing(*args, **kwargs)
- module.routing = types.MethodType(_routing_with_replay_target, module)
+ original_routing_name = getattr(
+ getattr(original_routing, "__func__", None), "__name__", ""
+ )
+ if original_routing_name == "_hash_routing":
+ routing_wrapper = _hash_routing_with_replay_target
+ elif original_routing_name == "_moe_routing":
+ routing_wrapper = _moe_routing_with_replay_target
+ else:
+ routing_wrapper = _routing_with_replay_target
+ module.routing = types.MethodType(routing_wrapper, module)
setattr(module, "_art_routing_replay_target_patched", True)
self._router_bindings[router_key] = {
"module": module,
diff --git a/src/art/megatron/runtime/bridge_runtime.py b/src/art/megatron/runtime/bridge_runtime.py
index 4b6b89d7f..8aabed47b 100644
--- a/src/art/megatron/runtime/bridge_runtime.py
+++ b/src/art/megatron/runtime/bridge_runtime.py
@@ -1,6 +1,6 @@
from __future__ import annotations
-from collections.abc import Iterable, Mapping
+from collections.abc import Callable, Iterable, Mapping
import contextlib
import fnmatch
from typing import Any, cast
@@ -23,6 +23,8 @@
from art.megatron.model_support.spec import HfWeightSource
+_Fp32PreservedTensor = tuple[torch.nn.Module, str, torch.Tensor, bool]
+
class ExpertTensorSlice:
__slots__ = ("global_start", "global_stop", "tensor")
@@ -240,6 +242,8 @@ def load_unique_hf_keys_once(
hf_state_dict: Mapping[str, torch.Tensor],
*,
bridge: MegatronModelBridge | None = None,
+ extra_keys: Callable[[Iterable[str], Mapping[str, torch.Tensor]], Iterable[str]]
+ | None = None,
) -> dict[str, torch.Tensor | ExpertTensorSlice]:
task_list = list(tasks)
prefetch_task_by_key: dict[str, Any] = {}
@@ -248,6 +252,9 @@ def load_unique_hf_keys_once(
continue
for key in _iter_hf_param_names(task.mapping.hf_param):
prefetch_task_by_key.setdefault(key, task)
+ if extra_keys is not None:
+ for key in extra_keys(tuple(sorted(prefetch_task_by_key)), hf_state_dict):
+ prefetch_task_by_key.setdefault(key, None)
keys = sorted(prefetch_task_by_key)
expert_slice_ranges: dict[str, tuple[int, int]] = {}
expert_slice_task_by_key: dict[str, Any] = {}
@@ -415,21 +422,62 @@ def _wrap_with_mp_wrapper(
) -> list[MegatronModule]:
if not (model_config.fp16 or model_config.bf16) or mixed_precision_wrapper is None:
return model
- keep_in_fp32: list[tuple[Any, torch.Tensor]] = []
- for model_module in model:
- for submodule in model_module.modules():
- if hasattr(submodule, "_maintain_float32_expert_bias"):
- expert_bias = getattr(submodule, "expert_bias", None)
- if expert_bias is not None:
- keep_in_fp32.append((submodule, expert_bias.data.clone()))
+ keep_in_fp32 = _collect_fp32_preserved_tensors(model)
wrapped = [
mixed_precision_wrapper(model_config, model_module) for model_module in model
]
- for submodule, fp32_data in keep_in_fp32:
- submodule.expert_bias.data = fp32_data
+ _restore_fp32_preserved_tensors(keep_in_fp32)
return wrapped
+def _collect_fp32_preserved_tensors(
+ model: list[MegatronModule],
+) -> list[_Fp32PreservedTensor]:
+ """Snapshot tensors explicitly marked to survive Megatron fp16/bf16 casts."""
+
+ keep_in_fp32: list[_Fp32PreservedTensor] = []
+ for model_module in model:
+ for submodule in model_module.modules():
+ fp32_parameter_names = set(getattr(submodule, "_keep_fp32_parameters", ()))
+ fp32_buffer_names = set(getattr(submodule, "_keep_fp32_buffers", ()))
+ explicit_names = fp32_parameter_names | fp32_buffer_names
+ seen: set[str] = set()
+ if hasattr(submodule, "_maintain_float32_expert_bias"):
+ expert_bias = getattr(submodule, "expert_bias", None)
+ if isinstance(expert_bias, torch.nn.Parameter):
+ keep_in_fp32.append(
+ (submodule, "expert_bias", expert_bias.data.clone(), True)
+ )
+ seen.add("expert_bias")
+ for name in explicit_names:
+ tensor = getattr(submodule, name, None)
+ if isinstance(tensor, torch.nn.Parameter):
+ keep_in_fp32.append((submodule, name, tensor.data.clone(), True))
+ seen.add(name)
+ elif isinstance(tensor, torch.Tensor):
+ keep_in_fp32.append((submodule, name, tensor.data.clone(), False))
+ seen.add(name)
+ for name, param in submodule.named_parameters(recurse=False):
+ if name not in seen and getattr(param, "_keep_fp32", False):
+ keep_in_fp32.append((submodule, name, param.data.clone(), True))
+ seen.add(name)
+ for name, buffer in submodule.named_buffers(recurse=False):
+ if name not in seen and getattr(buffer, "_keep_fp32", False):
+ keep_in_fp32.append((submodule, name, buffer.data.clone(), False))
+ seen.add(name)
+ return keep_in_fp32
+
+
+def _restore_fp32_preserved_tensors(
+ keep_in_fp32: list[_Fp32PreservedTensor],
+) -> None:
+ for submodule, name, fp32_data, is_parameter in keep_in_fp32:
+ if is_parameter:
+ getattr(submodule, name).data = fp32_data
+ else:
+ submodule._buffers[name] = fp32_data
+
+
def _art_get_model(
model_provider: ModelProviderMixin,
ddp_config: DistributedDataParallelConfig,
@@ -527,11 +575,13 @@ def _column_parallel_hf_to_megatron(
splits = list(torch.chunk(hf_weights, self.tp_size, dim=0))
else:
splits = None
- return self.scatter_to_tp_ranks(
+ return _scatter_to_tp_ranks(
+ self,
splits,
target_param.shape,
target_param.dtype,
target_param.device,
+ output_tensor=target_param.data,
)
@@ -542,20 +592,50 @@ def _scatter_to_tp_ranks(
dtype: torch.dtype,
device: torch.device,
src_rank: int = 0,
+ output_tensor: torch.Tensor | None = None,
) -> torch.Tensor:
if self.tp_size == 1:
- return cast(list[torch.Tensor], splits)[0].to(
- device=device, dtype=dtype, non_blocking=True
- )
- output = torch.empty(output_shape, dtype=dtype, device=device)
+ shard = cast(list[torch.Tensor], splits)[0]
+ if output_tensor is None:
+ return shard.to(device=device, dtype=dtype, non_blocking=True)
+ output_tensor.copy_(shard, non_blocking=True)
+ if output_tensor.device.type == "cuda":
+ torch.cuda.synchronize(output_tensor.device)
+ return output_tensor
+ output = (
+ torch.empty(output_shape, dtype=dtype, device=device)
+ if output_tensor is None
+ else output_tensor
+ )
dist = cast(Any, torch.distributed)
global_src = dist.get_global_rank(group=self.tp_group, group_rank=src_rank)
- scatter_list = None
- if self.tp_rank == src_rank and splits:
- scatter_list = [
- shard.to(device=device, dtype=dtype, non_blocking=True) for shard in splits
- ]
- dist.scatter(output, scatter_list, src=global_src, group=self.tp_group)
+ if self.tp_rank == src_rank:
+ if not splits:
+ raise RuntimeError("source TP rank must provide tensor splits")
+ if len(splits) != self.tp_size:
+ raise RuntimeError(
+ f"source TP rank got {len(splits)} tensor splits for TP size "
+ f"{self.tp_size}"
+ )
+ for peer_rank, shard in enumerate(splits):
+ if peer_rank == src_rank:
+ output.copy_(shard, non_blocking=True)
+ continue
+ send_buffer = torch.empty(output_shape, dtype=dtype, device=device)
+ send_buffer.copy_(shard, non_blocking=True)
+ if send_buffer.device.type == "cuda":
+ torch.cuda.current_stream(send_buffer.device).synchronize()
+ dist.send(
+ send_buffer,
+ dst=dist.get_global_rank(group=self.tp_group, group_rank=peer_rank),
+ group=self.tp_group,
+ )
+ if send_buffer.device.type == "cuda":
+ torch.cuda.synchronize(send_buffer.device)
+ else:
+ dist.recv(output, src=global_src, group=self.tp_group)
+ if output.device.type == "cuda":
+ torch.cuda.synchronize(output.device)
return output
@@ -598,7 +678,12 @@ def _optimized_load_weights_hf_to_megatron(
stack.enter_context(megatron_model[0].hide_loss_modules())
tasks = self.build_conversion_tasks(hf_pretrained, megatron_model)
hf_state_dict = hf_pretrained.state
- raw_cache = load_unique_hf_keys_once(tasks, hf_state_dict, bridge=self)
+ raw_cache = load_unique_hf_keys_once(
+ tasks,
+ hf_state_dict,
+ bridge=self,
+ extra_keys=getattr(self, "art_extra_hf_prefetch_keys", None),
+ )
cached_state = _CachedStateLookup(cache=raw_cache, source=hf_state_dict)
description = f"Loading from {hf_pretrained.model_name_or_path}"
pending_device_copy = False
@@ -642,7 +727,8 @@ def _optimized_load_weights_hf_to_megatron(
f" Bridge type: {type(task.mapping).__name__}\n"
f" HF mapping: {task.mapping.hf_param}"
)
- task.param_weight.data.copy_(converted_weights, non_blocking=True)
+ if converted_weights.data_ptr() != task.param_weight.data.data_ptr():
+ task.param_weight.data.copy_(converted_weights, non_blocking=True)
if task.param_weight.device.type == "cuda":
pending_device_copy = True
if pending_device_copy and torch.cuda.is_available():
diff --git a/src/art/megatron/service.py b/src/art/megatron/service.py
index 53d62d6a4..7535a9c76 100644
--- a/src/art/megatron/service.py
+++ b/src/art/megatron/service.py
@@ -8,6 +8,7 @@
import socket
import sys
from typing import Any, AsyncIterator, Literal, TypedDict, cast
+from urllib.parse import urlparse
import warnings
from peft.tuners.lora.config import LoraConfig
@@ -31,6 +32,10 @@
from ..vllm_runtime import (
ManagedVllmRuntime,
VllmRuntimeLaunchConfig,
+ get_external_vllm_runtime_config,
+ map_checkpoint_path_for_vllm,
+ normalize_vllm_server_url,
+ wait_for_vllm_http_runtime,
)
from .lora import (
LORA_ALPHA,
@@ -242,14 +247,22 @@ def rollout_weights_mode(self) -> Literal["lora", "merged"]:
@property
def _vllm_base_url(self) -> str:
+ if external_runtime := get_external_vllm_runtime_config(self.config):
+ return normalize_vllm_server_url(external_runtime.server_url)
return self._vllm_runtime.base_url
@property
def _vllm_host(self) -> str:
+ if external_runtime := get_external_vllm_runtime_config(self.config):
+ parsed = urlparse(normalize_vllm_server_url(external_runtime.server_url))
+ return parsed.hostname or self._vllm_runtime.host
return self._vllm_runtime.host
@property
def _vllm_port(self) -> int:
+ if external_runtime := get_external_vllm_runtime_config(self.config):
+ parsed = urlparse(normalize_vllm_server_url(external_runtime.server_url))
+ return parsed.port or (443 if parsed.scheme == "https" else 80)
return self._vllm_runtime.port
@_vllm_port.setter
@@ -258,6 +271,8 @@ def _vllm_port(self, port: int) -> None:
@property
def _vllm_api_key(self) -> str | None:
+ if external_runtime := get_external_vllm_runtime_config(self.config):
+ return external_runtime.api_key
return self._vllm_runtime.api_key
@property
@@ -446,6 +461,9 @@ def _runtime_request_kwargs(self) -> _RuntimeRequestKwargs:
headers = self._runtime_headers()
return {"headers": headers} if headers else {}
+ def _vllm_checkpoint_path(self, checkpoint_path: str) -> str:
+ return map_checkpoint_path_for_vllm(self.config, checkpoint_path)
+
def _sleep_mode_enabled(self) -> bool:
return bool(self.config.get("engine_args", {}).get("enable_sleep_mode", True))
@@ -624,7 +642,7 @@ async def _reload_adapter(self, checkpoint_path: str, step: int) -> None:
f"{self._vllm_base_url}/v1/load_lora_adapter",
json={
"lora_name": f"{self.model_name}@{step}",
- "lora_path": checkpoint_path,
+ "lora_path": self._vllm_checkpoint_path(checkpoint_path),
"load_inplace": True,
},
**self._runtime_request_kwargs(),
@@ -909,6 +927,7 @@ async def start_openai_server(
) -> tuple[str, int]:
self._raise_if_child_failed()
lora_path = self._resolve_active_lora_path()
+ external_runtime = get_external_vllm_runtime_config(self.config)
if not self.is_dedicated and not self._sleep_mode_enabled():
raise ValueError(
@@ -916,6 +935,21 @@ async def start_openai_server(
"for the external vLLM runtime"
)
+ if external_runtime is not None:
+ if self.rollout_weights_mode != "lora":
+ raise RuntimeError(
+ "External vLLM runtime requires LoRA rollout weights"
+ )
+ await wait_for_vllm_http_runtime(
+ base_url=self._vllm_base_url,
+ timeout=external_runtime.health_timeout_s,
+ headers=self._runtime_headers(),
+ )
+ await self._reload_adapter(lora_path, self._latest_step)
+ self._loaded_adapter_steps.add(self._latest_step)
+ self._status(f"External vLLM runtime is ready at {self._vllm_base_url}")
+ return self._vllm_host, self._vllm_port
+
port = (config or {}).get("server_args", {}).get("port", 8000)
location = await self._start_vllm_subprocess(lora_path, port, config)
if self.rollout_weights_mode == "lora":
diff --git a/src/art/megatron/shared_prefix_state.py b/src/art/megatron/shared_prefix_state.py
index e6ecccfc8..97dbb6d93 100644
--- a/src/art/megatron/shared_prefix_state.py
+++ b/src/art/megatron/shared_prefix_state.py
@@ -31,6 +31,7 @@
move_gdn_rank_execution_plan_to_device,
parse_gdn_shared_prefix_segments,
)
+from art.megatron.model_support.spec import SharedPrefixModelStateContext
class SharedPrefixAttentionState(FlexSharedPrefixAttentionState):
@@ -38,6 +39,7 @@ class SharedPrefixAttentionState(FlexSharedPrefixAttentionState):
group_ids: Tensor
parent_ids: Tensor
+ model_state: dict[str, Any] = Field(default_factory=dict)
gdn_execution_spec: GdnPackedExecutionSpec | None = None
gdn_execution_plan: GdnRankExecutionPlan | None = None
gdn_hidden_layout: str = "attention"
@@ -59,6 +61,7 @@ def create_shared_prefix_state(
input_pos: Tensor | None = None,
sliding_windows: tuple[int, ...] = (),
build_gdn_execution_spec: bool = False,
+ model_support_handler: Any | None = None,
attention_token_layout_index: TokenLayoutIndex | None = None,
attention_head_dim: int | None = None,
attention_value_head_dim: int | None = None,
@@ -106,6 +109,16 @@ def create_shared_prefix_state(
sliding_block_masks=sliding_block_masks,
group_ids=group_ids_cpu,
parent_ids=parent_ids_cpu,
+ model_state=_build_model_state_once(
+ model_support_handler,
+ input_pos=input_pos_cpu,
+ group_ids=group_ids_cpu,
+ parent_ids=parent_ids_cpu,
+ device=device,
+ attention_token_layout_index=attention_token_layout_index,
+ attention_head_dim=attention_head_dim,
+ attention_value_head_dim=attention_value_head_dim,
+ ),
gdn_execution_spec=gdn_execution_spec,
gdn_execution_plan=_build_gdn_execution_plan_once(
gdn_execution_spec,
@@ -118,6 +131,34 @@ def create_shared_prefix_state(
)
+def _build_model_state_once(
+ model_support_handler: Any | None,
+ *,
+ input_pos: Tensor | None,
+ group_ids: Tensor,
+ parent_ids: Tensor,
+ device: torch.device,
+ attention_token_layout_index: TokenLayoutIndex | None,
+ attention_head_dim: int | None,
+ attention_value_head_dim: int | None,
+) -> dict[str, Any]:
+ if model_support_handler is None:
+ return {}
+ return dict(
+ model_support_handler.build_shared_prefix_model_state(
+ SharedPrefixModelStateContext(
+ input_pos=input_pos,
+ group_ids=group_ids,
+ parent_ids=parent_ids,
+ device=device,
+ attention_token_layout_index=attention_token_layout_index,
+ attention_head_dim=attention_head_dim,
+ attention_value_head_dim=attention_value_head_dim,
+ )
+ )
+ )
+
+
def _metadata_cpu(tensor: Tensor) -> Tensor:
tensor = tensor.detach()
if tensor.device.type != "cpu" or tensor.dtype != torch.int64:
diff --git a/src/art/megatron/training/microbatches.py b/src/art/megatron/training/microbatches.py
index 8bf65017f..de6650cba 100644
--- a/src/art/megatron/training/microbatches.py
+++ b/src/art/megatron/training/microbatches.py
@@ -315,6 +315,7 @@ def _causal_attention_state(
*,
sliding_windows: tuple[int, ...] = (),
build_gdn_execution_spec: bool,
+ model_support_handler: Any,
attention_head_dim: int | None = None,
attention_value_head_dim: int | None = None,
) -> Any:
@@ -327,6 +328,7 @@ def _causal_attention_state(
input_pos=torch.arange(seq_len, dtype=torch.int64).unsqueeze(0),
sliding_windows=sliding_windows,
build_gdn_execution_spec=build_gdn_execution_spec,
+ model_support_handler=model_support_handler,
attention_head_dim=attention_head_dim,
attention_value_head_dim=attention_value_head_dim,
)
@@ -360,6 +362,7 @@ def _prepare_dense_rl_micro(
build_gdn_execution_spec=bool(
getattr(model_support_handler, "build_gdn_execution_spec", False)
),
+ model_support_handler=model_support_handler,
attention_head_dim=getattr(provider, "kv_channels", None),
attention_value_head_dim=getattr(provider, "kv_channels", None),
)
@@ -561,6 +564,7 @@ def _prepare_dense_sft_micro(
build_gdn_execution_spec=bool(
getattr(model_support_handler, "build_gdn_execution_spec", False)
),
+ model_support_handler=model_support_handler,
attention_head_dim=getattr(provider, "kv_channels", None),
attention_value_head_dim=getattr(provider, "kv_channels", None),
),
diff --git a/src/art/megatron/training/streaming_weight_offload.py b/src/art/megatron/training/streaming_weight_offload.py
index c7da2d702..a36ca5414 100644
--- a/src/art/megatron/training/streaming_weight_offload.py
+++ b/src/art/megatron/training/streaming_weight_offload.py
@@ -27,6 +27,8 @@
STREAMING_COMPILED_LAYERS_MESSAGE = (
"Streaming weight offload managing compiled transformer layers"
)
+# Quantized/custom kernels may vector-load weights and scales from streamed views.
+STREAMED_PARAM_ALIGNMENT_BYTES = 256
def _rank0_info(rank: int, message: str, *args: object) -> None:
@@ -199,7 +201,7 @@ def _pre_forward(self, layer_state: _LayerState) -> None:
)
def _post_forward(self, layer_state: _LayerState) -> None:
- if is_checkpointing() and not torch.is_grad_enabled():
+ if not torch.is_grad_enabled():
self._start_offload(layer_state)
self._prefetch_window(layer_state.index + self.config.resident_layers, 1, 1)
@@ -490,13 +492,16 @@ def _build_tensor_groups(
grouped.setdefault(param.dtype, []).append((name, param))
groups: list[_TensorGroup] = []
for dtype, dtype_params in grouped.items():
- total_numel = sum(param.numel() for _name, param in dtype_params)
- cpu_flat = torch.empty(total_numel, dtype=dtype, device="cpu")
+ element_size = dtype_params[0][1].element_size()
+ alignment_numel = max(
+ 1,
+ (STREAMED_PARAM_ALIGNMENT_BYTES + element_size - 1) // element_size,
+ )
specs: list[_ParamSpec] = []
offset = 0
for name, param in dtype_params:
+ offset = _align_numel(offset, alignment_numel)
numel = param.numel()
- cpu_flat[offset : offset + numel].copy_(param.detach().view(-1).cpu())
specs.append(
_ParamSpec(
name=name,
@@ -507,10 +512,19 @@ def _build_tensor_groups(
)
)
offset += numel
+ cpu_flat = torch.empty(offset, dtype=dtype, device="cpu")
+ for spec in specs:
+ cpu_flat[spec.offset : spec.offset + spec.numel].copy_(
+ spec.param.detach().view(-1).cpu()
+ )
groups.append(_TensorGroup(dtype=dtype, cpu_flat=cpu_flat, specs=specs))
return groups
+def _align_numel(offset: int, alignment_numel: int) -> int:
+ return ((offset + alignment_numel - 1) // alignment_numel) * alignment_numel
+
+
def _validate_streamed_param(spec: _ParamSpec) -> None:
if spec.param.requires_grad:
raise RuntimeError(
diff --git a/src/art/transformers/patches.py b/src/art/transformers/patches.py
index 6eb28b007..fe3509cb8 100644
--- a/src/art/transformers/patches.py
+++ b/src/art/transformers/patches.py
@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
import torch
from transformers import masking_utils
@@ -19,15 +19,7 @@ def _patched_preprocess_mask_arguments(
position_ids: Optional[torch.Tensor],
layer_idx: Optional[int],
encoder_hidden_states: Optional[torch.Tensor] = None,
-) -> tuple[
- bool,
- Optional[Union[torch.Tensor, "BlockMask"]],
- Optional[torch.Tensor],
- int,
- int,
- int,
- int,
-]:
+) -> tuple[Any, ...]:
if position_ids is not None and len(position_ids.shape) == 3:
position_ids = position_ids[0]
return _preprocess_mask_arguments(
diff --git a/src/art/vllm_runtime.py b/src/art/vllm_runtime.py
index 31e4b98b8..7445758d1 100644
--- a/src/art/vllm_runtime.py
+++ b/src/art/vllm_runtime.py
@@ -10,7 +10,8 @@
import shutil
import subprocess
import tempfile
-from typing import Any, Callable, Literal, TypedDict
+from typing import Any, Callable, Literal, Mapping, TypedDict
+from urllib.parse import urlparse
import httpx
from pydantic import BaseModel, ConfigDict, Field
@@ -25,6 +26,17 @@
RUNTIME_PACKAGE = "art-vllm-runtime"
RUNTIME_PROTOCOL_VERSION = 1
RUNTIME_INSTALL_MARKER = "openpipe-art-vllm-runtime"
+_TILELANG_ENV_KEYS = (
+ "PYTHONPATH",
+ "TVM_IMPORT_PYTHON_PATH",
+ "TVM_LIBRARY_PATH",
+ "TL_CUTLASS_PATH",
+ "TL_TEMPLATE_PATH",
+ "TL_COMPOSABLE_KERNEL_PATH",
+)
+_TILELANG_PATH_MARKERS = ("/site-packages/tilelang/", "\\site-packages\\tilelang\\")
+_FLASHINFER_WORKSPACE_ENV = "FLASHINFER_WORKSPACE_BASE"
+_ART_FLASHINFER_WORKSPACE_ENV = "ART_VLLM_RUNTIME_FLASHINFER_WORKSPACE_BASE"
VLLM_RUNTIME_CLOSE_TIMEOUT = 15.0
@@ -35,13 +47,24 @@ class VllmRuntimeLaunchConfig(BaseModel):
port: int
host: str = "127.0.0.1"
cuda_visible_devices: str
- lora_path: str
+ lora_path: str | None = None
served_model_name: str
rollout_weights_mode: Literal["lora", "merged"]
engine_args: dict[str, object] = Field(default_factory=dict)
server_args: dict[str, object] = Field(default_factory=dict)
+class ExternalVllmRuntimeConfig(BaseModel):
+ model_config = ConfigDict(extra="forbid")
+
+ mode: Literal["external"]
+ server_url: str
+ api_key: str | None = None
+ local_checkpoint_root: str | None = None
+ server_checkpoint_root: str | None = None
+ health_timeout_s: float = Field(default=120.0, gt=0)
+
+
class VllmRuntimeManifest(BaseModel):
model_config = ConfigDict(extra="forbid")
@@ -75,6 +98,123 @@ class VllmRuntimeRequestKwargs(TypedDict, total=False):
headers: dict[str, str]
+def is_external_vllm_runtime(config: Mapping[str, Any]) -> bool:
+ runtime_config = config.get("vllm_runtime")
+ return (
+ isinstance(runtime_config, Mapping) and runtime_config.get("mode") == "external"
+ )
+
+
+def get_external_vllm_runtime_config(
+ config: Mapping[str, Any],
+) -> ExternalVllmRuntimeConfig | None:
+ runtime_config = config.get("vllm_runtime")
+ if not isinstance(runtime_config, Mapping):
+ return None
+ if runtime_config.get("mode", "managed") != "external":
+ return None
+ return ExternalVllmRuntimeConfig.model_validate(runtime_config)
+
+
+def normalize_vllm_server_url(server_url: str) -> str:
+ normalized = server_url.rstrip("/")
+ if normalized.endswith("/v1"):
+ normalized = normalized[:-3].rstrip("/")
+ parsed = urlparse(normalized)
+ if parsed.scheme not in {"http", "https"} or not parsed.netloc:
+ raise ValueError(
+ f"External vLLM server_url must be an HTTP URL: {server_url!r}"
+ )
+ return normalized
+
+
+def openai_base_url_from_vllm_server_url(server_url: str) -> str:
+ return f"{normalize_vllm_server_url(server_url)}/v1"
+
+
+def map_checkpoint_path_for_vllm(
+ config: Mapping[str, Any],
+ checkpoint_path: str,
+) -> str:
+ runtime_config = get_external_vllm_runtime_config(config)
+ if runtime_config is None:
+ return checkpoint_path
+ local_root = runtime_config.local_checkpoint_root
+ server_root = runtime_config.server_checkpoint_root
+ if local_root is None and server_root is None:
+ return checkpoint_path
+ if not local_root or not server_root:
+ raise ValueError(
+ "Set both vllm_runtime.local_checkpoint_root and "
+ "vllm_runtime.server_checkpoint_root, or neither."
+ )
+ checkpoint_abs = os.path.abspath(checkpoint_path)
+ local_root_abs = os.path.abspath(local_root)
+ rel_path = os.path.relpath(checkpoint_abs, local_root_abs)
+ if rel_path == os.pardir or rel_path.startswith(os.pardir + os.sep):
+ raise ValueError(
+ f"Checkpoint path {checkpoint_path!r} is not under "
+ f"vllm_runtime.local_checkpoint_root {local_root!r}"
+ )
+ return os.path.join(server_root, rel_path)
+
+
+async def wait_for_vllm_http_runtime(
+ *,
+ base_url: str,
+ timeout: float,
+ headers: dict[str, str] | None = None,
+) -> None:
+ deadline = asyncio.get_running_loop().time() + timeout
+ url = f"{base_url.rstrip('/')}/health"
+ async with httpx.AsyncClient() as client:
+ while True:
+ try:
+ response = await client.get(url, headers=headers, timeout=5.0)
+ if response.status_code == 200:
+ return
+ except httpx.HTTPError:
+ pass
+ if asyncio.get_running_loop().time() >= deadline:
+ raise TimeoutError(
+ f"vLLM runtime did not become ready within {math.ceil(timeout)}s"
+ )
+ await asyncio.sleep(0.5)
+
+
+def _drop_tilelang_env_paths(value: str | None) -> str | None:
+ if value is None:
+ return None
+ kept = [
+ part
+ for part in value.split(os.pathsep)
+ if not any(marker in part for marker in _TILELANG_PATH_MARKERS)
+ ]
+ return os.pathsep.join(kept) if kept else None
+
+
+def _vllm_runtime_subprocess_env() -> dict[str, str]:
+ """Build a child env isolated from runtime-specific JIT path leaks.
+
+ TileLang mutates process env during import. If a vLLM runtime child inherits
+ those paths, spawn workers can load two TVM FFI libraries and abort on
+ duplicate global registration during model load.
+
+ FlashInfer writes absolute source/include paths into generated build.ninja
+ files. Sharing its default ~/.cache/flashinfer across source worktrees can
+ make one runtime compile kernels from another runtime's venv.
+ """
+ env = os.environ.copy()
+ for key in _TILELANG_ENV_KEYS:
+ value = _drop_tilelang_env_paths(env.get(key))
+ if value is None:
+ env.pop(key, None)
+ else:
+ env[key] = value
+ env[_FLASHINFER_WORKSPACE_ENV] = str(_vllm_runtime_flashinfer_workspace_base())
+ return env
+
+
class ManagedVllmRuntime:
def __init__(self, *, host: str = "127.0.0.1") -> None:
self.host = host
@@ -123,7 +263,7 @@ async def start(
self.process = subprocess.Popen(
managed_process_cmd(cmd),
cwd=str(get_vllm_runtime_working_dir()),
- env=os.environ.copy(),
+ env=_vllm_runtime_subprocess_env(),
stdout=self.log_file,
stderr=subprocess.STDOUT,
bufsize=1,
@@ -186,7 +326,13 @@ async def start(
def close(self) -> None:
if self.process is not None:
terminate_popen_process_group(
- self.process, timeout=VLLM_RUNTIME_CLOSE_TIMEOUT
+ self.process,
+ timeout=float(
+ os.environ.get(
+ "ART_VLLM_RUNTIME_CLOSE_TIMEOUT",
+ VLLM_RUNTIME_CLOSE_TIMEOUT,
+ )
+ ),
)
self.process = None
if self.log_file is not None:
@@ -227,6 +373,16 @@ def get_vllm_runtime_cache_root() -> Path:
return Path.home() / ".cache" / "art" / "vllm_runtime"
+def _vllm_runtime_flashinfer_workspace_base() -> Path:
+ override = os.environ.get(_ART_FLASHINFER_WORKSPACE_ENV)
+ if override:
+ return Path(override).expanduser()
+ runtime_root = get_vllm_runtime_project_root()
+ if runtime_root.exists():
+ return runtime_root.resolve().parent / "scratch" / "vllm_runtime_flashinfer"
+ return get_vllm_runtime_cache_root().expanduser() / "flashinfer_workspace"
+
+
def _bundled_runtime_dir() -> Path:
return Path(__file__).resolve().parent / "_vllm_runtime"
@@ -578,18 +734,24 @@ def _runtime_command_prefix() -> list[str]:
def build_vllm_runtime_server_cmd(config: VllmRuntimeLaunchConfig) -> list[str]:
- return [
+ command = [
*_runtime_command_prefix(),
f"--model={config.base_model}",
f"--port={config.port}",
f"--host={config.host}",
f"--cuda-visible-devices={config.cuda_visible_devices}",
- f"--lora-path={config.lora_path}",
- f"--served-model-name={config.served_model_name}",
- f"--rollout-weights-mode={config.rollout_weights_mode}",
- f"--engine-args-json={json.dumps(config.engine_args)}",
- f"--server-args-json={json.dumps(config.server_args)}",
]
+ if config.lora_path is not None:
+ command.append(f"--lora-path={config.lora_path}")
+ command.extend(
+ [
+ f"--served-model-name={config.served_model_name}",
+ f"--rollout-weights-mode={config.rollout_weights_mode}",
+ f"--engine-args-json={json.dumps(config.engine_args)}",
+ f"--server-args-json={json.dumps(config.server_args)}",
+ ]
+ )
+ return command
async def wait_for_vllm_runtime(
diff --git a/tests/integration/megatron/lora/merged_vllm_serving.py b/tests/integration/megatron/lora/merged_vllm_serving.py
index 7a9a1fd8a..6909ca461 100644
--- a/tests/integration/megatron/lora/merged_vllm_serving.py
+++ b/tests/integration/megatron/lora/merged_vllm_serving.py
@@ -1,9 +1,11 @@
from __future__ import annotations
import asyncio
+from contextlib import contextmanager
import os
from pathlib import Path
import socket
+from typing import Any, Iterator, cast
from pydantic import BaseModel, Field
import torch
@@ -15,9 +17,15 @@
from ..model_support.oracle_harness import (
ORACLE_TOPOLOGY,
OracleCaseConfig,
+ Topology,
ensure_case_artifacts,
)
from ..model_support.oracle_worker import provider_topology_env
+from ..model_support.workflow_resources import (
+ handler_workflow_resources_for_base_model,
+ resolve_stage_resources_for_visible_gpus,
+ validate_dedicated_test_resources,
+)
_TRAINER_GPU_IDS_ENV = "ART_MODEL_SUPPORT_TRAINER_GPU_IDS"
_INFERENCE_GPU_IDS_ENV = "ART_MODEL_SUPPORT_INFERENCE_GPU_IDS"
@@ -66,14 +74,28 @@ def _resolve_dedicated_gpu_ids() -> tuple[list[int], list[int]]:
return [0], [1]
-def _init_runtime_config(case_config: OracleCaseConfig) -> None:
+@contextmanager
+def _temporary_env(updates: dict[str, str]) -> Iterator[None]:
+ previous = {name: os.environ.get(name) for name in updates}
+ os.environ.update(updates)
+ try:
+ yield
+ finally:
+ for name, value in previous.items():
+ if value is None:
+ os.environ.pop(name, None)
+ else:
+ os.environ[name] = value
+
+
+def _init_runtime_config(case_config: OracleCaseConfig, topology: Any) -> None:
art.init_megatron_runtime_config(
topology=art.MegatronTopologyConfig(
- tp=ORACLE_TOPOLOGY.tp,
- cp=ORACLE_TOPOLOGY.cp,
- ep=ORACLE_TOPOLOGY.ep,
- pp=ORACLE_TOPOLOGY.pp,
- etp=ORACLE_TOPOLOGY.etp,
+ tp=topology.tp,
+ cp=topology.cp,
+ ep=topology.ep,
+ pp=topology.pp,
+ etp=topology.etp,
),
packed_sequence_length=case_config.packed_tensors.sequence_length,
)
@@ -82,7 +104,50 @@ def _init_runtime_config(case_config: OracleCaseConfig) -> None:
async def _run_merged_vllm_serving(
case_config: OracleCaseConfig,
) -> MergedVllmServingReport:
- trainer_gpu_ids, inference_gpu_ids = _resolve_dedicated_gpu_ids()
+ workflow_resources = handler_workflow_resources_for_base_model(
+ case_config.base_model,
+ allow_unvalidated_arch=case_config.allow_unvalidated_arch,
+ )
+ stage_resources = (
+ workflow_resources.merged_vllm_serving
+ if workflow_resources is not None
+ else None
+ )
+ topology: Topology = ORACLE_TOPOLOGY
+ megatron_env: dict[str, str] = {}
+ engine_args: dev.EngineArgs = dev.EngineArgs()
+ if stage_resources is not None:
+ stage_resources = resolve_stage_resources_for_visible_gpus(
+ "merged_vllm_serving",
+ stage_resources,
+ visible_gpu_count=int(torch.cuda.device_count()),
+ )
+ if stage_resources.megatron is None or stage_resources.vllm is None:
+ raise RuntimeError(
+ "merged_vllm_serving resources require Megatron and vLLM"
+ )
+ trainer_gpu_ids = list(stage_resources.megatron.gpu_ids)
+ inference_gpu_ids = list(stage_resources.vllm.gpu_ids)
+ validate_dedicated_test_resources(
+ stage_name="merged_vllm_serving",
+ trainer_gpu_ids=trainer_gpu_ids,
+ inference_gpu_ids=inference_gpu_ids,
+ allow_overlap=stage_resources.allow_gpu_overlap,
+ )
+ resource_topology = stage_resources.megatron.topology
+ topology = Topology(
+ tp=resource_topology.tp,
+ ep=resource_topology.ep,
+ etp=resource_topology.etp,
+ dp=resource_topology.dp,
+ cp=resource_topology.cp,
+ pp=resource_topology.pp,
+ sp=resource_topology.sp,
+ )
+ megatron_env = dict(stage_resources.megatron_env)
+ engine_args = cast(dev.EngineArgs, stage_resources.vllm.engine_args())
+ else:
+ trainer_gpu_ids, inference_gpu_ids = _resolve_dedicated_gpu_ids()
service_name = "model_support_merged_validation"
case_artifacts = ensure_case_artifacts(case_config)
output_dir = str(Path(case_artifacts.case_dir) / "merged_vllm_serving")
@@ -92,10 +157,12 @@ async def _run_merged_vllm_serving(
inference_gpu_ids=inference_gpu_ids,
rollout_weights_mode="merged",
allow_unvalidated_arch=case_config.allow_unvalidated_arch,
+ engine_args=engine_args,
)
- dev.validate_dedicated_config(internal_config)
- with provider_topology_env(ORACLE_TOPOLOGY):
- _init_runtime_config(case_config)
+ if stage_resources is None:
+ dev.validate_dedicated_config(internal_config)
+ with _temporary_env(megatron_env), provider_topology_env(topology):
+ _init_runtime_config(case_config, topology)
service = MegatronService(
model_name=service_name,
base_model=case_config.base_model,
diff --git a/tests/integration/megatron/lora/native_vllm_lora.py b/tests/integration/megatron/lora/native_vllm_lora.py
index a5689275b..d0040e860 100644
--- a/tests/integration/megatron/lora/native_vllm_lora.py
+++ b/tests/integration/megatron/lora/native_vllm_lora.py
@@ -6,6 +6,7 @@
import shutil
import socket
import tempfile
+from typing import cast
from pydantic import BaseModel, Field
import torch
@@ -21,6 +22,11 @@
ensure_case_artifacts,
)
from ..model_support.oracle_worker import provider_topology_env
+from ..model_support.workflow_resources import (
+ handler_workflow_resources_for_base_model,
+ resolve_stage_resources_for_visible_gpus,
+ validate_dedicated_test_resources,
+)
_TRAINER_GPU_IDS_ENV = "ART_MODEL_SUPPORT_TRAINER_GPU_IDS"
_INFERENCE_GPU_IDS_ENV = "ART_MODEL_SUPPORT_INFERENCE_GPU_IDS"
@@ -122,7 +128,33 @@ def _init_runtime_config(case_config: OracleCaseConfig) -> None:
async def _run_native_vllm_lora(
case_config: OracleCaseConfig,
) -> NativeVllmLoraServingReport:
- trainer_gpu_ids, inference_gpu_ids = _resolve_dedicated_gpu_ids()
+ workflow_resources = handler_workflow_resources_for_base_model(
+ case_config.base_model,
+ allow_unvalidated_arch=case_config.allow_unvalidated_arch,
+ )
+ stage_resources = (
+ workflow_resources.native_vllm_lora if workflow_resources is not None else None
+ )
+ if stage_resources is not None:
+ stage_resources = resolve_stage_resources_for_visible_gpus(
+ "native_vllm_lora",
+ stage_resources,
+ visible_gpu_count=int(torch.cuda.device_count()),
+ )
+ if stage_resources.vllm is None:
+ raise RuntimeError("native_vllm_lora resources require vLLM")
+ trainer_gpu_ids = [0]
+ inference_gpu_ids = list(stage_resources.vllm.gpu_ids)
+ validate_dedicated_test_resources(
+ stage_name="native_vllm_lora",
+ trainer_gpu_ids=trainer_gpu_ids,
+ inference_gpu_ids=inference_gpu_ids,
+ allow_overlap=True,
+ )
+ engine_args = cast(dev.EngineArgs, stage_resources.vllm.engine_args())
+ else:
+ trainer_gpu_ids, inference_gpu_ids = _resolve_dedicated_gpu_ids()
+ engine_args = dev.EngineArgs()
service_name = "model_support_native_lora_validation"
case_artifacts = ensure_case_artifacts(case_config)
output_root = Path(case_artifacts.case_dir) / "native_vllm_lora"
@@ -133,8 +165,10 @@ async def _run_native_vllm_lora(
inference_gpu_ids=inference_gpu_ids,
rollout_weights_mode="lora",
allow_unvalidated_arch=case_config.allow_unvalidated_arch,
+ engine_args=engine_args,
)
- dev.validate_dedicated_config(internal_config)
+ if stage_resources is None:
+ dev.validate_dedicated_config(internal_config)
with provider_topology_env(ORACLE_TOPOLOGY):
_init_runtime_config(case_config)
service = MegatronService(
diff --git a/tests/integration/megatron/lora/test_merged_weight_export.py b/tests/integration/megatron/lora/test_merged_weight_export.py
index 54a155f1f..fc95cfa42 100644
--- a/tests/integration/megatron/lora/test_merged_weight_export.py
+++ b/tests/integration/megatron/lora/test_merged_weight_export.py
@@ -1,3 +1,4 @@
+from types import SimpleNamespace
from typing import Any, cast
import httpx
diff --git a/tests/integration/megatron/model_support/chat_template_rollout.py b/tests/integration/megatron/model_support/chat_template_rollout.py
index 4067eedbf..f126fb3ee 100644
--- a/tests/integration/megatron/model_support/chat_template_rollout.py
+++ b/tests/integration/megatron/model_support/chat_template_rollout.py
@@ -6,6 +6,7 @@
import art
from art.local import LocalBackend
+from art.local.backend import _tokenizer_cache_key
from art.preprocessing.pack import PackedTensors
from art.preprocessing.tokenize import (
TokenizedResult,
@@ -101,12 +102,18 @@ def run_chat_template_rollout(base_model: str) -> ChatTemplateRolloutReport:
base_model=base_model,
_internal_config={"init_args": {"max_seq_length": 2048}},
)
- tokenizer_key = (base_model, None)
+ internal_config = model._internal_config
+ assert internal_config is not None
+ tokenizer_key = _tokenizer_cache_key(base_model, internal_config)
tokenizer = backend._tokenizers.get(tokenizer_key)
if tokenizer is None:
from transformers import AutoTokenizer
- tokenizer = AutoTokenizer.from_pretrained(base_model)
+ tokenizer = backend._configure_training_tokenizer(
+ AutoTokenizer.from_pretrained(base_model),
+ model=model,
+ internal_config=internal_config,
+ )
backend._tokenizers[tokenizer_key] = tokenizer
inputs = build_chat_template_conformance_inputs(tokenizer)
diff --git a/tests/integration/megatron/model_support/forward_trace.py b/tests/integration/megatron/model_support/forward_trace.py
index cf3aca284..4a0410c6f 100644
--- a/tests/integration/megatron/model_support/forward_trace.py
+++ b/tests/integration/megatron/model_support/forward_trace.py
@@ -7,8 +7,11 @@
import torch
from .trace_uids import (
+ TRACE_ROW_TOKEN_UIDS_ATTR,
+ TRACE_UID_SPAN_ATTR,
expand_token_uids_for_heads,
extract_tensor_attr,
+ normalize_row_token_uids,
row_token_uids_from_trace_sources,
)
@@ -54,6 +57,15 @@
)
+def _module_layer_index(module_name: str) -> int | None:
+ marker = "decoder.layers."
+ marker_index = module_name.find(marker)
+ if marker_index < 0:
+ return None
+ raw = module_name[marker_index + len(marker) :].split(".", 1)[0]
+ return int(raw) if raw.isdigit() else None
+
+
def _trace_hook(fn: Callable[..., Any]) -> Callable[..., Any]:
return torch.compiler.disable(fn)
@@ -71,6 +83,15 @@ def _safe_int(value: Any, default: int = 0) -> int:
return default
+def _call_cp_world_size(call: dict[str, Any]) -> int:
+ rank_meta = call.get("rank_meta")
+ if isinstance(rank_meta, list) and rank_meta:
+ rank_meta = rank_meta[0]
+ if not isinstance(rank_meta, dict):
+ return 1
+ return _safe_int(rank_meta.get("cp_world_size"), 1)
+
+
def _safe_ps_stat(name: str, default: int) -> int:
"""Reads one Megatron parallel-state integer when available."""
try:
@@ -269,11 +290,15 @@ def __init__(
*,
enabled: bool,
capture_name_tokens: tuple[str, ...] = CAPTURE_NAME_TOKENS,
+ capture_layer_outputs: bool = True,
+ max_layer_index: int | None = None,
micro_start_callback: Callable[[int | None, int], None] | None = None,
strict_output_match: bool = True,
) -> None:
self.enabled = enabled
self.capture_name_tokens = capture_name_tokens
+ self.capture_layer_outputs = capture_layer_outputs
+ self.max_layer_index = max_layer_index
self.micro_start_callback = micro_start_callback
self.strict_output_match = strict_output_match
self.current_step_index: int | None = None
@@ -307,6 +332,13 @@ def _register_hooks(self, model_chunks: list[Any]) -> None:
named_modules = list(chunk.named_modules())
module_by_name = dict(named_modules)
for module_name, module in named_modules:
+ layer_index = _module_layer_index(module_name)
+ if (
+ self.max_layer_index is not None
+ and layer_index is not None
+ and layer_index > self.max_layer_index
+ ):
+ continue
trace_module_name = _normalize_trace_module_name(
f"chunk{chunk_index}.{module_name}"
)
@@ -318,7 +350,8 @@ def _register_hooks(self, model_chunks: list[Any]) -> None:
if metadata:
self._trace_metadata_by_name[trace_module_name] = metadata
is_layer_output = (
- ".decoder.layers." in module_name
+ self.capture_layer_outputs
+ and ".decoder.layers." in module_name
and module_name.rsplit(".", 1)[-1].isdigit()
)
if not is_layer_output and not any(
@@ -593,7 +626,8 @@ def _hook(_module: Any, inputs: Any, output: Any) -> None:
if isinstance(primary_output, torch.Tensor) and primary_output.ndim > 0
else None
)
- row_token_uids, _uid_span = self._row_token_uids_for_trace(
+ row_token_uids, _uid_span = self._row_token_uids_for_capture(
+ module_name=name,
inputs=inputs,
output=output,
module=module,
@@ -704,6 +738,47 @@ def _row_token_uids_for_trace(
prefer_uid_span=prefer_uid_span,
)
+ @staticmethod
+ def _module_row_token_uids(
+ module: Any,
+ *,
+ row_count: int,
+ ) -> tuple[torch.Tensor | None, int | None]:
+ row_token_uids = normalize_row_token_uids(
+ getattr(module, TRACE_ROW_TOKEN_UIDS_ATTR, None)
+ )
+ if row_token_uids is None or int(row_token_uids.numel()) != int(row_count):
+ return None, None
+ uid_span = getattr(module, TRACE_UID_SPAN_ATTR, None)
+ return (
+ row_token_uids,
+ uid_span if isinstance(uid_span, int) and uid_span > 0 else None,
+ )
+
+ @classmethod
+ def _row_token_uids_for_capture(
+ cls,
+ *,
+ module_name: str,
+ inputs: Any,
+ output: Any,
+ module: Any,
+ row_count: int | None,
+ ) -> tuple[torch.Tensor | None, int | None]:
+ if row_count is not None and not cls._is_moe_expert_forward_module(module_name):
+ row_token_uids, uid_span = cls._module_row_token_uids(
+ module,
+ row_count=row_count,
+ )
+ if row_token_uids is not None:
+ return row_token_uids, uid_span
+ return cls._row_token_uids_for_trace(
+ inputs=inputs,
+ output=output,
+ module=module,
+ row_count=row_count,
+ )
+
@classmethod
def _slice_row_aligned_value(
cls,
@@ -1404,6 +1479,39 @@ def _propagate_attention_output_row_token_uids(
head_count=head_count,
)
+ @classmethod
+ def _restore_non_cp_sequence_row_token_uids(
+ cls,
+ trace: dict[str, list[dict[str, Any]]],
+ ) -> None:
+ for module_name, calls in trace.items():
+ if cls._is_moe_expert_forward_module(module_name):
+ continue
+ for call in calls:
+ if (
+ "rank_meta" not in call
+ or _call_cp_world_size(call) > 1
+ or "row_uid_span" in call
+ ):
+ continue
+ tensor = call.get("primary_output")
+ row_token_uids = call.get("row_token_uids")
+ if (
+ not isinstance(tensor, torch.Tensor)
+ or tensor.ndim == 0
+ or not isinstance(row_token_uids, torch.Tensor)
+ or row_token_uids.ndim != 1
+ or int(row_token_uids.numel()) != int(tensor.shape[0])
+ or not bool((row_token_uids >= 0).all().item())
+ or int(row_token_uids.unique().numel())
+ != int(row_token_uids.numel())
+ ):
+ continue
+ call["row_token_uids"] = torch.arange(
+ row_token_uids.numel(),
+ dtype=torch.int64,
+ )
+
@classmethod
def canonicalize_trace(
cls,
@@ -1426,6 +1534,7 @@ def canonicalize_trace(
call[PRIMARY_OUTPUT_CANONICAL_KEY] = True
cls._propagate_decoder_row_token_uids(trace)
cls._propagate_attention_output_row_token_uids(trace)
+ cls._restore_non_cp_sequence_row_token_uids(trace)
for calls in trace.values():
for call in calls:
cls._canonicalize_call_row_token_order(call)
diff --git a/tests/integration/megatron/model_support/hf_parity_worker.py b/tests/integration/megatron/model_support/hf_parity_worker.py
index 6e7944f67..1bbc2aaaf 100644
--- a/tests/integration/megatron/model_support/hf_parity_worker.py
+++ b/tests/integration/megatron/model_support/hf_parity_worker.py
@@ -306,6 +306,39 @@ def _timed_load_weights(*args: Any, **kwargs: Any) -> Any:
model_bridge._art_hf_parity_timing_wrapped = True
+def _is_bridge_hf_load_hook(hook: Any) -> bool:
+ fn = getattr(hook, "func", hook)
+ name = getattr(fn, "__name__", "")
+ qualname = getattr(fn, "__qualname__", "")
+ return name in {
+ "load_weights_hf_to_megatron",
+ "_optimized_load_weights_hf_to_megatron",
+ } or qualname.endswith(".load_weights_hf_to_megatron")
+
+
+def _remove_bridge_hf_load_hook(provider_bundle: Any) -> None:
+ """Disable raw checkpoint load when parity seeds from HF oracle state."""
+
+ provider = provider_bundle.provider
+ hooks = list(getattr(provider, "_pre_wrap_hooks", []))
+ kept = [hook for hook in hooks if not _is_bridge_hf_load_hook(hook)]
+ if len(kept) == len(hooks):
+ raise RuntimeError(
+ "HF parity expected a Bridge HF-load pre-wrap hook to remove"
+ )
+ provider._pre_wrap_hooks = kept
+
+
+def _configure_hf_parity_provider_bundle(
+ provider_bundle: Any,
+ *,
+ use_hf_reference_state: bool,
+) -> None:
+ if use_hf_reference_state:
+ _remove_bridge_hf_load_hook(provider_bundle)
+ _install_bridge_timing_debug(provider_bundle)
+
+
def _load_hf_model(
*,
base_model: str,
@@ -315,15 +348,35 @@ def _load_hf_model(
) -> Any:
from transformers import AutoConfig, AutoModelForCausalLM
+ from art.megatron.model_support.registry import get_model_support_handler
+
+ handler = get_model_support_handler(base_model)
+ ensure_hf_reference_registered = getattr(
+ handler, "ensure_hf_reference_registered", None
+ )
+ if ensure_hf_reference_registered is not None:
+ ensure_hf_reference_registered()
config = AutoConfig.from_pretrained(base_model, trust_remote_code=True)
set_hf_config_num_layers(config, num_layers)
zero_hf_dropout_config(config)
+ prepare_hf_reference_config = getattr(handler, "prepare_hf_reference_config", None)
+ if prepare_hf_reference_config is not None:
+ prepare_hf_reference_config(config)
+ hf_reference_from_pretrained_kwargs = getattr(
+ handler, "hf_reference_from_pretrained_kwargs", None
+ )
+ extra_kwargs = (
+ hf_reference_from_pretrained_kwargs(config=config, dtype=dtype)
+ if hf_reference_from_pretrained_kwargs is not None
+ else {}
+ )
model = AutoModelForCausalLM.from_pretrained(
base_model,
config=config,
trust_remote_code=True,
torch_dtype=dtype,
low_cpu_mem_usage=True,
+ **extra_kwargs,
)
model.train()
return cast(Any, model).to(device)
@@ -339,6 +392,37 @@ def _collect_hf_grads(model: Any) -> dict[str, torch.Tensor]:
return grads
+def _collect_hf_state_dict(model: Any) -> dict[str, torch.Tensor]:
+ return {
+ key: value.detach().cpu().clone()
+ for key, value in model.state_dict().items()
+ if _is_language_hf_param_name(key)
+ }
+
+
+def _normalize_hf_reference_state_for_hf_parity(
+ *,
+ base_model: str,
+ model: Any,
+ state: dict[str, torch.Tensor],
+) -> dict[str, torch.Tensor]:
+ from art.megatron.model_support.registry import get_model_support_handler
+
+ handler = get_model_support_handler(base_model)
+ normalize = getattr(handler, "normalize_hf_reference_state_for_hf_parity", None)
+ if normalize is not None:
+ normalize(state, config=model.config)
+ return state
+
+
+def _use_hf_reference_state_for_hf_parity(base_model: str) -> bool:
+ from art.megatron.model_support.registry import get_model_support_handler
+
+ handler = get_model_support_handler(base_model)
+ enabled = getattr(handler, "use_hf_reference_state_for_hf_parity", None)
+ return bool(enabled()) if enabled is not None else False
+
+
def _bridge_compatible_hf_key(key: str, expected_keys: set[str]) -> str:
if key in expected_keys:
return key
@@ -454,11 +538,21 @@ def _focus_derivative_tensor_map(
loss_active_last_layer_experts: set[int],
) -> dict[str, torch.Tensor]:
focused: dict[str, torch.Tensor] = {}
+ active_router_expert_sets = {
+ layer_index: set(int(row) for row in rows.reshape(-1).tolist())
+ for layer_index, rows in active_router_rows.items()
+ if rows.numel() > 0
+ }
for key, value in tensor_map.items():
if match := _EXPERT_WEIGHT_PATTERN.match(key):
+ layer_index = int(match.group("layer"))
+ expert_index = int(match.group("expert"))
+ active_experts = active_router_expert_sets.get(layer_index)
+ if active_experts is not None and expert_index not in active_experts:
+ continue
if (
- int(match.group("layer")) == last_layer_index
- and int(match.group("expert")) not in loss_active_last_layer_experts
+ layer_index == last_layer_index
+ and expert_index not in loss_active_last_layer_experts
):
continue
focused_value = value
@@ -489,6 +583,7 @@ def _run_hf_sft_step(
torch.Tensor,
dict[str, torch.Tensor],
MoeRoutingReplayBundle | None,
+ dict[str, torch.Tensor] | None,
]:
_debug("loading HF model")
model = _load_hf_model(
@@ -540,6 +635,15 @@ def _run_hf_sft_step(
token_count += int(mask.sum().item())
(masked_losses.sum() / total_token_count).backward()
grads = _collect_hf_grads(model)
+ hf_reference_state_dict = (
+ _normalize_hf_reference_state_for_hf_parity(
+ base_model=base_model,
+ model=model,
+ state=_collect_hf_state_dict(model),
+ )
+ if _use_hf_reference_state_for_hf_parity(base_model)
+ else None
+ )
routing_replay_bundle = route_capture.build_replay_bundle(topology=topology)
scalar_loss = (loss_sum / max(token_count, 1)).detach().cpu().reshape(1)
output_vector = torch.cat(trainable_losses, dim=0).to(dtype=torch.float32)
@@ -548,7 +652,13 @@ def _run_hf_sft_step(
if torch.cuda.is_available():
torch.cuda.empty_cache()
_debug("finished HF step")
- return output_vector, scalar_loss, grads, routing_replay_bundle
+ return (
+ output_vector,
+ scalar_loss,
+ grads,
+ routing_replay_bundle,
+ hf_reference_state_dict,
+ )
def _install_hf_qwen35_gdn_fp32_reference(model: Any, *, base_model: str) -> None:
@@ -572,10 +682,18 @@ def _build_megatron_runtime(
*,
moe_routing_replay_bundle: MoeRoutingReplayBundle | None = None,
) -> megatron_train.TrainingRuntime:
+ use_hf_reference_state = _use_hf_reference_state_for_hf_parity(
+ request.case_config.base_model
+ )
return megatron_train.build_training_runtime(
model_identifier=request.case_config.base_model,
provider_torch_dtype=_dtype_for_precision(request.case_config.precision),
- provider_bundle_configure=_install_bridge_timing_debug,
+ provider_bundle_configure=lambda provider_bundle: (
+ _configure_hf_parity_provider_bundle(
+ provider_bundle,
+ use_hf_reference_state=use_hf_reference_state,
+ )
+ ),
provider_configure=lambda provider: _configure_provider(
provider, ORACLE_TOPOLOGY, request.case_config
),
@@ -644,6 +762,53 @@ def _mapping_targets_language_only(mapping: Any) -> bool:
return all(_is_language_hf_param_name(name) for name in names)
+def _hf_param_names_for_mapping(mapping: Any) -> set[str]:
+ names = _language_hf_param_names(mapping)
+ if not names:
+ return set()
+ return set(names)
+
+
+def _build_hf_parity_conversion_tasks(
+ *,
+ bridge: Any,
+ model: list[Any],
+ hf_keys: set[str],
+) -> list[Any]:
+ tasks = []
+ for task in build_art_conversion_tasks(bridge=bridge, model=model):
+ mapping_names = _hf_param_names_for_mapping(task.mapping)
+ if not mapping_names:
+ tasks.append(task)
+ continue
+ if mapping_names & hf_keys:
+ tasks.append(task)
+ return tasks
+
+
+def _seed_megatron_from_hf_reference_state(
+ runtime: megatron_train.TrainingRuntime,
+ *,
+ tasks: list[Any],
+ hf_reference_state_dict: dict[str, torch.Tensor],
+) -> None:
+ model_bridge = runtime.bridge._model_bridge
+ for task in tasks:
+ if task.mapping is None:
+ continue
+ hf_weights = model_bridge.maybe_modify_loaded_hf_weight(
+ task.mapping.hf_param,
+ hf_reference_state_dict,
+ )
+ converted_weights = task.mapping.hf_to_megatron(
+ hf_weights, task.megatron_module
+ )
+ if isinstance(task.param_weight, torch.nn.Parameter):
+ task.param_weight.data.copy_(converted_weights.to(task.param_weight.device))
+ elif isinstance(task.param_weight, torch.Tensor):
+ task.param_weight.copy_(converted_weights.to(task.param_weight.device))
+
+
def _filter_language_only_tensor_map(
tensor_map: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
@@ -772,6 +937,7 @@ def _convert_megatron_tasks_to_hf(
*,
mode: str,
tasks: list[Any] | None = None,
+ hf_state_dict_override: dict[str, torch.Tensor] | None = None,
) -> dict[str, torch.Tensor]:
if tasks is None:
tasks = [
@@ -783,7 +949,11 @@ def _convert_megatron_tasks_to_hf(
if isinstance(task.param_weight, torch.nn.Parameter)
]
model_bridge = runtime.bridge._model_bridge
- hf_state_dict = runtime.bridge.hf_pretrained.state
+ hf_state_dict = (
+ hf_state_dict_override
+ if hf_state_dict_override is not None
+ else runtime.bridge.hf_pretrained.state
+ )
grouped_buffers: dict[str, dict[int, torch.Tensor]] = {}
converted: dict[str, torch.Tensor] = {}
additive_grad_keys: set[str] = set()
@@ -847,6 +1017,7 @@ def _run_megatron_sft_step(
sample_indices: list[int | None],
device: torch.device,
moe_routing_replay_bundle: MoeRoutingReplayBundle | None = None,
+ hf_reference_state_dict: dict[str, torch.Tensor] | None = None,
) -> tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]:
runtime = _build_megatron_runtime(
request,
@@ -865,16 +1036,34 @@ def _run_megatron_sft_step(
sample_index=sample_indices,
global_grad_accumulation_sequences=request.case_config.grad_accumulation_sequences,
)
- _debug("initializing Megatron optimizer state")
- megatron_train._eager_initialize_optimizer_state(runtime.optimizer)
- tasks = [
- task
- for task in build_art_conversion_tasks(
+ if hf_reference_state_dict is None:
+ tasks = [
+ task
+ for task in build_art_conversion_tasks(
+ bridge=runtime.bridge,
+ model=runtime.model,
+ )
+ if isinstance(task.param_weight, torch.nn.Parameter)
+ ]
+ else:
+ seed_tasks = _build_hf_parity_conversion_tasks(
bridge=runtime.bridge,
model=runtime.model,
+ hf_keys=set(hf_reference_state_dict),
)
- if isinstance(task.param_weight, torch.nn.Parameter)
- ]
+ tasks = [
+ task
+ for task in seed_tasks
+ if isinstance(task.param_weight, torch.nn.Parameter)
+ ]
+ _debug("seeding Megatron weights from HF oracle state")
+ _seed_megatron_from_hf_reference_state(
+ runtime,
+ tasks=seed_tasks,
+ hf_reference_state_dict=hf_reference_state_dict,
+ )
+ _debug("initializing Megatron optimizer state")
+ megatron_train._eager_initialize_optimizer_state(runtime.optimizer)
_debug(f"built {len(tasks)} Megatron conversion tasks")
for chunk in runtime.model:
if hasattr(chunk, "zero_grad_buffer"):
@@ -940,6 +1129,7 @@ def _run_megatron_sft_step(
runtime,
mode="grad",
tasks=derivative_tasks,
+ hf_state_dict_override=hf_reference_state_dict,
)
_debug("exported Megatron grads")
if runtime.moe_routing_replay_controller is not None:
@@ -1030,7 +1220,13 @@ def _worker_run(request: HfParityRunRequest) -> None:
)
try:
_debug("starting HF parity worker")
- hf_outputs, hf_loss, hf_grads, moe_routing_replay_bundle = _run_hf_sft_step(
+ (
+ hf_outputs,
+ hf_loss,
+ hf_grads,
+ moe_routing_replay_bundle,
+ hf_reference_state_dict,
+ ) = _run_hf_sft_step(
base_model=request.case_config.base_model,
num_layers=request.case_config.num_layers,
micro_inputs=micro_inputs,
@@ -1045,6 +1241,7 @@ def _worker_run(request: HfParityRunRequest) -> None:
sample_indices=sample_indices,
device=device,
moe_routing_replay_bundle=moe_routing_replay_bundle,
+ hf_reference_state_dict=hf_reference_state_dict,
)
_debug("finished HF and Megatron steps, building report")
normalized_hf_grads = _normalize_hf_grads_for_bridge(
diff --git a/tests/integration/megatron/model_support/lora_coverage.py b/tests/integration/megatron/model_support/lora_coverage.py
index 07097c4d4..7d998de06 100644
--- a/tests/integration/megatron/model_support/lora_coverage.py
+++ b/tests/integration/megatron/model_support/lora_coverage.py
@@ -22,6 +22,13 @@
from .oracle_worker import _configure_provider, provider_topology_env
_WRAPPED_TARGET_SUFFIXES: dict[str, tuple[str, ...]] = {
+ "q_a_proj": (".self_attn.q_a_proj",),
+ "q_b_proj": (".self_attn.q_b_proj",),
+ "kv_proj": (".self_attn.kv_proj",),
+ "o_a_proj": (".self_attn.o_a_proj",),
+ "o_b_proj": (".self_attn.o_b_proj",),
+ "compressor.kv_proj": (".self_attn.compressor.kv_proj",),
+ "compressor.gate_proj": (".self_attn.compressor.gate_proj",),
"q_proj": (".self_attn.q_proj",),
"k_proj": (".self_attn.k_proj",),
"v_proj": (".self_attn.v_proj",),
@@ -107,6 +114,27 @@ def _covered_exported_target_modules(
) -> set[str]:
covered: set[str] = set()
for base_name, adapter_weights in adapter_weights_by_base.items():
+ if base_name.endswith(".self_attention.wq_a.weight"):
+ covered.add("q_a_proj")
+ continue
+ if base_name.endswith(".self_attention.wq_b.weight"):
+ covered.add("q_b_proj")
+ continue
+ if base_name.endswith(".self_attention.wkv.weight"):
+ covered.add("kv_proj")
+ continue
+ if base_name.endswith(".self_attention.wo_a.weight"):
+ covered.add("o_a_proj")
+ continue
+ if base_name.endswith(".self_attention.wo_b.weight"):
+ covered.add("o_b_proj")
+ continue
+ if base_name.endswith(".self_attention.compressor.wkv.weight"):
+ covered.add("compressor.kv_proj")
+ continue
+ if base_name.endswith(".self_attention.compressor.wgate.weight"):
+ covered.add("compressor.gate_proj")
+ continue
if base_name.endswith(".self_attention.linear_qkv.weight"):
for adapter_weight in adapter_weights:
adapter_key = getattr(adapter_weight, "adapter_key", None)
diff --git a/tests/integration/megatron/model_support/oracle_harness.py b/tests/integration/megatron/model_support/oracle_harness.py
index d5bf7c581..2ea15447e 100644
--- a/tests/integration/megatron/model_support/oracle_harness.py
+++ b/tests/integration/megatron/model_support/oracle_harness.py
@@ -220,6 +220,17 @@ def world_size(self) -> int:
Topology(tp=2, ep=2, etp=1, dp=1, cp=2, sp=True),
Topology(tp=2, ep=4, etp=2, dp=2, cp=2, sp=True),
]
+
+
+def _without_context_parallel(topology: Topology) -> Topology:
+ return topology.model_copy(update={"dp": topology.dp * topology.cp, "cp": 1})
+
+
+CP_UNSUPPORTED_MOE_TOPOLOGIES = [
+ _without_context_parallel(topology) for topology in TOPOLOGIES[:-1]
+] + [
+ Topology(tp=2, ep=2, etp=2, dp=2, cp=1, sp=True),
+]
DENSE_TOPOLOGIES = [
Topology(tp=1, ep=1, etp=1, dp=1, sp=False),
Topology(tp=2, ep=1, etp=1, dp=1, sp=True),
@@ -751,9 +762,30 @@ def oracle_topology(*, is_moe: bool = True) -> Topology:
return ORACLE_TOPOLOGY if is_moe else DENSE_ORACLE_TOPOLOGY
-def selected_suite_topologies(*, is_moe: bool = True) -> list[Topology]:
+def _filter_context_parallel_support(
+ topologies: list[Topology],
+ *,
+ is_moe: bool,
+ cp_supported: bool,
+) -> list[Topology]:
+ if cp_supported:
+ return topologies
+ if is_moe:
+ return list(CP_UNSUPPORTED_MOE_TOPOLOGIES)
+ return [_without_context_parallel(topology) for topology in topologies]
+
+
+def selected_suite_topologies(
+ *,
+ is_moe: bool = True,
+ cp_supported: bool = True,
+) -> list[Topology]:
"""Returns the correctness topology list for a model family."""
- return list(TOPOLOGIES if is_moe else DENSE_TOPOLOGIES)
+ return _filter_context_parallel_support(
+ list(TOPOLOGIES if is_moe else DENSE_TOPOLOGIES),
+ is_moe=is_moe,
+ cp_supported=cp_supported,
+ )
def stable_case_id(case_config: OracleCaseConfig) -> str:
@@ -1204,7 +1236,11 @@ def _stacked_layers(
if len(reference_shapes) != 1 or len(candidate_shapes) != 1:
original_names = original_names_by_group[normalized]
for original_name, (reference, candidate) in zip(original_names, group):
- stacked_pairs.append((original_name, reference, candidate))
+ # Keep one synthetic layer axis so layer-averaged comparison
+ # does not treat tensor rows/features as layer entries.
+ stacked_pairs.append(
+ (original_name, reference.unsqueeze(0), candidate.unsqueeze(0))
+ )
continue
stacked_pairs.append(
(
@@ -2063,8 +2099,17 @@ def run_variant(
def run_suite(
self,
variants: list[VariantSpec],
+ *,
+ prune_reference_artifacts: bool = True,
+ prune_case_artifacts: bool = True,
) -> list[VariantReport]:
- """Runs variants in order and stops at the first unexpected signal."""
+ """Runs variants in order and stops at the first unexpected signal.
+
+ Reference and case artifacts are normally pruned when the suite exits.
+ Callers that immediately run another comparison suite against the same
+ reference can defer that pruning so the second suite does not have to
+ regenerate or fail on missing forward traces.
+ """
reports: list[VariantReport] = []
try:
for variant in variants:
@@ -2078,8 +2123,10 @@ def run_suite(
/ "variant_report.json",
)
finally:
- self._prune_reference_artifacts()
- _prune_case_artifacts(self.case_dir)
+ if prune_reference_artifacts:
+ self._prune_reference_artifacts()
+ if prune_case_artifacts:
+ _prune_case_artifacts(self.case_dir)
return reports
@@ -2125,13 +2172,18 @@ def _suite_variants(
objective: OracleObjective,
*,
is_moe: bool = True,
+ cp_supported: bool = True,
max_world_size: int | None = None,
variant_flex_backend: FlexBackend | None = None,
+ phase_pass_fns: dict[str, PhasePassFn] | None = None,
) -> list[VariantSpec]:
"""Builds the standard oracle suite variant ordering."""
- phase_pass = _default_phase_pass_fns()
+ phase_pass = phase_pass_fns or _default_phase_pass_fns()
variants: list[VariantSpec] = []
- for topology in selected_suite_topologies(is_moe=is_moe)[1:]:
+ for topology in selected_suite_topologies(
+ is_moe=is_moe,
+ cp_supported=cp_supported,
+ )[1:]:
if max_world_size is not None and topology.world_size() > max_world_size:
continue
variants.append(
@@ -2152,6 +2204,11 @@ def run_suite(
max_world_size: int | None = None,
oracle_flex_backend: FlexBackend | None = None,
variant_flex_backend: FlexBackend | None = None,
+ cp_supported: bool = True,
+ phase_pass_fns: dict[str, PhasePassFn] | None = None,
+ use_fp32_lora_reference: bool = True,
+ prune_reference_artifacts: bool = True,
+ prune_case_artifacts: bool = True,
) -> list[VariantReport]:
"""Runs non-oracle topologies against the canonical replay-backed oracle."""
reports: list[VariantReport] = []
@@ -2161,15 +2218,20 @@ def run_suite(
case_config=case_config,
oracle_flex_backend=oracle_flex_backend,
variant_flex_backend=variant_flex_backend,
+ use_fp32_lora_reference=use_fp32_lora_reference,
)
reports.extend(
runner.run_suite(
_suite_variants(
objective,
is_moe=case_config.is_moe,
+ cp_supported=cp_supported,
max_world_size=max_world_size,
variant_flex_backend=variant_flex_backend,
- )
+ phase_pass_fns=phase_pass_fns,
+ ),
+ prune_reference_artifacts=prune_reference_artifacts,
+ prune_case_artifacts=prune_case_artifacts,
)
)
return reports
diff --git a/tests/integration/megatron/model_support/oracle_worker.py b/tests/integration/megatron/model_support/oracle_worker.py
index 29a537469..42da1eab5 100644
--- a/tests/integration/megatron/model_support/oracle_worker.py
+++ b/tests/integration/megatron/model_support/oracle_worker.py
@@ -20,6 +20,7 @@
ParallelTopology as ReplayParallelTopology,
)
from art.preprocessing.pack import PackedTensors
+from art.utils.lifecycle import terminate_popen_process_group
from ..routing_replay.bundle import build_bundle_from_forward_trace_dir
from ..routing_replay.trace import install_moe_routing_trace_hooks
@@ -103,6 +104,7 @@ def run_worker_subprocess(
worker_log_path = topology_dir / "worker.log"
live_log_raw = os.environ.get("ART_ORACLE_LIVE_TRAINING_LOG")
live_log_path = None if not live_log_raw else Path(live_log_raw)
+ run: subprocess.Popen[str] | None = None
worker_log_path.parent.mkdir(parents=True, exist_ok=True)
with worker_log_path.open("w", encoding="utf-8") as worker_log:
live_log = None
@@ -129,6 +131,7 @@ def run_worker_subprocess(
stderr=subprocess.STDOUT,
text=True,
bufsize=1,
+ start_new_session=True,
)
assert run.stdout is not None
for line in run.stdout:
@@ -140,6 +143,8 @@ def run_worker_subprocess(
live_log.flush()
run.returncode = run.wait()
finally:
+ if run is not None and run.poll() is None:
+ terminate_popen_process_group(run)
if live_log is not None:
live_log.close()
combined_output = "".join(combined_lines).strip()
@@ -389,7 +394,12 @@ def _configure_provider(
topology: Topology,
case_config: OracleCaseConfig,
) -> None:
- """Applies deterministic topology/model overrides to provider config."""
+ """Applies deterministic topology/model overrides to provider config.
+
+ Handler-specific oracle hooks are validation-only. They keep large model
+ families such as DSV4 fit-sized while preserving the layer families and
+ kernel-facing invariants under test.
+ """
del topology
provider.num_layers = case_config.num_layers
if case_config.precision == "fp32":
@@ -405,6 +415,15 @@ def _configure_provider(
provider.attention_dropout = 0.0
if hasattr(provider, "hidden_dropout"):
provider.hidden_dropout = 0.0
+ from art.megatron.model_support.registry import get_model_support_handler
+
+ handler = get_model_support_handler(
+ case_config.base_model,
+ allow_unvalidated_arch=case_config.allow_unvalidated_arch,
+ )
+ configure_oracle_provider = getattr(handler, "configure_oracle_provider", None)
+ if configure_oracle_provider is not None:
+ configure_oracle_provider(provider, case_config=case_config)
@contextmanager
diff --git a/tests/integration/megatron/model_support/packed_position_ids.py b/tests/integration/megatron/model_support/packed_position_ids.py
index 657fdc1ec..397e7ce12 100644
--- a/tests/integration/megatron/model_support/packed_position_ids.py
+++ b/tests/integration/megatron/model_support/packed_position_ids.py
@@ -52,6 +52,7 @@
"qwen3_moe",
"qwen3_5_dense",
"qwen3_5_moe",
+ "dsv4",
"gpt_oss_moe",
}
)
@@ -613,6 +614,7 @@ def _logits_equivalence_check(
build_gdn_execution_spec=bool(
getattr(handler, "build_gdn_execution_spec", False)
),
+ model_support_handler=handler,
attention_head_dim=getattr(provider, "kv_channels", None),
attention_value_head_dim=getattr(provider, "kv_channels", None),
)
@@ -663,6 +665,7 @@ def _logits_equivalence_check(
build_gdn_execution_spec=bool(
getattr(handler, "build_gdn_execution_spec", False)
),
+ model_support_handler=handler,
attention_head_dim=getattr(provider, "kv_channels", None),
attention_value_head_dim=getattr(provider, "kv_channels", None),
)
diff --git a/tests/integration/megatron/model_support/test_dsv4_real_path_correctness.py b/tests/integration/megatron/model_support/test_dsv4_real_path_correctness.py
new file mode 100644
index 000000000..d0eee8460
--- /dev/null
+++ b/tests/integration/megatron/model_support/test_dsv4_real_path_correctness.py
@@ -0,0 +1,170 @@
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import contextmanager, redirect_stderr, redirect_stdout
+import os
+from pathlib import Path
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from .oracle_harness import ( # noqa: E402
+ LIVE_TRAINING_LOG_PATH,
+ MetricThresholdRule,
+ OracleCaseConfig,
+ OracleObjective,
+ PhasePassFn,
+ VariantReport,
+ VariantRunner,
+ VariantSpec,
+ available_gpu_count,
+ selected_oracle_objectives,
+ selected_suite_topologies,
+)
+
+BASE_MODEL = "deepseek-ai/DeepSeek-V4-Flash"
+NUM_LAYERS = 4
+REPO_ROOT = Path(__file__).resolve().parents[4]
+CORRECTNESS_LOG_PATH = REPO_ROOT / ".local" / "correctness.log"
+ORACLE_LIVE_TRAINING_LOG_ENV = "ART_ORACLE_LIVE_TRAINING_LOG"
+_EXPECTED_COMPRESS_RATIOS = [0, 0, 4, 128]
+_EXPECTED_LAYER_TYPES = [
+ "sliding_attention",
+ "sliding_attention",
+ "compressed_sparse_attention",
+ "heavily_compressed_attention",
+]
+_EXPECTED_MLP_LAYER_TYPES = ["hash_moe", "hash_moe", "hash_moe", "moe"]
+
+
+def test_dsv4_real_path_bf16_correctness(
+ capsys: pytest.CaptureFixture[str],
+) -> None:
+ """Runs the real DSV4 ART Megatron path through the standard oracle harness."""
+ _assert_representative_dsv4_layers()
+ gpu_count = available_gpu_count()
+ reports: list[VariantReport] = []
+ with capsys.disabled():
+ print(
+ f"\nDSV4 real-path bf16 correctness log: {CORRECTNESS_LOG_PATH}",
+ flush=True,
+ )
+ print(
+ f"DSV4 real-path bf16 live training log: {LIVE_TRAINING_LOG_PATH}",
+ flush=True,
+ )
+ CORRECTNESS_LOG_PATH.parent.mkdir(parents=True, exist_ok=True)
+ LIVE_TRAINING_LOG_PATH.parent.mkdir(parents=True, exist_ok=True)
+ LIVE_TRAINING_LOG_PATH.write_text(
+ (
+ "DSV4 real-path bf16 live training log.\n"
+ "Topology worker output is appended below. If no topology sections "
+ "appear, complete cached artifacts were reused.\n"
+ ),
+ encoding="utf-8",
+ )
+ with _temporary_env(**{ORACLE_LIVE_TRAINING_LOG_ENV: str(LIVE_TRAINING_LOG_PATH)}):
+ with CORRECTNESS_LOG_PATH.open("w", encoding="utf-8") as log_file:
+ with redirect_stdout(log_file), redirect_stderr(log_file):
+ print("DSV4 real-path bf16 correctness")
+ print(f"base_model={BASE_MODEL}")
+ print(f"num_layers={NUM_LAYERS}")
+ print(f"precision=bf16")
+ print(f"visible_gpus={gpu_count}")
+ print(f"live_training_log={LIVE_TRAINING_LOG_PATH}")
+ for objective in selected_oracle_objectives():
+ runner = VariantRunner(
+ objective=objective,
+ case_config=OracleCaseConfig(
+ base_model=BASE_MODEL,
+ precision="bf16",
+ num_layers=NUM_LAYERS,
+ ),
+ use_fp32_lora_reference=False,
+ )
+ variants = _dsv4_bf16_variants(
+ objective=objective,
+ max_world_size=gpu_count,
+ )
+ if variants:
+ reports.extend(runner.run_suite(variants))
+ if not reports:
+ CORRECTNESS_LOG_PATH.write_text(
+ f"DSV4 real-path bf16 correctness skipped. Need at least 2 GPUs; found {gpu_count}.\n",
+ encoding="utf-8",
+ )
+ pytest.skip(f"Need at least 2 GPUs for DSV4 correctness; found {gpu_count}.")
+ assert all(report.signal == "pass" for report in reports)
+
+
+@contextmanager
+def _temporary_env(**updates: str) -> Iterator[None]:
+ previous = {key: os.environ.get(key) for key in updates}
+ os.environ.update(updates)
+ try:
+ yield
+ finally:
+ for key, value in previous.items():
+ if value is None:
+ os.environ.pop(key, None)
+ else:
+ os.environ[key] = value
+
+
+def _dsv4_bf16_variants(
+ *,
+ objective: OracleObjective,
+ max_world_size: int,
+) -> list[VariantSpec]:
+ phase_pass = _dsv4_bf16_phase_pass_fns()
+ variants: list[VariantSpec] = []
+ for topology in selected_suite_topologies(is_moe=True, cp_supported=False)[1:]:
+ if topology.world_size() > max_world_size:
+ continue
+ variants.append(
+ VariantSpec(
+ name=f"{objective}_dsv4_bf16_topology_{topology.slug()}",
+ objective=objective,
+ topology=topology,
+ pass_fn_by_phase=phase_pass,
+ )
+ )
+ return variants
+
+
+def _dsv4_bf16_phase_pass_fns() -> dict[str, PhasePassFn]:
+ non_zero_scales = {"typical_abs_scale": 0.0, "candidate_abs_scale": 0.0}
+ fwd = MetricThresholdRule(
+ limits={"mean_abs_pct": 3.0},
+ minimums=non_zero_scales,
+ )
+ loss = MetricThresholdRule(limits={"mean_abs_pct": 3.0})
+ grad = MetricThresholdRule(
+ limits={"mean_abs_pct": 5.0},
+ minimums=non_zero_scales,
+ )
+ router_topk = MetricThresholdRule(
+ limits={"topk_mismatch_fraction": 0.0, "top1_mismatch_fraction": 0.0}
+ )
+ return {
+ "forward": fwd,
+ "outputs": fwd,
+ "losses": loss,
+ "grads": grad,
+ "deltas": grad,
+ "router_scores": fwd,
+ "router_topk_ids": router_topk,
+ }
+
+
+def _assert_representative_dsv4_layers() -> None:
+ from transformers import AutoConfig
+
+ from art.megatron.dsv4.hf_config import ensure_dsv4_hf_model_registered
+
+ ensure_dsv4_hf_model_registered()
+ config = AutoConfig.from_pretrained(BASE_MODEL, trust_remote_code=True)
+ assert list(config.compress_ratios[:NUM_LAYERS]) == _EXPECTED_COMPRESS_RATIOS
+ assert list(config.layer_types[:NUM_LAYERS]) == _EXPECTED_LAYER_TYPES
+ assert list(config.mlp_layer_types[:NUM_LAYERS]) == _EXPECTED_MLP_LAYER_TYPES
diff --git a/tests/integration/megatron/model_support/test_oracle_harness_invariants.py b/tests/integration/megatron/model_support/test_oracle_harness_invariants.py
index 52dac4893..991e925ec 100644
--- a/tests/integration/megatron/model_support/test_oracle_harness_invariants.py
+++ b/tests/integration/megatron/model_support/test_oracle_harness_invariants.py
@@ -239,6 +239,65 @@ def test_forward_trace_prefers_local_tensor_uids_over_module_fallback() -> None:
assert torch.equal(row_uids, torch.tensor([4, 7]))
+def test_forward_trace_capture_prefers_dense_module_uids_for_sequence_modules() -> None:
+ module = type("ModuleWithDenseTraceUids", (), {})()
+ output = torch.zeros((2, 1), dtype=torch.float32)
+ setattr(module, "_art_trace_row_token_uids", torch.tensor([10, 11]))
+ setattr(output, "_art_trace_row_token_uids", torch.tensor([4, 7]))
+
+ row_uids, _uid_span = ForwardTraceCapture._row_token_uids_for_capture(
+ module_name="chunk0.module.decoder.layers.0.self_attention",
+ inputs=(),
+ output=output,
+ module=module,
+ row_count=2,
+ )
+
+ assert row_uids is not None
+ assert torch.equal(row_uids, torch.tensor([10, 11]))
+
+
+def test_forward_trace_capture_keeps_expert_uid_span_for_expert_modules() -> None:
+ module = type("ModuleWithDenseTraceUids", (), {})()
+ output = torch.zeros((2, 1), dtype=torch.float32)
+ setattr(module, "_art_trace_row_token_uids", torch.tensor([10, 11]))
+ setattr(output, "_art_trace_row_token_uids", torch.tensor([4, 7]))
+ setattr(output, "_art_trace_uid_span", 16)
+
+ row_uids, uid_span = ForwardTraceCapture._row_token_uids_for_capture(
+ module_name="chunk0.module.decoder.layers.0.mlp.experts.linear_fc1",
+ inputs=(),
+ output=output,
+ module=module,
+ row_count=2,
+ )
+
+ assert uid_span == 16
+ assert row_uids is not None
+ assert torch.equal(row_uids, torch.tensor([4, 7]))
+
+
+def test_forward_trace_restores_dense_non_cp_sequence_uids_before_sorting() -> None:
+ trace: dict[str, list[dict[str, Any]]] = {
+ "chunk0.module.decoder.layers.0.self_attention": [
+ {
+ "primary_output": torch.tensor([[1.0], [2.0], [3.0]]),
+ "row_token_uids": torch.tensor([10, 30, 20]),
+ "rank_meta": [
+ {"cp_world_size": 1, "tp_rank": 0},
+ {"cp_world_size": 1, "tp_rank": 1},
+ ],
+ }
+ ]
+ }
+
+ ForwardTraceCapture.canonicalize_trace(trace)
+
+ call = trace["chunk0.module.decoder.layers.0.self_attention"][0]
+ assert torch.equal(call["row_token_uids"], torch.tensor([0, 1, 2]))
+ assert torch.equal(call["primary_output"], torch.tensor([[1.0], [2.0], [3.0]]))
+
+
def test_forward_trace_extracts_empty_router_topk_with_config_hint() -> None:
topk = _extract_router_topk(
(
@@ -353,6 +412,44 @@ def test_forward_trace_canonicalizes_row_outputs_by_token_uid() -> None:
)
+def test_forward_trace_canonicalizes_router_outputs_without_output_attrs() -> None:
+ module = type("RouterWithDenseTraceUids", (), {})()
+ probs = torch.tensor([[3.0], [1.0], [2.0]])
+ routing_map = torch.tensor([[True], [False], [True]])
+ setattr(module, "_art_trace_row_token_uids", torch.tensor([3, 1, 2]))
+
+ row_uids, _uid_span = ForwardTraceCapture._row_token_uids_for_capture(
+ module_name="chunk0.module.decoder.layers.0.mlp.router",
+ inputs=(),
+ output=(probs, routing_map),
+ module=module,
+ row_count=3,
+ )
+ assert row_uids is not None
+
+ trace: dict[str, list[dict[str, Any]]] = {
+ "chunk0.module.decoder.layers.0.mlp.router": [
+ {
+ "primary_output": probs,
+ "router_topk_scores": probs,
+ "router_topk_ids": torch.tensor([[3], [1], [2]]),
+ "output": {"probs": probs, "routing_map": routing_map},
+ "row_token_uids": row_uids,
+ }
+ ]
+ }
+
+ ForwardTraceCapture.canonicalize_trace(trace)
+
+ call = trace["chunk0.module.decoder.layers.0.mlp.router"][0]
+ assert torch.equal(call["row_token_uids"], torch.tensor([1, 2, 3]))
+ assert torch.equal(call["output"]["probs"], torch.tensor([[1.0], [2.0], [3.0]]))
+ assert torch.equal(
+ call["output"]["routing_map"],
+ torch.tensor([[False], [True], [True]]),
+ )
+
+
def test_forward_trace_expands_attention_output_uids_for_out_norm_heads() -> None:
trace: dict[str, list[dict[str, Any]]] = {
"chunk0.module.decoder.layers.0.self_attention": [
diff --git a/tests/integration/megatron/model_support/test_provider_support.py b/tests/integration/megatron/model_support/test_provider_support.py
index ae1098829..f8cd668d9 100644
--- a/tests/integration/megatron/model_support/test_provider_support.py
+++ b/tests/integration/megatron/model_support/test_provider_support.py
@@ -17,6 +17,7 @@
UnsupportedModelArchitectureError,
get_model_support_handler,
get_model_support_spec,
+ model_requires_merged_rollout,
model_uses_expert_parallel,
)
import art.megatron.provider as provider_module
@@ -130,6 +131,37 @@ def test_model_support_specs_own_moe_metadata() -> None:
assert model_uses_expert_parallel("OpenPipe/Qwen3-14B-Instruct") is False
assert model_uses_expert_parallel("Qwen/Qwen3-30B-A3B-Instruct-2507") is True
assert model_uses_expert_parallel("Qwen/Qwen3.5-35B-A3B") is True
+ assert model_uses_expert_parallel("deepseek-ai/DeepSeek-V4-Flash") is True
+
+
+def test_dsv4_prefers_validated_native_lora_rollout() -> None:
+ spec = get_model_support_spec("deepseek-ai/DeepSeek-V4-Flash")
+
+ assert spec.native_vllm_lora_status == "validated"
+ assert model_requires_merged_rollout("deepseek-ai/DeepSeek-V4-Flash") is False
+
+
+def test_dsv4_provider_disables_shared_expert_overlap(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ provider = _FakeProvider()
+ provider.num_moe_experts = 256
+ provider.moe_shared_expert_overlap = True
+ fake_bridge = _FakeBridge(
+ model_bridge=object(),
+ provider=provider,
+ )
+ monkeypatch.setattr(
+ provider_module.AutoBridge,
+ "from_hf_pretrained",
+ lambda *args, **kwargs: fake_bridge,
+ )
+ monkeypatch.setattr(provider_module.torch.cuda, "device_count", lambda: 2)
+
+ resolved = provider_module.get_provider("deepseek-ai/DeepSeek-V4-Flash")
+
+ assert resolved.moe_shared_expert_overlap is False
+ assert resolved.context_parallel_size == 1
def test_megatron_lora_rank_defaults_by_architecture() -> None:
@@ -418,6 +450,52 @@ def test_get_provider_bundle_honors_context_parallel_env_topology(
)
+def test_cp_unsupported_handler_defaults_to_tensor_parallel(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ provider = _FakeProvider()
+ provider.num_moe_experts = 8
+ fake_bridge = _FakeBridge(
+ model_bridge=object(),
+ provider=provider,
+ )
+ monkeypatch.setattr(
+ provider_module.AutoBridge,
+ "from_hf_pretrained",
+ lambda *args, **kwargs: fake_bridge,
+ )
+ monkeypatch.setattr(provider_module.torch.cuda, "device_count", lambda: 4)
+
+ bundle = provider_module.prepare_provider_bundle("deepseek-ai/DeepSeek-V4-Flash")
+ resolved = bundle.provider
+
+ assert resolved.tensor_model_parallel_size == 4
+ assert resolved.context_parallel_size == 1
+ assert resolved.expert_model_parallel_size == 4
+ assert resolved.sequence_parallel is True
+
+
+def test_cp_unsupported_handler_rejects_context_parallel_env(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ provider = _FakeProvider()
+ provider.num_moe_experts = 8
+ fake_bridge = _FakeBridge(
+ model_bridge=object(),
+ provider=provider,
+ )
+ monkeypatch.setattr(
+ provider_module.AutoBridge,
+ "from_hf_pretrained",
+ lambda *args, **kwargs: fake_bridge,
+ )
+ monkeypatch.setattr(provider_module.torch.cuda, "device_count", lambda: 4)
+ monkeypatch.setenv("ART_MEGATRON_CONTEXT_PARALLEL_SIZE", "2")
+
+ with pytest.raises(RuntimeError, match="does not implement context parallelism"):
+ provider_module.prepare_provider_bundle("deepseek-ai/DeepSeek-V4-Flash")
+
+
def test_qwen35_handler_keeps_standard_attention_on_flex_under_cp(
monkeypatch: pytest.MonkeyPatch,
) -> None:
diff --git a/tests/integration/megatron/model_support/test_workflow.py b/tests/integration/megatron/model_support/test_workflow.py
index 4f918ffcf..118fef764 100644
--- a/tests/integration/megatron/model_support/test_workflow.py
+++ b/tests/integration/megatron/model_support/test_workflow.py
@@ -1,5 +1,6 @@
import os
from types import SimpleNamespace
+from typing import cast
import pytest
@@ -29,6 +30,11 @@
run_yes_no_trainability_stage,
validated_architecture_representative_models,
)
+from .workflow_resources import (
+ _h200_equivalent_slots_for_total_gib,
+ handler_workflow_resources_for_base_model,
+ resolve_stage_resources_for_visible_gpus,
+)
@pytest.fixture(autouse=True)
@@ -70,10 +76,88 @@ def test_validated_architecture_representative_models_are_fixed() -> None:
"Qwen/Qwen3.5-27B",
"google/gemma-4-26B-A4B-it",
"google/gemma-4-31B-it",
+ "deepseek-ai/DeepSeek-V4-Flash",
"openai/gpt-oss-20b",
]
+def test_dsv4_runtime_stages_use_full_model_resources() -> None:
+ resources = handler_workflow_resources_for_base_model(
+ "deepseek-ai/DeepSeek-V4-Flash"
+ )
+ assert resources is not None
+ for stage in (
+ resources.train_inf_mismatch,
+ resources.yes_no_trainability,
+ resources.length_trainability,
+ ):
+ assert stage is not None
+ assert stage.required_world_size == 8
+ assert stage.requires_external_vllm is True
+ assert stage.megatron is not None
+ assert stage.megatron.gpu_ids == [0, 1, 2, 3, 4, 5, 6, 7]
+ assert stage.megatron.topology.tp == 2
+ assert stage.megatron.topology.ep == 8
+ assert stage.megatron.topology.cp == 1
+ assert stage.vllm is not None
+ assert stage.vllm.gpu_ids == [4, 5, 6, 7]
+ engine_args = stage.vllm.engine_args()
+ assert "hf_overrides" not in engine_args
+ assert engine_args.get("load_format") != "dummy"
+ assert engine_args["moe_backend"] == "triton_unfused"
+ assert engine_args["kv_cache_dtype"] == "fp8"
+ assert stage.megatron_env == {"ART_MEGATRON_STREAMING_WEIGHT_OFFLOAD": "1"}
+
+ for stage in (resources.merged_vllm_serving, resources.native_vllm_lora):
+ assert stage is not None
+ assert stage.vllm is not None
+ engine_args = stage.vllm.engine_args()
+ assert engine_args["load_format"] == "dummy"
+ hf_overrides = cast(dict[str, object], engine_args["hf_overrides"])
+ assert hf_overrides["num_hidden_layers"] == 4
+ assert resources.merged_vllm_serving is not None
+ assert resources.merged_vllm_serving.vllm is not None
+ assert resources.merged_vllm_serving.vllm.engine_args()["kv_cache_dtype"] == "fp8"
+ assert resources.native_vllm_lora is not None
+ assert resources.native_vllm_lora.vllm is not None
+ assert resources.native_vllm_lora.vllm.engine_args().get("max_loras", 2) == 2
+
+
+def test_dsv4_resources_remap_to_four_high_vram_gpus(monkeypatch) -> None:
+ resources = handler_workflow_resources_for_base_model(
+ "deepseek-ai/DeepSeek-V4-Flash"
+ )
+ assert resources is not None
+ assert resources.train_inf_mismatch is not None
+ monkeypatch.setattr(
+ "tests.integration.megatron.model_support.workflow_resources."
+ "_visible_h200_equivalent_gpus",
+ lambda *, visible_gpu_count: 8,
+ )
+
+ stage = resolve_stage_resources_for_visible_gpus(
+ "train_inf_mismatch",
+ resources.train_inf_mismatch,
+ visible_gpu_count=4,
+ )
+
+ assert stage.megatron is not None
+ assert stage.vllm is not None
+ assert stage.megatron.gpu_ids == [0, 1]
+ assert stage.megatron.topology.tp == 2
+ assert stage.megatron.topology.ep == 2
+ assert stage.vllm.gpu_ids == [2, 3]
+ assert stage.vllm.tensor_parallel_size == 2
+ assert stage.vllm.engine_args()["moe_backend"] == "triton_unfused"
+ assert stage.vllm.engine_args()["kv_cache_dtype"] == "fp8"
+
+
+def test_h200_equivalent_slots_tolerate_reported_gb300_vram() -> None:
+ assert _h200_equivalent_slots_for_total_gib(80.0) == 0
+ assert _h200_equivalent_slots_for_total_gib(139.0) == 1
+ assert _h200_equivalent_slots_for_total_gib(276.6) == 2
+
+
def test_inspect_architecture_for_workflow_uses_minimal_topology(monkeypatch) -> None:
seen_env: dict[str, str | None] = {}
@@ -628,7 +712,7 @@ def test_run_correctness_sensitivity_stage_runs_dense_models(monkeypatch) -> Non
case_configs: list[SimpleNamespace] = []
oracle_module = SimpleNamespace(
OracleCaseConfig=lambda **kwargs: SimpleNamespace(**kwargs),
- selected_suite_topologies=lambda *, is_moe: [
+ selected_suite_topologies=lambda *, is_moe, cp_supported=True: [
SimpleNamespace(world_size=lambda: 1, slug=lambda: "tp1"),
SimpleNamespace(world_size=lambda: 2, slug=lambda: "tp2"),
SimpleNamespace(world_size=lambda: 2, slug=lambda: "dp2"),
@@ -643,7 +727,7 @@ def test_run_correctness_sensitivity_stage_runs_dense_models(monkeypatch) -> Non
world_size=lambda: 2
),
available_gpu_count=lambda: 4,
- run_suite=lambda case_config, max_world_size: (
+ run_suite=lambda case_config, max_world_size, cp_supported=True, **kwargs: (
case_configs.append(case_config)
or [
SimpleNamespace(
@@ -991,7 +1075,7 @@ def test_run_correctness_sensitivity_stage_summarizes_reports(monkeypatch) -> No
)
oracle_module = SimpleNamespace(
OracleCaseConfig=lambda **kwargs: SimpleNamespace(**kwargs),
- selected_suite_topologies=lambda *, is_moe: [
+ selected_suite_topologies=lambda *, is_moe, cp_supported=True: [
SimpleNamespace(world_size=lambda: 1, slug=lambda: "tp1"),
SimpleNamespace(world_size=lambda: 2, slug=lambda: "tp2"),
],
@@ -1004,7 +1088,7 @@ def test_run_correctness_sensitivity_stage_summarizes_reports(monkeypatch) -> No
world_size=lambda: 2
),
available_gpu_count=lambda: 2,
- run_suite=lambda case_config, max_world_size: [
+ run_suite=lambda case_config, max_world_size, cp_supported=True, **kwargs: [
SimpleNamespace(
variant="sft_topology_tp2",
topology="tp2",
@@ -1051,6 +1135,70 @@ def test_run_correctness_sensitivity_stage_summarizes_reports(monkeypatch) -> No
assert stage.artifact_dir == "/tmp/oracle"
+def test_run_correctness_sensitivity_stage_uses_dsv4_real_path_config(
+ monkeypatch,
+) -> None:
+ architecture = ArchitectureReport(
+ base_model="deepseek-ai/DeepSeek-V4-Flash",
+ model_key="dsv4",
+ handler_key="dsv4",
+ layer_families=[LayerFamilyInstance(key="dsv4_attention", layer_index=0)],
+ recommended_min_layers=4,
+ )
+ captured: dict[str, object] = {}
+ oracle_module = SimpleNamespace(
+ OracleCaseConfig=lambda **kwargs: SimpleNamespace(**kwargs),
+ MetricThresholdRule=lambda **kwargs: SimpleNamespace(**kwargs),
+ selected_suite_topologies=lambda *, is_moe, cp_supported=True: [
+ SimpleNamespace(world_size=lambda: 1, slug=lambda: "tp1"),
+ SimpleNamespace(world_size=lambda: 2, slug=lambda: "tp2"),
+ ],
+ oracle_topology=lambda *, is_moe: SimpleNamespace(world_size=lambda: 1),
+ selected_oracle_objectives=lambda: ["rl"],
+ supported_sensitivity_mutations_for_objective=lambda objective, *, is_moe: [],
+ sensitivity_topology_for_mutation=lambda mutation, *, is_moe: SimpleNamespace(
+ world_size=lambda: 2
+ ),
+ available_gpu_count=lambda: 2,
+ run_suite=lambda case_config, **kwargs: (
+ captured.update(case_config=case_config, suite_kwargs=kwargs)
+ or [
+ SimpleNamespace(
+ variant="rl_topology_tp2",
+ topology="tp2",
+ signal="pass",
+ fail_count=0,
+ )
+ ]
+ ),
+ run_sensitivity_suite=lambda case_config, mutations, max_world_size: [],
+ ensure_case_artifacts=lambda case_config: SimpleNamespace(
+ case_dir="/tmp/oracle"
+ ),
+ keep_topology_artifacts=lambda: False,
+ )
+ monkeypatch.setattr(
+ "tests.integration.megatron.model_support.workflow._import_integration_module",
+ lambda name: oracle_module,
+ )
+ monkeypatch.setenv(SKIP_SENSITIVITY_ENV, "1")
+
+ stage = run_correctness_sensitivity_stage(
+ base_model="deepseek-ai/DeepSeek-V4-Flash",
+ architecture=architecture,
+ )
+
+ case_config = captured["case_config"]
+ suite_kwargs = cast(dict[str, object], captured["suite_kwargs"])
+ phase_pass_fns = cast(dict[str, object], suite_kwargs["phase_pass_fns"])
+ assert getattr(case_config, "precision") == "bf16"
+ assert suite_kwargs["use_fp32_lora_reference"] is False
+ assert getattr(phase_pass_fns["forward"], "limits") == {"mean_abs_pct": 3.0}
+ assert getattr(phase_pass_fns["grads"], "limits") == {"mean_abs_pct": 5.0}
+ assert stage.metrics["precision"] == "bf16"
+ assert stage.metrics["use_fp32_lora_reference"] is False
+
+
def test_run_correctness_sensitivity_stage_can_skip_sensitivity_only(
monkeypatch,
) -> None:
@@ -1063,7 +1211,7 @@ def test_run_correctness_sensitivity_stage_can_skip_sensitivity_only(
)
oracle_module = SimpleNamespace(
OracleCaseConfig=lambda **kwargs: SimpleNamespace(**kwargs),
- selected_suite_topologies=lambda *, is_moe: [
+ selected_suite_topologies=lambda *, is_moe, cp_supported=True: [
SimpleNamespace(world_size=lambda: 1, slug=lambda: "tp1"),
SimpleNamespace(world_size=lambda: 2, slug=lambda: "tp2"),
],
@@ -1076,7 +1224,7 @@ def test_run_correctness_sensitivity_stage_can_skip_sensitivity_only(
world_size=lambda: 4
),
available_gpu_count=lambda: 2,
- run_suite=lambda case_config, max_world_size: [
+ run_suite=lambda case_config, max_world_size, cp_supported=True, **kwargs: [
SimpleNamespace(
variant="sft_topology_tp2",
topology="tp2",
@@ -1154,8 +1302,7 @@ def _import_integration_module(name: str):
assert stage.name == "merged_vllm_serving"
assert stage.passed is True
- assert stage.metrics == {
- "base_model": "Qwen/Qwen3.5-35B-A3B",
- "served_model_name": "validation@0",
- }
+ assert stage.metrics["base_model"] == "Qwen/Qwen3.5-35B-A3B"
+ assert stage.metrics["served_model_name"] == "validation@0"
+ assert "readable_summary" in stage.metrics
assert stage.artifact_dir == "/tmp/merged-serving"
diff --git a/tests/integration/megatron/model_support/workflow.py b/tests/integration/megatron/model_support/workflow.py
index 0c9d6977a..0550d92e3 100644
--- a/tests/integration/megatron/model_support/workflow.py
+++ b/tests/integration/megatron/model_support/workflow.py
@@ -65,6 +65,7 @@
"qwen3_5_dense": "Qwen/Qwen3.5-27B",
"gemma4_moe": "google/gemma-4-26B-A4B-it",
"gemma4_dense": "google/gemma-4-31B-it",
+ "dsv4": "deepseek-ai/DeepSeek-V4-Flash",
"gpt_oss_moe": "openai/gpt-oss-20b",
}
SUBPROCESS_VALIDATION_STAGES = frozenset(
@@ -473,16 +474,23 @@ def run_correctness_sensitivity_stage(
allow_unvalidated_arch=allow_unvalidated_arch,
)
handler = get_model_support_handler_for_spec(spec)
+ cp_supported = bool(handler.cp_supported)
+ correctness_precision = handler.correctness_precision()
+ correctness_use_fp32_lora_reference = handler.correctness_use_fp32_lora_reference()
+ correctness_phase_pass_fns = handler.correctness_phase_pass_fns(oracle_harness)
case_config = oracle_harness.OracleCaseConfig(
base_model=base_model,
is_moe=handler.is_moe,
- precision="fp32",
+ precision=correctness_precision,
num_layers=max(1, architecture.recommended_min_layers),
num_steps=1,
allow_unvalidated_arch=allow_unvalidated_arch,
)
suite_topologies = list(
- oracle_harness.selected_suite_topologies(is_moe=handler.is_moe)
+ oracle_harness.selected_suite_topologies(
+ is_moe=handler.is_moe,
+ cp_supported=cp_supported,
+ )
)
objectives = list(oracle_harness.selected_oracle_objectives())
skip_sensitivity = _truthy_env(SKIP_SENSITIVITY_ENV)
@@ -526,6 +534,14 @@ def run_correctness_sensitivity_stage(
is_moe=handler.is_moe,
).world_size()
> max_world_size
+ or (
+ not cp_supported
+ and oracle_harness.sensitivity_topology_for_mutation(
+ mutation,
+ is_moe=handler.is_moe,
+ ).cp
+ > 1
+ )
]
mutations = [
mutation
@@ -539,6 +555,11 @@ def run_correctness_sensitivity_stage(
suite_reports = oracle_harness.run_suite(
case_config=case_config,
max_world_size=max_world_size,
+ cp_supported=cp_supported,
+ phase_pass_fns=correctness_phase_pass_fns,
+ use_fp32_lora_reference=correctness_use_fp32_lora_reference,
+ prune_reference_artifacts=skip_sensitivity or not mutations,
+ prune_case_artifacts=skip_sensitivity or not mutations,
)
sensitivity_reports = []
if skip_sensitivity:
@@ -572,7 +593,10 @@ def run_correctness_sensitivity_stage(
passed=True,
metrics={
"requested_num_layers": case_config.num_layers,
+ "precision": correctness_precision,
+ "use_fp32_lora_reference": correctness_use_fp32_lora_reference,
"is_moe": handler.is_moe,
+ "cp_supported": cp_supported,
"allow_unvalidated_arch": allow_unvalidated_arch,
"objectives": objectives,
"sensitivity_mutations": mutations,
@@ -644,14 +668,42 @@ def run_merged_vllm_serving_stage(
allow_unvalidated_arch=allow_unvalidated_arch,
)
report = merged_vllm_serving.run_merged_vllm_serving(case_config)
+ metrics = report.model_dump(mode="json")
+ warning_lines = _read_vllm_reload_warnings(report.output_dir)
+ metrics["vllm_reload_warning_count"] = len(warning_lines)
+ metrics["vllm_reload_warnings"] = warning_lines
+ metrics["readable_summary"] = _merged_vllm_serving_summary(metrics)
return ValidationStageResult(
name="merged_vllm_serving",
passed=bool(report.model_ids),
- metrics=report.model_dump(mode="json"),
+ metrics=metrics,
artifact_dir=report.output_dir,
)
+def _read_vllm_reload_warnings(output_dir: str) -> list[str]:
+ log_path = Path(output_dir) / "logs" / "vllm-runtime.log"
+ if not log_path.exists():
+ return []
+ return [
+ line.strip()
+ for line in log_path.read_text(encoding="utf-8", errors="replace").splitlines()
+ if "Failed to load weights" in line
+ ]
+
+
+def _merged_vllm_serving_summary(metrics: dict[str, Any]) -> list[str]:
+ lines = [
+ f"served_model_name={metrics.get('served_model_name', '')}",
+ f"model_ids={metrics.get('model_ids', [])}",
+ f"completion_text={metrics.get('completion_text', '')!r}",
+ f"vllm_reload_warning_count={metrics.get('vllm_reload_warning_count', 0)}",
+ ]
+ for warning in metrics.get("vllm_reload_warnings", []):
+ lines.append(f"vllm_reload_warning={warning}")
+ return lines
+
+
def run_chat_template_rollout_stage(
*,
base_model: str,
@@ -981,6 +1033,20 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
return args
+def _print_stage_result(stage: ValidationStageResult, *, indent: str = "") -> None:
+ status = "PASS" if stage.passed else "FAIL"
+ print(f"{indent}{stage.name}: {status}", flush=True)
+ child_indent = f"{indent} "
+ if stage.artifact_dir:
+ print(f"{child_indent}artifact_dir={stage.artifact_dir}", flush=True)
+ summary = stage.metrics.get("readable_summary")
+ if isinstance(summary, list):
+ for line in summary:
+ print(f"{child_indent}{line}", flush=True)
+ if not stage.passed:
+ print(f"{child_indent}metrics={stage.metrics}", flush=True)
+
+
def main(argv: list[str] | None = None) -> int:
args = _parse_args(argv)
if args.all_architectures:
@@ -996,12 +1062,7 @@ def main(argv: list[str] | None = None) -> int:
for report in all_report.reports:
print(f"base_model={report.base_model}", flush=True)
for stage in report.stages:
- status = "PASS" if stage.passed else "FAIL"
- print(f" {stage.name}: {status}", flush=True)
- if stage.artifact_dir:
- print(f" artifact_dir={stage.artifact_dir}", flush=True)
- if not stage.passed:
- print(f" metrics={stage.metrics}", flush=True)
+ _print_stage_result(stage, indent=" ")
print(f"report_json={args.output_json}", flush=True)
return 0 if all_report.passed else 1
report = build_validation_report(
@@ -1015,12 +1076,7 @@ def main(argv: list[str] | None = None) -> int:
allow_unvalidated_arch=args.allow_unsupported_arch,
)
for stage in report.stages:
- status = "PASS" if stage.passed else "FAIL"
- print(f"{stage.name}: {status}", flush=True)
- if stage.artifact_dir:
- print(f" artifact_dir={stage.artifact_dir}", flush=True)
- if not stage.passed:
- print(f" metrics={stage.metrics}", flush=True)
+ _print_stage_result(stage)
print(f"report_json={args.output_json}", flush=True)
return 0 if all(stage.passed for stage in report.stages) else 1
diff --git a/tests/integration/megatron/model_support/workflow_resources.py b/tests/integration/megatron/model_support/workflow_resources.py
new file mode 100644
index 000000000..7bc3091c6
--- /dev/null
+++ b/tests/integration/megatron/model_support/workflow_resources.py
@@ -0,0 +1,441 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+_H200_REFERENCE_VRAM_GIB = 140.0
+_H200_SLOT_TOLERANCE = 0.05
+
+
+class MegatronWorkflowTopology(BaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ tp: int = 1
+ ep: int = 1
+ etp: int = 1
+ dp: int = 1
+ cp: int = 1
+ pp: int = 1
+ sp: bool = False
+
+ def to_megatron_config(self) -> dict[str, int | None]:
+ return {
+ "tp": self.tp,
+ "ep": self.ep,
+ "etp": self.etp,
+ "cp": self.cp,
+ "pp": self.pp,
+ }
+
+ def to_oracle_topology_kwargs(self) -> dict[str, int | bool]:
+ return self.model_dump()
+
+ def to_train_inf_topology_kwargs(self) -> dict[str, int]:
+ return {
+ "tp": self.tp,
+ "ep": self.ep,
+ "etp": self.etp,
+ "dp": self.dp,
+ "cp": self.cp,
+ "pp": self.pp,
+ }
+
+
+class MegatronWorkflowResources(BaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ gpu_ids: list[int]
+ topology: MegatronWorkflowTopology
+
+
+class VllmWorkflowResources(BaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ gpu_ids: list[int]
+ tensor_parallel_size: int
+ enable_expert_parallel: bool = False
+ hf_overrides: dict[str, object] = Field(default_factory=dict)
+ extra_engine_args: dict[str, object] = Field(default_factory=dict)
+
+ def engine_args(self) -> dict[str, object]:
+ engine_args: dict[str, object] = {
+ "tensor_parallel_size": self.tensor_parallel_size,
+ }
+ if self.enable_expert_parallel:
+ engine_args["enable_expert_parallel"] = True
+ if self.hf_overrides:
+ engine_args["hf_overrides"] = dict(self.hf_overrides)
+ engine_args.update(self.extra_engine_args)
+ return engine_args
+
+
+class WorkflowStageResources(BaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ required_world_size: int
+ required_h200_equivalent_gpus: int | None = None
+ allow_gpu_overlap: bool = False
+ requires_external_vllm: bool = False
+ megatron: MegatronWorkflowResources | None = None
+ vllm: VllmWorkflowResources | None = None
+ high_vram_megatron: MegatronWorkflowResources | None = None
+ high_vram_vllm: VllmWorkflowResources | None = None
+ megatron_env: dict[str, str] = Field(default_factory=dict)
+
+
+class HandlerWorkflowResources(BaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ train_inf_mismatch: WorkflowStageResources | None = None
+ merged_vllm_serving: WorkflowStageResources | None = None
+ native_vllm_lora: WorkflowStageResources | None = None
+ yes_no_trainability: WorkflowStageResources | None = None
+ length_trainability: WorkflowStageResources | None = None
+ yes_no_trainability_variant: (
+ Literal[
+ "megatron_shared",
+ "megatron_dedicated",
+ "unsloth_dedicated",
+ ]
+ | None
+ ) = None
+
+
+_DSV4_TP2_EP8 = MegatronWorkflowTopology(
+ tp=2,
+ ep=8,
+ etp=1,
+ dp=4,
+ cp=1,
+ pp=1,
+ sp=True,
+)
+_DSV4_TP2_EP4 = MegatronWorkflowTopology(
+ tp=2,
+ ep=4,
+ etp=1,
+ dp=2,
+ cp=1,
+ pp=1,
+ sp=True,
+)
+_DSV4_TP2_EP2 = MegatronWorkflowTopology(
+ tp=2,
+ ep=2,
+ etp=1,
+ dp=1,
+ cp=1,
+ pp=1,
+ sp=True,
+)
+_DSV4_REPRESENTATIVE_NUM_LAYERS = 4
+_DSV4_REPRESENTATIVE_COMPRESS_RATIOS = [0, 0, 4, 128]
+_DSV4_REPRESENTATIVE_LAYER_TYPES = [
+ "sliding_attention",
+ "sliding_attention",
+ "compressed_sparse_attention",
+ "heavily_compressed_attention",
+]
+_DSV4_REPRESENTATIVE_MLP_LAYER_TYPES = ["hash_moe", "hash_moe", "hash_moe", "moe"]
+_DSV4_MEGATRON_ENV = {
+ "ART_DSV4_VALIDATION_NUM_LAYERS": str(_DSV4_REPRESENTATIVE_NUM_LAYERS)
+}
+_DSV4_STREAMING_OFFLOAD_ENV = {"ART_MEGATRON_STREAMING_WEIGHT_OFFLOAD": "1"}
+_DSV4_HF_OVERRIDES = {
+ "num_hidden_layers": _DSV4_REPRESENTATIVE_NUM_LAYERS,
+ "compress_ratios": _DSV4_REPRESENTATIVE_COMPRESS_RATIOS,
+ "layer_types": _DSV4_REPRESENTATIVE_LAYER_TYPES,
+ "mlp_layer_types": _DSV4_REPRESENTATIVE_MLP_LAYER_TYPES,
+}
+_DSV4_COMMON_VLLM_ENGINE_ARGS = {
+ "compilation_config": {
+ "cudagraph_mode": "NONE",
+ "pass_config": {"fuse_allreduce_rms": False},
+ },
+ "disable_custom_all_reduce": True,
+ "enforce_eager": True,
+ "gpu_memory_utilization": 0.82,
+ "kv_cache_dtype": "fp8",
+ "max_num_batched_tokens": 1032,
+}
+_DSV4_MERGED_VLLM_ENGINE_ARGS = {
+ **_DSV4_COMMON_VLLM_ENGINE_ARGS,
+ "moe_backend": "triton_unfused",
+}
+_DSV4_LORA_VLLM_ENGINE_ARGS = {
+ **_DSV4_COMMON_VLLM_ENGINE_ARGS,
+ "moe_backend": "triton_unfused",
+}
+_DSV4_REDUCED_VLLM_ENGINE_ARGS = {
+ **_DSV4_MERGED_VLLM_ENGINE_ARGS,
+ # The quick DSV4 vLLM serving gates use a reduced 4-layer validation model and then
+ # sync Megatron weights into vLLM through merged-weight transfer. Loading
+ # the full public checkpoint before that sync is incompatible with the
+ # reduced hf_overrides because vLLM still streams layer-4+ tensors.
+ "load_format": "dummy",
+}
+_DSV4_NATIVE_LORA_VLLM_ENGINE_ARGS = {
+ **_DSV4_LORA_VLLM_ENGINE_ARGS,
+ "load_format": "dummy",
+}
+_DSV4_MEGATRON = MegatronWorkflowResources(
+ gpu_ids=[0, 1, 2, 3, 4, 5, 6, 7],
+ topology=_DSV4_TP2_EP8,
+)
+_DSV4_HIGH_VRAM_MEGATRON = MegatronWorkflowResources(
+ gpu_ids=[0, 1],
+ topology=_DSV4_TP2_EP2,
+)
+_DSV4_FULL_VLLM_EP4 = VllmWorkflowResources(
+ gpu_ids=[4, 5, 6, 7],
+ tensor_parallel_size=4,
+ enable_expert_parallel=True,
+ extra_engine_args=_DSV4_LORA_VLLM_ENGINE_ARGS,
+)
+_DSV4_FULL_VLLM_EP2 = VllmWorkflowResources(
+ gpu_ids=[2, 3],
+ tensor_parallel_size=2,
+ enable_expert_parallel=True,
+ extra_engine_args=_DSV4_LORA_VLLM_ENGINE_ARGS,
+)
+_DSV4_REDUCED_VLLM_EP4 = VllmWorkflowResources(
+ gpu_ids=[4, 5, 6, 7],
+ tensor_parallel_size=4,
+ enable_expert_parallel=True,
+ hf_overrides=_DSV4_HF_OVERRIDES,
+ extra_engine_args=_DSV4_REDUCED_VLLM_ENGINE_ARGS,
+)
+_DSV4_REDUCED_VLLM_EP2 = VllmWorkflowResources(
+ gpu_ids=[2, 3],
+ tensor_parallel_size=2,
+ enable_expert_parallel=True,
+ hf_overrides=_DSV4_HF_OVERRIDES,
+ extra_engine_args=_DSV4_REDUCED_VLLM_ENGINE_ARGS,
+)
+_DSV4_REDUCED_NATIVE_VLLM_EP4 = VllmWorkflowResources(
+ gpu_ids=[0, 1, 2, 3],
+ tensor_parallel_size=4,
+ enable_expert_parallel=True,
+ hf_overrides=_DSV4_HF_OVERRIDES,
+ extra_engine_args=_DSV4_NATIVE_LORA_VLLM_ENGINE_ARGS,
+)
+
+# Explicitly for large models which do not fit in the default topology.
+HANDLER_WORKFLOW_RESOURCES: dict[str, HandlerWorkflowResources] = {
+ "dsv4": HandlerWorkflowResources(
+ train_inf_mismatch=WorkflowStageResources(
+ required_world_size=8,
+ required_h200_equivalent_gpus=8,
+ requires_external_vllm=True,
+ megatron=_DSV4_MEGATRON,
+ vllm=_DSV4_FULL_VLLM_EP4,
+ high_vram_megatron=_DSV4_HIGH_VRAM_MEGATRON,
+ high_vram_vllm=_DSV4_FULL_VLLM_EP2,
+ megatron_env=_DSV4_STREAMING_OFFLOAD_ENV,
+ ),
+ merged_vllm_serving=WorkflowStageResources(
+ required_world_size=8,
+ required_h200_equivalent_gpus=8,
+ megatron=_DSV4_MEGATRON,
+ vllm=_DSV4_REDUCED_VLLM_EP4,
+ high_vram_megatron=_DSV4_HIGH_VRAM_MEGATRON,
+ high_vram_vllm=_DSV4_REDUCED_VLLM_EP2,
+ megatron_env=_DSV4_MEGATRON_ENV,
+ ),
+ native_vllm_lora=WorkflowStageResources(
+ required_world_size=4,
+ vllm=_DSV4_REDUCED_NATIVE_VLLM_EP4,
+ ),
+ yes_no_trainability=WorkflowStageResources(
+ required_world_size=8,
+ required_h200_equivalent_gpus=8,
+ requires_external_vllm=True,
+ megatron=_DSV4_MEGATRON,
+ vllm=_DSV4_FULL_VLLM_EP4,
+ high_vram_megatron=_DSV4_HIGH_VRAM_MEGATRON,
+ high_vram_vllm=_DSV4_FULL_VLLM_EP2,
+ megatron_env=_DSV4_STREAMING_OFFLOAD_ENV,
+ ),
+ length_trainability=WorkflowStageResources(
+ required_world_size=8,
+ required_h200_equivalent_gpus=8,
+ requires_external_vllm=True,
+ megatron=_DSV4_MEGATRON,
+ vllm=_DSV4_FULL_VLLM_EP4,
+ high_vram_megatron=_DSV4_HIGH_VRAM_MEGATRON,
+ high_vram_vllm=_DSV4_FULL_VLLM_EP2,
+ megatron_env=_DSV4_STREAMING_OFFLOAD_ENV,
+ ),
+ yes_no_trainability_variant="megatron_dedicated",
+ ),
+}
+
+
+def handler_workflow_resources_for_base_model(
+ base_model: str,
+ *,
+ allow_unvalidated_arch: bool = False,
+) -> HandlerWorkflowResources | None:
+ from art.megatron.model_support.registry import get_model_support_spec
+
+ spec = get_model_support_spec(
+ base_model,
+ allow_unvalidated_arch=allow_unvalidated_arch,
+ )
+ return HANDLER_WORKFLOW_RESOURCES.get(spec.handler_key)
+
+
+def _h200_equivalent_slots_for_total_gib(total_gib: float) -> int:
+ return max(0, int(total_gib / _H200_REFERENCE_VRAM_GIB + _H200_SLOT_TOLERANCE))
+
+
+def _visible_h200_equivalent_gpus(*, visible_gpu_count: int) -> int:
+ try:
+ import torch
+ except ImportError:
+ return 0
+ if not torch.cuda.is_available():
+ return 0
+ equivalent = 0
+ for device_index in range(visible_gpu_count):
+ props = torch.cuda.get_device_properties(device_index)
+ total_gib = float(props.total_memory) / (1024**3)
+ equivalent += _h200_equivalent_slots_for_total_gib(total_gib)
+ return equivalent
+
+
+def _remap_gpu_ids_to_visible(
+ gpu_ids: list[int], *, visible_gpu_count: int
+) -> list[int]:
+ if all(0 <= gpu_id < visible_gpu_count for gpu_id in gpu_ids):
+ return list(gpu_ids)
+ if len(gpu_ids) > visible_gpu_count:
+ raise RuntimeError(
+ "Cannot remap workflow GPU ids to visible high-VRAM devices: "
+ f"gpu_ids={gpu_ids}, visible_gpu_count={visible_gpu_count}"
+ )
+ return list(range(len(gpu_ids)))
+
+
+def _validate_gpu_ids_visible(gpu_ids: list[int], *, visible_gpu_count: int) -> None:
+ invalid = [
+ gpu_id for gpu_id in gpu_ids if gpu_id < 0 or gpu_id >= visible_gpu_count
+ ]
+ if invalid:
+ raise RuntimeError(
+ f"Workflow GPU ids {gpu_ids} are not visible on host with "
+ f"{visible_gpu_count} GPUs"
+ )
+
+
+def resolve_stage_resources_for_visible_gpus(
+ stage_name: str,
+ stage_resources: WorkflowStageResources,
+ *,
+ visible_gpu_count: int,
+) -> WorkflowStageResources:
+ if visible_gpu_count >= stage_resources.required_world_size:
+ return stage_resources
+ required_equivalent = stage_resources.required_h200_equivalent_gpus
+ available_equivalent = _visible_h200_equivalent_gpus(
+ visible_gpu_count=visible_gpu_count
+ )
+ if required_equivalent is None or available_equivalent < required_equivalent:
+ raise RuntimeError(
+ f"Need {stage_resources.required_world_size} visible GPUs for "
+ f"{stage_name}, found {visible_gpu_count}. High-VRAM remapping "
+ f"requires {required_equivalent or stage_resources.required_world_size} "
+ f"H200-equivalent GPUs, found {available_equivalent}."
+ )
+ if (
+ stage_resources.high_vram_megatron is not None
+ or stage_resources.high_vram_vllm is not None
+ ):
+ megatron = stage_resources.high_vram_megatron or stage_resources.megatron
+ vllm = stage_resources.high_vram_vllm or stage_resources.vllm
+ if megatron is not None:
+ _validate_gpu_ids_visible(
+ megatron.gpu_ids,
+ visible_gpu_count=visible_gpu_count,
+ )
+ if vllm is not None:
+ _validate_gpu_ids_visible(
+ vllm.gpu_ids,
+ visible_gpu_count=visible_gpu_count,
+ )
+ return stage_resources.model_copy(update={"megatron": megatron, "vllm": vllm})
+ if not stage_resources.allow_gpu_overlap:
+ raise RuntimeError(
+ f"Need {stage_resources.required_world_size} visible GPUs for "
+ f"{stage_name}, found {visible_gpu_count}. No high-VRAM resource "
+ "override is configured for this stage."
+ )
+ megatron = stage_resources.megatron
+ if megatron is not None:
+ megatron = megatron.model_copy(
+ update={
+ "gpu_ids": _remap_gpu_ids_to_visible(
+ megatron.gpu_ids,
+ visible_gpu_count=visible_gpu_count,
+ )
+ }
+ )
+ vllm = stage_resources.vllm
+ if vllm is not None:
+ vllm = vllm.model_copy(
+ update={
+ "gpu_ids": _remap_gpu_ids_to_visible(
+ vllm.gpu_ids,
+ visible_gpu_count=visible_gpu_count,
+ )
+ }
+ )
+ return stage_resources.model_copy(update={"megatron": megatron, "vllm": vllm})
+
+
+def resolve_stage_resources_for_current_host(
+ stage_name: str,
+ stage_resources: WorkflowStageResources,
+) -> WorkflowStageResources:
+ try:
+ import torch
+ except ImportError:
+ visible_gpu_count = 0
+ else:
+ visible_gpu_count = int(torch.cuda.device_count())
+ return resolve_stage_resources_for_visible_gpus(
+ stage_name,
+ stage_resources,
+ visible_gpu_count=visible_gpu_count,
+ )
+
+
+def validate_visible_gpu_count(
+ stage_name: str,
+ stage_resources: WorkflowStageResources,
+ *,
+ visible_gpu_count: int,
+) -> None:
+ if visible_gpu_count < stage_resources.required_world_size:
+ raise RuntimeError(
+ f"Need {stage_resources.required_world_size} visible GPUs for "
+ f"{stage_name}, found {visible_gpu_count}"
+ )
+
+
+def validate_dedicated_test_resources(
+ *,
+ stage_name: str,
+ trainer_gpu_ids: list[int],
+ inference_gpu_ids: list[int],
+ allow_overlap: bool = False,
+) -> None:
+ if not trainer_gpu_ids:
+ raise RuntimeError(f"{stage_name} trainer GPU ids must be non-empty")
+ if not inference_gpu_ids:
+ raise RuntimeError(f"{stage_name} inference GPU ids must be non-empty")
+ if not allow_overlap and set(trainer_gpu_ids) & set(inference_gpu_ids):
+ raise RuntimeError(
+ f"{stage_name} trainer and inference GPU ids must not overlap"
+ )
diff --git a/tests/integration/megatron/routing_replay/trace.py b/tests/integration/megatron/routing_replay/trace.py
index 4538df56a..815b99e61 100644
--- a/tests/integration/megatron/routing_replay/trace.py
+++ b/tests/integration/megatron/routing_replay/trace.py
@@ -1,6 +1,5 @@
from __future__ import annotations
-import types
from typing import Any, Callable
from megatron.core.tensor_parallel import (
@@ -21,6 +20,11 @@ def _active_controller() -> Any | None:
return _CONTROLLER_GETTER()
+def _trace_hook(fn: Callable[..., Any]) -> Callable[..., Any]:
+ """Keep correctness-only routing trace monkeypatches out of Dynamo graphs."""
+ return torch.compiler.disable(fn)
+
+
@torch._dynamo.disable
def _dispatcher_local_token_uids(
controller: Any,
@@ -137,93 +141,6 @@ def _attach_trace_row_uids(
setattr(target, TRACE_UID_SPAN_ATTR, uid_span)
-def _attach_trace_row_uids_recursive(
- target: Any,
- *,
- row_token_uids: torch.Tensor,
- uid_span: int | None,
-) -> None:
- if isinstance(target, torch.Tensor):
- _attach_trace_row_uids(
- target,
- row_token_uids=row_token_uids,
- uid_span=uid_span,
- )
- return
- if isinstance(target, dict):
- for value in target.values():
- _attach_trace_row_uids_recursive(
- value,
- row_token_uids=row_token_uids,
- uid_span=uid_span,
- )
- return
- if isinstance(target, (list, tuple)):
- for value in target:
- _attach_trace_row_uids_recursive(
- value,
- row_token_uids=row_token_uids,
- uid_span=uid_span,
- )
-
-
-@torch._dynamo.disable
-def _attach_router_output_trace_row_uids(
- *,
- controller: Any,
- router_key: str,
- output: Any,
-) -> None:
- target_key = getattr(controller, "_router_prepared_target_keys", {}).get(router_key)
- if target_key is None:
- return
- token_uid_key, _call_index = target_key
- token_uids = getattr(controller, "_prepared_uid_sets", {}).get(token_uid_key)
- binding = getattr(controller, "_router_bindings", {}).get(router_key)
- if token_uids is None or binding is None:
- return
- router_token_uids = controller._token_uids_for_router_binding(
- token_uids,
- sequence_parallel=bool(binding["sequence_parallel"]),
- )
- _attach_trace_row_uids_recursive(
- output,
- row_token_uids=router_token_uids,
- uid_span=getattr(controller, "_global_uid_count", None),
- )
-
-
-def _install_router_output_trace_hooks(controller: Any | None) -> None:
- if controller is None:
- return
- for router_key, binding in getattr(controller, "_router_bindings", {}).items():
- module = binding.get("module")
- if module is None or getattr(module, "_art_oracle_router_trace_patched", False):
- continue
- original_routing = module.routing
-
- def patched_routing(
- router_module: Any,
- *args: Any,
- _original_routing: Any = original_routing,
- _router_key: str = router_key,
- **kwargs: Any,
- ) -> Any:
- del router_module
- output = _original_routing(*args, **kwargs)
- controller = _active_controller()
- if controller is not None:
- _attach_router_output_trace_row_uids(
- controller=controller,
- router_key=_router_key,
- output=output,
- )
- return output
-
- module.routing = types.MethodType(patched_routing, module)
- setattr(module, "_art_oracle_router_trace_patched", True)
-
-
@torch._dynamo.disable
def _propagate_grouped_mlp_trace_row_uids(source: Any, linear_fc2: Any) -> None:
row_token_uids, uid_span = _trace_row_uids_from_source(source)
@@ -415,7 +332,6 @@ def install_moe_routing_trace_hooks(
) -> None:
global _CONTROLLER_GETTER
_CONTROLLER_GETTER = controller_getter
- _install_router_output_trace_hooks(_active_controller())
try:
from megatron.core.transformer.moe.experts import TEGroupedMLP
from megatron.core.transformer.moe.token_dispatcher import (
@@ -609,23 +525,27 @@ def patched_fc2_forward(
)
return original_fc2_forward(self, x, tokens_per_expert)
- setattr(MoEAlltoAllTokenDispatcher, "preprocess", patched_preprocess)
+ setattr(MoEAlltoAllTokenDispatcher, "preprocess", _trace_hook(patched_preprocess))
setattr(
MoEAlltoAllTokenDispatcher,
"dispatch_preprocess",
- patched_dispatch_preprocess,
+ _trace_hook(patched_dispatch_preprocess),
+ )
+ setattr(
+ MoEAlltoAllTokenDispatcher,
+ "token_dispatch",
+ _trace_hook(patched_token_dispatch),
)
- setattr(MoEAlltoAllTokenDispatcher, "token_dispatch", patched_token_dispatch)
setattr(
MoEAlltoAllTokenDispatcher,
"dispatch_postprocess",
- patched_dispatch_postprocess,
+ _trace_hook(patched_dispatch_postprocess),
)
setattr(
MoEAlltoAllTokenDispatcher,
"combine_preprocess",
- patched_combine_preprocess,
+ _trace_hook(patched_combine_preprocess),
)
- setattr(TEGroupedMLP, "forward", patched_te_grouped_mlp_forward)
- setattr(MLPExpertsLinearFC2LoRA, "forward", patched_fc2_forward)
+ setattr(TEGroupedMLP, "forward", _trace_hook(patched_te_grouped_mlp_forward))
+ setattr(MLPExpertsLinearFC2LoRA, "forward", _trace_hook(patched_fc2_forward))
setattr(MoEAlltoAllTokenDispatcher, "_art_oracle_trace_patched", True)
diff --git a/tests/integration/megatron/runtime_isolation/test_runtime_launcher.py b/tests/integration/megatron/runtime_isolation/test_runtime_launcher.py
index b9bca13fe..65103ce3f 100644
--- a/tests/integration/megatron/runtime_isolation/test_runtime_launcher.py
+++ b/tests/integration/megatron/runtime_isolation/test_runtime_launcher.py
@@ -68,6 +68,51 @@ def test_build_runtime_server_cmd_honors_runtime_bin_override(monkeypatch) -> No
assert command[:2] == ["/opt/art/bin/runtime", "--wrapped"]
+def test_build_runtime_server_cmd_allows_lora_without_initial_adapter(
+ monkeypatch,
+ tmp_path: Path,
+) -> None:
+ monkeypatch.delenv("ART_VLLM_RUNTIME_BIN", raising=False)
+ runtime_root = tmp_path / "custom-runtime"
+ runtime_bin = runtime_root / ".venv" / "bin" / "art-vllm-runtime-server"
+ runtime_bin.parent.mkdir(parents=True, exist_ok=True)
+ runtime_bin.write_text("#!/bin/sh\n", encoding="ascii")
+ monkeypatch.setenv("ART_VLLM_RUNTIME_PROJECT_ROOT", str(runtime_root))
+
+ command = runtime.build_vllm_runtime_server_cmd(
+ runtime.VllmRuntimeLaunchConfig(
+ base_model="Qwen/Qwen3-14B",
+ port=8000,
+ host="0.0.0.0",
+ cuda_visible_devices="0,1",
+ served_model_name="test@0",
+ rollout_weights_mode="lora",
+ )
+ )
+
+ assert command[0] == str(runtime_bin)
+ assert not any(arg.startswith("--lora-path=") for arg in command)
+ assert "--rollout-weights-mode=lora" in command
+
+
+def test_external_checkpoint_path_mapping() -> None:
+ config = {
+ "vllm_runtime": {
+ "mode": "external",
+ "server_url": "http://inference:8000",
+ "local_checkpoint_root": "/mnt/ws_pvc/ws",
+ "server_checkpoint_root": "/remote/ws",
+ }
+ }
+
+ mapped = runtime.map_checkpoint_path_for_vllm(
+ config,
+ "/mnt/ws_pvc/ws/projects/art/.art/models/model/0001",
+ )
+
+ assert mapped == "/remote/ws/projects/art/.art/models/model/0001"
+
+
def test_get_vllm_runtime_nccl_so_path_queries_runtime_python(
monkeypatch,
tmp_path: Path,
@@ -90,7 +135,7 @@ def fake_run(command, *, capture_output: bool, text: bool):
return SimpleNamespace(returncode=0, stdout=f"{nccl_so_path}\n", stderr="")
monkeypatch.setenv("ART_VLLM_RUNTIME_PROJECT_ROOT", str(runtime_root))
- monkeypatch.setattr(runtime.subprocess, "run", fake_run)
+ monkeypatch.setattr(runtime, "subprocess", SimpleNamespace(run=fake_run))
assert runtime.get_vllm_runtime_nccl_so_path() == nccl_so_path.resolve()
command = seen["command"]
@@ -100,6 +145,61 @@ def fake_run(command, *, capture_output: bool, text: bool):
assert seen["text"] is True
+def test_vllm_runtime_subprocess_env_isolates_flashinfer_for_source_runtime(
+ monkeypatch,
+ tmp_path: Path,
+) -> None:
+ runtime_root = tmp_path / "vllm_runtime"
+ runtime_root.mkdir()
+ monkeypatch.setenv("ART_VLLM_RUNTIME_PROJECT_ROOT", str(runtime_root))
+ monkeypatch.setenv("FLASHINFER_WORKSPACE_BASE", "/shared/flashinfer")
+ monkeypatch.setenv(
+ "PYTHONPATH",
+ os.pathsep.join(
+ [
+ "/keep",
+ "/venv/lib/python3.12/site-packages/tilelang/vendored",
+ ]
+ ),
+ )
+
+ env = runtime._vllm_runtime_subprocess_env()
+
+ assert env["PYTHONPATH"] == "/keep"
+ assert env["FLASHINFER_WORKSPACE_BASE"] == str(
+ tmp_path / "scratch" / "vllm_runtime_flashinfer"
+ )
+
+
+def test_vllm_runtime_subprocess_env_isolates_flashinfer_for_managed_runtime(
+ monkeypatch,
+ tmp_path: Path,
+) -> None:
+ cache_root = tmp_path / "runtime_cache"
+ monkeypatch.setenv("ART_VLLM_RUNTIME_PROJECT_ROOT", str(tmp_path / "missing"))
+ monkeypatch.setenv("ART_VLLM_RUNTIME_CACHE_DIR", str(cache_root))
+ monkeypatch.setenv("FLASHINFER_WORKSPACE_BASE", "/shared/flashinfer")
+
+ env = runtime._vllm_runtime_subprocess_env()
+
+ assert env["FLASHINFER_WORKSPACE_BASE"] == str(cache_root / "flashinfer_workspace")
+
+
+def test_vllm_runtime_subprocess_env_honors_flashinfer_workspace_override(
+ monkeypatch,
+ tmp_path: Path,
+) -> None:
+ override = tmp_path / "explicit_flashinfer"
+ monkeypatch.setenv(
+ "ART_VLLM_RUNTIME_FLASHINFER_WORKSPACE_BASE",
+ str(override),
+ )
+
+ env = runtime._vllm_runtime_subprocess_env()
+
+ assert env["FLASHINFER_WORKSPACE_BASE"] == str(override)
+
+
def test_cleanup_old_managed_runtimes_only_deletes_marked_venvs(
monkeypatch,
tmp_path: Path,
diff --git a/tests/integration/megatron/train_inf_mismatch/output_parity.py b/tests/integration/megatron/train_inf_mismatch/output_parity.py
index ed46cee9f..839404894 100644
--- a/tests/integration/megatron/train_inf_mismatch/output_parity.py
+++ b/tests/integration/megatron/train_inf_mismatch/output_parity.py
@@ -11,6 +11,11 @@
from pydantic import BaseModel, ConfigDict, Field, model_validator
+from ..model_support.workflow_resources import (
+ handler_workflow_resources_for_base_model,
+ resolve_stage_resources_for_current_host,
+)
+
# These gates are intentionally bf16-scale, not fp32 oracle-scale. A 2026-05-18
# Qwen/Qwen3.5-35B-A3B diagnostic on the exact same real generated tokens found:
# vLLM generation vs Megatron: 2.916% mean_abs_pct, 0.0123 MAE, 0.883 top1,
@@ -22,8 +27,14 @@
# prefix route-conflict behavior on the measured path. With the workflow's
# 16-token completions, Qwen3.5 MoE reruns on 2026-05-25 measured 4.169% and
# 4.606% mean_abs_pct while staying under the KL gate, so its gate is 5%.
+# DeepSeek-V4-Flash uses vLLM quantized DSV4 kernels on the serving side while
+# Megatron materializes train-time bf16/fp32 tensors. A 2026-06-18 diagnostic
+# measured non-QAT Megatron vs vLLM generation at 19.016% mean_abs_pct and
+# 0.02603 candidate->target top20 KL; vLLM generation vs exact vLLM prompt
+# rescore was already 15.176% mean_abs_pct and 0.04424 KL.
BF16_FWD_MEAN_ABS_PCT_LIMIT = 4.0
BF16_FWD_MEAN_ABS_PCT_LIMIT_BY_MODEL_KEY = {
+ "dsv4": 20.0,
# Gemma 4 MoE long-prompt SWA native-LoRA runs showed high variation, with
# repeated samples reaching 7.6% mean_abs_pct and 0.0076 KL.
"gemma4_dense": 8.0,
@@ -33,6 +44,7 @@
}
TOP20_KL_CANDIDATE_TO_TARGET_LIMIT = 0.002
TOP20_KL_CANDIDATE_TO_TARGET_LIMIT_BY_MODEL_KEY = {
+ "dsv4": 0.07,
"gemma4_dense": 0.003,
"gemma4_moe": 0.008,
# GPT OSS MXFP4/native-LoRA repeats on 2026-07-06 stayed under the 4%
@@ -107,6 +119,10 @@ class TrainInfOutputParityConfig(BaseModel):
lora_target_modules: list[str] | None = None
engine_args: dict[str, Any] = Field(default_factory=dict)
server_args: dict[str, Any] = Field(default_factory=dict)
+ megatron_env: dict[str, str] = Field(default_factory=dict)
+ replay_vllm_routing: bool = False
+ external_vllm_server_url: str | None = None
+ external_vllm_api_key: str | None = None
@model_validator(mode="after")
def _set_default_rollout_modes(self) -> "TrainInfOutputParityConfig":
@@ -304,7 +320,28 @@ def model_support_is_moe(
return get_model_support_handler_for_spec(spec).is_moe
+def model_supports_context_parallel(
+ base_model: str,
+ *,
+ allow_unvalidated_arch: bool = False,
+) -> bool:
+ from art.megatron.model_support.registry import (
+ get_model_support_handler_for_spec,
+ get_model_support_spec,
+ )
+
+ spec = get_model_support_spec(
+ base_model,
+ allow_unvalidated_arch=allow_unvalidated_arch,
+ )
+ return bool(getattr(get_model_support_handler_for_spec(spec), "cp_supported", True))
+
+
def config_from_env() -> TrainInfOutputParityConfig:
+ train_inf_external_url_env = "ART_TRAIN_INF_MISMATCH_EXTERNAL_VLLM_URL"
+ model_support_external_url_env = "ART_MODEL_SUPPORT_EXTERNAL_VLLM_URL"
+ train_inf_external_key_env = "ART_TRAIN_INF_MISMATCH_EXTERNAL_VLLM_API_KEY"
+ model_support_external_key_env = "ART_MODEL_SUPPORT_EXTERNAL_VLLM_API_KEY"
config = TrainInfOutputParityConfig(
base_model=os.environ.get(
"ART_TRAIN_INF_MISMATCH_BASE_MODEL",
@@ -323,6 +360,44 @@ def config_from_env() -> TrainInfOutputParityConfig:
)
== "1",
)
+ workflow_resources = handler_workflow_resources_for_base_model(
+ config.base_model,
+ allow_unvalidated_arch=config.allow_unvalidated_arch,
+ )
+ stage_resources = (
+ workflow_resources.train_inf_mismatch
+ if workflow_resources is not None
+ else None
+ )
+ if stage_resources is not None:
+ stage_resources = resolve_stage_resources_for_current_host(
+ "train_inf_mismatch",
+ stage_resources,
+ )
+ if stage_resources is not None:
+ if (
+ stage_resources.megatron is not None
+ and "ART_TRAIN_INF_MISMATCH_TRAINER_GPU_IDS" not in os.environ
+ ):
+ config.trainer_gpu_ids = list(stage_resources.megatron.gpu_ids)
+ if (
+ stage_resources.vllm is not None
+ and "ART_TRAIN_INF_MISMATCH_INFERENCE_GPU_IDS" not in os.environ
+ ):
+ config.inference_gpu_ids = list(stage_resources.vllm.gpu_ids)
+ if stage_resources.megatron is not None:
+ config.topology = config.topology.model_copy(
+ update=stage_resources.megatron.topology.to_train_inf_topology_kwargs()
+ )
+ if stage_resources.vllm is not None:
+ config.engine_args = {
+ **stage_resources.vllm.engine_args(),
+ **config.engine_args,
+ }
+ config.megatron_env = {
+ **stage_resources.megatron_env,
+ **config.megatron_env,
+ }
if raw_modes := os.environ.get("ART_TRAIN_INF_MISMATCH_ROLLOUT_MODES"):
config.rollout_modes = _parse_rollout_modes(raw_modes)
if raw_seq_len := os.environ.get("ART_TRAIN_INF_MISMATCH_SEQUENCE_LENGTH"):
@@ -335,18 +410,50 @@ def config_from_env() -> TrainInfOutputParityConfig:
("ART_TRAIN_INF_MISMATCH_TP", "tp"),
("ART_TRAIN_INF_MISMATCH_EP", "ep"),
("ART_TRAIN_INF_MISMATCH_ETP", "etp"),
+ ("ART_TRAIN_INF_MISMATCH_DP", "dp"),
("ART_TRAIN_INF_MISMATCH_CP", "cp"),
("ART_TRAIN_INF_MISMATCH_PP", "pp"),
):
if raw_value := os.environ.get(env_name):
config.topology = config.topology.model_copy(update={attr: int(raw_value)})
- if not model_support_is_moe(
+ is_moe = model_support_is_moe(
config.base_model,
allow_unvalidated_arch=config.allow_unvalidated_arch,
- ):
+ )
+ cp_supported = model_supports_context_parallel(
+ config.base_model,
+ allow_unvalidated_arch=config.allow_unvalidated_arch,
+ )
+ if not is_moe:
config.topology = config.topology.model_copy(update={"ep": 1, "etp": 1})
+ if not cp_supported and "ART_TRAIN_INF_MISMATCH_CP" not in os.environ:
+ updates = {"cp": 1}
+ if stage_resources is None and "ART_TRAIN_INF_MISMATCH_DP" not in os.environ:
+ updates["dp"] = config.topology.dp * config.topology.cp
+ config.topology = config.topology.model_copy(update=updates)
if raw_targets := os.environ.get("ART_TRAIN_INF_MISMATCH_LORA_TARGET_MODULES"):
config.lora_target_modules = _parse_str_list(raw_targets)
+ raw_url = os.environ.get(train_inf_external_url_env)
+ if raw_url is None and stage_resources is not None:
+ raw_url = os.environ.get(model_support_external_url_env)
+ if raw_url:
+ config.external_vllm_server_url = raw_url
+ config.external_vllm_api_key = os.environ.get(
+ train_inf_external_key_env,
+ os.environ.get(
+ model_support_external_key_env,
+ "art-external-vllm",
+ ),
+ )
+ if (
+ stage_resources is not None
+ and stage_resources.requires_external_vllm
+ and not config.external_vllm_server_url
+ ):
+ raise RuntimeError(
+ "train_inf_mismatch for this model requires an external vLLM server. "
+ f"Set {train_inf_external_url_env} or {model_support_external_url_env}."
+ )
return config
@@ -390,10 +497,21 @@ def build_logical_token_map(packed_tensors: dict[str, Any]) -> LogicalTokenMap:
tokens = packed_tensors["tokens"]
group_ids = packed_tensors["group_ids"]
parent_ids = packed_tensors["parent_ids"]
+ assistant_mask = packed_tensors.get("assistant_mask")
+ logprobs = packed_tensors.get("logprobs")
prompts: list[LogicalPrompt] = []
logical_tokens: list[LogicalToken] = []
prompt_id_by_tokens: dict[tuple[int, ...], int] = {}
+ def scored_token(sample_id: int, packed_i: int) -> bool:
+ if assistant_mask is not None and not bool(assistant_mask[sample_id, packed_i]):
+ return False
+ if logprobs is not None:
+ value = float(logprobs[sample_id, packed_i])
+ if math.isnan(value):
+ return False
+ return True
+
for sample_id in range(int(tokens.shape[0])):
families = _prompt_family_segments(group_ids[sample_id], parent_ids[sample_id])
for family_id, (prompt_segment, completion_segments) in enumerate(families):
@@ -402,15 +520,20 @@ def build_logical_token_map(packed_tensors: dict[str, Any]) -> LogicalTokenMap:
for completion_id, (completion_start, completion_end) in enumerate(
completion_segments
):
- if completion_end - completion_start < 2:
+ last_scored_i = None
+ for packed_i in range(completion_start + 1, completion_end):
+ if scored_token(sample_id, packed_i):
+ last_scored_i = packed_i
+ if last_scored_i is None:
continue
+ effective_completion_end = last_scored_i + 1
flat = [
int(value)
for value in tokens[sample_id, prompt_start:prompt_end].tolist()
] + [
int(value)
for value in tokens[
- sample_id, completion_start:completion_end
+ sample_id, completion_start:effective_completion_end
].tolist()
]
flat_key = tuple(flat)
@@ -429,7 +552,9 @@ def build_logical_token_map(packed_tensors: dict[str, Any]) -> LogicalTokenMap:
token_ids=flat,
)
)
- for packed_i in range(completion_start + 1, completion_end):
+ for packed_i in range(completion_start + 1, effective_completion_end):
+ if not scored_token(sample_id, packed_i):
+ continue
logical_tokens.append(
LogicalToken(
token_id=int(tokens[sample_id, packed_i].item()),
@@ -886,6 +1011,7 @@ def _run_logits(
build_gdn_execution_spec=bool(
getattr(runtime.model_support_handler, "build_gdn_execution_spec", False)
),
+ model_support_handler=runtime.model_support_handler,
attention_head_dim=getattr(runtime.provider, "kv_channels", None),
attention_value_head_dim=getattr(runtime.provider, "kv_channels", None),
)
diff --git a/tests/integration/megatron/train_inf_mismatch/real_path.py b/tests/integration/megatron/train_inf_mismatch/real_path.py
index b882160e8..18f5ebb05 100644
--- a/tests/integration/megatron/train_inf_mismatch/real_path.py
+++ b/tests/integration/megatron/train_inf_mismatch/real_path.py
@@ -2,7 +2,9 @@
import argparse
import asyncio
-from contextlib import asynccontextmanager
+from contextlib import asynccontextmanager, contextmanager
+import hashlib
+import json
import os
from pathlib import Path
import random
@@ -10,13 +12,13 @@
import socket
import subprocess
import sys
-from typing import Any, AsyncIterator, cast
+from typing import Any, AsyncIterator, Iterator, cast
import uuid
from openai.types.chat.chat_completion import Choice
from pydantic import BaseModel, ConfigDict, Field
-from art.dev.model import RolloutWeightsMode
+from art.dev.model import InternalModelConfig, RolloutWeightsMode
from art.preprocessing.moe_routing import (
MoeRoutingPackStats,
PackedMoeRoutingReplay,
@@ -71,6 +73,7 @@ class RealPathConfig(BaseModel):
diagnose_base: bool = False
trace_layers: bool = False
trace_enforce_eager: bool = False
+ adapter_cache_dir: str | None = None
class RealPathMegatronWorkerRequest(BaseModel):
@@ -106,6 +109,23 @@ class RealPathBaseDiagnosticBundle(BaseModel):
megatron_forward_trace_dir: str | None = None
+@contextmanager
+def _temporary_env(updates: dict[str, str] | None) -> Iterator[None]:
+ if not updates:
+ yield
+ return
+ previous = {key: os.environ.get(key) for key in updates}
+ os.environ.update(updates)
+ try:
+ yield
+ finally:
+ for key, value in previous.items():
+ if value is None:
+ os.environ.pop(key, None)
+ else:
+ os.environ[key] = value
+
+
class RealPathTrainInfReport(BaseModel):
base_model: str
artifact_dir: str
@@ -117,6 +137,8 @@ class RealPathTrainInfReport(BaseModel):
base_moe_routing_shared_prefix_conflict_rows: int | None = None
base_moe_routing_shared_prefix_conflict_slots: int | None = None
adapter_path: str
+ adapter_cache_key: str
+ adapter_cache_hit: bool
megatron_base_scores: str | None = None
vllm_base_scores: str | None = None
megatron_lora_scores: str
@@ -135,6 +157,12 @@ class RealPathTrainInfReport(BaseModel):
passed: bool
+class AdapterCacheResult(BaseModel):
+ path: str
+ cache_key: str
+ cache_hit: bool
+
+
def _real_path_rollout_mode(config: TrainInfOutputParityConfig) -> RolloutMode:
return config.rollout_modes[0]
@@ -196,6 +224,8 @@ def config_from_env() -> RealPathConfig:
config.diagnose_base = True
if raw := os.environ.get("ART_REAL_PATH_TRACE_ENFORCE_EAGER"):
config.trace_enforce_eager = raw == "1"
+ if raw := os.environ.get("ART_TRAIN_INF_MISMATCH_ADAPTER_CACHE_DIR"):
+ config.adapter_cache_dir = raw
return config
@@ -306,10 +336,19 @@ async def _collect_real_trajectory_groups(
from transformers import AutoTokenizer
import art
+ from art.megatron.model_support.tokenizer import (
+ configure_tokenizer_for_model_support,
+ )
if config.rollouts_per_prompt < 2:
raise ValueError("real-path mismatch requires at least two rollouts per prompt")
- tokenizer = AutoTokenizer.from_pretrained(config.output_parity.base_model)
+ tokenizer = configure_tokenizer_for_model_support(
+ AutoTokenizer.from_pretrained(config.output_parity.base_model),
+ base_model=config.output_parity.base_model,
+ internal_config={
+ "allow_unvalidated_arch": config.output_parity.allow_unvalidated_arch
+ },
+ )
chat_template_kwargs: dict[str, Any] = {}
if isinstance(tokenizer.chat_template, str):
if "enable_thinking" in tokenizer.chat_template:
@@ -819,6 +858,134 @@ def _make_nonzero_adapter(
return _run_real_path_megatron_worker(request, adapter_only=True).adapter_path or ""
+def _adapter_cache_key(config: TrainInfOutputParityConfig) -> str:
+ from art.megatron.model_support import vllm_lora_config_for_model
+
+ from .output_parity import _adapter_config
+
+ def jsonable(value: Any) -> Any:
+ if isinstance(value, dict):
+ return {key: jsonable(item) for key, item in value.items()}
+ if isinstance(value, set):
+ return [jsonable(item) for item in sorted(value, key=str)]
+ if isinstance(value, (list, tuple)):
+ return [jsonable(item) for item in value]
+ return value
+
+ adapter_config = _adapter_config(config)
+ published_adapter_config = vllm_lora_config_for_model(
+ config.base_model,
+ adapter_config,
+ allow_unvalidated_arch=config.allow_unvalidated_arch,
+ )
+ payload = {
+ "schema": 2,
+ "base_model": config.base_model,
+ "seed": config.seed,
+ "allow_unvalidated_arch": config.allow_unvalidated_arch,
+ "lora_target_modules": _lora_target_modules(config),
+ "adapter_config": adapter_config,
+ "published_adapter_config": published_adapter_config,
+ }
+ encoded = json.dumps(
+ jsonable(payload),
+ sort_keys=True,
+ separators=(",", ":"),
+ ).encode()
+ return hashlib.sha256(encoded).hexdigest()
+
+
+def _default_adapter_cache_dir() -> Path:
+ return REPO_ROOT / "scratch" / "train_inf_mismatch_adapters"
+
+
+def _adapter_cache_dir(config: RealPathConfig) -> Path:
+ if config.adapter_cache_dir:
+ return Path(config.adapter_cache_dir)
+ return _default_adapter_cache_dir()
+
+
+def _adapter_cache_manifest_path(adapter_path: Path) -> Path:
+ return adapter_path / "art_train_inf_mismatch_adapter_cache.json"
+
+
+def _cached_adapter_is_valid(adapter_path: Path, *, cache_key: str) -> bool:
+ manifest_path = _adapter_cache_manifest_path(adapter_path)
+ model_path = adapter_path / "adapter_model.safetensors"
+ config_path = adapter_path / "adapter_config.json"
+ if not (manifest_path.exists() and model_path.exists() and config_path.exists()):
+ return False
+ try:
+ manifest = _read_json(manifest_path)
+ except Exception:
+ return False
+ return (
+ manifest.get("cache_key") == cache_key
+ and manifest.get("non_identity") is True
+ and int(manifest.get("adapter_model_bytes", 0)) == model_path.stat().st_size
+ and model_path.stat().st_size > 0
+ )
+
+
+def _copy_adapter_dir(src: Path, dst: Path, *, cache_key: str) -> None:
+ dst.mkdir(parents=True, exist_ok=True)
+ for filename in ("adapter_model.safetensors", "adapter_config.json"):
+ shutil.copy2(src / filename, dst / filename)
+ _write_json(
+ _adapter_cache_manifest_path(dst),
+ {
+ "cache_key": cache_key,
+ "non_identity": True,
+ "adapter_model_bytes": (dst / "adapter_model.safetensors").stat().st_size,
+ },
+ )
+
+
+def _prune_adapter_cache_dir(cache_dir: Path, *, keep_key: str) -> None:
+ if not cache_dir.exists():
+ return
+ cache_root = cache_dir.resolve()
+ keep_path = (cache_root / keep_key).resolve()
+ for candidate in cache_root.iterdir():
+ if not candidate.is_dir() or candidate.is_symlink():
+ continue
+ resolved = candidate.resolve()
+ if resolved == keep_path:
+ continue
+ resolved.relative_to(cache_root)
+ if _adapter_cache_manifest_path(resolved).exists():
+ shutil.rmtree(resolved)
+
+
+def _make_or_reuse_nonzero_adapter(
+ *,
+ config: RealPathConfig,
+ artifact_dir: Path,
+) -> AdapterCacheResult:
+ cache_key = _adapter_cache_key(config.output_parity)
+ cache_dir = _adapter_cache_dir(config)
+ cache_path = cache_dir / cache_key
+ if _cached_adapter_is_valid(cache_path, cache_key=cache_key):
+ _prune_adapter_cache_dir(cache_dir, keep_key=cache_key)
+ return AdapterCacheResult(
+ path=str(cache_path),
+ cache_key=cache_key,
+ cache_hit=True,
+ )
+ adapter_path = Path(
+ _make_nonzero_adapter(config=config.output_parity, artifact_dir=artifact_dir)
+ )
+ if not adapter_path:
+ raise RuntimeError("Real-path adapter worker did not create an adapter")
+ _copy_adapter_dir(adapter_path, cache_path, cache_key=cache_key)
+ _prune_adapter_cache_dir(cache_dir, keep_key=cache_key)
+ return AdapterCacheResult(
+ path=str(cache_path),
+ cache_key=cache_key,
+ cache_hit=False,
+ )
+
+
def _run_logits_with_replay(
*,
runtime: Any,
@@ -923,6 +1090,7 @@ def _real_path_megatron_worker(
from art.megatron import train as megatron_train
from art.megatron.model_support.lora_disk import load_lora_tensors_for_megatron
+ from art.megatron.training.weight_offload import WeightOffloadManager
from art.preprocessing.pack import packed_tensors_from_dir
local_rank = int(os.environ["LOCAL_RANK"])
@@ -959,92 +1127,113 @@ def _configure_worker_bundle(bundle: Any) -> None:
for chunk in runtime.model:
chunk.eval()
+ weight_offload = None
+ if not adapter_only:
+ weight_offload = WeightOffloadManager.from_env(
+ model=runtime.model,
+ rank=torch.distributed.get_rank(), # type: ignore[possibly-missing-attribute]
+ compile_enabled=runtime.transformer_layers_compiled,
+ )
+ weight_offload.install()
+ weight_offload.after_job()
+ weight_offload.before_job()
+
artifact_dir = Path(request.artifact_dir)
adapter_path: Path | None = None
- if request.weight_state == "lora":
- if request.adapter_path is None:
- initial_state = _collect_full_lora_state(cast(list[Any], runtime.model))
- if torch.distributed.get_rank() == 0: # type: ignore[possibly-missing-attribute]
+ job_completed = False
+ try:
+ if request.weight_state == "lora":
+ if request.adapter_path is None:
+ initial_state = _collect_full_lora_state(cast(list[Any], runtime.model))
+ if torch.distributed.get_rank() == 0: # type: ignore[possibly-missing-attribute]
+ adapter_path = artifact_dir / "real_path_active_lora"
+ initialized = _build_deterministic_nonzero_lora(
+ initial_state or {},
+ seed=request.config.seed,
+ )
+ _save_vllm_lora_adapter(
+ lora_path=adapter_path,
+ state=initialized,
+ runtime=runtime,
+ config=request.config,
+ )
+ torch.distributed.barrier() # type: ignore[possibly-missing-attribute]
adapter_path = artifact_dir / "real_path_active_lora"
- initialized = _build_deterministic_nonzero_lora(
- initial_state or {},
- seed=request.config.seed,
+ else:
+ adapter_path = Path(request.adapter_path)
+ adapter_model = load_lora_tensors_for_megatron(
+ str(adapter_path),
+ handler=runtime.model_support_handler,
+ allow_unvalidated_arch=request.config.allow_unvalidated_arch,
+ )
+ megatron_train.load_adapter_into_model(runtime.model, adapter_model)
+
+ if adapter_only:
+ if torch.distributed.get_rank() == 0: # type: ignore[possibly-missing-attribute]
+ result = RealPathMegatronWorkerResult(
+ score_path="",
+ adapter_path=str(adapter_path)
+ if adapter_path is not None
+ else None,
)
- _save_vllm_lora_adapter(
- lora_path=adapter_path,
- state=initialized,
- runtime=runtime,
- config=request.config,
+ _write_json(
+ artifact_dir / "real_path_adapter_worker_result.json",
+ result.model_dump(mode="json"),
)
torch.distributed.barrier() # type: ignore[possibly-missing-attribute]
- adapter_path = artifact_dir / "real_path_active_lora"
- else:
- adapter_path = Path(request.adapter_path)
- adapter_model = load_lora_tensors_for_megatron(
- str(adapter_path),
- handler=runtime.model_support_handler,
- allow_unvalidated_arch=request.config.allow_unvalidated_arch,
+ torch.distributed.destroy_process_group() # type: ignore[possibly-missing-attribute]
+ return
+
+ packed_tensors = packed_tensors_from_dir(**request.disk_packed_tensors)
+ logical_map = LogicalTokenMap.model_validate(
+ _read_json(Path(request.logical_map_path))
+ )
+ forward_trace_capture = None
+ if request.forward_trace_dir is not None:
+ from ..model_support.forward_trace import (
+ CAPTURE_NAME_TOKENS,
+ ForwardTraceCapture,
+ )
+
+ forward_trace_capture = ForwardTraceCapture(
+ runtime.model,
+ enabled=True,
+ capture_name_tokens=(*CAPTURE_NAME_TOKENS, ".decoder.final_layernorm"),
+ strict_output_match=True,
+ )
+ forward_trace_capture.set_step(
+ 0,
+ list(range(int(packed_tensors["tokens"].shape[0]))),
+ )
+ score = _score_megatron_runtime(
+ runtime=runtime,
+ packed_tensors=cast(dict[str, Any], packed_tensors),
+ logical_map=logical_map,
+ weight_state=request.weight_state,
+ rollout_mode=_real_path_rollout_mode(request.config),
+ global_grad_accumulation_sequences=request.global_grad_accumulation_sequences,
+ forward_trace_capture=forward_trace_capture,
+ forward_trace_dir=request.forward_trace_dir,
)
- megatron_train.load_adapter_into_model(runtime.model, adapter_model)
- if adapter_only:
if torch.distributed.get_rank() == 0: # type: ignore[possibly-missing-attribute]
+ score_path = (
+ artifact_dir / f"real_path_megatron_{request.weight_state}.json"
+ )
+ _write_json(score_path, score.model_dump(mode="json"))
result = RealPathMegatronWorkerResult(
- score_path="",
+ score_path=str(score_path),
adapter_path=str(adapter_path) if adapter_path is not None else None,
)
_write_json(
- artifact_dir / "real_path_adapter_worker_result.json",
+ artifact_dir
+ / f"real_path_megatron_{request.weight_state}_worker_result.json",
result.model_dump(mode="json"),
)
- torch.distributed.barrier() # type: ignore[possibly-missing-attribute]
- torch.distributed.destroy_process_group() # type: ignore[possibly-missing-attribute]
- return
-
- packed_tensors = packed_tensors_from_dir(**request.disk_packed_tensors)
- logical_map = LogicalTokenMap.model_validate(
- _read_json(Path(request.logical_map_path))
- )
- forward_trace_capture = None
- if request.forward_trace_dir is not None:
- from ..model_support.forward_trace import (
- CAPTURE_NAME_TOKENS,
- ForwardTraceCapture,
- )
-
- forward_trace_capture = ForwardTraceCapture(
- runtime.model,
- enabled=True,
- capture_name_tokens=(*CAPTURE_NAME_TOKENS, ".decoder.final_layernorm"),
- strict_output_match=True,
- )
- forward_trace_capture.set_step(
- 0,
- list(range(int(packed_tensors["tokens"].shape[0]))),
- )
- score = _score_megatron_runtime(
- runtime=runtime,
- packed_tensors=cast(dict[str, Any], packed_tensors),
- logical_map=logical_map,
- weight_state=request.weight_state,
- rollout_mode=_real_path_rollout_mode(request.config),
- global_grad_accumulation_sequences=request.global_grad_accumulation_sequences,
- forward_trace_capture=forward_trace_capture,
- forward_trace_dir=request.forward_trace_dir,
- )
-
- if torch.distributed.get_rank() == 0: # type: ignore[possibly-missing-attribute]
- score_path = artifact_dir / f"real_path_megatron_{request.weight_state}.json"
- _write_json(score_path, score.model_dump(mode="json"))
- result = RealPathMegatronWorkerResult(
- score_path=str(score_path),
- adapter_path=str(adapter_path) if adapter_path is not None else None,
- )
- _write_json(
- artifact_dir
- / f"real_path_megatron_{request.weight_state}_worker_result.json",
- result.model_dump(mode="json"),
- )
+ job_completed = True
+ finally:
+ if weight_offload is not None and job_completed:
+ weight_offload.after_job()
torch.distributed.barrier() # type: ignore[possibly-missing-attribute]
torch.distributed.destroy_process_group() # type: ignore[possibly-missing-attribute]
@@ -1066,6 +1255,7 @@ def _run_real_path_megatron_worker(
env["CUDA_VISIBLE_DEVICES"] = ",".join(
str(value) for value in request.config.trainer_gpu_ids
)
+ env.update(request.config.megatron_env)
env["PYTHONUNBUFFERED"] = "1"
tests_dir = str(REPO_ROOT / "tests")
env["PYTHONPATH"] = (
@@ -1142,11 +1332,8 @@ async def run_real_path_train_inf_mismatch(
allow_unvalidated_arch=parity_config.allow_unvalidated_arch,
)
_write_json(artifact_dir / "real_path_config.json", config.model_dump(mode="json"))
- adapter_path = _make_nonzero_adapter(
- config=parity_config, artifact_dir=artifact_dir
- )
- if not adapter_path:
- raise RuntimeError("Real-path adapter worker did not create an adapter")
+ adapter = _make_or_reuse_nonzero_adapter(config=config, artifact_dir=artifact_dir)
+ adapter_path = adapter.path
_init_art_megatron_runtime_config(parity_config)
backend = MegatronBackend(
@@ -1154,16 +1341,9 @@ async def run_real_path_train_inf_mismatch(
enable_expert_replay=is_moe,
)
backend_open = False
- model = art.TrainableModel(
- name=f"train-inf-real-{uuid.uuid4().hex[:8]}",
- project="train_inf_mismatch",
- base_model=parity_config.base_model,
- lora_config=(
- {"target_modules": _lora_target_modules(parity_config)}
- if parity_config.lora_target_modules is not None
- else None
- ),
- _internal_config={
+ internal_config = cast(
+ InternalModelConfig,
+ {
"trainer_gpu_ids": parity_config.trainer_gpu_ids,
"inference_gpu_ids": parity_config.inference_gpu_ids,
"rollout_weights_mode": _real_path_rollout_weights_mode(parity_config),
@@ -1181,11 +1361,29 @@ async def run_real_path_train_inf_mismatch(
},
},
)
+ if parity_config.external_vllm_server_url is not None:
+ internal_config["vllm_runtime"] = {
+ "mode": "external",
+ "server_url": parity_config.external_vllm_server_url,
+ "api_key": parity_config.external_vllm_api_key,
+ }
+ model = art.TrainableModel(
+ name=f"train-inf-real-{uuid.uuid4().hex[:8]}",
+ project="train_inf_mismatch",
+ base_model=parity_config.base_model,
+ lora_config=(
+ {"target_modules": _lora_target_modules(parity_config)}
+ if parity_config.lora_target_modules is not None
+ else None
+ ),
+ _internal_config=internal_config,
+ )
_move_adapter_to_step_zero(adapter_path=adapter_path, model=model, backend=backend)
try:
- await model.register(backend)
- backend_open = True
+ with _temporary_env(parity_config.megatron_env):
+ await model.register(backend)
+ backend_open = True
trajectory_groups = await _collect_real_trajectory_groups(
model=model,
config=config,
@@ -1342,6 +1540,8 @@ async def run_real_path_train_inf_mismatch(
else None
),
adapter_path=adapter_path,
+ adapter_cache_key=adapter.cache_key,
+ adapter_cache_hit=adapter.cache_hit,
megatron_base_scores=(
base_diagnostic.megatron_score_path
if base_diagnostic is not None
diff --git a/tests/integration/megatron/train_inf_mismatch/test_config.py b/tests/integration/megatron/train_inf_mismatch/test_config.py
new file mode 100644
index 000000000..bce110550
--- /dev/null
+++ b/tests/integration/megatron/train_inf_mismatch/test_config.py
@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import torch
+
+from . import output_parity
+from .output_parity import config_from_env
+
+
+def test_cp_unsupported_default_converts_cp_to_dp_without_changing_tp(
+ monkeypatch,
+) -> None:
+ monkeypatch.setenv("BASE_MODEL", "Qwen/Qwen3.5-35B-A3B")
+ monkeypatch.delenv("ART_TRAIN_INF_MISMATCH_TP", raising=False)
+ monkeypatch.delenv("ART_TRAIN_INF_MISMATCH_CP", raising=False)
+ monkeypatch.delenv("ART_TRAIN_INF_MISMATCH_DP", raising=False)
+ monkeypatch.setattr(
+ output_parity,
+ "handler_workflow_resources_for_base_model",
+ lambda base_model, *, allow_unvalidated_arch=False: None,
+ )
+ monkeypatch.setattr(output_parity, "model_support_is_moe", lambda *_, **__: True)
+ monkeypatch.setattr(
+ output_parity,
+ "model_supports_context_parallel",
+ lambda *_, **__: False,
+ )
+
+ config = config_from_env()
+
+ assert config.topology.tp == 1
+ assert config.topology.cp == 1
+ assert config.topology.dp == 2
+ assert config.topology.ep == 2
+ assert config.topology.world_size() == 2
+
+
+def test_cp_unsupported_model_uses_non_cp_default_topology(monkeypatch) -> None:
+ monkeypatch.setattr(torch.cuda, "device_count", lambda: 4)
+ monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+ monkeypatch.setattr(
+ torch.cuda,
+ "get_device_properties",
+ lambda device: SimpleNamespace(total_memory=284 * 1024**3),
+ )
+ monkeypatch.setenv("BASE_MODEL", "deepseek-ai/DeepSeek-V4-Flash")
+ monkeypatch.delenv("ART_TRAIN_INF_MISMATCH_TRAINER_GPU_IDS", raising=False)
+ monkeypatch.delenv("ART_TRAIN_INF_MISMATCH_INFERENCE_GPU_IDS", raising=False)
+ monkeypatch.delenv("ART_TRAIN_INF_MISMATCH_TP", raising=False)
+ monkeypatch.delenv("ART_TRAIN_INF_MISMATCH_CP", raising=False)
+ monkeypatch.delenv("ART_TRAIN_INF_MISMATCH_EP", raising=False)
+ monkeypatch.setenv("ART_MODEL_SUPPORT_EXTERNAL_VLLM_URL", "http://127.0.0.1:8000")
+
+ config = config_from_env()
+
+ assert config.topology.cp == 1
+ assert config.topology.tp == 2
+ assert config.topology.ep == 2
+ assert config.topology.dp == 1
+ assert config.trainer_gpu_ids == [0, 1]
+ assert config.inference_gpu_ids == [2, 3]
+ assert config.engine_args["tensor_parallel_size"] == 2
+ assert config.engine_args["enable_expert_parallel"] is True
+ assert config.engine_args["kv_cache_dtype"] == "fp8"
+ assert config.engine_args["moe_backend"] == "triton_unfused"
+ assert config.megatron_env == {"ART_MEGATRON_STREAMING_WEIGHT_OFFLOAD": "1"}
+ assert config.external_vllm_server_url == "http://127.0.0.1:8000"
diff --git a/tests/integration/megatron/trainability/test_config.py b/tests/integration/megatron/trainability/test_config.py
index b265c2b3d..75e5cc37c 100644
--- a/tests/integration/megatron/trainability/test_config.py
+++ b/tests/integration/megatron/trainability/test_config.py
@@ -1,12 +1,15 @@
import asyncio
+from types import SimpleNamespace
from typing import cast
from openai.types.chat.chat_completion import ChatCompletion, Choice
from openai.types.chat.chat_completion_message import ChatCompletionMessage
import pytest
+import torch
import art
+from .test_live_length_trainability import _use_default_moe_dedicated_placement
from .yes_no_trainability import (
_build_internal_config,
_build_variant,
@@ -17,6 +20,7 @@
_variant_max_steps,
_variant_packed_sequence_length,
_variant_rollouts_per_prompt,
+ _variant_train_kwargs,
)
@@ -105,6 +109,7 @@ def test_megatron_variants_keep_short_packed_sequence_default(monkeypatch) -> No
)
assert _variant_packed_sequence_length(variant) == 1024
+ assert _variant_train_kwargs(variant) == {"packed_sequence_length": 1024}
config = _build_internal_config(
variant, base_model="Qwen/Qwen3-30B-A3B-Instruct-2507"
)
@@ -128,6 +133,7 @@ def test_unsloth_variant_uses_chunk_aligned_training_length(monkeypatch) -> None
)
assert _variant_packed_sequence_length(variant) == 1024
+ assert _variant_train_kwargs(variant) == {"packed_sequence_length": 1024}
assert _variant_init_args(variant) == {"max_seq_length": 1024}
assert _build_internal_config(
variant, base_model="Qwen/Qwen3-30B-A3B-Instruct-2507"
@@ -162,6 +168,7 @@ def test_validated_dense_model_uses_dense_shared_topology(
base_model="Qwen/Qwen3.5-4B",
)
assert built_variant.topology is not None
+ assert built_variant.topology.tp == 1
assert built_variant.topology.cp == 2
assert built_variant.topology.ep == 1
assert built_variant.topology.etp == 1
@@ -194,3 +201,80 @@ def test_qwen3_5_moe_shared_variant_enables_expert_parallel(monkeypatch) -> None
assert config["rollout_weights_mode"] == "lora"
assert config["engine_args"]["enable_expert_parallel"] is True
+
+
+def test_dsv4_trainability_uses_large_model_dedicated_resources(
+ monkeypatch,
+) -> None:
+ monkeypatch.setattr(torch.cuda, "device_count", lambda: 4)
+ monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+ monkeypatch.setattr(
+ torch.cuda,
+ "get_device_properties",
+ lambda device: SimpleNamespace(total_memory=284 * 1024**3),
+ )
+ monkeypatch.setattr(
+ "tests.integration.megatron.trainability.yes_no_trainability."
+ "_safe_gpu_memory_utilization",
+ lambda device_ids: 0.5,
+ )
+ monkeypatch.setenv("ART_MODEL_SUPPORT_EXTERNAL_VLLM_URL", "http://127.0.0.1:8000")
+
+ default_variant = _default_variant_name(
+ "deepseek-ai/DeepSeek-V4-Flash",
+ )
+ variant = _build_variant(
+ default_variant,
+ base_model="deepseek-ai/DeepSeek-V4-Flash",
+ )
+ config = _build_internal_config(
+ variant,
+ base_model="deepseek-ai/DeepSeek-V4-Flash",
+ )
+
+ assert default_variant == "megatron_dedicated"
+ assert variant.topology is not None
+ assert variant.topology.tp == 2
+ assert variant.topology.ep == 2
+ assert variant.topology.cp == 1
+ assert variant.topology.dp == 1
+ assert variant.topology.sp is True
+ assert variant.trainer_gpu_ids == [0, 1]
+ assert variant.inference_gpu_ids == [2, 3]
+ assert config["engine_args"]["tensor_parallel_size"] == 2
+ assert config["engine_args"]["enable_expert_parallel"] is True
+ assert config["engine_args"]["kv_cache_dtype"] == "fp8"
+ assert config["engine_args"].get("moe_backend") == "triton_unfused"
+ assert "megatron_topology" not in config
+ assert config["vllm_runtime"] == {
+ "mode": "external",
+ "server_url": "http://127.0.0.1:8000",
+ "api_key": "art-external-vllm",
+ }
+
+
+def test_dsv4_length_trainability_keeps_handler_resources(monkeypatch) -> None:
+ monkeypatch.setattr(torch.cuda, "device_count", lambda: 8)
+ monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+ monkeypatch.setattr(
+ torch.cuda,
+ "get_device_properties",
+ lambda device: SimpleNamespace(total_memory=140 * 1024**3),
+ )
+
+ variant = _build_variant(
+ "megatron_dedicated",
+ base_model="deepseek-ai/DeepSeek-V4-Flash",
+ resource_stage_name="length_trainability",
+ )
+ _use_default_moe_dedicated_placement(
+ variant,
+ base_model="deepseek-ai/DeepSeek-V4-Flash",
+ )
+
+ assert variant.topology is not None
+ assert variant.trainer_gpu_ids == [0, 1, 2, 3, 4, 5, 6, 7]
+ assert variant.inference_gpu_ids == [4, 5, 6, 7]
+ assert variant.topology.tp == 2
+ assert variant.topology.ep == 8
+ assert variant.topology.cp == 1
diff --git a/tests/integration/megatron/trainability/test_live_length_trainability.py b/tests/integration/megatron/trainability/test_live_length_trainability.py
index 8e8aa313b..2921c2667 100644
--- a/tests/integration/megatron/trainability/test_live_length_trainability.py
+++ b/tests/integration/megatron/trainability/test_live_length_trainability.py
@@ -31,6 +31,7 @@
_get_env_int,
_init_megatron_runtime_config,
_list_model_ids,
+ _trainability_stage_resources,
)
torch = pytest.importorskip("torch")
@@ -210,6 +211,13 @@ def _target_tokens() -> int:
def _use_default_moe_dedicated_placement(variant: Any, *, base_model: str) -> None:
if not model_uses_expert_parallel(base_model, allow_unvalidated_arch=True):
return
+ stage_resources = _trainability_stage_resources(
+ base_model,
+ stage_name="length_trainability",
+ allow_unvalidated_arch=True,
+ )
+ if stage_resources is not None:
+ return
if os.environ.get(TRAINER_GPU_IDS_ENV) or os.environ.get(INFERENCE_GPU_IDS_ENV):
return
if torch.cuda.device_count() < 3:
@@ -636,6 +644,7 @@ async def run_length_trainability_async(
"megatron_dedicated",
base_model=base_model,
allow_unvalidated_arch=allow_unvalidated_arch,
+ resource_stage_name="length_trainability",
)
_use_default_moe_dedicated_placement(variant, base_model=base_model)
max_steps = _length_max_steps()
@@ -667,6 +676,7 @@ async def run_length_trainability_async(
variant,
base_model=base_model,
allow_unvalidated_arch=allow_unvalidated_arch,
+ resource_stage_name="length_trainability",
)
internal_config["engine_args"]["max_model_len"] = _get_env_int(
"ART_MODEL_SUPPORT_LENGTH_MAX_MODEL_LEN",
@@ -682,8 +692,18 @@ async def run_length_trainability_async(
chat_template_kwargs = _length_chat_template_kwargs(base_model, tokenizer)
rollout_weights_mode = internal_config["rollout_weights_mode"]
_init_megatron_runtime_config(variant)
+ stage_resources = _trainability_stage_resources(
+ base_model,
+ stage_name="length_trainability",
+ allow_unvalidated_arch=allow_unvalidated_arch,
+ )
+ backend_env = stage_resources.megatron_env if stage_resources is not None else {}
- async with _backend_context(variant, backend_root=backend_root) as backend:
+ async with _backend_context(
+ variant,
+ backend_root=backend_root,
+ extra_env=backend_env,
+ ) as backend:
model = art.TrainableModel(
name=f"length-{uuid.uuid4().hex[:8]}",
project="integration-tests",
diff --git a/tests/integration/megatron/trainability/yes_no_trainability.py b/tests/integration/megatron/trainability/yes_no_trainability.py
index a3f7918ae..67aba7078 100644
--- a/tests/integration/megatron/trainability/yes_no_trainability.py
+++ b/tests/integration/megatron/trainability/yes_no_trainability.py
@@ -8,7 +8,7 @@
from pathlib import Path
import re
import time
-from typing import Any, AsyncIterator, Iterator, Literal, cast
+from typing import Any, AsyncIterator, Iterator, Literal, TypedDict, cast
import uuid
from pydantic import BaseModel, Field
@@ -20,16 +20,23 @@
from art.megatron.backend import MegatronBackend
from art.megatron.model_support.registry import (
get_model_support_spec,
+ model_supports_context_parallel,
model_uses_expert_parallel,
)
from art.megatron.model_support.spec import RolloutWeightsMode
from ..model_support.oracle_harness import Topology, oracle_topology
from ..model_support.oracle_worker import provider_topology_env
+from ..model_support.workflow_resources import (
+ handler_workflow_resources_for_base_model,
+ resolve_stage_resources_for_current_host,
+)
_TRAINER_GPU_IDS_ENV = "ART_MODEL_SUPPORT_TRAINER_GPU_IDS"
_INFERENCE_GPU_IDS_ENV = "ART_MODEL_SUPPORT_INFERENCE_GPU_IDS"
_SHARED_GPU_IDS_ENV = "ART_MODEL_SUPPORT_SHARED_GPU_IDS"
+_EXTERNAL_VLLM_URL_ENV = "ART_MODEL_SUPPORT_EXTERNAL_VLLM_URL"
+_EXTERNAL_VLLM_API_KEY_ENV = "ART_MODEL_SUPPORT_EXTERNAL_VLLM_API_KEY"
_TRAINABILITY_ROOT = (
Path(__file__).resolve().parents[4] / ".local" / "model_support_validation"
)
@@ -40,6 +47,11 @@
"megatron_dedicated",
"unsloth_dedicated",
]
+_RESOURCE_STAGE_NAME = Literal["yes_no_trainability", "length_trainability"]
+
+
+class _TrainKwargs(TypedDict):
+ packed_sequence_length: int
class TrainabilityStepReport(BaseModel):
@@ -96,9 +108,11 @@ def build_prompts() -> list[str]:
for length in (3, 2)
for words in permutations(("yes", "no", "maybe"), length)
for body in [
- ", ".join(f"'{word}'" if use_quotes else word for word in words)
- if length == 3
- else " or ".join(f"'{word}'" if use_quotes else word for word in words)
+ (
+ ", ".join(f"'{word}'" if use_quotes else word for word in words)
+ if length == 3
+ else " or ".join(f"'{word}'" if use_quotes else word for word in words)
+ )
]
]
if prompt_count <= len(prompts):
@@ -117,6 +131,50 @@ def _parse_gpu_id_env(name: str) -> list[int] | None:
return [int(part.strip()) for part in raw.split(",") if part.strip()]
+def _external_vllm_runtime_config() -> dev.VllmRuntimeArgs | None:
+ server_url = os.environ.get(_EXTERNAL_VLLM_URL_ENV)
+ if server_url is None or server_url.strip() == "":
+ return None
+ return {
+ "mode": "external",
+ "server_url": server_url,
+ "api_key": os.environ.get(_EXTERNAL_VLLM_API_KEY_ENV, "art-external-vllm"),
+ }
+
+
+def _topology_with_env_overrides(topology: Topology) -> Topology:
+ updates: dict[str, int | bool] = {}
+ for env_name, attr in (
+ ("ART_MODEL_SUPPORT_TP", "tp"),
+ ("ART_MODEL_SUPPORT_EP", "ep"),
+ ("ART_MODEL_SUPPORT_ETP", "etp"),
+ ("ART_MODEL_SUPPORT_DP", "dp"),
+ ("ART_MODEL_SUPPORT_CP", "cp"),
+ ("ART_MODEL_SUPPORT_PP", "pp"),
+ ("ART_MODEL_SUPPORT_VPP", "vpp"),
+ ):
+ if raw_value := os.environ.get(env_name):
+ updates[attr] = int(raw_value)
+ if raw_sp := os.environ.get("ART_MODEL_SUPPORT_SP"):
+ updates["sp"] = raw_sp.strip().lower() in {"1", "true", "yes", "on"}
+ return topology.model_copy(update=updates) if updates else topology
+
+
+def _variant_with_env_overrides(
+ variant: _TrainabilityVariant,
+) -> _TrainabilityVariant:
+ trainer_gpu_ids = _parse_gpu_id_env(_TRAINER_GPU_IDS_ENV)
+ inference_gpu_ids = _parse_gpu_id_env(_INFERENCE_GPU_IDS_ENV)
+ updates: dict[str, object] = {}
+ if trainer_gpu_ids is not None:
+ updates["trainer_gpu_ids"] = trainer_gpu_ids
+ if inference_gpu_ids is not None:
+ updates["inference_gpu_ids"] = inference_gpu_ids
+ if variant.topology is not None:
+ updates["topology"] = _topology_with_env_overrides(variant.topology)
+ return variant.model_copy(update=updates) if updates else variant
+
+
def _resolve_shared_gpu_ids() -> list[int]:
if shared_gpu_ids := _parse_gpu_id_env(_SHARED_GPU_IDS_ENV):
return shared_gpu_ids
@@ -331,44 +389,133 @@ def _artifact_dir(base_model: str, variant_name: _VARIANT_NAME) -> Path:
return path
+def _trainability_stage_resources(
+ base_model: str,
+ *,
+ stage_name: _RESOURCE_STAGE_NAME,
+ allow_unvalidated_arch: bool = False,
+):
+ workflow_resources = handler_workflow_resources_for_base_model(
+ base_model,
+ allow_unvalidated_arch=allow_unvalidated_arch,
+ )
+ if workflow_resources is None:
+ return None
+ stage_resources = getattr(workflow_resources, stage_name)
+ if stage_resources is None:
+ return None
+ return resolve_stage_resources_for_current_host(stage_name, stage_resources)
+
+
def _build_variant(
variant_name: _VARIANT_NAME,
*,
base_model: str,
allow_unvalidated_arch: bool = False,
+ resource_stage_name: _RESOURCE_STAGE_NAME = "yes_no_trainability",
) -> _TrainabilityVariant:
+ stage_resources = _trainability_stage_resources(
+ base_model,
+ stage_name=resource_stage_name,
+ allow_unvalidated_arch=allow_unvalidated_arch,
+ )
is_moe = model_uses_expert_parallel(
base_model,
allow_unvalidated_arch=allow_unvalidated_arch,
)
+ cp_supported = model_supports_context_parallel(
+ base_model,
+ allow_unvalidated_arch=allow_unvalidated_arch,
+ )
if variant_name == "megatron_shared":
- shared_gpu_ids = _resolve_shared_gpu_ids()
- return _TrainabilityVariant(
- name=variant_name,
- backend_name="megatron",
- placement_mode="shared",
- topology=_SHARED_MEGATRON_TOPOLOGY
- if is_moe
- else _DENSE_SHARED_MEGATRON_TOPOLOGY,
- trainer_gpu_ids=shared_gpu_ids,
- inference_gpu_ids=shared_gpu_ids,
+ if (
+ stage_resources is not None
+ and stage_resources.megatron is not None
+ and stage_resources.vllm is not None
+ ):
+ shared_gpu_ids = sorted(
+ {*stage_resources.megatron.gpu_ids, *stage_resources.vllm.gpu_ids}
+ )
+ else:
+ shared_gpu_ids = _resolve_shared_gpu_ids()
+ if not cp_supported:
+ shared_world_size = len(shared_gpu_ids)
+ return _variant_with_env_overrides(
+ _TrainabilityVariant(
+ name=variant_name,
+ backend_name="megatron",
+ placement_mode="shared",
+ topology=Topology(
+ tp=shared_world_size,
+ ep=shared_world_size if is_moe else 1,
+ etp=1,
+ dp=1,
+ cp=1,
+ sp=shared_world_size > 1,
+ ),
+ trainer_gpu_ids=shared_gpu_ids,
+ inference_gpu_ids=shared_gpu_ids,
+ )
+ )
+ return _variant_with_env_overrides(
+ _TrainabilityVariant(
+ name=variant_name,
+ backend_name="megatron",
+ placement_mode="shared",
+ topology=(
+ _SHARED_MEGATRON_TOPOLOGY
+ if is_moe
+ else _DENSE_SHARED_MEGATRON_TOPOLOGY
+ ),
+ trainer_gpu_ids=shared_gpu_ids,
+ inference_gpu_ids=shared_gpu_ids,
+ )
+ )
+ if (
+ variant_name == "megatron_dedicated"
+ and stage_resources is not None
+ and stage_resources.megatron is not None
+ and stage_resources.vllm is not None
+ ):
+ workflow_topology = stage_resources.megatron.topology
+ return _variant_with_env_overrides(
+ _TrainabilityVariant(
+ name=variant_name,
+ backend_name="megatron",
+ placement_mode="dedicated",
+ topology=Topology(
+ tp=workflow_topology.tp,
+ ep=workflow_topology.ep,
+ etp=workflow_topology.etp,
+ dp=workflow_topology.dp,
+ sp=workflow_topology.sp,
+ cp=workflow_topology.cp,
+ pp=workflow_topology.pp,
+ ),
+ trainer_gpu_ids=list(stage_resources.megatron.gpu_ids),
+ inference_gpu_ids=list(stage_resources.vllm.gpu_ids),
+ )
)
trainer_gpu_ids, inference_gpu_ids = _resolve_dedicated_gpu_ids()
if variant_name == "megatron_dedicated":
- return _TrainabilityVariant(
+ return _variant_with_env_overrides(
+ _TrainabilityVariant(
+ name=variant_name,
+ backend_name="megatron",
+ placement_mode="dedicated",
+ topology=oracle_topology(is_moe=is_moe),
+ trainer_gpu_ids=trainer_gpu_ids,
+ inference_gpu_ids=inference_gpu_ids,
+ )
+ )
+ return _variant_with_env_overrides(
+ _TrainabilityVariant(
name=variant_name,
- backend_name="megatron",
+ backend_name="local",
placement_mode="dedicated",
- topology=oracle_topology(is_moe=is_moe),
trainer_gpu_ids=trainer_gpu_ids,
inference_gpu_ids=inference_gpu_ids,
)
- return _TrainabilityVariant(
- name=variant_name,
- backend_name="local",
- placement_mode="dedicated",
- trainer_gpu_ids=trainer_gpu_ids,
- inference_gpu_ids=inference_gpu_ids,
)
@@ -376,6 +523,10 @@ def _variant_packed_sequence_length(variant: _TrainabilityVariant) -> int:
return _get_env_int("ART_MODEL_SUPPORT_YES_NO_PACKED_SEQUENCE_LENGTH", 1024)
+def _variant_train_kwargs(variant: _TrainabilityVariant) -> _TrainKwargs:
+ return {"packed_sequence_length": _variant_packed_sequence_length(variant)}
+
+
def _variant_init_args(variant: _TrainabilityVariant) -> dev.InitArgs:
return {"max_seq_length": _variant_packed_sequence_length(variant)}
@@ -383,7 +534,10 @@ def _variant_init_args(variant: _TrainabilityVariant) -> dev.InitArgs:
def _init_megatron_runtime_config(variant: _TrainabilityVariant) -> None:
if variant.topology is None:
return
- art.init_megatron_runtime_config(
+ init_runtime_config = getattr(art, "init_megatron_runtime_config", None)
+ if init_runtime_config is None:
+ return
+ init_runtime_config(
topology=art.MegatronTopologyConfig(
tp=variant.topology.tp,
cp=variant.topology.cp,
@@ -420,6 +574,15 @@ def _default_variant_name(
*,
allow_unvalidated_arch: bool = False,
) -> _VARIANT_NAME:
+ workflow_resources = handler_workflow_resources_for_base_model(
+ base_model,
+ allow_unvalidated_arch=allow_unvalidated_arch,
+ )
+ if (
+ workflow_resources is not None
+ and workflow_resources.yes_no_trainability_variant is not None
+ ):
+ return workflow_resources.yes_no_trainability_variant
if (
_rollout_weights_mode(
base_model,
@@ -437,16 +600,43 @@ def _build_internal_config(
base_model: str,
rollout_weights_mode: RolloutWeightsMode | None = None,
allow_unvalidated_arch: bool = False,
+ resource_stage_name: _RESOURCE_STAGE_NAME = "yes_no_trainability",
) -> dev.InternalModelConfig:
shared = variant.placement_mode == "shared"
- inference_gpu_ids = (
- variant.inference_gpu_ids if not shared else _resolve_shared_gpu_ids()
+ inference_gpu_ids = variant.inference_gpu_ids
+ stage_resources = _trainability_stage_resources(
+ base_model,
+ stage_name=resource_stage_name,
+ allow_unvalidated_arch=allow_unvalidated_arch,
+ )
+ stage_resources_apply = (
+ not shared
+ and variant.backend_name == "megatron"
+ and stage_resources is not None
+ and stage_resources.megatron is not None
+ and stage_resources.vllm is not None
+ and variant.trainer_gpu_ids == stage_resources.megatron.gpu_ids
+ and variant.inference_gpu_ids == stage_resources.vllm.gpu_ids
)
+ if stage_resources_apply:
+ assert stage_resources is not None
+ assert stage_resources.vllm is not None
+ vllm_resources = stage_resources.vllm
+ else:
+ vllm_resources = None
engine_args = _engine_args_for_yes_no_trainability(
inference_gpu_ids=inference_gpu_ids,
- tensor_parallel_size=len(inference_gpu_ids) if shared else 1,
+ tensor_parallel_size=(
+ vllm_resources.tensor_parallel_size
+ if vllm_resources is not None
+ else len(inference_gpu_ids)
+ if shared
+ else 1
+ ),
enable_expert_parallel=(
- shared
+ vllm_resources.enable_expert_parallel
+ if vllm_resources is not None
+ else shared
and variant.backend_name == "megatron"
and model_uses_expert_parallel(
base_model,
@@ -455,6 +645,10 @@ def _build_internal_config(
),
enable_sleep_mode=True if shared else None,
)
+ if vllm_resources is not None:
+ engine_args.update(vllm_resources.engine_args())
+ elif shared and stage_resources is not None and stage_resources.vllm is not None:
+ engine_args.update(stage_resources.vllm.extra_engine_args)
engine_args["model"] = base_model
internal_config = dev.InternalModelConfig(
rollout_weights_mode=rollout_weights_mode
@@ -466,10 +660,23 @@ def _build_internal_config(
init_args=_variant_init_args(variant),
allow_unvalidated_arch=allow_unvalidated_arch,
)
+ external_runtime = _external_vllm_runtime_config()
+ if (
+ stage_resources is not None
+ and stage_resources.requires_external_vllm
+ and external_runtime is None
+ ):
+ raise RuntimeError(
+ f"{resource_stage_name} for this model requires an external vLLM server. "
+ f"Set {_EXTERNAL_VLLM_URL_ENV}."
+ )
+ if external_runtime is not None:
+ internal_config["vllm_runtime"] = external_runtime
if not shared:
internal_config["trainer_gpu_ids"] = variant.trainer_gpu_ids
internal_config["inference_gpu_ids"] = variant.inference_gpu_ids
- dev.validate_dedicated_config(internal_config)
+ if not stage_resources_apply:
+ dev.validate_dedicated_config(internal_config)
return internal_config
@@ -704,8 +911,28 @@ async def run_yes_no_trainability_async(
_internal_config=internal_config,
report_metrics=[],
)
+ train_kwargs = _variant_train_kwargs(variant)
+ workflow_resources = handler_workflow_resources_for_base_model(
+ base_model,
+ allow_unvalidated_arch=allow_unvalidated_arch,
+ )
+ stage_resources = (
+ workflow_resources.yes_no_trainability
+ if workflow_resources is not None
+ else None
+ )
+ if stage_resources is not None:
+ stage_resources = resolve_stage_resources_for_current_host(
+ "yes_no_trainability",
+ stage_resources,
+ )
+ backend_env = {
+ **(stage_resources.megatron_env if stage_resources is not None else {}),
+ **(extra_env or {}),
+ }
+
async with _backend_context(
- variant, backend_root=backend_root, extra_env=extra_env
+ variant, backend_root=backend_root, extra_env=backend_env
) as backend:
await model.register(backend)
output_dir = Path(model.base_path) / model.project / "models" / model.name
@@ -756,6 +983,7 @@ async def run_yes_no_trainability_async(
1e-4,
),
loss_fn="cispo",
+ packed_sequence_length=train_kwargs["packed_sequence_length"],
)
await model.log(
train_groups,
diff --git a/tests/unit/test_chat_template_config.py b/tests/unit/test_chat_template_config.py
deleted file mode 100644
index 9ba17ba09..000000000
--- a/tests/unit/test_chat_template_config.py
+++ /dev/null
@@ -1,69 +0,0 @@
-import pytest
-
-from art.local.backend import (
- _apply_configured_chat_template,
- _apply_configured_chat_template_server_args,
- _tokenizer_cache_key,
-)
-
-
-class _Tokenizer:
- chat_template = "base"
-
-
-def test_apply_configured_chat_template_reads_path(tmp_path) -> None:
- path = tmp_path / "chat_template.jinja"
- path.write_text("custom template", encoding="utf-8")
- tokenizer = _Tokenizer()
-
- _apply_configured_chat_template(
- tokenizer, # type: ignore[arg-type]
- {"chat_template_path": str(path)},
- )
-
- assert tokenizer.chat_template == "custom template"
-
-
-def test_apply_configured_chat_template_server_args_preserves_explicit_override() -> (
- None
-):
- config_dict = {"server_args": {"chat_template": "explicit"}}
-
- _apply_configured_chat_template_server_args(
- config_dict,
- {
- "chat_template": "default",
- "chat_template_content_format": "string",
- },
- )
-
- assert config_dict["server_args"] == {
- "chat_template": "explicit",
- "chat_template_content_format": "string",
- }
-
-
-def test_configured_chat_template_rejects_ambiguous_config(tmp_path) -> None:
- path = tmp_path / "chat_template.jinja"
- path.write_text("custom template", encoding="utf-8")
-
- with pytest.raises(ValueError, match="Set only one"):
- _apply_configured_chat_template(
- _Tokenizer(), # type: ignore[arg-type]
- {
- "chat_template": "raw",
- "chat_template_path": str(path),
- },
- )
-
-
-def test_tokenizer_cache_key_includes_chat_template_content(tmp_path) -> None:
- path = tmp_path / "chat_template.jinja"
- path.write_text("first template", encoding="utf-8")
- first_key = _tokenizer_cache_key("base-model", {"chat_template_path": str(path)})
-
- path.write_text("second template", encoding="utf-8")
- second_key = _tokenizer_cache_key("base-model", {"chat_template_path": str(path)})
-
- assert _tokenizer_cache_key("base-model", {}) == ("base-model", None)
- assert first_key != second_key
diff --git a/tests/unit/test_dsv4_vllm_runtime_patches.py b/tests/unit/test_dsv4_vllm_runtime_patches.py
new file mode 100644
index 000000000..1ed497f8d
--- /dev/null
+++ b/tests/unit/test_dsv4_vllm_runtime_patches.py
@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+import sys
+import types
+from types import SimpleNamespace
+
+import torch
+
+
+def _load_dsv4_patches_module():
+ path = (
+ Path(__file__).resolve().parents[2]
+ / "vllm_runtime/src/art_vllm_runtime/dsv4_patches.py"
+ )
+ spec = importlib.util.spec_from_file_location(
+ "_art_vllm_runtime_dsv4_patches",
+ path,
+ )
+ assert spec is not None
+ module = importlib.util.module_from_spec(spec)
+ assert spec.loader is not None
+ spec.loader.exec_module(module)
+ return module
+
+
+def test_dsv4_compressor_helper_uses_punica_metadata_without_full_batch_lora(
+ monkeypatch,
+) -> None:
+ patches = _load_dsv4_patches_module()
+ fake_vllm = types.ModuleType("vllm")
+ fake_platforms = types.ModuleType("vllm.platforms")
+ setattr(
+ fake_platforms,
+ "current_platform",
+ SimpleNamespace(can_update_inplace=lambda: True),
+ )
+ setattr(fake_vllm, "platforms", fake_platforms)
+ monkeypatch.setitem(sys.modules, "vllm", fake_vllm)
+ monkeypatch.setitem(
+ sys.modules,
+ "vllm.platforms",
+ fake_platforms,
+ )
+ monkeypatch.setattr(
+ patches, "_register_dsv4_lora_expand_fp32_output_op", lambda: None
+ )
+
+ expand_calls: list[tuple[tuple[int, ...], int]] = []
+
+ def fake_expand(inputs, lora_b, output, *args):
+ offset = args[-1]
+ width = lora_b.shape[2]
+ expand_calls.append((tuple(lora_b.shape), offset))
+ output[:, offset : offset + width].add_(inputs.sum(dim=-1, keepdim=True))
+
+ monkeypatch.setattr(
+ torch.ops.vllm,
+ "art_dsv4_lora_expand_fp32_output",
+ fake_expand,
+ raising=False,
+ )
+
+ class FakeTokenMappingMeta:
+ def meta_args(self, token_count, specialize_active_lora):
+ assert token_count == 4
+ assert specialize_active_lora is False
+ return (
+ torch.tensor([0, 0, 1, 1], dtype=torch.int32),
+ torch.tensor([0, 1, 2, 3], dtype=torch.int32),
+ torch.tensor([2, 2, 0], dtype=torch.int32),
+ torch.tensor([0, 2, 4, 4], dtype=torch.int32),
+ torch.tensor([0, 1, -1], dtype=torch.int32),
+ torch.tensor([False]),
+ torch.tensor([2], dtype=torch.int32),
+ )
+
+ class FakeWrapper:
+ no_lora = False
+ indices_len = [4]
+ lora_config = SimpleNamespace(specialize_active_lora=False)
+ token_mapping_meta = FakeTokenMappingMeta()
+
+ def add_shrink(self, buffers, x, lora_a_stacked, scale):
+ assert buffers.shape == (2, 4, 2)
+ assert scale == 1.0
+ buffers.copy_(
+ torch.arange(buffers.numel(), dtype=torch.float32).view_as(buffers)
+ )
+ return None
+
+ module = SimpleNamespace(
+ lora_a_stacked=(
+ torch.zeros(2, 1, 2, 4, dtype=torch.bfloat16),
+ torch.zeros(2, 1, 2, 4, dtype=torch.bfloat16),
+ ),
+ lora_b_stacked=(
+ torch.zeros(2, 1, 3, 2, dtype=torch.bfloat16),
+ torch.zeros(2, 1, 5, 2, dtype=torch.bfloat16),
+ ),
+ output_slices=(3, 5),
+ punica_wrapper=FakeWrapper(),
+ tp_size=1,
+ )
+ output = torch.zeros(4, 8, dtype=torch.float32)
+
+ result = patches._apply_dsv4_compressor_lora_to_existing_output(
+ module, torch.zeros(4, 4, dtype=torch.bfloat16), output
+ )
+
+ assert result is output
+ assert expand_calls == [((2, 1, 3, 2), 0), ((2, 1, 5, 2), 3)]
+ assert output.abs().sum() > 0
diff --git a/tests/unit/test_tinker_renderers.py b/tests/unit/test_tinker_renderers.py
index c03f87129..59b01961b 100644
--- a/tests/unit/test_tinker_renderers.py
+++ b/tests/unit/test_tinker_renderers.py
@@ -111,7 +111,7 @@ def test_qwen3_5_parse_response_handles_xml_tool_calls() -> None:
message, success = renderer.parse_response(response)
- assert success == renderers.ParseTermination.STOP_SEQUENCE
+ assert success is True or getattr(success, "value", None) == "stop_sequence"
assert message["content"] == [
{"type": "thinking", "thinking": "reasoning"},
{"type": "text", "text": "Answer first.\n\n"},
diff --git a/tests/unit/test_train_inf_mismatch_logical_map.py b/tests/unit/test_train_inf_mismatch_logical_map.py
new file mode 100644
index 000000000..34987a6e3
--- /dev/null
+++ b/tests/unit/test_train_inf_mismatch_logical_map.py
@@ -0,0 +1,37 @@
+import math
+from pathlib import Path
+import sys
+
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+
+from tests.integration.megatron.train_inf_mismatch.output_parity import (
+ build_logical_token_map,
+)
+
+
+def test_logical_map_excludes_masked_template_tail_token() -> None:
+ packed = {
+ "tokens": torch.tensor([[10, 20, 30, 99, 4, 5, 1, 0]], dtype=torch.long),
+ "group_ids": torch.tensor([[7, 7, 7, 8, 8, 8, 8, -1]], dtype=torch.long),
+ "parent_ids": torch.tensor([[7, 7, 7, 7, 7, 7, 7, -1]], dtype=torch.long),
+ "assistant_mask": torch.tensor(
+ [[False, False, False, False, True, True, False, False]]
+ ),
+ "logprobs": torch.tensor(
+ [[math.nan, math.nan, math.nan, math.nan, -0.1, -0.2, math.nan, math.nan]],
+ dtype=torch.float32,
+ ),
+ }
+
+ logical_map = build_logical_token_map(packed)
+
+ assert len(logical_map.prompts) == 1
+ prompt = logical_map.prompts[0]
+ assert prompt.packed_prompt_length == 3
+ assert prompt.scored_token_start_index == 4
+ assert prompt.token_ids == [10, 20, 30, 99, 4, 5]
+ assert [token.token_id for token in logical_map.tokens] == [4, 5]
+ assert [token.vllm_prompt_token_index for token in logical_map.tokens] == [4, 5]
+ assert [token.art_packed_token_index for token in logical_map.tokens] == [4, 5]
diff --git a/tests/unit/test_vllm_lora_delta.py b/tests/unit/test_vllm_lora_delta.py
new file mode 100644
index 000000000..7b6930a3d
--- /dev/null
+++ b/tests/unit/test_vllm_lora_delta.py
@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+from types import SimpleNamespace
+
+import torch
+
+
+def _load_lora_delta_module():
+ path = (
+ Path(__file__).resolve().parents[2]
+ / "vllm_runtime/src/art_vllm_runtime/lora_delta.py"
+ )
+ spec = importlib.util.spec_from_file_location("_art_vllm_runtime_lora_delta", path)
+ assert spec is not None
+ module = importlib.util.module_from_spec(spec)
+ assert spec.loader is not None
+ spec.loader.exec_module(module)
+ return module
+
+
+def test_additive_weight_loader_uses_legacy_loader_for_plain_merged_column_param():
+ lora_delta = _load_lora_delta_module()
+ param = torch.nn.Parameter(torch.zeros(2, 4))
+ loaded = torch.arange(8, dtype=torch.float32).view(2, 4)
+ calls = []
+
+ class Owner:
+ def weight_loader_v2(self, loader_param, loaded_weight, shard_id):
+ del shard_id
+ loader_param.load_merged_column_weight(loaded_weight=loaded_weight)
+
+ def weight_loader(self, loader_param, loaded_weight, shard_id):
+ calls.append((loader_param, shard_id))
+ loader_param.data.copy_(loaded_weight)
+
+ owner = Owner()
+ loader = lora_delta._additive_weight_loader(param, owner.weight_loader_v2)
+ result = loader(param, loaded, 0)
+
+ assert result is None
+ assert calls == [(param, 0)]
+ assert torch.equal(param, loaded)
+
+
+def test_additive_weight_loader_keeps_v2_for_vllm_parameter_like_param():
+ lora_delta = _load_lora_delta_module()
+ param = torch.nn.Parameter(torch.zeros(2, 4))
+ loaded = torch.arange(8, dtype=torch.float32).view(2, 4)
+ calls = []
+
+ def load_merged_column_weight(*, loaded_weight, **_kwargs):
+ calls.append("v2")
+ param.data.copy_(loaded_weight)
+
+ setattr(param, "load_merged_column_weight", load_merged_column_weight)
+ owner = SimpleNamespace(
+ weight_loader_v2=lambda loader_param, loaded_weight, shard_id: (
+ loader_param.load_merged_column_weight(loaded_weight=loaded_weight)
+ ),
+ weight_loader=lambda *_args, **_kwargs: calls.append("legacy"),
+ )
+ loader = lora_delta._additive_weight_loader(param, owner.weight_loader_v2)
+ loader(param, loaded, 0)
+
+ assert calls == ["v2"]
+ assert torch.equal(param, loaded)
diff --git a/uv.lock b/uv.lock
index 36ad513d8..149ff5f9e 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1220,7 +1220,7 @@ name = "cuda-bindings"
version = "12.9.7"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "cuda-pathfinder", marker = "sys_platform == 'linux' or sys_platform == 'win32' or extra == 'extra-12-openpipe-art-megatron'" },
+ { name = "cuda-pathfinder", marker = "sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-12-openpipe-art-megatron') or (sys_platform == 'win32' and extra != 'extra-12-openpipe-art-backend') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron')" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/32/45/557d4ed1fa54f0c7db8aee083229f624990d69f7d00f55477eed5c7e169a/cuda_bindings-12.9.7-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0666d3c082ef8f4b2d670950589373550e9f3bf564d635dd883f24a0b40402ff", size = 7071026, upload-time = "2026-05-27T18:44:13.356Z" },
@@ -1291,37 +1291,37 @@ wheels = [
[package.optional-dependencies]
cublas = [
- { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux' or (sys_platform == 'win32' and extra != 'extra-12-openpipe-art-backend' and extra != 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
]
cudart = [
- { name = "nvidia-cuda-runtime-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "nvidia-cuda-runtime-cu12", marker = "sys_platform == 'linux' or (sys_platform == 'win32' and extra != 'extra-12-openpipe-art-backend' and extra != 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
]
cufft = [
- { name = "nvidia-cufft-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "nvidia-cufft-cu12", marker = "sys_platform == 'linux' or (sys_platform == 'win32' and extra != 'extra-12-openpipe-art-backend' and extra != 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
]
cufile = [
- { name = "nvidia-cufile-cu12", marker = "sys_platform == 'linux' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "nvidia-cufile-cu12", marker = "sys_platform == 'linux'" },
]
cupti = [
- { name = "nvidia-cuda-cupti-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "nvidia-cuda-cupti-cu12", marker = "sys_platform == 'linux' or (sys_platform == 'win32' and extra != 'extra-12-openpipe-art-backend' and extra != 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
]
curand = [
- { name = "nvidia-curand-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "nvidia-curand-cu12", marker = "sys_platform == 'linux' or (sys_platform == 'win32' and extra != 'extra-12-openpipe-art-backend' and extra != 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
]
cusolver = [
- { name = "nvidia-cusolver-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "nvidia-cusolver-cu12", marker = "sys_platform == 'linux' or (sys_platform == 'win32' and extra != 'extra-12-openpipe-art-backend' and extra != 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
]
cusparse = [
- { name = "nvidia-cusparse-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "nvidia-cusparse-cu12", marker = "sys_platform == 'linux' or (sys_platform == 'win32' and extra != 'extra-12-openpipe-art-backend' and extra != 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
]
nvjitlink = [
- { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux' or (sys_platform == 'win32' and extra != 'extra-12-openpipe-art-backend' and extra != 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
]
nvrtc = [
- { name = "nvidia-cuda-nvrtc-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "nvidia-cuda-nvrtc-cu12", marker = "sys_platform == 'linux' or (sys_platform == 'win32' and extra != 'extra-12-openpipe-art-backend' and extra != 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
]
nvtx = [
- { name = "nvidia-nvtx-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "nvidia-nvtx-cu12", marker = "sys_platform == 'linux' or (sys_platform == 'win32' and extra != 'extra-12-openpipe-art-backend' and extra != 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
]
[[package]]
@@ -1938,7 +1938,7 @@ version = "0.5.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "fla-core" },
- { name = "transformers", version = "5.6.2", source = { registry = "https://pypi.org/simple" } },
+ { name = "transformers", version = "5.12.1", source = { registry = "https://pypi.org/simple" } },
]
sdist = { url = "https://files.pythonhosted.org/packages/79/5c/1db76cc829c951117a3112f306d50333bd71399d2e35807fe7c99ffc2007/flash_linear_attention-0.5.0.tar.gz", hash = "sha256:22b789a47f07738b4382ecdf775d7bb40e0d803c467c34f8e2ecd6a1dc780938", size = 160419, upload-time = "2026-04-21T20:25:42.344Z" }
wheels = [
@@ -2502,7 +2502,7 @@ name = "gunicorn"
version = "25.3.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "packaging" },
+ { name = "packaging", marker = "sys_platform != 'win32'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/c4/f4/e78fa054248fab913e2eab0332c6c2cb07421fca1ce56d8fe43b6aef57a4/gunicorn-25.3.0.tar.gz", hash = "sha256:f74e1b2f9f76f6cd1ca01198968bd2dd65830edc24b6e8e4d78de8320e2fe889", size = 634883, upload-time = "2026-03-27T00:00:26.092Z" }
wheels = [
@@ -3844,7 +3844,7 @@ dependencies = [
{ name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-12-openpipe-art-megatron') or (sys_platform == 'linux' and extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (sys_platform == 'linux' and extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker') or (sys_platform == 'win32' and extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (sys_platform == 'win32' and extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
{ name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "(sys_platform == 'linux' and extra == 'extra-12-openpipe-art-megatron') or (sys_platform == 'win32' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
{ name = "tqdm" },
- { name = "transformers", version = "5.6.2", source = { registry = "https://pypi.org/simple" } },
+ { name = "transformers", version = "5.12.1", source = { registry = "https://pypi.org/simple" } },
{ name = "typing-extensions" },
{ name = "wandb" },
]
@@ -4394,7 +4394,7 @@ name = "nvidia-cudnn-cu12"
version = "9.19.0.56"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux' or (sys_platform == 'win32' and extra != 'extra-12-openpipe-art-backend' and extra != 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/09/b8/277c51962ee46fa3e5b203ac5f76107c650f781d6891e681e28e6f3e9fe6/nvidia_cudnn_cu12-9.19.0.56-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:08caaf27fe556aca82a3ee3b5aa49a77e7de0cfcb7ff4e5c29da426387a8267e", size = 656910700, upload-time = "2026-02-03T20:40:25.508Z" },
@@ -4423,7 +4423,7 @@ name = "nvidia-cufft-cu12"
version = "11.3.3.83"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux' or (sys_platform == 'win32' and extra != 'extra-12-openpipe-art-backend' and extra != 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/60/bc/7771846d3a0272026c416fbb7e5f4c1f146d6d80704534d0b187dd6f4800/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:848ef7224d6305cdb2a4df928759dca7b1201874787083b6e7550dd6765ce69a", size = 193109211, upload-time = "2025-03-07T01:44:56.873Z" },
@@ -4455,9 +4455,9 @@ name = "nvidia-cusolver-cu12"
version = "11.7.3.90"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
- { name = "nvidia-cusparse-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
- { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux' or (sys_platform == 'win32' and extra != 'extra-12-openpipe-art-backend' and extra != 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "nvidia-cusparse-cu12", marker = "sys_platform == 'linux' or (sys_platform == 'win32' and extra != 'extra-12-openpipe-art-backend' and extra != 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux' or (sys_platform == 'win32' and extra != 'extra-12-openpipe-art-backend' and extra != 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/c8/32/f7cd6ce8a7690544d084ea21c26e910a97e077c9b7f07bf5de623ee19981/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:db9ed69dbef9715071232caa9b69c52ac7de3a95773c2db65bdba85916e4e5c0", size = 267229841, upload-time = "2025-03-07T01:46:54.356Z" },
@@ -4470,7 +4470,7 @@ name = "nvidia-cusparse-cu12"
version = "12.5.8.93"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux' or (sys_platform == 'win32' and extra != 'extra-12-openpipe-art-backend' and extra != 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/bc/f7/cd777c4109681367721b00a106f491e0d0d15cfa1fd59672ce580ce42a97/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:9b6c161cb130be1a07a27ea6923df8141f3c295852f4b260c65f18f3e0a091dc", size = 288117129, upload-time = "2025-03-07T01:47:40.407Z" },
@@ -4799,7 +4799,7 @@ megatron = [
{ name = "transformer-engine" },
{ name = "transformer-engine-cu12" },
{ name = "transformer-engine-torch" },
- { name = "transformers", version = "5.6.2", source = { registry = "https://pypi.org/simple" } },
+ { name = "transformers", version = "5.12.1", source = { registry = "https://pypi.org/simple" } },
]
plotting = [
{ name = "matplotlib" },
@@ -4905,7 +4905,7 @@ requires-dist = [
{ name = "transformer-engine-cu12", marker = "extra == 'megatron'", specifier = "==2.11.0" },
{ name = "transformer-engine-torch", marker = "extra == 'megatron'", git = "https://github.com/NVIDIA/TransformerEngine.git?subdirectory=transformer_engine%2Fpytorch&rev=v2.11" },
{ name = "transformers", marker = "extra == 'backend'", specifier = "==5.2.0" },
- { name = "transformers", marker = "extra == 'megatron'", specifier = "==5.6.2" },
+ { name = "transformers", marker = "extra == 'megatron'", specifier = "==5.12.1" },
{ name = "transformers", marker = "extra == 'tinker'", specifier = ">=5.2.0,<=5.5.3" },
{ name = "trl", marker = "extra == 'backend'", specifier = "==0.20.0" },
{ name = "typer", specifier = ">=0.15.2" },
@@ -5236,7 +5236,7 @@ dependencies = [
{ name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "(sys_platform == 'linux' and extra == 'extra-12-openpipe-art-backend') or (sys_platform == 'linux' and extra == 'extra-12-openpipe-art-megatron') or (sys_platform == 'win32' and extra == 'extra-12-openpipe-art-backend') or (sys_platform == 'win32' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
{ name = "tqdm" },
{ name = "transformers", version = "5.2.0", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-12-openpipe-art-backend' or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
- { name = "transformers", version = "5.6.2", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-12-openpipe-art-megatron'" },
+ { name = "transformers", version = "5.12.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-12-openpipe-art-megatron'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/86/cf/037f1e3d5186496c05513a6754639e2dab3038a05f384284d49a9bd06a2d/peft-0.19.1.tar.gz", hash = "sha256:0d97542fe96dcdaa20d3b81c06f26f988618f416a73544ab23c3618ccb674a40", size = 763738, upload-time = "2026-04-16T15:46:45.105Z" }
wheels = [
@@ -7696,16 +7696,16 @@ name = "tilelang"
version = "0.1.10"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "apache-tvm-ffi", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
- { name = "cloudpickle", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
- { name = "ml-dtypes", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
- { name = "numpy", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
- { name = "psutil", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
- { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
- { name = "torch-c-dlpack-ext", marker = "(python_full_version < '3.14' and sys_platform == 'linux') or (python_full_version < '3.14' and sys_platform == 'win32')" },
- { name = "tqdm", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
- { name = "typing-extensions", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
- { name = "z3-solver", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+ { name = "apache-tvm-ffi", marker = "platform_machine != 's390x' and sys_platform == 'linux'" },
+ { name = "cloudpickle", marker = "platform_machine != 's390x' and sys_platform == 'linux'" },
+ { name = "ml-dtypes", marker = "platform_machine != 's390x' and sys_platform == 'linux'" },
+ { name = "numpy", marker = "platform_machine != 's390x' and sys_platform == 'linux'" },
+ { name = "psutil", marker = "platform_machine != 's390x' and sys_platform == 'linux'" },
+ { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "platform_machine != 's390x' and sys_platform == 'linux'" },
+ { name = "torch-c-dlpack-ext", marker = "python_full_version < '3.14' and platform_machine != 's390x' and sys_platform == 'linux'" },
+ { name = "tqdm", marker = "platform_machine != 's390x' and sys_platform == 'linux'" },
+ { name = "typing-extensions", marker = "platform_machine != 's390x' and sys_platform == 'linux'" },
+ { name = "z3-solver", marker = "platform_machine != 's390x' and sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/77/5c/07146b4527656102e48d21c2599aa80477e83ea3f149ac0df3b15a247bd4/tilelang-0.1.10.tar.gz", hash = "sha256:d8813e668fcf75843bc2d68c633c352b419c1e292895a6038a4aadd943e56c2b", size = 93184128, upload-time = "2026-05-25T03:58:57.006Z" }
wheels = [
@@ -7953,20 +7953,20 @@ resolution-markers = [
"(python_full_version < '3.13' and sys_platform == 'linux' and extra != 'extra-12-openpipe-art-backend' and extra != 'extra-12-openpipe-art-megatron') or (python_full_version < '3.13' and sys_platform == 'win32' and extra != 'extra-12-openpipe-art-backend' and extra != 'extra-12-openpipe-art-megatron')",
]
dependencies = [
- { name = "cuda-bindings", marker = "sys_platform == 'linux' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "cuda-bindings", marker = "sys_platform == 'linux'" },
{ name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
- { name = "filelock", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
- { name = "fsspec", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
- { name = "jinja2", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
- { name = "networkx", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
- { name = "nvidia-cudnn-cu12", marker = "sys_platform == 'linux' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
- { name = "nvidia-cusparselt-cu12", marker = "sys_platform == 'linux' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
- { name = "nvidia-nccl-cu12", marker = "sys_platform == 'linux' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
- { name = "nvidia-nvshmem-cu12", marker = "sys_platform == 'linux' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
- { name = "setuptools", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
- { name = "sympy", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
- { name = "triton", marker = "sys_platform == 'linux' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
- { name = "typing-extensions", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" },
+ { name = "filelock", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+ { name = "fsspec", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+ { name = "jinja2", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+ { name = "networkx", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+ { name = "nvidia-cudnn-cu12", marker = "sys_platform == 'linux'" },
+ { name = "nvidia-cusparselt-cu12", marker = "sys_platform == 'linux'" },
+ { name = "nvidia-nccl-cu12", marker = "sys_platform == 'linux'" },
+ { name = "nvidia-nvshmem-cu12", marker = "sys_platform == 'linux'" },
+ { name = "setuptools", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+ { name = "sympy", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+ { name = "triton", marker = "sys_platform == 'linux'" },
+ { name = "typing-extensions", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
wheels = [
{ url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:9c8f38efee365cb9d334de8a83ce52fc7e5fc9e5a7b0853285efa1b69e00b0f2", upload-time = "2026-04-27T17:41:30Z" },
@@ -8178,7 +8178,7 @@ wheels = [
[[package]]
name = "transformers"
-version = "5.6.2"
+version = "5.12.1"
source = { registry = "https://pypi.org/simple" }
resolution-markers = [
"python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'linux'",
@@ -8211,9 +8211,9 @@ dependencies = [
{ name = "tqdm", marker = "extra == 'extra-12-openpipe-art-megatron'" },
{ name = "typer", marker = "extra == 'extra-12-openpipe-art-megatron'" },
]
-sdist = { url = "https://files.pythonhosted.org/packages/a4/e9/c6c80a07690142a7d05444271f47b9f3c8aac7dea01d52e1137ee480ad78/transformers-5.6.2.tar.gz", hash = "sha256:e657134c3e5a6bc00a3c35f4e2674bb51adfcd89898495b788a18552bac2b91a", size = 8311867, upload-time = "2026-04-23T18:33:29.332Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/aa/7c/8240f612819718100a9346dc28dea6a11370c3ca9c8c6eabadd3dea4ef29/transformers-5.12.1.tar.gz", hash = "sha256:679ee731c8225347889ad4fb3b2c926a62e9da3b7d284e9d12c791da7272466b", size = 8924054, upload-time = "2026-06-15T17:27:50.604Z" }
wheels = [
- { url = "https://files.pythonhosted.org/packages/5d/95/0b0218149b0d6f14df35f5b8f676fa83df4f19ed253c3cc447107ef86eca/transformers-5.6.2-py3-none-any.whl", hash = "sha256:f8d3a1bb96778fed9b8aabfd0dd6e19843e4b0f2bb6b59f32b8a92051b0f348f", size = 10364898, upload-time = "2026-04-23T18:33:26.081Z" },
+ { url = "https://files.pythonhosted.org/packages/df/56/bbd60dd8668055803bf8ba55a81f9b8a8b31497f620109a9671d26a2076d/transformers-5.12.1-py3-none-any.whl", hash = "sha256:2a5e109d2021265df7098ffbb738295acaf5ad256f12cbc586db2ea4dcbb1a8a", size = 11150587, upload-time = "2026-06-15T17:27:46.679Z" },
]
[[package]]
diff --git a/vllm_runtime/pyproject.toml b/vllm_runtime/pyproject.toml
index 673f66585..412a9c1c1 100644
--- a/vllm_runtime/pyproject.toml
+++ b/vllm_runtime/pyproject.toml
@@ -5,7 +5,7 @@ description = "Tiny ART-owned vLLM runtime package"
requires-python = ">=3.12,<3.13"
dependencies = [
"nvidia-nccl-cu12==2.28.9 ; sys_platform == 'linux'",
- "transformers==5.6.2",
+ "transformers==5.12.1",
"vllm @ https://github.com/vllm-project/vllm/releases/download/v0.23.0/vllm-0.23.0%2Bcu129-cp38-abi3-manylinux_2_28_x86_64.whl ; sys_platform == 'linux'",
]
@@ -37,5 +37,5 @@ override-dependencies = [
"torch @ https://download.pytorch.org/whl/test/cu128/torch-2.11.0%2Bcu128-cp312-cp312-manylinux_2_28_x86_64.whl",
"torchaudio @ https://download.pytorch.org/whl/test/cu128/torchaudio-2.11.0%2Bcu128-cp312-cp312-manylinux_2_28_x86_64.whl",
"torchvision @ https://download.pytorch.org/whl/test/cu128/torchvision-0.26.0%2Bcu128-cp312-cp312-manylinux_2_28_x86_64.whl",
- "transformers==5.6.2",
+ "transformers==5.12.1",
]
diff --git a/vllm_runtime/src/art_vllm_runtime/dedicated_server.py b/vllm_runtime/src/art_vllm_runtime/dedicated_server.py
index 856c055b9..535960338 100644
--- a/vllm_runtime/src/art_vllm_runtime/dedicated_server.py
+++ b/vllm_runtime/src/art_vllm_runtime/dedicated_server.py
@@ -15,7 +15,7 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
parser.add_argument("--port", type=int, required=True)
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--cuda-visible-devices", required=True)
- parser.add_argument("--lora-path", required=True, help="Initial checkpoint path")
+ parser.add_argument("--lora-path", help="Optional initial checkpoint path")
parser.add_argument("--served-model-name", required=True)
parser.add_argument(
"--rollout-weights-mode",
@@ -139,6 +139,8 @@ def _append_cli_arg(vllm_args: list[str], key: str, value: object) -> None:
def main(argv: list[str] | None = None) -> None:
args = parse_args(argv)
+ if args.rollout_weights_mode == "merged" and not args.lora_path:
+ raise SystemExit("--lora-path is required for merged rollout weights")
engine_args = json.loads(args.engine_args_json)
server_args = json.loads(args.server_args_json)
@@ -164,12 +166,11 @@ def main(argv: list[str] | None = None) -> None:
f"--served-model-name={args.served_model_name}",
]
if args.rollout_weights_mode == "lora":
- vllm_args.extend(
- [
- "--enable-lora",
- f"--lora-modules={args.served_model_name}={args.lora_path}",
- ]
- )
+ vllm_args.append("--enable-lora")
+ if args.lora_path:
+ vllm_args.append(
+ f"--lora-modules={args.served_model_name}={args.lora_path}"
+ )
for extra_args in (engine_args, server_args):
for key, value in extra_args.items():
_append_cli_arg(vllm_args, key, value)
diff --git a/vllm_runtime/src/art_vllm_runtime/dsv4_patches.py b/vllm_runtime/src/art_vllm_runtime/dsv4_patches.py
new file mode 100644
index 000000000..383bfe6ca
--- /dev/null
+++ b/vllm_runtime/src/art_vllm_runtime/dsv4_patches.py
@@ -0,0 +1,1603 @@
+"""DSV4-specific monkey patches for the ART-owned vLLM runtime."""
+
+import importlib
+from typing import Any
+
+
+def apply_dsv4_vllm_runtime_patches() -> None:
+ patch_layerwise_reload_shadow_attrs()
+ patch_dsv4_attn_sink_layerwise_reload()
+ patch_dsv4_mhc_pre_fixed_split()
+ patch_dsv4_lora_support()
+ patch_dsv4_mla_lora_aliases()
+ patch_dsv4_fast_path_lora()
+ patch_dsv4_triton_moe_topk6_routing()
+ patch_lora_linear_base_attr_proxy()
+ patch_marlin_lora_swiglu_limit()
+
+
+def _drop_reload_shadow_attrs(layer: Any, names: Any) -> None:
+ for name in names:
+ if (
+ name in getattr(layer, "__dict__", {})
+ and name not in layer._parameters
+ and name not in layer._buffers
+ and name not in layer._modules
+ ):
+ delattr(layer, name)
+
+
+def patch_layerwise_reload_shadow_attrs() -> None:
+ """Allow vLLM layerwise reload to restore processed DSV4 MegaMoE params.
+
+ DeepSeek V4 MegaMoE drops loader-side Parameters after transforming them for
+ DeepGEMM. Some vLLM builds leave same-name plain attributes behind; PyTorch
+ then rejects register_parameter during the next checkpoint-format reload.
+ """
+ from vllm.model_executor.model_loader.reload import layerwise, meta
+
+ if getattr(meta, "_art_reload_shadow_attrs_patched", False):
+ return
+
+ original_restore_layer_on_meta = meta.restore_layer_on_meta
+ original_place_kernel_tensors = layerwise._place_kernel_tensors
+
+ def restore_layer_on_meta(layer: Any, info: Any) -> None:
+ restore_params, restore_buffers = info.restore_metadata
+ _drop_reload_shadow_attrs(layer, tuple(restore_params) + tuple(restore_buffers))
+ return original_restore_layer_on_meta(layer, info)
+
+ def _place_kernel_tensors(layer: Any, info: Any) -> None:
+ assert info.kernel_tensors is not None
+ parameters, buffers = info.kernel_tensors
+ _drop_reload_shadow_attrs(layer, tuple(parameters) + tuple(buffers))
+ return original_place_kernel_tensors(layer, info)
+
+ restore_layer_on_meta.__art_patched__ = True # type: ignore[attr-defined]
+ _place_kernel_tensors.__art_patched__ = True # type: ignore[attr-defined]
+ meta.restore_layer_on_meta = restore_layer_on_meta # type: ignore[method-assign]
+ layerwise.restore_layer_on_meta = restore_layer_on_meta # type: ignore[method-assign]
+ layerwise._place_kernel_tensors = _place_kernel_tensors # type: ignore[method-assign]
+ setattr(meta, "_art_reload_shadow_attrs_patched", True)
+
+
+def _import_dsv4_model_module() -> Any | None:
+ for module_name in (
+ "vllm.model_executor.models.deepseek_v4",
+ "vllm.models.deepseek_v4.nvidia.model",
+ ):
+ try:
+ return importlib.import_module(module_name)
+ except ImportError:
+ continue
+ return None
+
+
+def patch_dsv4_attn_sink_layerwise_reload() -> None:
+ """Route DSV4 attention-sink loads through vLLM's layerwise loader.
+
+ Merged-weight transfer uses vLLM checkpoint-format reload. During that path,
+ every loadable parameter must be applied through its `weight_loader`; direct
+ `copy_` into `attn_sink` bypasses layerwise accounting and finalize restores
+ the old kernel tensor. With `load_format=dummy`, that old tensor is the
+ initialized sink, not the checkpoint sink.
+ """
+ dsv4_model = _import_dsv4_model_module()
+ if dsv4_model is None:
+ return
+ from vllm.model_executor.models.utils import is_pp_missing_parameter
+
+ model_cls = getattr(dsv4_model, "DeepseekV4Model", None)
+ if model_cls is None:
+ return
+ original = model_cls.load_weights
+ if getattr(original, "__art_patched__", False):
+ return
+
+ def load_weights(self: Any, weights: Any) -> set[str]:
+ stacked_params_mapping = [
+ ("gate_up_proj", "w1", 0),
+ ("gate_up_proj", "w3", 1),
+ ("attn.fused_wqa_wkv", "attn.wq_a", 0),
+ ("attn.fused_wqa_wkv", "attn.wkv", 1),
+ ("compressor.fused_wkv_wgate", "compressor.wkv", 0),
+ ("compressor.fused_wkv_wgate", "compressor.wgate", 1),
+ ]
+ params_dict = dict(self.named_parameters())
+ loaded_params: set[str] = set()
+
+ tp_size = dsv4_model.get_tensor_model_parallel_world_size()
+ tp_rank = dsv4_model.get_tensor_model_parallel_rank()
+ n_head = self.config.num_attention_heads
+ n_local_head = n_head // tp_size
+ head_rank_start = n_local_head * tp_rank
+ head_rank_end = n_local_head * (tp_rank + 1)
+ expert_mapping = self.get_expert_mapping()
+
+ for name, loaded_weight in weights:
+ for param_name, weight_name, shard_id in stacked_params_mapping:
+ if ".experts." in name:
+ continue
+ if weight_name not in name:
+ continue
+ name = name.replace(weight_name, param_name)
+
+ if is_pp_missing_parameter(name, self):
+ break
+ param = params_dict[name]
+ param.weight_loader(param, loaded_weight, shard_id)
+ loaded_params.add(name)
+ break
+ else:
+ if ".experts." in name:
+ if (
+ "weight_scale" in name
+ and loaded_weight.dtype == dsv4_model.torch.float8_e8m0fnu
+ ):
+ loaded_weight = loaded_weight.view(dsv4_model.torch.uint8)
+ for mapping in expert_mapping:
+ param_name, weight_name, expert_id, expert_shard_id = mapping
+ if weight_name not in name:
+ continue
+ name_mapped = name.replace(weight_name, param_name)
+ if is_pp_missing_parameter(name_mapped, self):
+ continue
+ param = params_dict[name_mapped]
+ success = param.weight_loader(
+ param,
+ loaded_weight,
+ name_mapped,
+ shard_id=expert_shard_id,
+ expert_id=expert_id,
+ return_success=True,
+ )
+ if success:
+ name = name_mapped
+ break
+ loaded_params.add(name_mapped)
+ continue
+ if "attn_sink" in name:
+ if is_pp_missing_parameter(name, self):
+ continue
+ param = params_dict[name]
+ narrow_weight = loaded_weight[head_rank_start:head_rank_end]
+ padded_weight = loaded_weight.new_full(
+ tuple(param.shape), -float("inf")
+ )
+ padded_weight[: narrow_weight.shape[0]].copy_(narrow_weight)
+ weight_loader = getattr(
+ param, "weight_loader", dsv4_model.default_weight_loader
+ )
+ weight_loader(param, padded_weight)
+ loaded_params.add(name)
+ continue
+
+ if is_pp_missing_parameter(name, self):
+ continue
+ param = params_dict[name]
+ weight_loader = getattr(
+ param, "weight_loader", dsv4_model.default_weight_loader
+ )
+ weight_loader(param, loaded_weight)
+ loaded_params.add(name)
+
+ return loaded_params
+
+ load_weights.__art_patched__ = True # type: ignore[attr-defined]
+ model_cls.load_weights = load_weights # type: ignore[method-assign]
+
+
+def patch_dsv4_mhc_pre_fixed_split() -> None:
+ """Make DSV4 mHC pre reductions invariant to total prefill length.
+
+ vLLM's current DSV4 mHC pre path chooses split-K from ``num_tokens``. The
+ same prompt prefix can therefore use a different reduction tree when a
+ suffix is appended, producing float32-level post/comb mix drift that becomes
+ bf16 output drift and later MoE route changes. Pinning the DSV4 prenorm
+ shape to the TileLang kernel default split count keeps the reduction plan
+ stable without changing model math.
+ """
+ try:
+ mhc = importlib.import_module("vllm.model_executor.layers.mhc")
+ except ImportError:
+ return
+ original = getattr(mhc, "compute_num_split", None)
+ if original is None or getattr(original, "__art_dsv4_fixed_split_patched__", False):
+ return
+
+ def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:
+ if block_k == 64 and k == 16_384:
+ return 16
+ return original(block_k, k, grid_size)
+
+ compute_num_split.__art_dsv4_fixed_split_patched__ = True # type: ignore[attr-defined]
+ mhc.compute_num_split = compute_num_split
+
+
+def patch_dsv4_lora_support() -> None:
+ """Enable vLLM's existing LoRA manager for ART-served DSV4.
+
+ DSV4 itself does not need a custom LoRA executor here. Once the model
+ advertises packed MLA/shared-expert modules and MoE expert children, vLLM
+ wraps the same FusedMoE module it already uses for serving. With LoRA
+ enabled, vLLM's modular MoE selector picks Marlin, whose expert backend
+ supports fused MoE LoRA. Do not point this patch at the FlashInfer TRTLLM
+ MXFP4 backend; that backend currently has no LoRA hooks.
+ """
+ dsv4_model = _import_dsv4_model_module()
+ if dsv4_model is None:
+ return
+ model_cls = getattr(dsv4_model, "DeepseekV4ForCausalLM", None)
+ if model_cls is None or getattr(model_cls, "_art_dsv4_lora_patched", False):
+ return
+ model_cls.supports_lora = True
+ model_cls.embedding_modules = {}
+ model_cls.packed_modules_mapping = {
+ "fused_wqa_wkv": ["wq_a", "wkv"],
+ "fused_wkv_wgate": ["wkv", "wgate"],
+ "gate_up_proj": ["gate_proj", "up_proj"],
+ }
+ model_cls.is_3d_moe_weight = False
+ model_cls.is_non_gated_moe = False
+ model_cls.lora_skip_prefixes = ["mtp", "indexer"]
+ model_cls._art_dsv4_lora_patched = True
+ _patch_dsv4_lora_manager_indexer_skip(model_cls)
+
+
+def _patch_dsv4_lora_manager_indexer_skip(model_cls: type) -> None:
+ from vllm.lora.model_manager import LoRAModelManager
+
+ original = LoRAModelManager._match_target_modules
+ if getattr(original, "__art_dsv4_indexer_skip_patched__", False):
+ return
+
+ duplicate_mla_aliases = {"fused_wqa_wkv", "wq_b", "wo_a", "wo_b"}
+
+ def _match_target_modules(self: Any, module_name: str) -> bool:
+ if isinstance(self.model, model_cls) and ".indexer." in module_name:
+ return False
+ if (
+ isinstance(self.model, model_cls)
+ and ".attn.mla_attn." in module_name
+ and module_name.rsplit(".", 1)[-1] in duplicate_mla_aliases
+ ):
+ return False
+ return original(self, module_name)
+
+ _match_target_modules.__art_dsv4_indexer_skip_patched__ = True # type: ignore[attr-defined]
+ LoRAModelManager._match_target_modules = _match_target_modules # type: ignore[method-assign]
+
+
+def patch_dsv4_mla_lora_aliases() -> None:
+ """Keep DSV4 MLA wrapper references aligned with vLLM LoRA replacements.
+
+ DeepSeek V4 stores direct references to several attention linears inside
+ ``mla_attn`` during model construction. vLLM's LoRA manager later replaces
+ the canonical modules on ``attn``. Without refreshing the aliases, the DSV4
+ custom attention path keeps calling the original unwrapped linears.
+ """
+ from vllm.lora.model_manager import LoRAModelManager
+
+ original = LoRAModelManager.register_module
+ if getattr(original, "__art_dsv4_mla_alias_patched__", False):
+ return
+
+ alias_names = {"fused_wqa_wkv", "wq_b", "wo_a", "wo_b"}
+
+ def register_module(self: Any, module_name: str, module: Any) -> Any:
+ result = original(self, module_name, module)
+ leaf = module_name.rsplit(".", 1)[-1]
+ if leaf not in alias_names:
+ return result
+ parent_name, _, _ = module_name.rpartition(".")
+ if parent_name.endswith(".mla_attn"):
+ mla_name = parent_name
+ else:
+ mla_name = f"{parent_name}.mla_attn"
+ try:
+ mla_attn = self.model.get_submodule(mla_name)
+ except AttributeError:
+ return result
+ if mla_attn is not None and hasattr(mla_attn, leaf):
+ setattr(mla_attn, leaf, module)
+ return result
+
+ register_module.__art_dsv4_mla_alias_patched__ = True # type: ignore[attr-defined]
+ LoRAModelManager.register_module = register_module # type: ignore[method-assign]
+
+
+def _is_lora_wrapped_linear(module: Any) -> bool:
+ return all(
+ hasattr(module, name)
+ for name in ("lora_a_stacked", "lora_b_stacked", "punica_wrapper")
+ )
+
+
+def _apply_lora_to_existing_linear_output(
+ module: Any,
+ x: Any,
+ output: Any,
+) -> Any:
+ if not _is_lora_wrapped_linear(module):
+ return output
+ wrapper = module.punica_wrapper
+ if getattr(wrapper, "no_lora", False):
+ return output
+ if getattr(wrapper, "indices_len", [None])[0] is None:
+ return output
+ return module._apply_lora_to_output(x, output)
+
+
+def _register_dsv4_lora_expand_fp32_output_op() -> None:
+ if getattr(_register_dsv4_lora_expand_fp32_output_op, "_registered", False):
+ return
+
+ import torch
+ from vllm import envs
+ from vllm.lora.ops.triton_ops.utils import get_lora_op_configs, supports_pdl
+ from vllm.triton_utils import tl, triton
+ from vllm.utils.torch_utils import direct_register_custom_op
+
+ @triton.jit
+ def _kernel(
+ input_ptr,
+ lora_b_ptr,
+ output_ptr,
+ token_indices_sorted_by_lora_ids,
+ num_tokens_per_lora,
+ lora_token_start_loc,
+ lora_ids,
+ M,
+ slice_offset,
+ input_stride_m,
+ input_stride_k,
+ lora_stride_lora,
+ lora_stride_out,
+ lora_stride_k,
+ output_stride_m,
+ output_stride_n,
+ RANK: tl.constexpr,
+ OUT_WIDTH: tl.constexpr,
+ BLOCK_M: tl.constexpr,
+ BLOCK_N: tl.constexpr,
+ BLOCK_K: tl.constexpr,
+ USE_GDC: tl.constexpr,
+ launch_pdl: tl.constexpr,
+ ):
+ cta_n_num = tl.cdiv(OUT_WIDTH, BLOCK_N)
+ cta_m_num = tl.cdiv(M, BLOCK_M)
+ pid_mn = tl.program_id(0)
+ pid_m = pid_mn % cta_m_num
+ pid_n = (pid_mn // cta_m_num) % cta_n_num
+ lora_idx = tl.program_id(1)
+
+ lora_id = tl.load(lora_ids + lora_idx)
+ if lora_id == -1:
+ return
+
+ lora_m_size = tl.load(num_tokens_per_lora + lora_idx)
+ cta_m_offset = pid_m * BLOCK_M
+ if cta_m_offset >= lora_m_size:
+ return
+
+ cta_m_len = min(BLOCK_M, lora_m_size - cta_m_offset)
+ lora_m_start = tl.load(lora_token_start_loc + lora_idx)
+ row_ptr = token_indices_sorted_by_lora_ids + lora_m_start + cta_m_offset
+ row_offsets = tl.arange(0, BLOCK_M) % cta_m_len
+ rows = tl.load(row_ptr + row_offsets)
+
+ offs_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+ accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+ for k_start in range(0, RANK, BLOCK_K):
+ offs_k = tl.arange(0, BLOCK_K) + k_start
+ x = tl.load(
+ input_ptr
+ + rows[:, None] * input_stride_m
+ + offs_k[None, :] * input_stride_k,
+ mask=(row_offsets[:, None] < cta_m_len) & (offs_k[None, :] < RANK),
+ other=0.0,
+ )
+ if USE_GDC:
+ tl.extra.cuda.gdc_wait()
+ b = tl.load(
+ lora_b_ptr
+ + lora_id * lora_stride_lora
+ + offs_n[None, :] * lora_stride_out
+ + offs_k[:, None] * lora_stride_k,
+ mask=(offs_k[:, None] < RANK) & (offs_n[None, :] < OUT_WIDTH),
+ other=0.0,
+ ).to(tl.float32)
+ accumulator += tl.dot(x, b)
+
+ out_cols = slice_offset + offs_n
+ out_ptrs = (
+ output_ptr
+ + rows[:, None] * output_stride_m
+ + out_cols[None, :] * output_stride_n
+ )
+ mask = (row_offsets[:, None] < cta_m_len) & (offs_n[None, :] < OUT_WIDTH)
+ old = tl.load(out_ptrs, mask=mask)
+ tl.store(out_ptrs, old + accumulator, mask=mask)
+
+ @torch.inference_mode()
+ def _impl(
+ inputs: torch.Tensor,
+ lora_b_weight: torch.Tensor,
+ output_tensor: torch.Tensor,
+ token_indices_sorted_by_lora_ids: torch.Tensor,
+ num_tokens_per_lora: torch.Tensor,
+ lora_token_start_loc: torch.Tensor,
+ lora_ids: torch.Tensor,
+ no_lora_flag_cpu: torch.Tensor,
+ num_active_loras: torch.Tensor,
+ slice_offset: int,
+ ) -> None:
+ assert no_lora_flag_cpu.numel() == 1
+ if no_lora_flag_cpu.item():
+ return
+ assert inputs.dtype == torch.float32
+ assert output_tensor.dtype == torch.float32
+ assert output_tensor.is_contiguous()
+ assert lora_b_weight.ndim == 4
+ assert lora_b_weight.size(1) == 1
+ assert lora_b_weight.is_contiguous()
+
+ m = inputs.size(0)
+ out_width = lora_b_weight.size(2)
+ rank = lora_b_weight.size(3)
+ kernel_config = get_lora_op_configs(
+ op_type="expand",
+ max_loras=lora_ids.size(0),
+ batch=m,
+ hidden_size=out_width,
+ rank=rank,
+ num_slices=1,
+ add_inputs=True,
+ )
+ block_m = kernel_config["block_m"]
+ block_n = kernel_config["block_n"]
+ block_k = kernel_config["block_k"]
+ use_gdc = supports_pdl(inputs.device) and envs.VLLM_LORA_ENABLE_DUAL_STREAM
+ grid = (
+ triton.cdiv(m, block_m) * triton.cdiv(out_width, block_n),
+ num_active_loras.item(),
+ )
+ _kernel[grid](
+ inputs,
+ lora_b_weight,
+ output_tensor,
+ token_indices_sorted_by_lora_ids,
+ num_tokens_per_lora,
+ lora_token_start_loc,
+ lora_ids,
+ m,
+ slice_offset,
+ inputs.stride(0),
+ inputs.stride(1),
+ lora_b_weight.stride(0),
+ lora_b_weight.stride(2),
+ lora_b_weight.stride(3),
+ output_tensor.stride(0),
+ output_tensor.stride(1),
+ RANK=rank,
+ OUT_WIDTH=out_width,
+ BLOCK_M=block_m,
+ BLOCK_N=block_n,
+ BLOCK_K=block_k,
+ USE_GDC=use_gdc,
+ num_warps=kernel_config["num_warps"],
+ num_ctas=kernel_config["num_ctas"],
+ num_stages=kernel_config["num_stages"],
+ launch_pdl=use_gdc,
+ )
+
+ def _fake(
+ inputs: torch.Tensor,
+ lora_b_weight: torch.Tensor,
+ output_tensor: torch.Tensor,
+ token_indices_sorted_by_lora_ids: torch.Tensor,
+ num_tokens_per_lora: torch.Tensor,
+ lora_token_start_loc: torch.Tensor,
+ lora_ids: torch.Tensor,
+ no_lora_flag_cpu: torch.Tensor,
+ num_active_loras: torch.Tensor,
+ slice_offset: int,
+ ) -> None:
+ return None
+
+ direct_register_custom_op(
+ op_name="art_dsv4_lora_expand_fp32_output",
+ op_func=_impl,
+ mutates_args=["output_tensor"],
+ fake_impl=_fake,
+ )
+ _register_dsv4_lora_expand_fp32_output_op._registered = True # type: ignore[attr-defined]
+
+
+def _apply_dsv4_compressor_lora_to_existing_output(
+ module: Any,
+ x: Any,
+ output: Any,
+) -> Any:
+ if not _is_active_lora_wrapped_linear(module):
+ return output
+
+ import torch
+ from vllm.platforms import current_platform
+
+ _register_dsv4_lora_expand_fp32_output_op()
+ original_shape = output.shape if output.ndim == 3 else None
+ if x.ndim == 3 and output.ndim == 3:
+ x = x.flatten(0, 1)
+ output = output.flatten(0, 1)
+ if x.ndim != 2 or output.ndim != 2:
+ raise RuntimeError(
+ "DSV4 compressor LoRA expects 2D hidden/output tensors, got "
+ f"x={tuple(x.shape)}, output={tuple(output.shape)}."
+ )
+ if output.dtype != torch.float32:
+ raise RuntimeError(
+ "DSV4 compressor LoRA must preserve the fp32 compressor projection, "
+ f"got output dtype {output.dtype}."
+ )
+ if getattr(module, "tp_size", 1) != 1:
+ raise RuntimeError(
+ "DSV4 compressor LoRA expects disable_tp compressor linears."
+ )
+
+ wrapper = module.punica_wrapper
+ local_rank = module.lora_a_stacked[0].shape[2]
+ buffers = torch.empty_strided(
+ (len(module.output_slices), x.shape[0], local_rank),
+ (x.shape[0] * local_rank, local_rank, 1),
+ dtype=torch.float32,
+ device=x.device,
+ )
+ shrunk = wrapper.add_shrink(buffers, x, module.lora_a_stacked, 1.0)
+ if not current_platform.can_update_inplace():
+ buffers = shrunk
+
+ (
+ _,
+ token_indices_sorted_by_lora_ids,
+ num_tokens_per_lora,
+ lora_token_start_loc,
+ lora_ids,
+ no_lora_flag_cpu,
+ num_active_loras,
+ ) = wrapper.token_mapping_meta.meta_args(
+ x.shape[0], wrapper.lora_config.specialize_active_lora
+ )
+ offset = 0
+ for idx, width in enumerate(module.output_slices):
+ torch.ops.vllm.art_dsv4_lora_expand_fp32_output(
+ buffers[idx],
+ module.lora_b_stacked[idx],
+ output,
+ token_indices_sorted_by_lora_ids,
+ num_tokens_per_lora,
+ lora_token_start_loc,
+ lora_ids,
+ no_lora_flag_cpu,
+ num_active_loras,
+ offset,
+ )
+ offset += width
+ if offset != output.shape[-1]:
+ raise RuntimeError(
+ "DSV4 compressor LoRA output slice width mismatch: "
+ f"slices={module.output_slices}, output={tuple(output.shape)}."
+ )
+ return output.reshape(original_shape) if original_shape is not None else output
+
+
+def _is_active_lora_wrapped_linear(module: Any) -> bool:
+ if not _is_lora_wrapped_linear(module):
+ return False
+ wrapper = module.punica_wrapper
+ return not getattr(wrapper, "no_lora", False) and (
+ getattr(wrapper, "indices_len", [None])[0] is not None
+ )
+
+
+def _register_dsv4_inv_rope_lora_input_op() -> None:
+ if getattr(_register_dsv4_inv_rope_lora_input_op, "_registered", False):
+ return
+
+ import torch
+ from vllm.platforms import current_platform
+ from vllm.triton_utils import tl, triton
+ from vllm.utils.torch_utils import direct_register_custom_op
+
+ @triton.jit
+ def _kernel(
+ o_ptr,
+ positions_ptr,
+ cos_sin_cache_ptr,
+ fp8_ptr,
+ scale_ptr,
+ lora_input_ptr,
+ num_tokens,
+ heads_per_group: tl.constexpr,
+ o_stride_token,
+ o_stride_head,
+ cache_stride_pos,
+ fp8_stride_group,
+ fp8_stride_token,
+ scale_stride_group,
+ scale_stride_k,
+ lora_stride_group,
+ lora_stride_token,
+ fp8_max: tl.constexpr,
+ eps: tl.constexpr,
+ QUANT_GROUP_SIZE: tl.constexpr,
+ CHUNKS_PER_HEAD: tl.constexpr,
+ ROPE_START: tl.constexpr,
+ HALF_ROPE: tl.constexpr,
+ TMA_ALIGNED_SCALES: tl.constexpr,
+ ):
+ pid_token = tl.program_id(0).to(tl.int64)
+ pid_gh = tl.program_id(1).to(tl.int64)
+ g = pid_gh // heads_per_group
+ head_in_group = pid_gh % heads_per_group
+ qb_start = head_in_group * CHUNKS_PER_HEAD
+
+ if pid_token >= num_tokens:
+ if TMA_ALIGNED_SCALES:
+ scale_addr = (
+ scale_ptr
+ + g * scale_stride_group
+ + pid_token
+ + head_in_group * scale_stride_k
+ )
+ tl.store(scale_addr, tl.zeros((), dtype=tl.int32))
+ else:
+ block_offsets = tl.arange(0, CHUNKS_PER_HEAD)
+ qb_indices = qb_start + block_offsets
+ scale_addrs = (
+ scale_ptr
+ + g * scale_stride_group
+ + pid_token
+ + qb_indices * scale_stride_k
+ )
+ tl.store(scale_addrs, tl.zeros((CHUNKS_PER_HEAD,), dtype=tl.float32))
+ return
+
+ head_dim: tl.constexpr = CHUNKS_PER_HEAD * QUANT_GROUP_SIZE
+ offsets = tl.arange(0, head_dim)
+ input_base = o_ptr + pid_token * o_stride_token + pid_gh * o_stride_head
+ x = tl.load(input_base + offsets).to(tl.float32)
+
+ rope_abs_start: tl.constexpr = (
+ CHUNKS_PER_HEAD - 1
+ ) * QUANT_GROUP_SIZE + ROPE_START
+ pos = tl.load(positions_ptr + pid_token)
+ cache_base = cos_sin_cache_ptr + pos * cache_stride_pos
+ is_rope = offsets >= rope_abs_start
+ rope_local = offsets - rope_abs_start
+
+ x_partner = tl.load(input_base + (offsets ^ 1), mask=is_rope, other=0.0).to(
+ tl.float32
+ )
+ cs_idx = tl.maximum(rope_local >> 1, 0)
+ cos_v = tl.load(cache_base + cs_idx, mask=is_rope, other=1.0)
+ sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope, other=0.0)
+ x_add = x * cos_v + x_partner * sin_v
+ x_sub = x * cos_v - x_partner * sin_v
+ is_even = (rope_local & 1) == 0
+ x = tl.where(is_rope, tl.where(is_even, x_add, x_sub), x)
+
+ group_head_offset = head_in_group * head_dim
+ lora_base = (
+ lora_input_ptr
+ + g * lora_stride_group
+ + pid_token * lora_stride_token
+ + group_head_offset
+ )
+ tl.store(lora_base + offsets, x)
+
+ x_2d = tl.reshape(tl.abs(x), (CHUNKS_PER_HEAD, QUANT_GROUP_SIZE))
+ block_absmax = tl.maximum(tl.max(x_2d, axis=1), eps)
+ scale_raw = block_absmax * (1.0 / fp8_max)
+ scales = tl.math.exp2(tl.ceil(tl.log2(scale_raw)))
+ scales_exp = tl.reshape(
+ tl.broadcast_to(
+ tl.reshape(scales, (CHUNKS_PER_HEAD, 1)),
+ (CHUNKS_PER_HEAD, QUANT_GROUP_SIZE),
+ ),
+ (head_dim,),
+ )
+ x_quant = tl.clamp(x / scales_exp, -fp8_max, fp8_max).to(tl.float8e4nv)
+
+ fp8_base = (
+ fp8_ptr
+ + g * fp8_stride_group
+ + pid_token * fp8_stride_token
+ + qb_start * QUANT_GROUP_SIZE
+ )
+ tl.store(fp8_base + offsets, x_quant)
+
+ block_offsets = tl.arange(0, CHUNKS_PER_HEAD)
+ qb_indices = qb_start + block_offsets
+ if TMA_ALIGNED_SCALES:
+ scale_bits = scales.to(tl.int32, bitcast=True)
+ ue8m0_bytes = (scale_bits >> 23) & 0xFF
+ packed_val = tl.sum(ue8m0_bytes << (block_offsets * 8))
+ scale_addr = (
+ scale_ptr
+ + g * scale_stride_group
+ + pid_token
+ + head_in_group * scale_stride_k
+ )
+ tl.store(scale_addr, packed_val)
+ else:
+ scale_addrs = (
+ scale_ptr
+ + g * scale_stride_group
+ + pid_token
+ + qb_indices * scale_stride_k
+ )
+ tl.store(scale_addrs, scales)
+
+ def _impl(
+ o: torch.Tensor,
+ positions: torch.Tensor,
+ cos_sin_cache: torch.Tensor,
+ fp8_buf: torch.Tensor,
+ scale_buf: torch.Tensor,
+ lora_input: torch.Tensor,
+ heads_per_group: int,
+ quant_group_size: int,
+ chunks_per_head: int,
+ rope_start: int,
+ half_rope: int,
+ tma_aligned_scales: bool,
+ fp8_max: float,
+ tma_aligned_t: int,
+ num_tokens: int,
+ ) -> None:
+ grid = (tma_aligned_t, fp8_buf.shape[0] * heads_per_group)
+ pdl_kwargs = {} if current_platform.is_rocm() else {"launch_pdl": False}
+ _kernel[grid](
+ o,
+ positions,
+ cos_sin_cache,
+ fp8_buf,
+ scale_buf,
+ lora_input,
+ num_tokens,
+ heads_per_group=heads_per_group,
+ o_stride_token=o.stride(0),
+ o_stride_head=o.stride(1),
+ cache_stride_pos=cos_sin_cache.stride(0),
+ fp8_stride_group=fp8_buf.stride(0),
+ fp8_stride_token=fp8_buf.stride(1),
+ scale_stride_group=scale_buf.stride(0),
+ scale_stride_k=scale_buf.stride(2),
+ lora_stride_group=lora_input.stride(0),
+ lora_stride_token=lora_input.stride(1),
+ fp8_max=fp8_max,
+ eps=1e-10,
+ QUANT_GROUP_SIZE=quant_group_size,
+ CHUNKS_PER_HEAD=chunks_per_head,
+ ROPE_START=rope_start,
+ HALF_ROPE=half_rope,
+ TMA_ALIGNED_SCALES=tma_aligned_scales,
+ num_stages=1,
+ num_warps=1,
+ **pdl_kwargs,
+ )
+
+ def _fake(
+ o: torch.Tensor,
+ positions: torch.Tensor,
+ cos_sin_cache: torch.Tensor,
+ fp8_buf: torch.Tensor,
+ scale_buf: torch.Tensor,
+ lora_input: torch.Tensor,
+ heads_per_group: int,
+ quant_group_size: int,
+ chunks_per_head: int,
+ rope_start: int,
+ half_rope: int,
+ tma_aligned_scales: bool,
+ fp8_max: float,
+ tma_aligned_t: int,
+ num_tokens: int,
+ ) -> None:
+ return None
+
+ direct_register_custom_op(
+ op_name="art_dsv4_inv_rope_fp8_quant_lora_input",
+ op_func=_impl,
+ mutates_args=["fp8_buf", "scale_buf", "lora_input"],
+ fake_impl=_fake,
+ )
+ _register_dsv4_inv_rope_lora_input_op._registered = True # type: ignore[attr-defined]
+
+
+def _dsv4_fused_inv_rope_fp8_quant_with_lora_input(
+ dsv4_attn: Any,
+ o: Any,
+ positions: Any,
+ cos_sin_cache: Any,
+ *,
+ n_groups: int,
+ heads_per_group: int,
+ lora_dtype: Any,
+ nope_dim: int = 448,
+ rope_dim: int = 64,
+ quant_group_size: int = 128,
+ tma_aligned_scales: bool = False,
+) -> tuple[Any, Any, Any]:
+ import torch
+ from vllm.utils.deep_gemm import get_tma_aligned_size
+
+ num_tokens, num_heads, head_dim = o.shape
+ assert num_heads == n_groups * heads_per_group
+ assert head_dim == nope_dim + rope_dim
+ assert head_dim % quant_group_size == 0
+ assert nope_dim % quant_group_size == (quant_group_size - rope_dim)
+ assert rope_dim % 2 == 0
+ assert cos_sin_cache.shape[-1] == rope_dim
+ assert cos_sin_cache.dtype == torch.float32
+
+ d = heads_per_group * head_dim
+ num_scale_blocks = d // quant_group_size
+ chunks_per_head = head_dim // quant_group_size
+ fp8_dtype = torch.float8_e4m3fn
+ tma_aligned_t = get_tma_aligned_size(num_tokens, 4)
+ scale_inner = (
+ (num_scale_blocks + 3) // 4 if tma_aligned_scales else num_scale_blocks
+ )
+
+ fp8_buf = torch.empty((n_groups, num_tokens, d), dtype=fp8_dtype, device=o.device)
+ scale_dtype = torch.int32 if tma_aligned_scales else torch.float32
+ scale_buf = torch.empty(
+ n_groups * scale_inner * tma_aligned_t,
+ dtype=scale_dtype,
+ device=o.device,
+ ).as_strided(
+ (n_groups, num_tokens, scale_inner),
+ (scale_inner * tma_aligned_t, 1, tma_aligned_t),
+ )
+ lora_input = torch.empty(
+ (n_groups, num_tokens, d),
+ dtype=lora_dtype,
+ device=o.device,
+ )
+ dsv4_attn.torch.ops.vllm.art_dsv4_inv_rope_fp8_quant_lora_input(
+ o,
+ positions,
+ cos_sin_cache,
+ fp8_buf,
+ scale_buf,
+ lora_input,
+ heads_per_group,
+ quant_group_size,
+ chunks_per_head,
+ nope_dim % quant_group_size,
+ rope_dim // 2,
+ tma_aligned_scales,
+ torch.finfo(fp8_dtype).max,
+ tma_aligned_t,
+ num_tokens,
+ )
+ return fp8_buf.transpose(0, 1), scale_buf.transpose(0, 1), lora_input
+
+
+def _dsv4_wo_a_lora_b_group_cache(
+ wo_a: Any,
+ *,
+ n_local_groups: int,
+ out_per_group: int,
+) -> tuple[Any, ...]:
+ source = wo_a.lora_b_stacked[0]
+ key = (
+ source.data_ptr(),
+ getattr(source, "_version", None),
+ n_local_groups,
+ out_per_group,
+ source.shape[-1],
+ )
+ if getattr(wo_a, "_art_wo_a_lora_b_group_cache_key", None) != key:
+ wo_a._art_wo_a_lora_b_group_cache = tuple(
+ source[
+ :, :, group * out_per_group : (group + 1) * out_per_group, :
+ ].contiguous()
+ for group in range(n_local_groups)
+ )
+ wo_a._art_wo_a_lora_b_group_cache_key = key
+ return wo_a._art_wo_a_lora_b_group_cache
+
+
+def _apply_dsv4_wo_a_lora_fast(
+ wo_a: Any,
+ z: Any,
+ *,
+ lora_input: Any,
+ n_local_groups: int,
+) -> Any:
+ if not _is_active_lora_wrapped_linear(wo_a):
+ return z
+
+ import torch
+ from vllm.distributed import tensor_model_parallel_all_gather
+ from vllm.platforms import current_platform
+
+ wrapper = wo_a.punica_wrapper
+ out_per_group = z.shape[-1]
+ z_flat = z.view(z.shape[0], n_local_groups * out_per_group)
+ group_b = _dsv4_wo_a_lora_b_group_cache(
+ wo_a,
+ n_local_groups=n_local_groups,
+ out_per_group=out_per_group,
+ )
+ local_rank = wo_a.lora_a_stacked[0].shape[2]
+ buffers = torch.empty_strided(
+ (n_local_groups, z.shape[0], local_rank),
+ (z.shape[0] * local_rank, local_rank, 1),
+ dtype=torch.float32,
+ device=z.device,
+ )
+ for group, lora_b in enumerate(group_b):
+ buffer = buffers[group : group + 1]
+ buffer.zero_()
+ shrunk = wrapper.add_shrink(buffer, lora_input[group], wo_a.lora_a_stacked, 1.0)
+ if not current_platform.can_update_inplace():
+ buffer = shrunk
+ buffer = tensor_model_parallel_all_gather(buffer)
+ expanded = wrapper.add_expand(
+ z_flat,
+ buffer,
+ (lora_b,),
+ (out_per_group,),
+ offset_start=group * out_per_group,
+ add_inputs=True,
+ )
+ if not current_platform.can_update_inplace():
+ z_flat = expanded
+ z = z_flat.view_as(z)
+ return z
+
+
+def _patch_dsv4_compressor_fast_path_lora(attention_cls: Any) -> None:
+ if getattr(attention_cls, "_art_compressor_fast_path_lora_patched", False):
+ return
+
+ original_attn_gemm_parallel_execute = attention_cls.attn_gemm_parallel_execute
+
+ def attn_gemm_parallel_execute(self: Any, hidden_states: Any) -> tuple[Any, ...]:
+ qr_kv, kv_score, indexer_kv_score, indexer_weights = (
+ original_attn_gemm_parallel_execute(self, hidden_states)
+ )
+ if self.compressor is not None:
+ kv_score = _apply_dsv4_compressor_lora_to_existing_output(
+ self.compressor.fused_wkv_wgate,
+ hidden_states,
+ kv_score,
+ )
+ if self.indexer is not None:
+ indexer_kv_score = _apply_dsv4_compressor_lora_to_existing_output(
+ self.indexer.compressor.fused_wkv_wgate,
+ hidden_states,
+ indexer_kv_score,
+ )
+ return qr_kv, kv_score, indexer_kv_score, indexer_weights
+
+ attn_gemm_parallel_execute.__art_patched__ = True # type: ignore[attr-defined]
+ attention_cls.attn_gemm_parallel_execute = attn_gemm_parallel_execute
+ attention_cls._art_compressor_fast_path_lora_patched = True
+
+
+def _dsv4_deep_gemm_fp8_o_proj_with_lora(
+ o_proj_mod: Any,
+ o: Any,
+ positions: Any,
+ cos_sin_cache: Any,
+ wo_a: Any,
+ wo_b: Any,
+ *,
+ n_groups: int,
+ heads_per_group: int,
+ nope_dim: int,
+ rope_dim: int,
+ o_lora_rank: int,
+ einsum_recipe: tuple[int, int, int],
+ tma_aligned_scales: bool,
+) -> Any:
+ wo_a_lora_input = None
+ if _is_active_lora_wrapped_linear(wo_a):
+ o_fp8, o_scale, wo_a_lora_input = (
+ _dsv4_fused_inv_rope_fp8_quant_with_lora_input(
+ o_proj_mod,
+ o,
+ positions,
+ cos_sin_cache,
+ n_groups=n_groups,
+ heads_per_group=heads_per_group,
+ lora_dtype=wo_a.lora_a_stacked[0].dtype,
+ nope_dim=nope_dim,
+ rope_dim=rope_dim,
+ tma_aligned_scales=tma_aligned_scales,
+ )
+ )
+ else:
+ o_fp8, o_scale = o_proj_mod.fused_inv_rope_fp8_quant(
+ o,
+ positions,
+ cos_sin_cache,
+ n_groups=n_groups,
+ heads_per_group=heads_per_group,
+ nope_dim=nope_dim,
+ rope_dim=rope_dim,
+ tma_aligned_scales=tma_aligned_scales,
+ )
+
+ z = o_proj_mod.torch.empty(
+ (o.shape[0], n_groups, o_lora_rank),
+ device=o.device,
+ dtype=o_proj_mod.torch.bfloat16,
+ )
+ o_proj_mod.fp8_einsum(
+ "bhr,hdr->bhd",
+ (o_fp8, o_scale),
+ (wo_a.weight, wo_a.weight_scale_inv),
+ z,
+ recipe=einsum_recipe,
+ )
+ if wo_a_lora_input is not None:
+ z = _apply_dsv4_wo_a_lora_fast(
+ wo_a,
+ z,
+ lora_input=wo_a_lora_input,
+ n_local_groups=n_groups,
+ )
+ return wo_b(z.flatten(1))
+
+
+def _patch_dsv4_cuda_o_proj_lora(attn_cls: Any, o_proj_mod: Any) -> None:
+ if getattr(attn_cls, "_art_wo_a_fast_path_lora_patched", False):
+ return
+
+ def _o_proj(self: Any, o: Any, positions: Any) -> Any:
+ return _dsv4_deep_gemm_fp8_o_proj_with_lora(
+ o_proj_mod,
+ o,
+ positions,
+ self.rotary_emb.cos_sin_cache,
+ self.wo_a,
+ self.wo_b,
+ n_groups=self.n_local_groups,
+ heads_per_group=self.n_local_heads // self.n_local_groups,
+ nope_dim=self.nope_head_dim,
+ rope_dim=self.rope_head_dim,
+ o_lora_rank=self.o_lora_rank,
+ einsum_recipe=self._einsum_recipe,
+ tma_aligned_scales=self._tma_aligned_scales,
+ )
+
+ _o_proj.__art_patched__ = True # type: ignore[attr-defined]
+ attn_cls._o_proj = _o_proj
+ attn_cls._art_wo_a_fast_path_lora_patched = True
+
+
+def _patch_current_dsv4_fast_path_lora() -> bool:
+ try:
+ dsv4_attention = importlib.import_module("vllm.models.deepseek_v4.attention")
+ except ModuleNotFoundError:
+ return False
+
+ attention_cls = getattr(dsv4_attention, "DeepseekV4Attention", None)
+ if attention_cls is None:
+ return False
+
+ _patch_dsv4_compressor_fast_path_lora(attention_cls)
+
+ try:
+ o_proj_mod = importlib.import_module(
+ "vllm.models.deepseek_v4.nvidia.ops.o_proj"
+ )
+ except ModuleNotFoundError:
+ return True
+
+ for module_name, class_name in (
+ (
+ "vllm.models.deepseek_v4.nvidia.flashmla",
+ "DeepseekV4FlashMLAAttention",
+ ),
+ (
+ "vllm.models.deepseek_v4.nvidia.flashinfer_sparse",
+ "DeepseekV4FlashInferMLAAttention",
+ ),
+ ):
+ try:
+ module = importlib.import_module(module_name)
+ except ModuleNotFoundError:
+ continue
+ attn_cls = getattr(module, class_name, None)
+ if attn_cls is not None:
+ _patch_dsv4_cuda_o_proj_lora(attn_cls, o_proj_mod)
+ return True
+
+
+def patch_dsv4_fast_path_lora() -> None:
+ """Apply LoRA deltas on DSV4 paths that read base weights directly.
+
+ vLLM's generic LoRA manager can wrap DSV4 linear modules, but the DSV4
+ Flash runtime bypasses some wrapped forwards for performance: compressor
+ projections are direct ``hidden @ fused_wkv_wgate.weight.T`` calls, and
+ ``wo_a`` is a custom inverse-RoPE/FP8/einsum path. Without this patch vLLM
+ accepts and activates these adapter tensors while silently omitting their
+ deltas from generation.
+ """
+ _register_dsv4_inv_rope_lora_input_op()
+ _register_dsv4_lora_expand_fp32_output_op()
+ if _patch_current_dsv4_fast_path_lora():
+ return
+
+ dsv4_attn = importlib.import_module(
+ "vllm.model_executor.layers.deepseek_v4_attention"
+ )
+ wrapper_cls = getattr(dsv4_attn, "DeepseekV4MultiHeadLatentAttentionWrapper", None)
+ if wrapper_cls is None:
+ return
+ if getattr(wrapper_cls, "_art_fast_path_lora_patched", False):
+ return
+
+ original_attn_gemm_parallel_execute = wrapper_cls.attn_gemm_parallel_execute
+ original_forward = wrapper_cls.forward
+
+ def attn_gemm_parallel_execute(self: Any, hidden_states: Any) -> tuple[Any, ...]:
+ qr_kv, kv_score, indexer_kv_score, indexer_weights = (
+ original_attn_gemm_parallel_execute(self, hidden_states)
+ )
+ if self.compressor is not None:
+ kv_score = _apply_dsv4_compressor_lora_to_existing_output(
+ self.compressor.fused_wkv_wgate,
+ hidden_states,
+ kv_score,
+ )
+ if self.indexer is not None:
+ indexer_kv_score = _apply_dsv4_compressor_lora_to_existing_output(
+ self.indexer.compressor.fused_wkv_wgate,
+ hidden_states,
+ indexer_kv_score,
+ )
+ return qr_kv, kv_score, indexer_kv_score, indexer_weights
+
+ def forward(
+ self: Any,
+ positions: Any,
+ hidden_states: Any,
+ llama_4_scaling: Any | None = None,
+ ) -> Any:
+ if dsv4_attn.current_platform.is_rocm():
+ return original_forward(self, positions, hidden_states, llama_4_scaling)
+
+ num_tokens = hidden_states.shape[0]
+ o_padded = dsv4_attn.torch.empty(
+ (num_tokens, self.padded_heads, self.head_dim),
+ dtype=hidden_states.dtype,
+ device=hidden_states.device,
+ )
+
+ dsv4_attn.torch.ops.vllm.deepseek_v4_attention(
+ hidden_states,
+ positions,
+ o_padded,
+ self.layer_name,
+ )
+ o = o_padded[:, : self.n_local_heads, :]
+
+ wo_a_lora_input = None
+ if _is_active_lora_wrapped_linear(self.wo_a):
+ o_fp8, o_scale, wo_a_lora_input = (
+ _dsv4_fused_inv_rope_fp8_quant_with_lora_input(
+ dsv4_attn,
+ o,
+ positions,
+ self.rotary_emb.cos_sin_cache,
+ n_groups=self.n_local_groups,
+ heads_per_group=self.n_local_heads // self.n_local_groups,
+ lora_dtype=self.wo_a.lora_a_stacked[0].dtype,
+ nope_dim=self.nope_head_dim,
+ rope_dim=self.rope_head_dim,
+ tma_aligned_scales=self._tma_aligned_scales,
+ )
+ )
+ else:
+ o_fp8, o_scale = dsv4_attn.fused_inv_rope_fp8_quant(
+ o,
+ positions,
+ self.rotary_emb.cos_sin_cache,
+ n_groups=self.n_local_groups,
+ heads_per_group=self.n_local_heads // self.n_local_groups,
+ nope_dim=self.nope_head_dim,
+ rope_dim=self.rope_head_dim,
+ tma_aligned_scales=self._tma_aligned_scales,
+ )
+
+ z = dsv4_attn.torch.empty(
+ (num_tokens, self.n_local_groups, self.o_lora_rank),
+ device=o.device,
+ dtype=dsv4_attn.torch.bfloat16,
+ )
+ dsv4_attn.torch.ops.vllm.deepseek_v4_fp8_einsum(
+ o_fp8,
+ o_scale,
+ self.wo_a.weight,
+ self.wo_a.weight_scale_inv,
+ z,
+ "bhr,hdr->bhd",
+ list(self._einsum_recipe),
+ )
+ if wo_a_lora_input is not None:
+ z = _apply_dsv4_wo_a_lora_fast(
+ self.wo_a,
+ z,
+ lora_input=wo_a_lora_input,
+ n_local_groups=self.n_local_groups,
+ )
+ return self.wo_b(z.flatten(1))
+
+ attn_gemm_parallel_execute.__art_patched__ = True # type: ignore[attr-defined]
+ forward.__art_patched__ = True # type: ignore[attr-defined]
+ wrapper_cls.attn_gemm_parallel_execute = attn_gemm_parallel_execute
+ wrapper_cls.forward = forward
+ wrapper_cls._art_fast_path_lora_patched = True
+
+
+def _next_power_of_two(value: int) -> int:
+ return 1 << (max(value, 1) - 1).bit_length()
+
+
+def patch_dsv4_triton_moe_topk6_routing() -> None:
+ """Make vLLM's DSV4 Triton MoE routing sort compile for top-k 6.
+
+ Triton's ``tl.arange`` requires a power-of-two range. Current vLLM's
+ DSV4/MXFP4 routing path sorts ``BLOCK_M * num_experts_per_tok`` entries;
+ DSV4 uses top-k 6, so the first profile run tries ``tl.arange(0, 192)`` and
+ the engine exits before serving starts. Keep the original indexing stride at
+ 192, but sort over a padded power-of-two vector and mask padded lanes.
+ """
+ try:
+ import torch
+ import triton
+ import triton.language as tl
+ from vllm.third_party.triton_kernels.routing_details._expt_data import (
+ _expt_data_compute,
+ )
+ except ImportError:
+ return
+
+ @triton.jit
+ def _routing_compute_indx_pow2(
+ pid_m,
+ GatherIndx,
+ ScatterIndx,
+ GateScal,
+ ExptScal,
+ ExptIndx,
+ PartialOffs,
+ stride_pm,
+ stride_pn,
+ TokensStart,
+ n_tokens,
+ BLOCK_M: tl.constexpr,
+ N_EXPTS_ACT: tl.constexpr,
+ BLOCK_SIZE_PADDED: tl.constexpr,
+ ):
+ if isinstance(n_tokens, tl.tensor) and n_tokens.dtype.is_ptr():
+ n_tokens = tl.load(n_tokens)
+ n_gates = n_tokens * N_EXPTS_ACT
+ block_size: tl.constexpr = N_EXPTS_ACT * BLOCK_M
+ tl.static_assert(BLOCK_SIZE_PADDED >= block_size)
+ tl.static_assert(BLOCK_SIZE_PADDED <= 32768)
+
+ local_offs = tl.arange(0, BLOCK_SIZE_PADDED)
+ valid_local = local_offs < block_size
+ offs = pid_m * block_size + local_offs
+ expert = tl.load(
+ ExptIndx + offs,
+ mask=valid_local & (offs < n_gates),
+ other=-1,
+ ).to(tl.uint32)
+
+ kv_pairs = ((expert << 16) | local_offs).to(tl.uint32)
+ kv_pairs = tl.sort(kv_pairs, 0)
+ expert = kv_pairs >> 16
+ offs = pid_m * block_size + (kv_pairs & 0xFFFF)
+ mask = expert != 0xFFFF
+ gate_scal = tl.load(ExptScal + offs, mask=mask)
+
+ x = kv_pairs & 0xFFFF0000 | 0x00000001
+ expts_and_inclusive_run_lengths = tl.associative_scan(x, 0, _keyed_add_pow2)
+ exclusive_run_lengths = (expts_and_inclusive_run_lengths - 1) & 0xFFFF
+
+ gates = tl.load(PartialOffs + pid_m * stride_pm + expert * stride_pn, mask=mask)
+ gates += tl.load(TokensStart + expert, mask=mask)
+ gates += exclusive_run_lengths
+
+ tl.store(ScatterIndx + offs, gates, mask=mask)
+ tl.store(GatherIndx + gates, offs, mask=mask)
+ tl.store(GateScal + gates, gate_scal, mask=mask)
+
+ @triton.jit
+ def _keyed_add_pow2(x, y):
+ key_mask: tl.constexpr = 0xFFFF0000
+ kx = x & key_mask
+ ky = y & key_mask
+ return tl.where(kx == ky, x + y - kx, y)
+
+ @triton.jit
+ def _combined_routing_compute_pow2(
+ GatherIndx,
+ ScatterIndx,
+ GateScal,
+ ExptScal,
+ ExptIndx,
+ PartialOffs,
+ stride_pm,
+ stride_pn,
+ TokensStart,
+ n_tokens,
+ BLOCK_M: tl.constexpr,
+ N_EXPTS_ACT: tl.constexpr,
+ Hist,
+ MDTileStarts,
+ tile_starts_stridem,
+ MDTileInfo,
+ tile_info_stridem,
+ first_tile_dim_log2,
+ SIZES: tl.constexpr,
+ BLOCK: tl.constexpr,
+ blocks2a,
+ BLOCK_SIZE_PADDED: tl.constexpr,
+ ):
+ pid = tl.program_id(0)
+ if pid < blocks2a:
+ _expt_data_compute(
+ Hist,
+ MDTileStarts,
+ tile_starts_stridem,
+ MDTileInfo,
+ tile_info_stridem,
+ first_tile_dim_log2,
+ SIZES,
+ BLOCK,
+ )
+ else:
+ pid -= blocks2a
+ _routing_compute_indx_pow2(
+ pid,
+ GatherIndx,
+ ScatterIndx,
+ GateScal,
+ ExptScal,
+ ExptIndx,
+ PartialOffs,
+ stride_pm,
+ stride_pn,
+ TokensStart,
+ n_tokens,
+ BLOCK_M,
+ N_EXPTS_ACT,
+ BLOCK_SIZE_PADDED,
+ )
+
+ for module_name in (
+ "vllm.third_party.triton_kernels.routing",
+ "triton_kernels.routing",
+ ):
+ try:
+ routing = importlib.import_module(module_name)
+ except ImportError:
+ continue
+ original_forward = routing.SortTokens.forward
+ if getattr(original_forward, "__art_dsv4_topk6_pow2_patched__", False):
+ continue
+
+ def forward(
+ ctx: Any,
+ expt_scal: Any,
+ expt_indx: Any,
+ n_expts_tot: int,
+ bitmatrix: Any,
+ _routing: Any = routing,
+ ) -> Any:
+ hist_block_m = 32
+ indx_offs_block_m = 512
+ memset_block = 1024
+ cdiv = triton.cdiv
+
+ device = expt_scal.device
+ dtype = expt_scal.dtype
+ n_tokens_raw, _ = bitmatrix.shape
+ n_tokens_pad, n_expts_act = expt_scal.shape
+ n_gates_pad = n_tokens_pad * n_expts_act
+
+ hist, partial_hist = bitmatrix.sum(partials_block_size=hist_block_m)
+ hist = hist[:n_expts_tot]
+ assert hist.dtype == torch.int32
+ expt_offs = torch.empty(n_expts_tot, dtype=torch.int32, device=device)
+ combined_indx = torch.empty(
+ n_gates_pad * 2, dtype=torch.int32, device=device
+ )
+ topk_indx = combined_indx[:n_gates_pad]
+ gate_indx = combined_indx[n_gates_pad:]
+ gate_scal = torch.empty(n_gates_pad, dtype=dtype, device=device)
+
+ (
+ token_offs_combined,
+ token_offs_raw,
+ token_offs_pad,
+ block_pid_map,
+ blocks1a,
+ blocks2a,
+ memset_block_a,
+ hist2_block_m,
+ block_m_log2_start,
+ block_m_num,
+ ) = _routing._compute_expt_data_internal(hist, n_expts_tot, n_gates_pad)
+
+ blocks1b = cdiv(n_gates_pad * 2, memset_block) + n_expts_tot + 1
+ blocks2b = cdiv(n_tokens_pad, hist_block_m)
+
+ _routing._combined_routing_memset[(blocks1a + blocks1b,)](
+ combined_indx,
+ n_gates_pad * 2,
+ -1,
+ memset_block,
+ hist,
+ expt_offs,
+ hist.shape[0],
+ n_expts_tot,
+ partial_hist,
+ partial_hist.shape[0],
+ partial_hist.stride(0),
+ partial_hist.stride(1),
+ token_offs_combined,
+ token_offs_combined.stride(0),
+ blocks1a,
+ block_pid_map,
+ block_m_log2_start,
+ SIZES=block_m_num,
+ BLOCK_A=memset_block_a,
+ BLOCK_N=512,
+ BLOCK_M=indx_offs_block_m,
+ )
+
+ indx_offs = partial_hist
+ _combined_routing_compute_pow2[(blocks2a + blocks2b,)](
+ topk_indx,
+ gate_indx,
+ gate_scal,
+ expt_scal,
+ expt_indx,
+ indx_offs,
+ indx_offs.stride(0),
+ indx_offs.stride(1),
+ expt_offs,
+ n_tokens_raw,
+ hist_block_m,
+ n_expts_act,
+ hist,
+ token_offs_pad,
+ token_offs_pad.stride(0),
+ block_pid_map,
+ block_pid_map.stride(0),
+ block_m_log2_start,
+ block_m_num,
+ hist2_block_m,
+ blocks2a,
+ BLOCK_SIZE_PADDED=_next_power_of_two(hist_block_m * n_expts_act),
+ )
+
+ ctx.n_tokens_raw = n_tokens_raw
+ ctx.n_tokens_pad = n_tokens_pad
+ ctx.n_expts_act = n_expts_act
+ ctx.save_for_backward(gate_indx)
+ return (
+ hist,
+ topk_indx,
+ gate_indx,
+ gate_scal,
+ token_offs_raw,
+ token_offs_pad,
+ block_pid_map,
+ )
+
+ forward.__art_dsv4_topk6_pow2_patched__ = True # type: ignore[attr-defined]
+ forward.__art_original__ = original_forward # type: ignore[attr-defined]
+ routing.SortTokens.forward = staticmethod(forward)
+ routing._combined_routing_compute = _combined_routing_compute_pow2
+
+
+def _base_layer_attr_proxy(name: str) -> property:
+ def attr(self: Any) -> Any:
+ return getattr(self.base_layer, name)
+
+ return property(attr)
+
+
+def patch_lora_linear_base_attr_proxy() -> None:
+ """Expose DSV4 base metadata through vLLM linear LoRA wrappers.
+
+ DeepSeek V4's output attention path calls a custom FP8 einsum directly and
+ reads ``wo_a.weight_scale_inv`` next to ``wo_a.weight``. vLLM's linear LoRA
+ wrappers already proxy ``weight`` but not the quant scale. Its router also
+ reads dynamic gate metadata from ``self.gate`` after that gate can be LoRA
+ wrapped. Keep these tensors owned by the base layer instead of copying or
+ re-registering them on every wrapper.
+ """
+ from vllm.lora.layers.base_linear import BaseLinearLayerWithLoRA
+
+ if getattr(BaseLinearLayerWithLoRA, "_art_base_attr_proxy_patched", False):
+ return
+
+ for name in ("weight_scale_inv", "tid2eid", "e_score_correction_bias"):
+ if not hasattr(BaseLinearLayerWithLoRA, name):
+ setattr(BaseLinearLayerWithLoRA, name, _base_layer_attr_proxy(name))
+ BaseLinearLayerWithLoRA._art_base_attr_proxy_patched = True
+
+
+def patch_marlin_lora_swiglu_limit() -> None:
+ """Keep Marlin MoE LoRA active when DSV4 uses a SwiGLU clamp limit.
+
+ vLLM's Marlin LoRA path injects W13 LoRA inside the activation callback and
+ stores that activated cache for W2 LoRA. DSV4 sets ``gemm1_clamp_limit``;
+ upstream Marlin bypasses the callback in that case and calls the clamp op
+ directly, so W13 LoRA is skipped and W2 LoRA later misses ``cache2``. Route
+ the callback through the same clamp op while preserving Marlin execution.
+ """
+ try:
+ marlin_moe = importlib.import_module(
+ "vllm.model_executor.layers.fused_moe.fused_marlin_moe"
+ )
+ except ModuleNotFoundError:
+ return
+
+ from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+ from vllm.model_executor.layers.fused_moe.utils import swiglu_limit_func
+
+ MarlinExperts = marlin_moe.MarlinExperts
+
+ original_apply = MarlinExperts.apply
+ if getattr(original_apply, "__art_patched__", False):
+ return
+
+ sentinel = object()
+
+ def apply(self: Any, *args: Any, **kwargs: Any) -> Any:
+ clamp_limit = getattr(self, "gemm1_clamp_limit", None)
+ if getattr(self, "_lora_context", None) is None or clamp_limit is None:
+ return original_apply(self, *args, **kwargs)
+
+ original_activation = self.activation
+ previous_activation = self.__dict__.get("activation", sentinel)
+ previous_clamp_limit = self.gemm1_clamp_limit
+
+ def activation_with_clamp(
+ activation: Any,
+ output: Any,
+ input: Any,
+ ) -> None:
+ if activation == MoEActivation.SILU:
+ swiglu_limit_func(output, input, clamp_limit)
+ else:
+ original_activation(activation, output, input)
+
+ self.activation = activation_with_clamp
+ self.gemm1_clamp_limit = None
+ try:
+ return original_apply(self, *args, **kwargs)
+ finally:
+ self.gemm1_clamp_limit = previous_clamp_limit
+ if previous_activation is sentinel:
+ delattr(self, "activation")
+ else:
+ self.activation = previous_activation
+
+ apply.__art_patched__ = True # type: ignore[attr-defined]
+ MarlinExperts.apply = apply # type: ignore[method-assign]
diff --git a/vllm_runtime/src/art_vllm_runtime/lora_delta.py b/vllm_runtime/src/art_vllm_runtime/lora_delta.py
index 30952f93e..8952bb5d2 100644
--- a/vllm_runtime/src/art_vllm_runtime/lora_delta.py
+++ b/vllm_runtime/src/art_vllm_runtime/lora_delta.py
@@ -87,6 +87,7 @@ def _iter_lora_checkpoint_deltas(
expert_a = a_tensor[expert_id * rank : (expert_id + 1) * rank]
delta = b_expert.float().matmul(expert_a.float()).mul_(scaling)
if previous_b is not None:
+ assert previous_lora_tensors is not None
previous_a = previous_lora_tensors[a_key][
expert_id * rank : (expert_id + 1) * rank
]
@@ -120,6 +121,7 @@ def _iter_lora_checkpoint_deltas(
expert_a = a_tensor[expert_id * rank : (expert_id + 1) * rank]
delta = b_expert.float().matmul(expert_a.float()).mul_(scaling)
if previous_b is not None:
+ assert previous_lora_tensors is not None
previous_a = previous_lora_tensors[a_key][
expert_id * rank : (expert_id + 1) * rank
]
@@ -157,6 +159,25 @@ def _default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) ->
param.data.copy_(loaded_weight)
+def _call_weight_loader(
+ loader: Any,
+ loader_param: torch.Tensor,
+ loaded_weight: torch.Tensor,
+ *args: Any,
+ **kwargs: Any,
+) -> Any:
+ if not hasattr(loader_param, "load_merged_column_weight"):
+ owner = getattr(loader, "__self__", None)
+ legacy_loader = getattr(owner, "weight_loader", None)
+ if (
+ legacy_loader is not None
+ and legacy_loader is not loader
+ and getattr(loader, "__name__", "") == "weight_loader_v2"
+ ):
+ return legacy_loader(loader_param, loaded_weight, *args, **kwargs)
+ return loader(loader_param, loaded_weight, *args, **kwargs)
+
+
def _additive_weight_loader(param: torch.Tensor, original_loader: Any) -> Any:
def load_delta(
loader_param: torch.Tensor,
@@ -168,7 +189,13 @@ def load_delta(
scratch = torch.zeros_like(real_data)
loader_param.data = scratch
try:
- result = original_loader(loader_param, loaded_weight, *args, **kwargs)
+ result = _call_weight_loader(
+ original_loader,
+ loader_param,
+ loaded_weight,
+ *args,
+ **kwargs,
+ )
finally:
loader_param.data = real_data
if result is not False:
diff --git a/vllm_runtime/src/art_vllm_runtime/patches.py b/vllm_runtime/src/art_vllm_runtime/patches.py
index 91f0298b2..5077613dc 100644
--- a/vllm_runtime/src/art_vllm_runtime/patches.py
+++ b/vllm_runtime/src/art_vllm_runtime/patches.py
@@ -1,6 +1,7 @@
"""Monkey patches and bootstrap contract for the ART-owned vLLM runtime."""
import ctypes
+import importlib
import inspect
import logging
from typing import Any
@@ -11,6 +12,7 @@
def apply_vllm_runtime_patches() -> None:
+ from art_vllm_runtime.dsv4_patches import apply_dsv4_vllm_runtime_patches
from art_vllm_runtime.gemma4_moe_lora_patch import (
patch_gemma4_moe_lora_support,
)
@@ -21,6 +23,7 @@ def apply_vllm_runtime_patches() -> None:
patch_listen_for_disconnect()
patch_tool_parser_manager()
patch_nccl_unique_id_bootstrap()
+ apply_dsv4_vllm_runtime_patches()
patch_art_lora_delta_weight_update()
patch_gemma4_checkpoint_weight_update_reload()
patch_routed_experts_prefix_cache_sidecar()
@@ -65,8 +68,6 @@ def _patch_gemma4_moe_experts_per_tok_alias() -> None:
return
def num_experts_per_tok(self: Any) -> Any:
- # vLLM's routed-expert sidecar uses the Mistral MoE field name, while
- # HF Gemma4 stores the same router top-k value as top_k_experts.
return self.top_k_experts
Gemma4TextConfig.num_experts_per_tok = property(num_experts_per_tok) # type: ignore[attr-defined]
@@ -91,7 +92,10 @@ def __init__(self, *args: object, **kwargs: object) -> None:
def patch_listen_for_disconnect() -> None:
- from vllm.entrypoints.serve.utils import api_utils
+ try:
+ api_utils = importlib.import_module("vllm.entrypoints.serve.utils.api_utils")
+ except ModuleNotFoundError:
+ api_utils = importlib.import_module("vllm.entrypoints.utils")
if getattr(api_utils, "_art_listen_for_disconnect_patched", False):
return
@@ -224,10 +228,6 @@ def start_weight_update(
"start_weight_update called while a weight update is "
"already active. Call finish_weight_update first."
)
- # vLLM's layerwise checkpoint reload corrupts Gemma4 after reloading
- # the original checkpoint. Direct model.load_weights keeps the update
- # path identical to initial checkpoint loading while preserving the
- # streaming NCCL transfer used by ART merged serving.
self._is_checkpoint_format = True
self._weight_update_active = True
@@ -372,16 +372,11 @@ def patch_routed_experts_prefix_cache_sidecar() -> None:
if getattr(routed_experts_capturer, "_art_prefix_route_sidecar_patched", False):
return
- if hasattr(routed_experts_capturer, "RoutedExpertsManager"):
- # vLLM 0.23 stores routed experts by physical KV-cache slot, so prefix
- # cache hits recover routes from the shared slot buffer without ART's
- # old per-request host-cache sidecar.
- setattr(routed_experts_capturer, "_art_prefix_route_sidecar_patched", True)
+ host_cls = getattr(routed_experts_capturer, "_RoutedExpertsHostCache", None)
+ capturer_cls = getattr(routed_experts_capturer, "_RoutedExpertsCapturerReal", None)
+ if host_cls is None or capturer_cls is None:
return
- host_cls = routed_experts_capturer._RoutedExpertsHostCache
- capturer_cls = routed_experts_capturer._RoutedExpertsCapturerReal
-
original_host_init = host_cls.__init__
original_get_or_grow_buffer = host_cls.get_or_grow_buffer
original_free_request = host_cls.free_request
diff --git a/vllm_runtime/uv.lock b/vllm_runtime/uv.lock
index f647be079..0c697f2ec 100644
--- a/vllm_runtime/uv.lock
+++ b/vllm_runtime/uv.lock
@@ -10,7 +10,7 @@ overrides = [
{ name = "torch", url = "https://download.pytorch.org/whl/test/cu128/torch-2.11.0%2Bcu128-cp312-cp312-manylinux_2_28_x86_64.whl" },
{ name = "torchaudio", url = "https://download.pytorch.org/whl/test/cu128/torchaudio-2.11.0%2Bcu128-cp312-cp312-manylinux_2_28_x86_64.whl" },
{ name = "torchvision", url = "https://download.pytorch.org/whl/test/cu128/torchvision-0.26.0%2Bcu128-cp312-cp312-manylinux_2_28_x86_64.whl" },
- { name = "transformers", specifier = "==5.6.2" },
+ { name = "transformers", specifier = "==5.12.1" },
]
[[package]]
@@ -142,7 +142,7 @@ dependencies = [
[package.metadata]
requires-dist = [
{ name = "nvidia-nccl-cu12", marker = "sys_platform == 'linux'", specifier = "==2.28.9" },
- { name = "transformers", specifier = "==5.6.2" },
+ { name = "transformers", specifier = "==5.12.1" },
{ name = "vllm", marker = "sys_platform == 'linux'", url = "https://github.com/vllm-project/vllm/releases/download/v0.23.0/vllm-0.23.0%2Bcu129-cp38-abi3-manylinux_2_28_x86_64.whl" },
]
@@ -2464,7 +2464,7 @@ wheels = [
[[package]]
name = "transformers"
-version = "5.6.2"
+version = "5.12.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "huggingface-hub" },
@@ -2477,9 +2477,9 @@ dependencies = [
{ name = "tqdm" },
{ name = "typer" },
]
-sdist = { url = "https://files.pythonhosted.org/packages/a4/e9/c6c80a07690142a7d05444271f47b9f3c8aac7dea01d52e1137ee480ad78/transformers-5.6.2.tar.gz", hash = "sha256:e657134c3e5a6bc00a3c35f4e2674bb51adfcd89898495b788a18552bac2b91a", size = 8311867, upload-time = "2026-04-23T18:33:29.332Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/aa/7c/8240f612819718100a9346dc28dea6a11370c3ca9c8c6eabadd3dea4ef29/transformers-5.12.1.tar.gz", hash = "sha256:679ee731c8225347889ad4fb3b2c926a62e9da3b7d284e9d12c791da7272466b", size = 8924054, upload-time = "2026-06-15T17:27:50.604Z" }
wheels = [
- { url = "https://files.pythonhosted.org/packages/5d/95/0b0218149b0d6f14df35f5b8f676fa83df4f19ed253c3cc447107ef86eca/transformers-5.6.2-py3-none-any.whl", hash = "sha256:f8d3a1bb96778fed9b8aabfd0dd6e19843e4b0f2bb6b59f32b8a92051b0f348f", size = 10364898, upload-time = "2026-04-23T18:33:26.081Z" },
+ { url = "https://files.pythonhosted.org/packages/df/56/bbd60dd8668055803bf8ba55a81f9b8a8b31497f620109a9671d26a2076d/transformers-5.12.1-py3-none-any.whl", hash = "sha256:2a5e109d2021265df7098ffbb738295acaf5ad256f12cbc586db2ea4dcbb1a8a", size = 11150587, upload-time = "2026-06-15T17:27:46.679Z" },
]
[[package]]