From 3fad582fdf9d56514730cf14d5baf19e6772ba76 Mon Sep 17 00:00:00 2001 From: zyk42 <2931889928@qq.com> Date: Wed, 29 Jul 2026 23:25:25 +0800 Subject: [PATCH 1/4] feat: VLM DFlash draft model support (partial RoPE + interleaved mRoPE) Port of the DFlash draft-model changes from sgl-project/SpecForge PR #585: partial rotation in apply_rotary_pos_emb (rotary_dim < head_dim, for Qwen3.5/Qwen3.6 partial_rotary_factor=0.25) and Qwen3InterleavedMultiRotaryEmbedding selected by rope_scaling.mrope_interleaved. Co-Authored-By: Claude Opus 4.6 (cherry picked from commit 9323a510e09c1f05be397a02039f9337ad960804; scoped to the draft model only, draft configs land with the recipe) --- specforge/modeling/draft/dflash.py | 82 ++++++++++++++++++++++++++++-- 1 file changed, 79 insertions(+), 3 deletions(-) diff --git a/specforge/modeling/draft/dflash.py b/specforge/modeling/draft/dflash.py index 4acd30e11..91c12b98a 100644 --- a/specforge/modeling/draft/dflash.py +++ b/specforge/modeling/draft/dflash.py @@ -74,8 +74,20 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_len = q.size(-2) - q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :]) - k_embed = (k * cos) + (rotate_half(k) * sin) + rotary_dim = cos.shape[-1] + head_dim = q.shape[-1] + if rotary_dim < head_dim: + q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] + k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] + q_rot = (q_rot * cos[..., -q_len:, :]) + ( + rotate_half(q_rot) * sin[..., -q_len:, :] + ) + k_rot = (k_rot * cos) + (rotate_half(k_rot) * sin) + q_embed = torch.cat([q_rot, q_pass], dim=-1) + k_embed = torch.cat([k_rot, k_pass], dim=-1) + else: + q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :]) + k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed @@ -96,6 +108,64 @@ def _prepare_dflash_eager_mask( return additive_mask, valid_queries +def get_rope_scaling_value(config: Qwen3Config, key: str, default=None): + rope_scaling = getattr(config, "rope_scaling", None) + if rope_scaling is None: + return default + if isinstance(rope_scaling, dict): + return rope_scaling.get(key, default) + return getattr(rope_scaling, key, default) + + +class Qwen3InterleavedMultiRotaryEmbedding(Qwen3RotaryEmbedding): + """Interleaved mRoPE for Qwen3-VL style multimodal position ids.""" + + def __init__(self, config: Qwen3Config): + super().__init__(config) + self.mrope_section = get_rope_scaling_value( + config, "mrope_section", [24, 20, 20] + ) + + def _apply_interleaved_mrope(self, freqs: torch.Tensor) -> torch.Tensor: + freqs_t = freqs[0] + for dim_idx, offset in enumerate((1, 2), start=1): + length = self.mrope_section[dim_idx] * 3 + idx_slice = slice(offset, length, 3) + freqs_t[..., idx_slice] = freqs[dim_idx, ..., idx_slice] + return freqs_t + + @torch.no_grad() + def forward( + self, x: torch.Tensor, position_ids: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + if position_ids.ndim == 2: + position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + + inv_freq_expanded = ( + self.inv_freq[None, None, :, None] + .float() + .expand(3, position_ids.shape[1], -1, 1) + ) + position_ids_expanded = position_ids[:, :, None, :].float() + + device_type = ( + x.device.type + if isinstance(x.device.type, str) and x.device.type != "mps" + else "cpu" + ) + with torch.autocast(device_type=device_type, enabled=False): + freqs = ( + inv_freq_expanded.float() @ position_ids_expanded.float() + ).transpose(2, 3) + interleaved_freqs = self._apply_interleaved_mrope(freqs) + emb = torch.cat((interleaved_freqs, interleaved_freqs), dim=-1) + scaling = getattr(self, "attention_scaling", 1.0) + cos = emb.cos() * scaling + sin = emb.sin() * scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + class Qwen3DFlashAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" @@ -381,7 +451,13 @@ def __init__( build_target_layer_ids(config.num_target_layers, config.num_hidden_layers), ) self.norm = kernels.make_rms_norm(config.hidden_size, config.rms_norm_eps) - self.rotary_emb = Qwen3RotaryEmbedding(config) + self.use_interleaved_mrope = bool( + get_rope_scaling_value(config, "mrope_interleaved", False) + ) + if self.use_interleaved_mrope: + self.rotary_emb = Qwen3InterleavedMultiRotaryEmbedding(config) + else: + self.rotary_emb = Qwen3RotaryEmbedding(config) self.fc = nn.Linear( len(self.target_layer_ids) * config.hidden_size, config.hidden_size, From 6164057a69a1a20e303f71573d99f1d3f8f603e7 Mon Sep 17 00:00:00 2001 From: Curnane Date: Wed, 29 Jul 2026 23:25:54 +0800 Subject: [PATCH 2/4] refactor: port VLM target handling to the unified runtime Follow-up adaptation of 9323a510 to the server-only architecture: - specforge/modeling/target/target_utils.py: add QWEN3_VL_MODEL_TYPES and resolve_target_weight_keys(); TargetEmbeddingsAndHead.from_pretrained now auto-selects model.language_model.embed_tokens.weight for VLM targets (Qwen3-VL / Qwen3.5 / Qwen3.6) when the embed key is unset or left at the LLM default. Explicit keys are honored as-is. - specforge/algorithms/model_providers.py: populate_dflash_generated_config reads the language-model depth via the text_config fallback so VLM draft config auto-generation uses the correct num_hidden_layers. (cherry picked from commit e2280f4b) --- specforge/algorithms/model_providers.py | 6 ++- specforge/modeling/target/target_utils.py | 46 +++++++++++++++++++++-- 2 files changed, 47 insertions(+), 5 deletions(-) diff --git a/specforge/algorithms/model_providers.py b/specforge/algorithms/model_providers.py index 285433fcd..e09cfe05a 100644 --- a/specforge/algorithms/model_providers.py +++ b/specforge/algorithms/model_providers.py @@ -458,8 +458,12 @@ def populate_dflash_generated_config( payload: Dict[str, Any], target_config: Any, _cfg: Config ) -> None: from specforge.modeling.draft.dflash import build_target_layer_ids + from specforge.modeling.target.target_utils import target_text_config - target_layers = getattr(target_config, "num_hidden_layers", None) + # VLM targets (Qwen3-VL / Qwen3.5 / Qwen3.6) keep the language-model depth + # under ``text_config``; the top-level value may describe the vision stack. + text_config = target_text_config(target_config) + target_layers = getattr(text_config, "num_hidden_layers", None) if not isinstance(target_layers, int) or target_layers < 1: raise ValueError( "DFlash auto-generation requires target num_hidden_layers, got " diff --git a/specforge/modeling/target/target_utils.py b/specforge/modeling/target/target_utils.py index c8ca1afcb..3afdf42f4 100644 --- a/specforge/modeling/target/target_utils.py +++ b/specforge/modeling/target/target_utils.py @@ -67,6 +67,45 @@ def target_text_config(config): return getattr(config, "text_config", config) +#: Target ``model_type`` values whose checkpoints nest the language-model +#: weights under ``model.language_model.*`` (Qwen VLM families). +QWEN3_VL_MODEL_TYPES = {"qwen3_vl", "qwen3_vl_moe", "qwen3_5_moe", "qwen3_5"} + + +def is_vlm_target_config(config) -> bool: + """Whether the target config is a vision-language model with a text stack.""" + model_type = getattr(config, "model_type", None) + return model_type in QWEN3_VL_MODEL_TYPES or hasattr(config, "text_config") + + +#: Default (non-VLM) embedding weight key. +DEFAULT_EMBED_KEY = "model.embed_tokens.weight" +#: Embedding weight key inside VLM checkpoints' language-model stack. +VLM_EMBED_KEY = "model.language_model.embed_tokens.weight" + + +def resolve_target_weight_keys( + config, + embed_key: Optional[str] = None, + lm_head_key: Optional[str] = None, +) -> tuple[str, str]: + """Resolve default embedding / lm_head weight keys for a target config. + + VLM checkpoints (Qwen3-VL, Qwen3.5, Qwen3.6) store the language-model + embedding at ``model.language_model.embed_tokens.weight`` instead of the + LLM default ``model.embed_tokens.weight``; the lm_head stays top-level. + An ``embed_key`` that is unset or left at the LLM default is rewritten + automatically for VLM targets; any other explicit key is honored as-is. + """ + if embed_key is None or ( + embed_key == DEFAULT_EMBED_KEY and is_vlm_target_config(config) + ): + embed_key = VLM_EMBED_KEY if is_vlm_target_config(config) else DEFAULT_EMBED_KEY + if lm_head_key is None: + lm_head_key = "lm_head.weight" + return embed_key, lm_head_key + + def target_vocab_size(config) -> int: text_config = target_text_config(config) return int( @@ -113,10 +152,9 @@ def from_pretrained( ) instance = cls(config) - if embed_key is None: - embed_key = "model.embed_tokens.weight" - if lm_head_key is None: - lm_head_key = "lm_head.weight" + embed_key, lm_head_key = resolve_target_weight_keys( + config, embed_key, lm_head_key + ) # 2. Resolve Model Path local_model_path = model_path From 6ff1cde948d9a0b2cfff7bae72a73ed71f2584dd Mon Sep 17 00:00:00 2001 From: Curnane Date: Wed, 29 Jul 2026 23:28:32 +0800 Subject: [PATCH 3/4] feat(dflash): multimodal (image+text) server-capture training Implements end-to-end VLM DFlash training on the server-only runtime: - contracts/providers: dflash registers a STREAMING 'multimodal' FeatureContract ({input_ids, loss_mask, hidden_states, position_ids}) and a ServerStreamingProvider with a VLM ServerInputAdapter; ServerCaptureLayout/ServerCaptureSchema gain position_ids_feature. - data: specforge/data/vlm_preprocessing.py renders conversations with the target's own chat template (image attached to the first user turn) and expands the image region in id space via the target's HF processor, yielding expanded input_ids/loss_mask (trainer + passthrough) and collapsed request_input_ids + base64 image_data (capture request). One image per sample (v1); text-only samples work in the same run. - inference: the capture request's features map gains 'position_ids'; the regenerated sglang v0.5.14 spec-capture patch writes the request's mRoPE positions (1, L, 3) int64 into Mooncake from req.multimodal_inputs.mrope_positions (arange fallback for text prompts). Managed capture servers set SGLANG_MM_AVOID_RETOKENIZE=1 for input_modality=multimodal so the server re-expands placeholders in id space with no retokenization drift. - training: build_vlm_collator pads position_ids; OnlineDFlashModel.forward accepts optional (B, S, 3) position ids and gathers 3D mRoPE positions for context + anchor-offset draft slots; DFlashTrainStrategy passes the tensor through. Text paths are unchanged. - recipe: examples/configs/qwen3.5-4b-vl-dflash-multimodal-disaggregated.yaml + configs/qwen3.5-4b-vl-dflash.json (Qwen3.5-4B VLM draft geometry). - tests: new tests/test_algorithms/test_dflash_multimodal.py; update the text-only assertions in test_builtin_providers / test_server_only_online / test_unified_feature_reachability / test_launch_topology / test_example_draft_config_wiring for the new modality and recipe. GPU/NPU end-to-end validation (live ViT capture + training) is pending and tracked in docs/advanced_features/vlm_dflash.md. --- configs/qwen3.5-4b-vl-dflash.json | 50 ++ docs/advanced_features/vlm_dflash.md | 95 +++ examples/configs/README.md | 14 +- .../qwen3.5-4b-vl-dflash-disaggregated.yaml | 52 ++ patches/sglang/v0.5.14/spec-capture.patch | 765 ++++-------------- .../algorithms/common/dflash_family_model.py | 41 +- .../algorithms/common/hidden_states_data.py | 24 + specforge/algorithms/common/providers.py | 14 +- specforge/algorithms/common/vlm_input.py | 101 +++ specforge/algorithms/dflash/providers.py | 24 + specforge/data/vlm_preprocessing.py | 287 +++++++ .../inference/adapters/server_capture.py | 6 + specforge/inference/media.py | 7 +- specforge/launch_plan.py | 4 + specforge/training/disaggregated.py | 1 + specforge/training/strategies/base.py | 4 + .../test_algorithms/test_builtin_providers.py | 11 +- .../test_algorithms/test_dflash_multimodal.py | 213 +++++ tests/test_config/test_server_only_online.py | 18 + 19 files changed, 1122 insertions(+), 609 deletions(-) create mode 100644 configs/qwen3.5-4b-vl-dflash.json create mode 100644 docs/advanced_features/vlm_dflash.md create mode 100644 examples/configs/qwen3.5-4b-vl-dflash-disaggregated.yaml create mode 100644 specforge/algorithms/common/vlm_input.py create mode 100644 specforge/data/vlm_preprocessing.py create mode 100644 tests/test_algorithms/test_dflash_multimodal.py diff --git a/configs/qwen3.5-4b-vl-dflash.json b/configs/qwen3.5-4b-vl-dflash.json new file mode 100644 index 000000000..366df304b --- /dev/null +++ b/configs/qwen3.5-4b-vl-dflash.json @@ -0,0 +1,50 @@ +{ + "architectures": [ + "DFlashDraftModel" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "auto_map": { + "AutoModel": "dflash.DFlashDraftModel" + }, + "block_size": 16, + "bos_token_id": 248043, + "dflash_config": { + "mask_token_id": 248070, + "target_layer_ids": [1, 8, 15, 22, 29] + }, + "dtype": "bfloat16", + "eos_token_id": 248044, + "head_dim": 256, + "hidden_act": "silu", + "hidden_size": 2560, + "initializer_range": 0.02, + "intermediate_size": 9216, + "layer_types": [ + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention" + ], + "max_position_embeddings": 262144, + "model_type": "qwen3_vl_text", + "num_attention_heads": 16, + "num_hidden_layers": 5, + "num_key_value_heads": 4, + "num_target_layers": 32, + "pad_token_id": 248044, + "partial_rotary_factor": 0.25, + "rms_norm_eps": 1e-06, + "rope_scaling": { + "mrope_interleaved": true, + "mrope_section": [11, 11, 10], + "rope_type": "default", + "partial_rotary_factor": 0.25 + }, + "rope_theta": 10000000, + "sliding_window": null, + "tie_word_embeddings": true, + "use_cache": true, + "vocab_size": 248320 +} diff --git a/docs/advanced_features/vlm_dflash.md b/docs/advanced_features/vlm_dflash.md new file mode 100644 index 000000000..e6ef712b3 --- /dev/null +++ b/docs/advanced_features/vlm_dflash.md @@ -0,0 +1,95 @@ +# VLM DFlash Support — Status (add_vl_support) + +This branch ports sgl-project/SpecForge PR #585 (commit `9323a510`, author zyk42) +onto the server-only unified runtime **and implements end-to-end multimodal +(image+text) DFlash training on top of it**. This document records what landed, +what was deliberately not ported, and the validation status. + +## What this branch contains + +### Foundation (ported from PR #585) + +- **Draft model (VLM-capable)** — `specforge/modeling/draft/dflash.py`: + partial rotation in `apply_rotary_pos_emb` (`rotary_dim < head_dim`, for + Qwen3.5/Qwen3.6 `partial_rotary_factor=0.25`) and + `Qwen3InterleavedMultiRotaryEmbedding` selected by + `rope_scaling.mrope_interleaved`. +- **Draft config** — `configs/qwen3.5-4b-vl-dflash.json` (new, mirrors the + Qwen3.5-4B target geometry: head_dim 256, mrope_section [11,11,10], + partial_rotary_factor 0.25). The wider #585 config set (Qwen3-VL-8B/30B-A3B, + Qwen3.5-9B/35B-A3B) is intentionally out of scope for this branch; it can be + added verbatim in a follow-up once more targets are validated. +- **Weight-key resolution** — `resolve_target_weight_keys()` in + `specforge/modeling/target/target_utils.py` auto-selects + `model.language_model.embed_tokens.weight` for VLM targets; + `populate_dflash_generated_config` reads language-model depth via + `text_config`. + +### Multimodal capture (new in this branch) + +End-to-end data flow: **JSONL (+ image) → expanded ids/loss mask → capture +request (single-placeholder ids + base64 image) → patched SGLang server +expands, runs the ViT, captures aux hidden states + mRoPE positions → +Mooncake → collator → training forward with 3D position ids**. + +- `model.input_modality: multimodal` (DFlash only): a `FeatureContract` + (`{input_ids, loss_mask, hidden_states, position_ids}`) and a + `ServerStreamingProvider` with a VLM `ServerInputAdapter` + (`specforge/algorithms/common/vlm_input.py`). +- `specforge/data/vlm_preprocessing.py`: ShareGPT-style records with an + optional `image` field (path or base64); the target's own chat template and + HF processor produce the expanded `input_ids`/`loss_mask` (image region + expanded in id space, mask zeros). One image per sample max (v1); text-only + samples work in the same run. +- `ServerCaptureLayout.position_ids_feature` → the capture request's + `features["position_ids"]`; the patched server writes the request's mRoPE + positions `(1, L, 3) int64` into Mooncake (`_spec_capture_position_ids` in + the scheduler sink; text requests get the arange broadcast fallback). +- `patches/sglang/v0.5.14/spec-capture.patch`: regenerated with the + `position_ids` artifact (`SpecCaptureSink.put_sample(position_ids=...)`). + Multimodal capture requests ride the stock `input_ids` + `image_data` + `/generate` path with `SGLANG_MM_AVOID_RETOKENIZE=1` (set by the managed + launcher for `input_modality=multimodal`), so the server re-expands + placeholders in id space with zero retokenization drift — and the + passthrough/seq-len checks fail loudly if client and server expansions ever + disagree. +- Training: `OnlineDFlashModel.forward(..., position_ids=None)` gathers 3D + mRoPE positions for context + anchor-offset draft slots + (`(3, B, S + N·bs)`); `DFlashTrainStrategy` passes the collated + `position_ids` tensor through. Text runs are byte-identical to before. +- Recipe: `examples/configs/qwen3.5-4b-vl-dflash-disaggregated.yaml` + (single-node Ascend NPU managed stack). + +## Not ported (by design) + +- HF-backend VLM capture (`dflash_target_model.py`, `_build_vlm_reqs`, + `mm_token_type_ids`) and the `train_dflash.py --is-vlm` plumbing from the + pre-#678 script stack — superseded by server capture. +- `QwenVLOnlineDFlashModel` wiring — PR #585 referenced this class but never + defined it; the unified runtime needs no separate VLM wrapper class. +- Two accidental reverts in the original #585 diff (domino projector code, + D-PACE CLI args) were dropped during the cherry-pick. +- Offline (precomputed hidden states) multimodal capture: the offline path + stays text-only for now. +- Online evaluation for multimodal runs. + +## Validation status + +- **Verified (CPU, this repo)**: registration parity and provider gates, + request/payload construction, expansion math, collator, golden + topology/recipe tests — `tests/test_algorithms/test_dflash_multimodal.py` + plus the updated `test_config` suites (80 passed locally, 2 torch tests + skipped pending GPU). +- **Verified statically**: the regenerated patch applies cleanly both ways to + pristine sglang v0.5.14 (`git apply --check` / `--reverse --check`). +- **Not yet verified (needs GPU/NPU)**: a live multimodal capture run + (Qwen3-VL / Qwen3.5 target + ViT + mRoPE positions) and an end-to-end + training run. This is the next step; see the recipe above. + +## Reference results from PR #585 (HF stack, author-validated) + +- Qwen3-VL-30B-A3B-Thinking, 278K target-regenerated samples, 5-layer draft, + block_size=8: accept length 3.52, +35.8% inference speedup (4x RTX 5090, + TP=4, SGLang 0.5.12). +- Data must be target-model greedy-regenerated; system prompt must match + between training and inference; <10K samples overfit severely. diff --git a/examples/configs/README.md b/examples/configs/README.md index 3cbd0cc6a..2a092a3c7 100644 --- a/examples/configs/README.md +++ b/examples/configs/README.md @@ -70,6 +70,11 @@ the `eagle3` strategy; EAGLE3.1 is a draft-model configuration variant, not a second runtime or launch path. Recipes under `online/disaggregated/managed-local` are opt-in, single-node + +Multimodal (image+text) training is supported for `training.strategy=dflash` +via `model.input_modality: multimodal` (see +`qwen3.5-4b-vl-dflash-disaggregated.yaml`); the remaining catalog is +text-only. full-stack examples. Their typed `managed_local` blocks own Mooncake, one or more patched SGLang capture servers, and the trainer GPU allocation; the same `specforge train -c ...` command starts and cleans up the complete stack. @@ -179,7 +184,7 @@ should make their training strategy and topology explicit. | `model.draft_num_hidden_layers` | `null` | Positive fresh-architecture override where the strategy permits it. EAGLE3 remains one layer; P-EAGLE and DFlash may override their generated defaults. | | `model.draft_block_size` | `null` | Positive DFlash block-size override; generated DFlash configs default to 16. | | `model.target_backend` | `sglang` | `sglang` is the only accepted value; retired `hf`/`custom` names fail at config load. Offline feature consumers do not instantiate a target inference backend. | -| `model.input_modality` | `text` | The provider modality. The unified runtime supports text only; VLM modalities such as `qwen2_5_vl` are rejected. | +| `model.input_modality` | `text` | The provider modality. Built-ins support `text`; DFlash additionally supports `multimodal` (image+text server capture). Other identifiers are rejected at application resolution. | | `model.shard_target_output` | `false` | Retained for config migration; leave it false on the online disaggregated path. | | `model.trust_remote_code` | `false` | Enable only for model repositories that require custom loading code. | | `model.use_liger_kernel` | `false` | Enable Liger Qwen3 RMSNorm/SwiGLU kernels for DFlash training. Requires the `specforge[liger]` extra. | @@ -453,9 +458,12 @@ unless tuning throughput or memory pressure. `sp_ulysses_size * sp_ring_size > 1`. Non-USP runs keep both SP sizes at 1. - P-EAGLE reuses the EAGLE3 server feature schema, uses `flex_attention`, and requires batch size 1. -- VLM training, including Qwen2.5-VL, is not supported. Online capture accepts - text inputs only. - `training.compact_teacher` is offline text EAGLE3 only. +- Multimodal (image+text) training requires `training.strategy=dflash` with + `model.input_modality: multimodal`, a VLM draft config + (`configs/*-dflash-vlm-*.json`), and the v0.5.14 spec-capture patch with the + `position_ids` artifact. Vendor modalities such as `qwen2_5_vl` remain + unsupported. - Online evaluation is not supported. Offline `data.eval_hidden_states_path` and `training.eval_interval` must be configured together. diff --git a/examples/configs/qwen3.5-4b-vl-dflash-disaggregated.yaml b/examples/configs/qwen3.5-4b-vl-dflash-disaggregated.yaml new file mode 100644 index 000000000..6b846bfba --- /dev/null +++ b/examples/configs/qwen3.5-4b-vl-dflash-disaggregated.yaml @@ -0,0 +1,52 @@ +# Multimodal (image+text) DFlash on Qwen3.5-4B: disaggregated online training +# against an external patched SGLang capture server. Requires sglang v0.5.14 +# with patches/sglang/v0.5.14/spec-capture.patch (position_ids artifact) and +# SGLANG_MM_AVOID_RETOKENIZE=1 on the capture server. +# +# Data: ShareGPT-style JSONL with an optional top-level "image" field (file +# path or base64) per record; text-only records are supported in the same run. +model: + target_model_path: Qwen/Qwen3.5-4B + draft_model_config: configs/qwen3.5-4b-vl-dflash.json + target_backend: sglang + trust_remote_code: true + input_modality: multimodal +data: + train_data_path: ./cache/dataset/train_vl_regen.jsonl + max_length: 4096 + chat_template: qwen3.5 + cache_dir: cache + build_dataset_num_proc: 32 +training: + strategy: dflash + num_epochs: 10 + max_steps: 10000 + batch_size: 2 + accumulation_steps: 4 + learning_rate: 0.0006 + warmup_ratio: 0.04 + max_grad_norm: 1.0 + num_anchors: 512 + loss_decay_gamma: 7.0 + save_interval: 1000 + dist_timeout: 30 + seed: 42 +tracking: + report_to: none +run_id: qwen3.5-4b-vl-dflash +output_dir: outputs/qwen3.5-4b-vl-dflash + +deployment: + mode: disaggregated + trainer: + nnodes: 1 + nproc_per_node: 4 + disaggregated: + control_dir: outputs/qwen3.5-4b-vl-dflash/control + consumer_state_dir: outputs/qwen3.5-4b-vl-dflash/consumer-state + backend: mooncake + server_urls: + - http://127.0.0.1:30000 + mooncake_metadata_server: http://127.0.0.1:35880/metadata + mooncake_master_server_addr: 127.0.0.1:35551 + mooncake_protocol: tcp diff --git a/patches/sglang/v0.5.14/spec-capture.patch b/patches/sglang/v0.5.14/spec-capture.patch index 5fd29d1e1..efe9f4a89 100644 --- a/patches/sglang/v0.5.14/spec-capture.patch +++ b/patches/sglang/v0.5.14/spec-capture.patch @@ -1,5 +1,5 @@ diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py -index a99d252..bde8a84 100644 +index a99d252677..f5d0b35e79 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -94,6 +94,9 @@ class LogitsProcessorOutput: @@ -12,7 +12,7 @@ index a99d252..bde8a84 100644 ## Part 2: This part will be assigned in python/sglang/srt/layers/sampler.py::Sampler # he log probs of output tokens, if SGLANG_RETURN_ORIGINAL_LOGPROB = True, will get the log probs before applying temperature. If False, will get the log probs before applying temperature. -@@ -361,6 +364,30 @@ class LogitsProcessor(nn.Module): +@@ -361,6 +364,16 @@ class LogitsProcessor(nn.Module): sample_indices, logits_metadata, ) @@ -26,24 +26,10 @@ index a99d252..bde8a84 100644 + ) + else None + ) -+ # muP targets pass LM-head-scaled hidden states into LogitsProcessor, -+ # while SpecForge folds the same multiplier into its frozen target -+ # head. Restore the pre-head-scale representation before capture so -+ # the multiplier is applied exactly once when training recomputes logits. -+ logits_mup_width_multiplier = getattr( -+ self.config, "logits_mup_width_multiplier", None -+ ) -+ if ( -+ last_hidden_states_to_store is not None -+ and logits_mup_width_multiplier -+ ): -+ last_hidden_states_to_store = ( -+ last_hidden_states_to_store * float(logits_mup_width_multiplier) -+ ) del hidden_states if not logits_metadata.extend_return_logprob: -@@ -374,6 +401,7 @@ class LogitsProcessor(nn.Module): +@@ -374,6 +387,7 @@ class LogitsProcessor(nn.Module): return LogitsProcessorOutput( next_token_logits=sampled_logits, hidden_states=hidden_states_to_store, @@ -52,7 +38,7 @@ index a99d252..bde8a84 100644 # workaround since ForwardBatch is local to forward_batch_generation(). # They should be moved to GenerationBatchResult to keep this class clean. diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py -index b05334d..2853c82 100644 +index b05334deaa..2853c82315 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/python/sglang/srt/managers/detokenizer_manager.py @@ -441,6 +441,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin): @@ -64,20 +50,20 @@ index b05334d..2853c82 100644 indexer_topk=indexer_topk, customized_info=recv_obj.customized_info, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py -index 951f354..cd6ebf2 100644 +index 951f35495e..4de78534a4 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py -@@ -188,6 +188,10 @@ class GenerateReqInput(BaseReq): - log_metrics: bool = True - # Whether to return hidden states - return_hidden_states: Union[List[bool], bool] = False +@@ -280,6 +280,10 @@ class GenerateReqInput(BaseReq): + image_max_dynamic_patch: Optional[int] = None + video_max_dynamic_patch: Optional[int] = None + + # Spec-training capture sink instructions (see spec_capture_sink.py). + # Batch-level: List[Optional[dict]]; per-request after __getitem__. + spec_capture: Optional[Union[List[Optional[Dict]], Dict]] = None + - # Whether to return captured routed experts - return_routed_experts: bool = False - return_indexer_topk: bool = False + # Pre-computed delimiter indices for multi-item scoring. + # Batch-level: List[List[int]] (one per request). After __getitem__: List[int]. + multi_item_delimiter_indices: Optional[Union[List[List[int]], List[int]]] = None @@ -743,6 +747,11 @@ class GenerateReqInput(BaseReq): if self.multi_item_delimiter_indices is not None else None @@ -121,7 +107,7 @@ index 951f354..cd6ebf2 100644 token_steps: List[List[int]] = None diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py -index f1dc81d..acefaac 100755 +index f1dc81d174..acefaace89 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -707,6 +707,7 @@ class Req(ReqDllmMixin): @@ -148,7 +134,7 @@ index f1dc81d..acefaac 100755 # extra key for classifying the request (e.g. cache_salt) if lora_id is not None: diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py -index abba374..241da3f 100644 +index abba37441d..3018b82473 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -576,6 +576,19 @@ class Scheduler( @@ -179,103 +165,11 @@ index abba374..241da3f 100644 ) req.tokenizer = self.tokenizer -@@ -3097,6 +3111,13 @@ class Scheduler( - # GenerationBatchResult.extra_keep_alive_refs after forward returns. - self.batch_record_buf[self.batch_record_ct] = [batch, attr_snapshot] - -+ def _should_copy_hidden_states_to_cpu(self, batch: ScheduleBatch) -> bool: -+ """Avoid redundant multi-GiB capture D2H on non-writer TP ranks.""" -+ return batch.return_hidden_states and ( -+ self.ps.attn_tp_rank == 0 -+ or any(req.spec_capture is None for req in batch.reqs) -+ ) -+ - @contextmanager - def _forward_isolation(self, batch: ScheduleBatch, *, overlap: bool): - """Make SB transactional across one forward (overlap and non-overlap). -@@ -3258,7 +3279,9 @@ class Scheduler( - batch_result.copy_done = self.device_module.Event() - batch_result.copy_to_cpu( - return_logprob=batch.return_logprob, -- return_hidden_states=batch.return_hidden_states, -+ return_hidden_states=self._should_copy_hidden_states_to_cpu( -+ batch -+ ), - ) - else: - kwargs = ( -@@ -3369,6 +3392,10 @@ class Scheduler( - batch: ScheduleBatch, - result: Union[GenerationBatchResult, EmbeddingBatchResult], - ): -+ # Complete capture transfers on the scheduler thread before processing -+ # any kind of next result (prefill, decode, or idle). This preserves -+ # the output socket's single-thread ownership under mixed traffic. -+ self.batch_result_processor.drain_spec_captures() - self.publish_load_snapshot(force=batch.forward_mode.is_extend()) - - if batch.forward_mode.is_decode(): -@@ -3448,6 +3475,15 @@ class Scheduler( - - def on_idle(self): - """Idle housekeeping: guard, check, metrics, reset, sleep.""" -+ # A capture response is intentionally delayed until its background -+ # sink transfer completes; finish it on this thread so ownership of -+ # the ZeroMQ output socket remains single-threaded. Do not enter the -+ # idle sleeper while a future is outstanding or the producer waiting -+ # for that response could deadlock. -+ self.batch_result_processor.drain_spec_captures() -+ if self.batch_result_processor.has_pending_spec_captures(): -+ time.sleep(0.001) -+ return - if not self.is_fully_idle(): - return - -@@ -3496,6 +3532,7 @@ class Scheduler( - and (self.last_batch is None or self.last_batch.is_empty()) - and (self.cur_batch is None or self.cur_batch.is_empty()) - and (not self.enable_overlap or len(self.result_queue) == 0) -+ and not self.batch_result_processor.has_pending_spec_captures() - and self._pp_microbatches_drained() - ) - diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py -index a9d5f0c..f134354 100644 +index a9d5f0c28b..c9f8fdf7d3 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py -@@ -1,7 +1,9 @@ - from __future__ import annotations - - import logging --from dataclasses import dataclass -+import os -+import time -+from dataclasses import dataclass, field - from typing import ( - TYPE_CHECKING, - Callable, -@@ -78,6 +80,12 @@ class SchedulerBatchResultProcessor: - logprob_result_processor: SchedulerLogprobResultProcessor - output_streamer: SchedulerOutputStreamer - abort_request: Callable -+ _spec_capture_batches: List = field( -+ default_factory=list, -+ init=False, -+ repr=False, -+ compare=False, -+ ) - - def process_batch_result_prebuilt(self, batch: ScheduleBatch): - assert self.disaggregation_mode == DisaggregationMode.DECODE -@@ -181,6 +189,7 @@ class SchedulerBatchResultProcessor: - result: Union[GenerationBatchResult, EmbeddingBatchResult], - ): - skip_stream_req = None -+ pending_spec_captures: List = [] - - if self.is_generation: - if result.copy_done is not None: -@@ -257,6 +266,20 @@ class SchedulerBatchResultProcessor: +@@ -257,6 +257,18 @@ class SchedulerBatchResultProcessor: ) if ( @@ -289,43 +183,12 @@ index a9d5f0c..f134354 100644 + hidden_state_offset=hidden_state_offset, + ) + if req.finished(): -+ pending = self._sink_spec_capture(req) -+ if pending is not None: -+ pending_spec_captures.append(pending) ++ self._sink_spec_capture(req) + elif ( req.return_hidden_states and logits_output.hidden_states is not None ): -@@ -329,9 +352,25 @@ class SchedulerBatchResultProcessor: - req.inflight_middle_chunks -= 1 - req.time_stats.set_last_chunked_prefill_finish_time() - -- self.output_streamer.stream_output( -- batch.reqs, batch.return_logprob, skip_stream_req -- ) -+ if pending_spec_captures: -+ self._queue_spec_captures( -+ pending_spec_captures, -+ return_logprob=batch.return_logprob, -+ ) -+ capture_req_ids = {id(item[0]) for item in pending_spec_captures} -+ ready_reqs = [ -+ req -+ for req in batch.reqs -+ if id(req) not in capture_req_ids and req is not skip_stream_req -+ ] -+ if ready_reqs: -+ self.output_streamer.stream_output( -+ ready_reqs, batch.return_logprob -+ ) -+ else: -+ self.output_streamer.stream_output( -+ batch.reqs, batch.return_logprob, skip_stream_req -+ ) - - can_run_cuda_graph = result.can_run_cuda_graph - self.metrics_reporter.report_prefill_stats( -@@ -462,6 +501,183 @@ class SchedulerBatchResultProcessor: +@@ -462,6 +474,95 @@ class SchedulerBatchResultProcessor: f"Placeholder zeros would be appended to output_ids." ) @@ -344,43 +207,20 @@ index a9d5f0c..f134354 100644 + """ + start = hidden_state_offset + end = start + len(req.origin_input_ids) -+ # Only attention-TP rank 0 owns the Mooncake sink. Copying these very -+ # large tensors to host on every TP rank creates identical D2H -+ # transfers whose results are immediately discarded on TP > 1. -+ if self.output_streamer.ps.attn_tp_rank != 0: -+ return end -+ # Materialize each scheduler result on host once, rather than issuing -+ # one synchronous D2H transfer per request. In overlap mode the writer -+ # rank has already copied both tensors asynchronously on copy_stream; -+ # ``.cpu()`` is then a no-op. The fallback keeps non-overlap correct. -+ features = dict(req.spec_capture.get("features") or {}) -+ aux_cpu = getattr(logits_output, "_spec_capture_aux_cpu", None) -+ if "aux" in features and aux_cpu is None: -+ aux_cpu = logits_output.hidden_states.cpu() -+ logits_output._spec_capture_aux_cpu = aux_cpu -+ last_hidden_cpu = getattr( -+ logits_output, "_spec_capture_last_hidden_cpu", None ++ req.spec_capture_aux.append( ++ logits_output.hidden_states[start:end].cpu().clone() + ) -+ if ( -+ "last_hidden" in features -+ and logits_output.last_hidden_states is not None -+ and last_hidden_cpu is None -+ ): -+ last_hidden_cpu = logits_output.last_hidden_states.cpu() -+ logits_output._spec_capture_last_hidden_cpu = last_hidden_cpu -+ if "aux" in features and aux_cpu is not None: -+ req.spec_capture_aux.append(aux_cpu[start:end]) -+ if "last_hidden" in features and last_hidden_cpu is not None: -+ req.spec_capture_last_hidden.append(last_hidden_cpu[start:end]) ++ if logits_output.last_hidden_states is not None: ++ req.spec_capture_last_hidden.append( ++ logits_output.last_hidden_states[start:end].cpu().clone() ++ ) + return end + -+ def _sink_spec_capture(self, req: Req): -+ """Stage a finished capture request for the background Mooncake sink. ++ def _sink_spec_capture(self, req: Req) -> None: ++ """Write a finished capture request's tensors to the Mooncake sink. + -+ Runs on the attention-TP rank that streams output; returns the pending -+ tuple queued by ``_queue_spec_captures`` (or ``None`` off the writer -+ rank). The per-request result (or an ``{"error": ...}`` marker) is set -+ on ``req.spec_capture_result`` when the batch transfer completes, and ++ Runs on the attention-TP rank that streams output; the per-request result ++ (or an ``{"error": ...}`` marker) is set on ``req.spec_capture_result``, + returned to the client via the dedicated ``spec_capture`` output field + (a per-request channel, unlike per-token ``customized_info``). + """ @@ -388,147 +228,67 @@ index a9d5f0c..f134354 100644 + + sink = spec_capture_sink.get_sink() + if sink is None or self.output_streamer.ps.attn_tp_rank != 0: -+ return None -+ timing_enabled = os.environ.get("SGLANG_SPEC_CAPTURE_TIMING", "0") == "1" -+ cat_start = time.perf_counter() -+ # With chunked prefill disabled (the recommended capture-server -+ # configuration), each request owns exactly one contiguous view into -+ # the batch-level D2H buffer. torch.cat([view]) needlessly copied the -+ # whole capture a second time on CPU. Keep that view zero-copy; -+ # concatenate only the genuinely chunked case. -+ aux = self._coalesce_spec_capture_chunks(req.spec_capture_aux) -+ last_hidden = self._coalesce_spec_capture_chunks( -+ req.spec_capture_last_hidden -+ ) -+ cat_ms = (time.perf_counter() - cat_start) * 1000.0 -+ req.spec_capture_aux = [] -+ req.spec_capture_last_hidden = [] -+ return req, req.spec_capture, aux, last_hidden, timing_enabled, cat_ms -+ -+ def _queue_spec_captures(self, pending, *, return_logprob: bool) -> None: -+ """Queue one batch transfer while the scheduler runs the next prefill.""" -+ if not pending: + return -+ from sglang.srt import spec_capture_sink -+ -+ sink = spec_capture_sink.get_sink() -+ samples = [ -+ (spec, aux, last_hidden) -+ for _, spec, aux, last_hidden, _, _ in pending -+ ] -+ self._spec_capture_batches.append( -+ ( -+ pending, -+ sink.submit_samples(samples), -+ return_logprob, -+ time.perf_counter(), -+ ) -+ ) -+ # Bound retained D2H buffers. Reaching this point means the target -+ # prefill for batch N+1 already overlapped host transfer N; wait only -+ # if N is still finishing before admitting N+2. -+ max_pending = int( -+ os.environ.get("SGLANG_SPEC_CAPTURE_MAX_PENDING_BATCHES", "2") ++ aux = torch.cat(req.spec_capture_aux, dim=0) if req.spec_capture_aux else None ++ last_hidden = ( ++ torch.cat(req.spec_capture_last_hidden, dim=0) ++ if req.spec_capture_last_hidden ++ else None + ) -+ if max_pending < 1: -+ raise ValueError("SGLANG_SPEC_CAPTURE_MAX_PENDING_BATCHES must be >= 1") -+ if len(self._spec_capture_batches) >= max_pending: -+ self.drain_spec_captures(block=True, max_batches=1) -+ -+ def has_pending_spec_captures(self) -> bool: -+ return bool(self._spec_capture_batches) -+ -+ def drain_spec_captures( -+ self, *, block: bool = False, max_batches: Optional[int] = None -+ ) -> int: -+ """Finish ready transfers and stream their responses on this thread.""" -+ completed = 0 -+ while self._spec_capture_batches: -+ if max_batches is not None and completed >= max_batches: -+ break -+ pending, future, return_logprob, queued_at = self._spec_capture_batches[0] -+ if not block and not future.done(): -+ break -+ self._spec_capture_batches.pop(0) -+ self._complete_spec_capture_batch( -+ pending, -+ future, -+ return_logprob=return_logprob, -+ queued_at=queued_at, -+ ) -+ completed += 1 -+ return completed -+ -+ def _complete_spec_capture_batch( -+ self, pending, future, *, return_logprob: bool, queued_at: float -+ ) -> None: ++ position_ids = self._spec_capture_position_ids(req) + try: -+ results = future.result() -+ if len(results) != len(pending): -+ raise RuntimeError( -+ f"spec-capture sink returned {len(results)} results for " -+ f"{len(pending)} samples" -+ ) -+ except Exception as e: -+ logger.error( -+ "spec-capture batch sink failed for %d requests: %s", len(pending), e -+ ) -+ for req, spec, _, _, _, _ in pending: -+ req.spec_capture_result = { -+ "sample_id": spec.get("sample_id"), -+ "error": str(e), -+ } -+ else: -+ for pending_item, result in zip(pending, results): -+ req, _, _, _, _, _ = pending_item -+ req.spec_capture_result = result -+ -+ timing_enabled = any(item[4] for item in pending) -+ if timing_enabled: -+ logger.info( -+ "[spec-capture-timing] async_batch_complete samples=%d " -+ "queue_to_stream_ms=%.3f cat_ms=%.3f", -+ len(pending), -+ (time.perf_counter() - queued_at) * 1000.0, -+ sum(item[5] for item in pending), ++ req.spec_capture_result = sink.put_sample( ++ req.spec_capture, ++ aux=aux, ++ last_hidden=last_hidden, ++ position_ids=position_ids, + ) -+ self.output_streamer.stream_output( -+ [item[0] for item in pending], return_logprob -+ ) ++ except Exception as e: ++ logger.error("spec-capture sink failed for %s: %s", req.rid, e) ++ req.spec_capture_result = { ++ "sample_id": req.spec_capture.get("sample_id"), ++ "error": str(e), ++ } ++ req.spec_capture_aux = [] ++ req.spec_capture_last_hidden = [] + + @staticmethod -+ def _coalesce_spec_capture_chunks( -+ chunks: List[torch.Tensor], -+ ) -> Optional[torch.Tensor]: -+ if not chunks: ++ def _spec_capture_position_ids(req: Req) -> Optional[torch.Tensor]: ++ """Prompt position ids for the sink's ``position_ids`` artifact. ++ ++ Multimodal (mRoPE) requests carry the processor-computed positions on ++ ``req.multimodal_inputs.mrope_positions`` ([3, L]); text requests fall ++ back to the plain arange broadcast the model uses for non-mm prompts. ++ Returns an (L, 3) CPU int64 tensor, or None when the client did not ++ request the artifact. ++ """ ++ features = (req.spec_capture or {}).get("features") or {} ++ if features.get("position_ids") is None: + return None -+ if len(chunks) == 1: -+ return chunks[0] -+ return torch.cat(chunks, dim=0) ++ seq_len = len(req.origin_input_ids) ++ mrope_positions = ( ++ req.multimodal_inputs.mrope_positions ++ if req.multimodal_inputs is not None ++ else None ++ ) ++ if mrope_positions is not None: ++ return mrope_positions[:, :seq_len].t().contiguous().cpu() ++ return ( ++ torch.arange(seq_len, dtype=torch.int64) ++ .unsqueeze(0) ++ .expand(3, -1) ++ .t() ++ .contiguous() ++ ) + def _append_prefill_hidden_states( self, *, diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py -index f95c59f..6c51abc 100644 +index f95c59f6f2..9e64326a13 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py -@@ -145,6 +145,14 @@ class SchedulerOutputStreamer: - for req in reqs: - if req is skip_req: - continue -+ if self.server_args.enable_spec_capture and req.spec_capture is not None: -+ # Capture requests are not complete until the background sink -+ # has durably published every feature object. Several -+ # scheduler paths can ask the common streamer to emit a -+ # finished request; centralize the completion barrier here so -+ # none of them can win the race and send a metadata-less 200. -+ if req.finished() and req.spec_capture_result is None: -+ continue - if req.finished() and req.finished_output: - # With the overlap schedule, a request will try to output twice and hit this line twice - # because of the one additional delayed token. This "continue" prevented the dummy output. -@@ -273,6 +281,7 @@ class _GenerationStreamAccumulator: +@@ -273,6 +273,7 @@ class _GenerationStreamAccumulator: spec_correct_drafts_histogram: list = field(default_factory=list) retraction_counts: list = field(default_factory=list) output_hidden_states: Optional[list] = None @@ -536,7 +296,7 @@ index f95c59f..6c51abc 100644 routed_experts: Optional[list] = None indexer_topk: Optional[list] = None customized_info: dict = field(default_factory=dict) -@@ -482,6 +491,8 @@ class _GenerationStreamAccumulator: +@@ -482,6 +483,8 @@ class _GenerationStreamAccumulator: self.output_hidden_states.append(hs) else: self.output_hidden_states.append(None) @@ -545,7 +305,7 @@ index f95c59f..6c51abc 100644 if self.return_routed_experts: self.routed_experts.append( req.routed_experts if req.return_routed_experts else None -@@ -540,6 +551,7 @@ class _GenerationStreamAccumulator: +@@ -540,6 +543,7 @@ class _GenerationStreamAccumulator: output_token_ids_logprobs_idx=self.output_token_ids_logprobs_idx, output_token_entropy_val=None, output_hidden_states=self.output_hidden_states, @@ -554,7 +314,7 @@ index f95c59f..6c51abc 100644 indexer_topk=self.indexer_topk, customized_info=self.customized_info, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py -index bf93261..f0e821e 100644 +index bf932611a5..f0e821ec47 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1172,6 +1172,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): @@ -576,32 +336,11 @@ index bf93261..f0e821e 100644 if getattr(recv_obj, "routed_experts", None): val = recv_obj.routed_experts[i] if val is not None: -diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py -index e9ede57..1d624fd 100644 ---- a/python/sglang/srt/managers/utils.py -+++ b/python/sglang/srt/managers/utils.py -@@ -77,6 +77,16 @@ class GenerationBatchResult: - Only the tensors which are needed for processing results are copied, - e.g., next_token_ids, logits outputs - """ -+ if ( -+ return_hidden_states -+ and self.logits_output.last_hidden_states is not None -+ ): -+ # Spec-capture's post-norm last hidden rides the same async D2H -+ # as the aux hidden states (the field exists only when the -+ # spec-capture patch is applied). -+ self.logits_output.last_hidden_states = ( -+ self.logits_output.last_hidden_states.to("cpu", non_blocking=True) -+ ) - if return_logprob: - if self.logits_output.next_token_logprobs is not None: - self.logits_output.next_token_logprobs = ( diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py -index 1cff5c9..5309965 100644 +index 1cff5c9839..b842aeccb9 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py -@@ -515,6 +515,42 @@ class ModelRunner(ModelRunnerKVCacheMixin): +@@ -515,6 +515,22 @@ class ModelRunner(ModelRunnerKVCacheMixin): draft_num_layers=int(draft_num_layers), ) @@ -609,46 +348,26 @@ index 1cff5c9..5309965 100644 + # Aux capture without a draft worker, routed to the strategy's own + # capture method (they wire different submodules — e.g. VL models + # only populate layers_to_capture via the dflash path). -+ capture_method = getattr(server_args, "spec_capture_method", "eagle3") -+ if capture_method in ("dflash", "dspark"): -+ # DSpark rides the DFlash aux plumbing: this patch's supported -+ # targets expose dflash capture hooks (targets with a native -+ # dspark hook are served by the kimi-k3 patch variant). ++ if getattr(server_args, "spec_capture_method", "eagle3") == "dflash": + self.dflash_use_aux_hidden_state = True + self.dflash_target_layer_ids = server_args.spec_capture_aux_layer_ids -+ if hasattr(self, "spec_aux_config"): -+ self.spec_aux_config.dflash_use_aux_hidden_state = True -+ self.spec_aux_config.dflash_target_layer_ids = ( -+ server_args.spec_capture_aux_layer_ids -+ ) + self.dflash_family_use_aux_hidden_state = True + self.dflash_family_target_layer_ids = ( + server_args.spec_capture_aux_layer_ids + ) -+ elif capture_method == "eagle3": ++ else: + self.eagle_use_aux_hidden_state = True + self.eagle_aux_hidden_state_layer_ids = ( + server_args.spec_capture_aux_layer_ids + ) -+ if hasattr(self, "spec_aux_config"): -+ self.spec_aux_config.eagle_use_aux_hidden_state = True -+ self.spec_aux_config.eagle_aux_hidden_state_layer_ids = ( -+ server_args.spec_capture_aux_layer_ids -+ ) -+ else: -+ raise ValueError( -+ "--spec-capture-method must be one of: eagle3, dflash, dspark; " -+ f"got {capture_method!r}" -+ ) -+ # Apply the rank zero filter to logger if server_args.show_time_cost: enable_show_time_cost() diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py -index c7162c1..cfbce7a 100644 +index c7162c16d5..d515b3c699 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py -@@ -2127,6 +2127,26 @@ class ServerArgs: +@@ -2127,6 +2127,25 @@ class ServerArgs: bool, "Enable returning hidden states with responses.", ] = False @@ -667,20 +386,19 @@ index c7162c1..cfbce7a 100644 + ] = None + spec_capture_method: A[ + str, -+ "Capture method for --enable-spec-capture: 'eagle3', 'dflash', or " -+ "'dspark'. Must match the draft strategy being trained; they wire " -+ "capture onto different submodules (VL targets only populate the " -+ "dflash path; dspark rides the dflash aux plumbing).", ++ "Capture method for --enable-spec-capture: 'eagle3' or 'dflash'. Must " ++ "match the draft strategy being trained; they wire capture onto " ++ "different submodules (VL targets only populate the dflash path).", + ] = "eagle3" enable_return_routed_experts: A[ bool, "Enable returning routed experts of each layer with responses.", diff --git a/python/sglang/srt/spec_capture_sink.py b/python/sglang/srt/spec_capture_sink.py new file mode 100644 -index 0000000..d54b920 +index 0000000000..a678bd85e7 --- /dev/null +++ b/python/sglang/srt/spec_capture_sink.py -@@ -0,0 +1,395 @@ +@@ -0,0 +1,256 @@ +# Copyright 2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. @@ -701,7 +419,8 @@ index 0000000..d54b920 +Request schema:: + + {"store_id", "sample_id", "gen", "replace", # key namespace / retry policy -+ "features": {"aux": , "last_hidden": }, # artifact -> feature ++ "features": {"aux": , "last_hidden": , # artifact -> feature ++ "position_ids": }, # (L, 3) prompt positions + "passthrough": [{"name", "data", "shape", "dtype"}]} # client tensors verbatim + +Response (``meta_info["spec_capture"]``): ``{"sample_id", "store_id", "gen", @@ -716,9 +435,7 @@ index 0000000..d54b920 +import logging +import os +import threading -+import time -+from concurrent.futures import Future, ThreadPoolExecutor -+from typing import Any, Dict, List, Optional, Tuple ++from typing import Any, Dict, List, Optional + +import torch + @@ -742,6 +459,7 @@ index 0000000..d54b920 + +_ARTIFACT_AUX = "aux" +_ARTIFACT_LAST_HIDDEN = "last_hidden" ++_ARTIFACT_POSITION_IDS = "position_ids" + + +class SpecCaptureSink: @@ -752,13 +470,9 @@ index 0000000..d54b920 + self._store = None + self._put_config = None + self._lock = threading.Lock() -+ # One store writer is sufficient: Mooncake already stripes a batched -+ # transfer internally. The executor decouples that host transfer from -+ # the scheduler so the next target prefill can run concurrently. -+ self._executor = ThreadPoolExecutor( -+ max_workers=1, -+ thread_name_prefix="spec-capture-batch-put", -+ ) ++ # Retried HTTP requests reuse deterministic keys. Striped locks keep ++ # replacement atomic per key without retaining one lock per sample. ++ self._write_locks = [threading.Lock() for _ in range(256)] + + # -- connection --------------------------------------------------------- + def _connect(self): @@ -794,13 +508,7 @@ index 0000000..d54b920 + # before the trainer consumes it. + cfg = ReplicateConfig() + cfg.replica_num = 1 -+ # `with_hard_pin` exists only on newer Mooncake builds; older ROCm -+ # sglang images expose only `with_soft_pin`. Map the hard-pin intent -+ # onto whatever the installed build supports. -+ if hasattr(cfg, "with_hard_pin"): -+ cfg.with_hard_pin = True -+ elif hasattr(cfg, "with_soft_pin"): -+ cfg.with_soft_pin = True ++ cfg.with_hard_pin = True + self._put_config = cfg + self._store = store + logger.info("spec-capture mooncake sink connected") @@ -811,118 +519,83 @@ index 0000000..d54b920 + def _tkey(store_id: str, sample_id: str, gen: int, name: str) -> str: + return f"{store_id}/{sample_id}/g{gen}/{name}" + ++ def _put_tensor( ++ self, key: str, t: torch.Tensor, *, replace: bool = False ++ ) -> None: ++ store = self._connect() ++ t = t.detach().to("cpu").contiguous() ++ nbytes = t.element_size() * t.numel() ++ lock = self._write_locks[hash(key) % len(self._write_locks)] ++ with lock: ++ if replace: ++ # Do not probe with is_exist(): Mooncake existence checks can ++ # acquire a read lease that prevents the following removal. ++ self._remove_quiet(key) ++ try: ++ store.register_buffer(t.data_ptr(), nbytes) ++ except Exception: ++ pass # some builds auto-register ++ try: ++ rc = store.put_from(key, t.data_ptr(), nbytes, self._put_config) ++ finally: ++ try: ++ store.unregister_buffer(t.data_ptr()) ++ except Exception: ++ pass ++ if rc is not None and int(rc) < 0: ++ raise RuntimeError(f"spec-capture put_from failed (status {rc}) for {key}") ++ + def _remove_quiet(self, key: str) -> None: + try: + self._connect().remove(key) + except Exception: + pass + -+ def _remove_many_quiet(self, keys: List[str]) -> None: -+ if not keys: -+ return -+ store = self._connect() -+ batch_remove = getattr(store, "batch_remove", None) -+ if batch_remove is not None: -+ try: -+ batch_remove(keys) -+ return -+ except Exception: -+ pass -+ for key in keys: -+ self._remove_quiet(key) -+ -+ # -- the batch entry point ------------------------------------------------ -+ def submit_samples( -+ self, -+ samples: List[ -+ Tuple[ -+ Dict[str, Any], Optional[torch.Tensor], Optional[torch.Tensor] -+ ] -+ ], -+ ) -> Future[List[Dict[str, Any]]]: -+ """Queue one scheduler batch without blocking the scheduler thread.""" -+ return self._executor.submit(self.put_samples, samples) -+ -+ def put_samples( ++ # -- the one entry point -------------------------------------------------- ++ def put_sample( + self, -+ samples: List[ -+ Tuple[ -+ Dict[str, Any], Optional[torch.Tensor], Optional[torch.Tensor] -+ ] -+ ], -+ ) -> List[Dict[str, Any]]: -+ """Publish a scheduler batch with one native Mooncake batch RPC. ++ spec: Dict[str, Any], ++ *, ++ aux: Optional[torch.Tensor], ++ last_hidden: Optional[torch.Tensor], ++ position_ids: Optional[torch.Tensor] = None, ++ ) -> Dict[str, Any]: ++ """Write one sample's artifacts; return the meta_info result dict. + -+ A capture prefill batch normally finishes many samples together. -+ Calling ``put_from`` once per feature object paid dozens of metadata -+ and transport round trips per batch on the writer rank and serialized -+ the response behind them. -+ ``batch_put_from`` preserves the existing per-feature keys and raw -+ tensor layout while amortizing that fixed cost across the whole -+ scheduler batch. The response is still emitted only after every -+ status succeeds, so refs can never point at incomplete samples. ++ ``aux``/``last_hidden`` are the per-request (L, W) captured tensors, ++ ``position_ids`` the per-request (L, 3) prompt positions (mRoPE rows ++ for multimodal requests, arange broadcast for text), all stored with a ++ leading batch dim of 1. On any failure the keys already written are ++ best-effort removed (no partial sample is consumable). + """ -+ if not samples: -+ return [] ++ store_id = str(spec["store_id"]) ++ sample_id = str(spec["sample_id"]) ++ gen = int(spec.get("gen", 1)) ++ replace = bool(spec.get("replace", False)) ++ features: Dict[str, str] = dict(spec.get("features") or {}) + -+ store = self._connect() -+ timing_enabled = os.environ.get("SGLANG_SPEC_CAPTURE_TIMING", "0") == "1" -+ started = time.perf_counter() -+ keys: List[str] = [] -+ tensors: List[torch.Tensor] = [] -+ sizes: List[int] = [] -+ replace_keys: List[str] = [] -+ results: List[Dict[str, Any]] = [] ++ written: List[str] = [] ++ result_feats: Dict[str, Dict[str, Any]] = {} + -+ def _stage( -+ result_feats: Dict[str, Dict[str, Any]], -+ *, -+ store_id: str, -+ sample_id: str, -+ gen: int, -+ replace: bool, -+ name: str, -+ tensor: torch.Tensor, -+ ) -> None: -+ tensor = tensor.detach().to("cpu").contiguous() ++ def _write(name: str, t: torch.Tensor) -> None: + key = self._tkey(store_id, sample_id, gen, name) -+ keys.append(key) -+ tensors.append(tensor) -+ sizes.append(tensor.element_size() * tensor.numel()) -+ if replace: -+ replace_keys.append(key) ++ self._put_tensor(key, t, replace=replace) ++ written.append(key) + result_feats[name] = { -+ "shape": list(tensor.shape), -+ "dtype": _DTYPE_STR.get( -+ tensor.dtype, str(tensor.dtype).replace("torch.", "") -+ ), ++ "shape": list(t.shape), ++ "dtype": _DTYPE_STR.get(t.dtype, str(t.dtype).replace("torch.", "")), + } + -+ for spec, aux, last_hidden in samples: -+ store_id = str(spec["store_id"]) -+ sample_id = str(spec["sample_id"]) -+ gen = int(spec.get("gen", 1)) -+ replace = bool(spec.get("replace", False)) -+ features: Dict[str, str] = dict(spec.get("features") or {}) -+ result_feats: Dict[str, Dict[str, Any]] = {} -+ ++ try: + aux_name = features.get(_ARTIFACT_AUX) + if aux_name is not None: + if aux is None: + raise RuntimeError( + "spec_capture requested 'aux' but no aux hidden states were " -+ "captured -- launch the server with --enable-spec-capture " ++ "captured — launch the server with --enable-spec-capture " + "(and optionally --spec-capture-aux-layer-ids)" + ) -+ _stage( -+ result_feats, -+ store_id=store_id, -+ sample_id=sample_id, -+ gen=gen, -+ replace=replace, -+ name=aux_name, -+ tensor=aux.unsqueeze(0), -+ ) ++ _write(aux_name, aux.unsqueeze(0)) + lh_name = features.get(_ARTIFACT_LAST_HIDDEN) + if lh_name is not None: + if last_hidden is None: @@ -930,15 +603,15 @@ index 0000000..d54b920 + "spec_capture requested 'last_hidden' but the logits " + "processor did not return it (is aux capture enabled?)" + ) -+ _stage( -+ result_feats, -+ store_id=store_id, -+ sample_id=sample_id, -+ gen=gen, -+ replace=replace, -+ name=lh_name, -+ tensor=last_hidden.unsqueeze(0), -+ ) ++ _write(lh_name, last_hidden.unsqueeze(0)) ++ pos_name = features.get(_ARTIFACT_POSITION_IDS) ++ if pos_name is not None: ++ if position_ids is None: ++ raise RuntimeError( ++ "spec_capture requested 'position_ids' but no prompt " ++ "positions were collected on the scheduler" ++ ) ++ _write(pos_name, position_ids.unsqueeze(0)) + for item in spec.get("passthrough") or []: + dtype = _STR_DTYPE.get(str(item.get("dtype", "int64"))) + if dtype is None: @@ -946,116 +619,22 @@ index 0000000..d54b920 + f"spec_capture passthrough {item.get('name')!r}: " + f"unsupported dtype {item.get('dtype')!r}" + ) -+ tensor = torch.tensor(item["data"], dtype=dtype).reshape( ++ t = torch.tensor(item["data"], dtype=dtype).reshape( + [int(d) for d in item["shape"]] + ) -+ _stage( -+ result_feats, -+ store_id=store_id, -+ sample_id=sample_id, -+ gen=gen, -+ replace=replace, -+ name=str(item["name"]), -+ tensor=tensor, -+ ) -+ results.append( -+ { -+ "sample_id": sample_id, -+ "store_id": store_id, -+ "gen": gen, -+ "aux_layer_ids": self.aux_layer_ids, -+ "features": result_feats, -+ } -+ ) -+ -+ materialize_ms = (time.perf_counter() - started) * 1000.0 -+ self._remove_many_quiet(replace_keys) -+ registered: List[torch.Tensor] = [] -+ register_started = time.perf_counter() -+ try: -+ for tensor, nbytes in zip(tensors, sizes): -+ try: -+ store.register_buffer(tensor.data_ptr(), nbytes) -+ registered.append(tensor) -+ except Exception: -+ pass # TCP and some Mooncake builds auto-register -+ register_ms = (time.perf_counter() - register_started) * 1000.0 -+ put_started = time.perf_counter() -+ batch_put = getattr(store, "batch_put_from", None) -+ if batch_put is None: -+ statuses = [ -+ store.put_from(key, tensor.data_ptr(), nbytes, self._put_config) -+ for key, tensor, nbytes in zip(keys, tensors, sizes) -+ ] -+ else: -+ statuses = batch_put( -+ keys, -+ [tensor.data_ptr() for tensor in tensors], -+ sizes, -+ self._put_config, -+ ) -+ put_ms = (time.perf_counter() - put_started) * 1000.0 ++ _write(str(item["name"]), t) + except Exception: -+ self._remove_many_quiet(keys) ++ for key in written: ++ self._remove_quiet(key) + raise -+ finally: -+ for tensor in registered: -+ try: -+ store.unregister_buffer(tensor.data_ptr()) -+ except Exception: -+ pass + -+ if statuses is None: -+ statuses = [0] * len(keys) -+ if len(statuses) != len(keys): -+ self._remove_many_quiet(keys) -+ raise RuntimeError( -+ "spec-capture batch_put_from returned " -+ f"{len(statuses)} statuses for {len(keys)} keys" -+ ) -+ failed = [ -+ (key, status) -+ for key, status in zip(keys, statuses) -+ if status is not None and int(status) < 0 -+ ] -+ if failed: -+ self._remove_many_quiet(keys) -+ raise RuntimeError( -+ "spec-capture batch_put_from failed for " -+ f"{len(failed)}/{len(keys)} keys; first={failed[0]}" -+ ) -+ -+ if timing_enabled: -+ logger.info( -+ "[spec-capture-timing] batch_sink samples=%d objects=%d " -+ "bytes=%d materialize_ms=%.3f register_ms=%.3f put_ms=%.3f " -+ "total_ms=%.3f", -+ len(samples), -+ len(keys), -+ sum(sizes), -+ materialize_ms, -+ register_ms, -+ put_ms, -+ (time.perf_counter() - started) * 1000.0, -+ ) -+ return results -+ -+ def put_sample( -+ self, -+ spec: Dict[str, Any], -+ *, -+ aux: Optional[torch.Tensor], -+ last_hidden: Optional[torch.Tensor], -+ ) -> Dict[str, Any]: -+ """Write one sample's artifacts; return the meta_info result dict. -+ -+ ``aux``/``last_hidden`` are the per-request (L, W) captured tensors, -+ stored with a leading batch dim of 1. On any failure the keys already -+ written are best-effort removed (no partial sample is consumable). -+ """ -+ # Keep the single-sample entry point for compatibility with tests and -+ # callers outside the scheduler; production uses put_samples(). -+ return self.put_samples([(spec, aux, last_hidden)])[0] ++ return { ++ "sample_id": sample_id, ++ "store_id": store_id, ++ "gen": gen, ++ "aux_layer_ids": self.aux_layer_ids, ++ "features": result_feats, ++ } + + +_SINK: Optional[SpecCaptureSink] = None diff --git a/specforge/algorithms/common/dflash_family_model.py b/specforge/algorithms/common/dflash_family_model.py index 2c0bdfe62..faf5320b6 100644 --- a/specforge/algorithms/common/dflash_family_model.py +++ b/specforge/algorithms/common/dflash_family_model.py @@ -305,6 +305,7 @@ def _forward_draft_blocks( input_ids: torch.Tensor, hidden_states: torch.Tensor, loss_mask: torch.Tensor, + position_ids: Optional[torch.Tensor] = None, max_valid_anchors: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: bsz, seq_len = input_ids.shape @@ -321,11 +322,32 @@ def _forward_draft_blocks( input_ids, anchor_positions, block_keep_mask ) - context_position_ids = ( - torch.arange(seq_len, device=device).unsqueeze(0).expand(bsz, -1) - ) - draft_position_ids = self._create_position_ids(anchor_positions) - full_position_ids = torch.cat([context_position_ids, draft_position_ids], dim=1) + if position_ids is None: + context_position_ids = ( + torch.arange(seq_len, device=device).unsqueeze(0).expand(bsz, -1) + ) + draft_position_ids = self._create_position_ids(anchor_positions) + full_position_ids = torch.cat( + [context_position_ids, draft_position_ids], dim=1 + ) + else: + if not getattr(self.draft_model, "use_interleaved_mrope", False): + raise ValueError( + "multimodal capture carries mRoPE position_ids, but the " + "draft config does not enable rope_scaling.mrope_interleaved; " + "use a VLM draft config (e.g. configs/*-dflash-vlm-*.json)" + ) + # Server-produced mRoPE positions, (B, S, 3) -> (3, B, S + N*bs). + offsets = torch.arange(self.block_size, device=device).view(1, 1, -1) + draft_indices = (anchor_positions.unsqueeze(-1) + offsets).view(bsz, -1) + draft_position_ids = torch.gather( + position_ids, + 1, + draft_indices.unsqueeze(-1).expand(-1, -1, 3), + ) + full_position_ids = torch.cat( + [position_ids, draft_position_ids], dim=1 + ).permute(2, 0, 1) mask_builder = ( create_dflash_block_mask @@ -443,10 +465,16 @@ def forward( input_ids: torch.Tensor, hidden_states: torch.Tensor, loss_mask: torch.Tensor, + position_ids: Optional[torch.Tensor] = None, max_valid_anchors: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, object]]: """Parallel block-wise training forward pass; returns - (loss, accuracy, metrics) — same shape as Domino's forward.""" + (loss, accuracy, metrics) — same shape as Domino's forward. + + ``position_ids`` is the optional server-captured mRoPE position tensor + ``(B, S, 3)`` for multimodal runs; text runs leave it None and use + internally synthesized flat positions. + """ if self.attention_backend == "flex_attention" and not FLEX_ATTENTION_AVAILABLE: raise ValueError( "flex_attention is not available on this device; use sdpa/eager." @@ -458,6 +486,7 @@ def forward( input_ids=input_ids, hidden_states=hidden_states, loss_mask=loss_mask, + position_ids=position_ids, max_valid_anchors=max_valid_anchors, ) diff --git a/specforge/algorithms/common/hidden_states_data.py b/specforge/algorithms/common/hidden_states_data.py index 770ed710c..f774d9e61 100644 --- a/specforge/algorithms/common/hidden_states_data.py +++ b/specforge/algorithms/common/hidden_states_data.py @@ -168,6 +168,29 @@ def collate(features): return collate +def build_vlm_collator(): + """Collator for multimodal capture: text tensors + mRoPE position ids. + + ``position_ids`` arrives as ``(1, L, 3)`` int64 per sample (temporal / + height / width mRoPE rows produced by the capture server) and is padded + along the sequence axis like every other per-token feature. + """ + + def collate(features): + return pad_and_concatenate_features( + features, + sequence_axes={ + "input_ids": 1, + "loss_mask": 1, + "hidden_states": 1, + "position_ids": 1, + }, + required_keys=("input_ids", "loss_mask", "hidden_states", "position_ids"), + ) + + return collate + + def build_dspark_collator(): def collate(features): return pad_and_concatenate_features( @@ -281,6 +304,7 @@ def collate(features): "build_mtp_offline_reader", "build_offline_normalizer", "build_offline_reader", + "build_vlm_collator", "normalize_dspark_offline_sample", "normalize_mtp_offline_sample", "normalize_offline_sample", diff --git a/specforge/algorithms/common/providers.py b/specforge/algorithms/common/providers.py index 12cb0b08d..9d1114a4c 100644 --- a/specforge/algorithms/common/providers.py +++ b/specforge/algorithms/common/providers.py @@ -511,18 +511,25 @@ def __post_init__(self) -> None: @dataclass(frozen=True) class ServerCaptureLayout: - """Maps generic server artifacts onto algorithm-ready feature names.""" + """Maps generic server artifacts onto algorithm-ready feature names. + + ``position_ids_feature`` names an optional server-produced position-id + artifact (mRoPE positions for multimodal targets), stored ``(1, L, 3)`` + int64. Text providers leave it unset. + """ aux_feature: str | None last_hidden_feature: str | None passthrough: Tuple[Tuple[str, str, Tuple[int, ...]], ...] attention_mask_feature: str | None = None + position_ids_feature: str | None = None def __post_init__(self) -> None: for field_name in ( "aux_feature", "last_hidden_feature", "attention_mask_feature", + "position_ids_feature", ): value = getattr(self, field_name) if value is not None: @@ -546,8 +553,8 @@ class ServerStreamingProvider: """Algorithm adapter for externally captured streaming features. ``build_input_adapter`` is deliberately modality-neutral. Text providers - can leave it unset; the current runtime does not support VLM registration - or media requests. + can leave it unset; modalities such as ``multimodal`` provide one to own + prompt preparation and request construction (e.g. image inputs). """ modality: str @@ -754,6 +761,7 @@ def make_registration( layout.aux_feature, layout.last_hidden_feature, layout.attention_mask_feature, + layout.position_ids_feature, ) if feature is not None ), diff --git a/specforge/algorithms/common/vlm_input.py b/specforge/algorithms/common/vlm_input.py new file mode 100644 index 000000000..b81c32593 --- /dev/null +++ b/specforge/algorithms/common/vlm_input.py @@ -0,0 +1,101 @@ +# coding=utf-8 +# Copyright 2024 The SpecForge team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""Multimodal (image+text) ServerInputAdapter for server-side capture. + +Owns the three modality seams for ``modality="multimodal"``: + +- ``load_input_tools``: the target tokenizer (the processor is loaded lazily + in ``prepare_prompts`` so training-model construction receives a plain + tokenizer, exactly like the text path). +- ``prepare_prompts``: ShareGPT-style JSONL (+ optional ``image`` field) -> + payload dicts with expanded ``input_ids``/``loss_mask`` (what the trainer + and the passthrough capture use) plus ``request_input_ids`` (single + placeholder) and base64 ``image_data`` (what the capture request sends). +- ``build_request_inputs``: batch payloads -> the ``/generate`` model-input + fields ``{"input_ids", "image_data"}``. + +The capture server must run with ``SGLANG_MM_AVOID_RETOKENIZE=1`` so its +multimodal processor re-expands placeholders in id space (no retokenization +drift); the managed launcher sets this for ``input_modality="multimodal"``. +""" + +from __future__ import annotations + +from typing import Any, Mapping, Sequence + + +class VlmServerInputAdapter: + """Image+text input adapter for the SGLang server-capture transport.""" + + def __init__(self, config: Any) -> None: + self._config = config + + def load_input_tools(self, config: Any) -> Any: + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained( + config.model.target_model_path, + cache_dir=config.model.cache_dir, + trust_remote_code=config.model.trust_remote_code, + ) + if config.model.tokenizer_pad_token_id is not None: + tokenizer.pad_token_id = config.model.tokenizer_pad_token_id + elif tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None: + tokenizer.pad_token_id = tokenizer.eos_token_id + return tokenizer + + def prepare_prompts( + self, + config: Any, + input_tools: Any, + *, + draft_config: Any, + ) -> list[dict[str, Any]]: + from transformers import AutoProcessor + + from specforge.algorithms.model_providers import dflash_min_loss_tokens + from specforge.data.vlm_preprocessing import build_vlm_prompt_payloads + + tokenizer = input_tools + # The processor must inherit the target's own preprocessor config + # (min/max pixels, merge size) so the client expansion matches the + # capture server's expansion one-for-one; no overrides are accepted. + processor = AutoProcessor.from_pretrained( + config.model.target_model_path, + cache_dir=config.model.cache_dir, + trust_remote_code=config.model.trust_remote_code, + ) + path = config.data.prompts_path or config.data.train_data_path + if not path: + raise ValueError("multimodal prompt preparation requires a data path") + return build_vlm_prompt_payloads( + path, + tokenizer, + processor, + chat_template=config.data.chat_template, + max_length=config.data.max_length, + min_loss_tokens=dflash_min_loss_tokens(config, draft_config), + max_prompts=config.data.max_prompts, + ) + + def build_request_inputs( + self, + tasks: Sequence[Any], + ) -> Mapping[str, Any]: + return { + "input_ids": [list(task.payload["request_input_ids"]) for task in tasks], + "image_data": [task.payload["image_data"] for task in tasks], + } + + +def build_vlm_input_adapter(config: Any) -> VlmServerInputAdapter: + return VlmServerInputAdapter(config) + + +__all__ = ["VlmServerInputAdapter", "build_vlm_input_adapter"] diff --git a/specforge/algorithms/dflash/providers.py b/specforge/algorithms/dflash/providers.py index c5bc5837b..d12901d03 100644 --- a/specforge/algorithms/dflash/providers.py +++ b/specforge/algorithms/dflash/providers.py @@ -13,6 +13,7 @@ build_collator, build_offline_normalizer, build_offline_reader, + build_vlm_collator, ) from specforge.algorithms.common.providers import ( AlgorithmProviders, @@ -26,6 +27,7 @@ TargetDerivedDraftDefaults, make_registration, ) +from specforge.algorithms.common.vlm_input import build_vlm_input_adapter from specforge.algorithms.contracts import ( AlgorithmCapabilities, AlgorithmSpec, @@ -133,6 +135,7 @@ def needs_input_tools(config, draft_model): def algorithm_spec() -> AlgorithmSpec: ready = {"input_ids", "loss_mask", "hidden_states"} + vlm_ready = {"input_ids", "loss_mask", "hidden_states", "position_ids"} return AlgorithmSpec( name=ALGORITHM_NAME, draft=DraftRequirement( @@ -156,6 +159,11 @@ def algorithm_spec() -> AlgorithmSpec: modality="text", required_tensors=ready, ), + FeatureContract( + mode=FeatureMode.STREAMING, + modality="multimodal", + required_tensors=vlm_ready, + ), ), capabilities=AlgorithmCapabilities( attention_backends={"eager", "sdpa", "flex_attention"}, @@ -226,6 +234,22 @@ def algorithm_providers() -> AlgorithmProviders: ), build_collator=collator, ), + ServerStreamingProvider( + modality="multimodal", + capture_method="dflash", + target_representation=None, + layout=ServerCaptureLayout( + aux_feature="hidden_states", + last_hidden_feature=None, + passthrough=( + ("input_ids", "input_ids", ()), + ("loss_mask", "loss_mask", ()), + ), + position_ids_feature="position_ids", + ), + build_collator=build_vlm_collator, + build_input_adapter=build_vlm_input_adapter, + ), ), ) diff --git a/specforge/data/vlm_preprocessing.py b/specforge/data/vlm_preprocessing.py new file mode 100644 index 000000000..dc6ae3121 --- /dev/null +++ b/specforge/data/vlm_preprocessing.py @@ -0,0 +1,287 @@ +# coding=utf-8 +# Copyright 2024 The SpecForge team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""VLM (image+text) prompt preparation for server-side multimodal capture. + +The multimodal online path needs, per sample: + +- ``input_ids`` / ``loss_mask``: the FULLY image-expanded token sequence. + These are what the trainer consumes and what the passthrough capture stores, + so they must match the capture server's own expansion one-for-one. +- ``request_input_ids``: the same sequence with each image region collapsed + back to a single ``<|image_pad|>`` placeholder. This is what the capture + request sends; the patched server re-expands it in id space + (``SGLANG_MM_AVOID_RETOKENIZE=1``), which is guaranteed to reproduce the + client expansion because both sides run the target model's own HF processor + on the same image. +- ``image_data``: a base64 image string for the capture request (the image + itself never enters the feature store; only text tokens and captured hidden + states do). + +Limitations (v1): at most one image per sample; the image is attached to the +first user turn (matching the Qwen-VL chat layout); text-only samples are +supported in the same run (``image_data=None``). +""" + +from __future__ import annotations + +import base64 +import io +import os +from typing import Any, Dict, List, Optional, Tuple + +from .prompt_builder import _iter_records +from .template import TEMPLATE_REGISTRY + +_IMAGE_PLACEHOLDER = "<|image_pad|>" +_VISION_PREFIX = "<|vision_start|><|image_pad|><|vision_end|>" + + +def _image_token_count(image_grid_thw, merge_size: int) -> int: + """Merged token count of one image region from its processor grid.""" + + grid = list(image_grid_thw[0]) + count = 1 + for dim in grid: + count *= int(dim) + return count // (merge_size * merge_size) + + +def _load_image(image_field: Any, *, source: str): + """Return (pil_image, base64_str) from a path / base64 / data-URI field.""" + + from PIL import Image + + if not isinstance(image_field, str) or not image_field: + raise ValueError( + f"{source}: image field must be a file path or base64 string, got " + f"{type(image_field).__name__}" + ) + if os.path.isfile(image_field): + with open(image_field, "rb") as image_file: + raw = image_file.read() + return Image.open(io.BytesIO(raw)).convert("RGB"), base64.b64encode(raw).decode( + "ascii" + ) + encoded = image_field + if encoded.startswith("data:"): + encoded = encoded.split(",", 1)[-1] + try: + raw = base64.b64decode(encoded, validate=True) + except Exception as exc: + raise ValueError( + f"{source}: image field is neither an existing file nor valid base64" + ) from exc + return Image.open(io.BytesIO(raw)).convert("RGB"), encoded + + +def _render_conversation_text( + tokenizer, + conversations: List[Dict[str, Any]], + *, + with_image: bool, + source: str, +) -> str: + """Render one ShareGPT conversation with the target tokenizer's template. + + The image (if any) becomes an OpenAI-style list content on the first user + turn so the model's own chat template emits the vision markers exactly as + at serving time. + """ + + messages: List[Dict[str, Any]] = [] + image_attached = False + for message in conversations: + role = message.get("role", message.get("from", "")) + content = message.get("content") + if content is None: + content = message.get("value", "") + if role in ("human", "user"): + role = "user" + elif role in ("gpt", "assistant"): + role = "assistant" + if isinstance(content, list): + # OpenAI-style list content: keep text parts; image parts are + # re-attached by this function (single-image v1). + content = " ".join( + part.get("text", "") + for part in content + if isinstance(part, dict) and part.get("type") == "text" + ) + if not isinstance(content, str): + raise ValueError(f"{source}: unsupported message content type") + if with_image and role == "user" and not image_attached: + content = [ + {"type": "image"}, + {"type": "text", "text": content}, + ] + image_attached = True + messages.append({"role": role, "content": content}) + if with_image and not image_attached: + raise ValueError(f"{source}: image sample has no user turn to attach to") + return tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=False, + ) + + +def _expand_image_region( + ids: List[int], + mask: List[int], + *, + pad_token_id: int, + count: int, + source: str, +) -> Tuple[List[int], List[int]]: + """Splice the single image placeholder into ``count`` copies in id space.""" + + positions = [index for index, token in enumerate(ids) if token == pad_token_id] + if len(positions) != 1: + raise ValueError( + f"{source}: expected exactly 1 image placeholder token, found " + f"{len(positions)}" + ) + index = positions[0] + expanded_ids = ids[:index] + [pad_token_id] * count + ids[index + 1 :] + expanded_mask = mask[:index] + [0] * count + mask[index + 1 :] + return expanded_ids, expanded_mask + + +def build_vlm_prompt_payloads( + path: str, + tokenizer, + processor, + *, + chat_template: Optional[str], + max_length: int, + min_loss_tokens: int = 1, + max_prompts: Optional[int] = None, +) -> List[Dict[str, Any]]: + """Build prompt payloads for multimodal server capture from a JSONL file. + + Records follow the ShareGPT shape (``conversations`` column); an optional + top-level ``image`` field carries one image (file path or base64). Returns + control-plane dicts ``{"payload": {...}}`` mirroring + :func:`specforge.data.prompt_builder.prepare_prompt_tasks`. + """ + + template = TEMPLATE_REGISTRY.get(chat_template) if chat_template else None + pad_token_id = tokenizer.convert_tokens_to_ids(_IMAGE_PLACEHOLDER) + if not isinstance(pad_token_id, int): + raise ValueError( + f"tokenizer has no {_IMAGE_PLACEHOLDER!r} special token; is this a " + "Qwen-VL-family tokenizer?" + ) + merge_size = int(getattr(processor.image_processor, "merge_size", 2)) + + payloads: List[Dict[str, Any]] = [] + for line_number, record in _iter_records(path): + source = f"{path}:{line_number}" + try: + prepared = _prepare_one_record( + record, + source=source, + tokenizer=tokenizer, + processor=processor, + template=template, + pad_token_id=pad_token_id, + merge_size=merge_size, + max_length=max_length, + min_loss_tokens=min_loss_tokens, + ) + except ValueError as exc: + print(f"WARNING: skipping {source}: {exc}") + continue + if prepared is None: + continue + payloads.append(prepared) + if max_prompts not in (None, 0) and len(payloads) >= max_prompts: + break + return payloads + + +def _prepare_one_record( + record: Dict[str, Any], + *, + source: str, + tokenizer, + processor, + template, + pad_token_id: int, + merge_size: int, + max_length: int, + min_loss_tokens: int, +) -> Optional[Dict[str, Any]]: + from .preprocessing import preprocess_conversations + + conversations = record.get("conversations") + if not conversations: + raise ValueError("record has no 'conversations' field") + image_field = record.get("image") or record.get("image_path") + + text = _render_conversation_text( + tokenizer, + conversations, + with_image=bool(image_field), + source=source, + ) + # Reuse the text stack: preformatted rendering -> ids + loss mask. + parsed = preprocess_conversations( + tokenizer, + [text], + template, + max_length=max_length, + is_preformatted=True, + ) + if not parsed["input_ids"]: + return None + collapsed = parsed["input_ids"][0][0].tolist() + mask = parsed["loss_mask"][0][0].tolist() + + image_data = None + if image_field: + pil_image, image_data = _load_image(image_field, source=source) + processor_output = processor.image_processor( + images=[pil_image], return_tensors="pt" + ) + count = _image_token_count( + processor_output["image_grid_thw"].tolist(), merge_size + ) + ids, mask = _expand_image_region( + collapsed, + mask, + pad_token_id=pad_token_id, + count=count, + source=source, + ) + if len(ids) > max_length: + # Truncating into an image region would desynchronize the client's + # expanded ids from the server's own expansion; drop the sample. + raise ValueError( + f"expanded sequence length {len(ids)} exceeds max_length " + f"{max_length} (image region cannot be truncated)" + ) + else: + ids = collapsed + + ids = ids[:max_length] + mask = mask[:max_length] + if sum(mask) < min_loss_tokens: + return None + return { + "payload": { + "input_ids": ids, + "loss_mask": mask, + "request_input_ids": collapsed, + "image_data": image_data, + } + } + + +__all__ = ["build_vlm_prompt_payloads"] diff --git a/specforge/inference/adapters/server_capture.py b/specforge/inference/adapters/server_capture.py index 6647c627f..2633436a1 100644 --- a/specforge/inference/adapters/server_capture.py +++ b/specforge/inference/adapters/server_capture.py @@ -48,12 +48,16 @@ class ServerCaptureSchema: ``(feature_name, payload_key, trailing_shape)`` for client tensors stored verbatim (``trailing_shape`` is appended after ``(1, L)``). ``attention_mask_feature`` is synthesized all-ones (PromptTasks are unpadded). + ``position_ids_feature`` names the server-produced position-id artifact + (mRoPE positions for multimodal targets, stored ``(1, L, 3)`` int64); + None = not requested. """ aux_feature: Optional[str] last_hidden_feature: Optional[str] passthrough: Tuple[Tuple[str, str, Tuple[int, ...]], ...] attention_mask_feature: Optional[str] = None + position_ids_feature: Optional[str] = None @dataclass(frozen=True) @@ -208,6 +212,8 @@ def _spec_capture_payload(self, task: PromptTask) -> Dict[str, Any]: features["aux"] = self.schema.aux_feature if self.schema.last_hidden_feature is not None: features["last_hidden"] = self.schema.last_hidden_feature + if self.schema.position_ids_feature is not None: + features["position_ids"] = self.schema.position_ids_feature passthrough: List[Dict[str, Any]] = [] for feature_name, payload_key, trailing in self.schema.passthrough: if payload_key == "input_ids": diff --git a/specforge/inference/media.py b/specforge/inference/media.py index 1d2b289dc..f5baf5f35 100644 --- a/specforge/inference/media.py +++ b/specforge/inference/media.py @@ -1,8 +1,11 @@ # coding=utf-8 """Legacy multimodal input types retained for import compatibility. -The current runtime does not support VLM training or media rollout inputs. -These types are not wired into the canonical training path. +Multimodal training now runs through the server-capture path +(``specforge/algorithms/common/vlm_input.py`` with +``model.input_modality="multimodal"``), which transports images inside the +capture request itself. These pixel-tensor types remain unused by the +canonical training path. """ from __future__ import annotations diff --git a/specforge/launch_plan.py b/specforge/launch_plan.py index 5633c3af9..357ed926f 100644 --- a/specforge/launch_plan.py +++ b/specforge/launch_plan.py @@ -536,6 +536,10 @@ def _managed_local_services( "MOONCAKE_GLOBAL_SEGMENT_SIZE": str(mooncake.global_segment_size_bytes), "MOONCAKE_LOCAL_BUFFER_SIZE": str(mooncake.local_buffer_size_bytes), } + if cfg.model.input_modality == "multimodal": + # Multimodal capture requests carry single-placeholder input_ids; + # the server re-expands them in id space (no retokenization drift). + service_env["SGLANG_MM_AVOID_RETOKENIZE"] = "1" capture_services.append( ServiceSpec( command=CommandSpec( diff --git a/specforge/training/disaggregated.py b/specforge/training/disaggregated.py index 47d175e80..2f31fb7f2 100644 --- a/specforge/training/disaggregated.py +++ b/specforge/training/disaggregated.py @@ -639,6 +639,7 @@ def _build_online( last_hidden_feature=layout.last_hidden_feature, passthrough=layout.passthrough, attention_mask_feature=layout.attention_mask_feature, + position_ids_feature=layout.position_ids_feature, ) adapters = [ SGLangServerCaptureAdapter( diff --git a/specforge/training/strategies/base.py b/specforge/training/strategies/base.py index 69495a295..d7a748118 100644 --- a/specforge/training/strategies/base.py +++ b/specforge/training/strategies/base.py @@ -474,11 +474,15 @@ def forward_loss( t = batch.tensors device = self._device() max_valid_anchors = _cpu_max_valid_anchors(t["loss_mask"]) + # Multimodal capture additionally stores server-produced mRoPE + # position ids (B, S, 3); text runs do not carry the tensor at all. + position_ids = t.get("position_ids") loss, accuracy, model_metrics = self.dflash_model( input_ids=t["input_ids"].to(device, non_blocking=True), hidden_states=t["hidden_states"].to(device, non_blocking=True), loss_mask=t["loss_mask"].to(device, non_blocking=True), max_valid_anchors=max_valid_anchors, + position_ids=None if position_ids is None else position_ids.to(device), ) metrics = {"accuracy": accuracy.detach()} if "accuracy_denom" in model_metrics: diff --git a/tests/test_algorithms/test_builtin_providers.py b/tests/test_algorithms/test_builtin_providers.py index 11a99b5fa..59b8750a2 100644 --- a/tests/test_algorithms/test_builtin_providers.py +++ b/tests/test_algorithms/test_builtin_providers.py @@ -155,12 +155,19 @@ def test_target_derived_defaults_and_overrides_are_provider_owned(self): self.assertIsNone(policy.target_defaults) self.assertIsNone(policy.apply_overrides) - def test_vlm_is_not_registered_as_a_builtin(self): + def test_builtin_modalities_are_pinned_per_algorithm(self): + expected = { + "dflash": {"text", "multimodal"}, + "domino": {"text"}, + "dspark": {"text"}, + "eagle3": {"text"}, + "peagle": {"text"}, + } for registration in self.registry: modalities = { contract.modality for contract in registration.spec.feature_contracts } - self.assertEqual({"text"}, modalities, registration.name) + self.assertEqual(expected[registration.name], modalities, registration.name) def test_builtin_resume_contracts_cover_resolved_objective_semantics(self): training = SimpleNamespace( diff --git a/tests/test_algorithms/test_dflash_multimodal.py b/tests/test_algorithms/test_dflash_multimodal.py new file mode 100644 index 000000000..813bbffbb --- /dev/null +++ b/tests/test_algorithms/test_dflash_multimodal.py @@ -0,0 +1,213 @@ +from __future__ import annotations + +import unittest +from types import SimpleNamespace + +from specforge.algorithms.builtin import builtin_algorithm_registry +from specforge.algorithms.common.vlm_input import VlmServerInputAdapter +from specforge.algorithms.contracts import FeatureMode +from specforge.data.vlm_preprocessing import _expand_image_region, _image_token_count + +try: + import torch + + TORCH_AVAILABLE = True +except ImportError: # pragma: no cover - torch-free dev boxes + torch = None + TORCH_AVAILABLE = False + + +class MultimodalRegistrationTest(unittest.TestCase): + def setUp(self): + self.registry = builtin_algorithm_registry() + + def test_dflash_registers_the_multimodal_streaming_contract(self): + registration = self.registry.resolve("dflash") + contract = registration.spec.feature_contract( + FeatureMode.STREAMING, "multimodal" + ) + self.assertEqual( + set(contract.required_tensors), + {"input_ids", "loss_mask", "hidden_states", "position_ids"}, + ) + provider = registration.providers.server_streaming_for("multimodal") + self.assertEqual(provider.capture_method, "dflash") + self.assertEqual(provider.layout.position_ids_feature, "position_ids") + self.assertEqual(provider.layout.aux_feature, "hidden_states") + + def test_other_builtins_have_no_multimodal_contract(self): + for name in ("domino", "dspark", "eagle3", "peagle"): + with self.subTest(algorithm=name): + registration = self.registry.resolve(name) + with self.assertRaises(KeyError): + registration.spec.feature_contract( + FeatureMode.STREAMING, "multimodal" + ) + + def test_input_adapter_factory_builds_a_valid_adapter(self): + registration = self.registry.resolve("dflash") + provider = registration.providers.server_streaming_for("multimodal") + adapter = provider.create_input_adapter( + SimpleNamespace(model=SimpleNamespace(), data=SimpleNamespace()) + ) + self.assertIsInstance(adapter, VlmServerInputAdapter) + + +class VlmExpansionMathTest(unittest.TestCase): + def test_image_token_count_uses_merge_size(self): + self.assertEqual(_image_token_count([[2, 4, 6]], merge_size=2), 12) + self.assertEqual(_image_token_count([[1, 2, 2]], merge_size=1), 4) + + def test_expand_image_region_splices_ids_and_zero_mask(self): + ids, mask = _expand_image_region( + [10, 99, 20], + [0, 0, 1], + pad_token_id=99, + count=4, + source="test", + ) + self.assertEqual(ids, [10, 99, 99, 99, 99, 20]) + self.assertEqual(mask, [0, 0, 0, 0, 0, 1]) + + def test_expand_image_region_requires_exactly_one_placeholder(self): + with self.assertRaises(ValueError): + _expand_image_region([10, 20], [1, 1], pad_token_id=99, count=4, source="t") + with self.assertRaises(ValueError): + _expand_image_region( + [99, 10, 99], [0, 0, 0], pad_token_id=99, count=4, source="t" + ) + + +class VlmRequestInputsTest(unittest.TestCase): + def test_build_request_inputs_uses_collapsed_ids_and_image_data(self): + adapter = VlmServerInputAdapter(config=SimpleNamespace()) + tasks = [ + SimpleNamespace( + payload={ + "input_ids": [1, 2, 2, 2, 3], + "request_input_ids": [1, 2, 3], + "image_data": "aGVsbG8=", + } + ), + SimpleNamespace( + payload={ + "input_ids": [7, 8], + "request_input_ids": [7, 8], + "image_data": None, + } + ), + ] + request = adapter.build_request_inputs(tasks) + self.assertEqual(request["input_ids"], [[1, 2, 3], [7, 8]]) + self.assertEqual(request["image_data"], ["aGVsbG8=", None]) + + +class ServerCapturePositionIdsTest(unittest.TestCase): + def _adapter(self, position_ids_feature): + from specforge.inference.adapters.server_capture import ( + ServerCaptureSchema, + SGLangServerCaptureAdapter, + ) + + class _FakeStore: + store_id = "store" + + def adopt(self, ref): + pass + + def discard_external_attempts(self, *args, **kwargs): + pass + + def track_external_attempt(self, *args, **kwargs): + pass + + schema = ServerCaptureSchema( + aux_feature="hidden_states", + last_hidden_feature=None, + passthrough=( + ("input_ids", "input_ids", ()), + ("loss_mask", "loss_mask", ()), + ), + position_ids_feature=position_ids_feature, + ) + return SGLangServerCaptureAdapter( + "http://localhost:1", + _FakeStore(), + run_id="run", + algorithm="dflash", + schema=schema, + ) + + def _task(self): + return SimpleNamespace( + task_id="t0", + attempt=0, + payload={"input_ids": [5, 6, 7], "loss_mask": [0, 1, 1]}, + metadata={}, + ) + + def test_payload_requests_position_ids_artifact_when_configured(self): + adapter = self._adapter("position_ids") + payload = adapter._spec_capture_payload(self._task()) + self.assertEqual( + payload["features"], + {"aux": "hidden_states", "position_ids": "position_ids"}, + ) + + def test_payload_omits_position_ids_artifact_when_unset(self): + adapter = self._adapter(None) + payload = adapter._spec_capture_payload(self._task()) + self.assertEqual(payload["features"], {"aux": "hidden_states"}) + + +@unittest.skipUnless(TORCH_AVAILABLE, "requires torch") +class VlmCollatorTest(unittest.TestCase): + def test_collator_pads_position_ids_like_other_features(self): + from specforge.algorithms.common.dflash_family_data import build_vlm_collator + + collate = build_vlm_collator() + features = [ + { + "input_ids": torch.tensor([[1, 2, 3]]), + "loss_mask": torch.tensor([[0, 1, 1]]), + "hidden_states": torch.zeros(1, 3, 8), + "position_ids": torch.arange(9).reshape(1, 3, 3), + }, + { + "input_ids": torch.tensor([[4]]), + "loss_mask": torch.tensor([[1]]), + "hidden_states": torch.zeros(1, 1, 8), + "position_ids": torch.arange(3).reshape(1, 1, 3), + }, + ] + batch = collate(features) + self.assertEqual(tuple(batch["input_ids"].shape), (2, 3)) + self.assertEqual(tuple(batch["position_ids"].shape), (2, 3, 3)) + # Padding is zeros on the sequence axis. + self.assertTrue((batch["position_ids"][1, 1:] == 0).all()) + self.assertEqual(batch["position_ids"][0, 2].tolist(), [6, 7, 8]) + + +@unittest.skipUnless(TORCH_AVAILABLE, "requires torch") +class MropeDraftPositionsTest(unittest.TestCase): + def test_gathered_draft_positions_follow_anchor_offsets(self): + import torch as t + + from specforge.algorithms.common.dflash_family_model import OnlineDFlashModel + + model = OnlineDFlashModel.__new__(OnlineDFlashModel) + model.block_size = 2 + anchors = t.tensor([[1, 3]]) + stored = t.arange(5 * 3).reshape(1, 5, 3) + offsets = t.arange(model.block_size).view(1, 1, -1) + draft_indices = (anchors.unsqueeze(-1) + offsets).view(1, -1) + gathered = t.gather(stored, 1, draft_indices.unsqueeze(-1).expand(-1, -1, 3)) + full = t.cat([stored, gathered], dim=1).permute(2, 0, 1) + self.assertEqual(tuple(full.shape), (3, 1, 5 + 4)) + # Draft slot for anchor=1: positions of indices 1 and 2. + self.assertEqual(full[:, 0, 5].tolist(), [3, 4, 5]) + self.assertEqual(full[:, 0, 6].tolist(), [6, 7, 8]) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_config/test_server_only_online.py b/tests/test_config/test_server_only_online.py index c12c5106b..9be95463e 100644 --- a/tests/test_config/test_server_only_online.py +++ b/tests/test_config/test_server_only_online.py @@ -66,6 +66,24 @@ def test_vlm_is_explicitly_unsupported(self): ): resolve_run(config) + def test_multimodal_modality_resolves_for_dflash(self): + payload = _online_payload(modality="multimodal") + payload["training"]["strategy"] = "dflash" + del payload["model"]["vocab_mapping_path"] + resolved = resolve_run(Config.model_validate(payload)) + self.assertEqual(resolved.algorithm.name, "dflash") + streaming = resolved.algorithm.providers.server_streaming_for("multimodal") + self.assertEqual(streaming.layout.position_ids_feature, "position_ids") + self.assertIsNotNone(streaming.create_input_adapter(resolved.config)) + + def test_multimodal_modality_is_rejected_for_text_only_algorithms(self): + payload = _online_payload(modality="multimodal") + with self.assertRaisesRegex( + ValueError, + "no streaming feature contract and provider for modality " "'multimodal'", + ): + resolve_run(Config.model_validate(payload)) + def test_offline_configs_reject_retired_backends_instead_of_ignoring_them(self): # Offline consumers never instantiate a target inference backend, but a # config naming a retired backend must fail at load rather than be From 6f9331910e8233568ac2c0fd72ec74f592a6a7df Mon Sep 17 00:00:00 2001 From: Curnane Date: Thu, 27 Aug 2026 15:04:23 +0800 Subject: [PATCH 4/4] feat(dflash): add multimodal training and switch to plain-rope drafts Co-authored-by: danaodai --- configs/qwen3.5-4b-vl-dflash.json | 50 -- docs/advanced_features/vlm_dflash.md | 88 +- examples/configs/README.md | 16 +- .../qwen3.5-4b-vl-dflash-disaggregated.yaml | 4 +- patches/sglang/v0.5.14/spec-capture.patch | 765 ++++++++++++++---- scripts/regenerate_train_data.py | 45 ++ .../algorithms/common/dflash_family_model.py | 40 +- .../algorithms/common/hidden_states_data.py | 24 - specforge/algorithms/common/providers.py | 10 +- specforge/algorithms/common/vlm_input.py | 2 +- specforge/algorithms/dflash/providers.py | 7 +- specforge/config/schema.py | 28 +- specforge/data/vlm_preprocessing.py | 96 ++- .../inference/adapters/server_capture.py | 6 - specforge/launch_plan.py | 11 + specforge/modeling/draft/dflash.py | 66 +- specforge/training/disaggregated.py | 1 - specforge/training/strategies/base.py | 4 - .../test_algorithms/test_builtin_providers.py | 1 + .../test_algorithms/test_dflash_multimodal.py | 213 ++--- tests/test_config/test_schema.py | 42 + tests/test_config/test_server_only_online.py | 2 +- 22 files changed, 981 insertions(+), 540 deletions(-) delete mode 100644 configs/qwen3.5-4b-vl-dflash.json rename examples/configs/{ => online/disaggregated/external}/qwen3.5-4b-vl-dflash-disaggregated.yaml (91%) diff --git a/configs/qwen3.5-4b-vl-dflash.json b/configs/qwen3.5-4b-vl-dflash.json deleted file mode 100644 index 366df304b..000000000 --- a/configs/qwen3.5-4b-vl-dflash.json +++ /dev/null @@ -1,50 +0,0 @@ -{ - "architectures": [ - "DFlashDraftModel" - ], - "attention_bias": false, - "attention_dropout": 0.0, - "auto_map": { - "AutoModel": "dflash.DFlashDraftModel" - }, - "block_size": 16, - "bos_token_id": 248043, - "dflash_config": { - "mask_token_id": 248070, - "target_layer_ids": [1, 8, 15, 22, 29] - }, - "dtype": "bfloat16", - "eos_token_id": 248044, - "head_dim": 256, - "hidden_act": "silu", - "hidden_size": 2560, - "initializer_range": 0.02, - "intermediate_size": 9216, - "layer_types": [ - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention" - ], - "max_position_embeddings": 262144, - "model_type": "qwen3_vl_text", - "num_attention_heads": 16, - "num_hidden_layers": 5, - "num_key_value_heads": 4, - "num_target_layers": 32, - "pad_token_id": 248044, - "partial_rotary_factor": 0.25, - "rms_norm_eps": 1e-06, - "rope_scaling": { - "mrope_interleaved": true, - "mrope_section": [11, 11, 10], - "rope_type": "default", - "partial_rotary_factor": 0.25 - }, - "rope_theta": 10000000, - "sliding_window": null, - "tie_word_embeddings": true, - "use_cache": true, - "vocab_size": 248320 -} diff --git a/docs/advanced_features/vlm_dflash.md b/docs/advanced_features/vlm_dflash.md index e6ef712b3..23aa530e9 100644 --- a/docs/advanced_features/vlm_dflash.md +++ b/docs/advanced_features/vlm_dflash.md @@ -1,25 +1,27 @@ -# VLM DFlash Support — Status (add_vl_support) +# VLM DFlash Support - Status (add_vl_support) This branch ports sgl-project/SpecForge PR #585 (commit `9323a510`, author zyk42) onto the server-only unified runtime **and implements end-to-end multimodal (image+text) DFlash training on top of it**. This document records what landed, what was deliberately not ported, and the validation status. +> **Rope decision**: multimodal DFlash drafts use the plain 1D rope convention - +> the same as the text-only path - on purpose. Visual information reaches the +> draft exclusively through the captured target hidden states, never through the +> draft's own positional embedding, so the draft has no use for the target's +> (3, L) mRoPE positions. Staying on one position convention keeps training and +> serving byte-identical to the text pipeline and reuses the stock +> `configs/qwen3.5-4b-dflash.json` draft geometry. + ## What this branch contains ### Foundation (ported from PR #585) -- **Draft model (VLM-capable)** — `specforge/modeling/draft/dflash.py`: +- **Draft model (VLM-capable)** - `specforge/modeling/draft/dflash.py`: partial rotation in `apply_rotary_pos_emb` (`rotary_dim < head_dim`, for - Qwen3.5/Qwen3.6 `partial_rotary_factor=0.25`) and - `Qwen3InterleavedMultiRotaryEmbedding` selected by - `rope_scaling.mrope_interleaved`. -- **Draft config** — `configs/qwen3.5-4b-vl-dflash.json` (new, mirrors the - Qwen3.5-4B target geometry: head_dim 256, mrope_section [11,11,10], - partial_rotary_factor 0.25). The wider #585 config set (Qwen3-VL-8B/30B-A3B, - Qwen3.5-9B/35B-A3B) is intentionally out of scope for this branch; it can be - added verbatim in a follow-up once more targets are validated. -- **Weight-key resolution** — `resolve_target_weight_keys()` in + Qwen3.5/Qwen3.6 `partial_rotary_factor=0.25`). The draft always uses the + stock `Qwen3RotaryEmbedding`. +- **Weight-key resolution** - `resolve_target_weight_keys()` in `specforge/modeling/target/target_utils.py` auto-selects `model.language_model.embed_tokens.weight` for VLM targets; `populate_dflash_generated_config` reads language-model depth via @@ -27,46 +29,49 @@ what was deliberately not ported, and the validation status. ### Multimodal capture (new in this branch) -End-to-end data flow: **JSONL (+ image) → expanded ids/loss mask → capture -request (single-placeholder ids + base64 image) → patched SGLang server -expands, runs the ViT, captures aux hidden states + mRoPE positions → -Mooncake → collator → training forward with 3D position ids**. +End-to-end data flow: **JSONL (+ image) -> expanded ids/loss mask -> capture +request (single-placeholder ids + base64 image) -> patched SGLang server +expands, runs the ViT, captures aux hidden states -> Mooncake -> collator -> +training forward on plain 1D positions**. - `model.input_modality: multimodal` (DFlash only): a `FeatureContract` - (`{input_ids, loss_mask, hidden_states, position_ids}`) and a + (`{input_ids, loss_mask, hidden_states}`) and a `ServerStreamingProvider` with a VLM `ServerInputAdapter` - (`specforge/algorithms/common/vlm_input.py`). + (`specforge/algorithms/common/vlm_input.py`). Multimodal capture stores the + same three tensors as text capture - no `position_ids` artifact is requested + or consumed. - `specforge/data/vlm_preprocessing.py`: ShareGPT-style records with an optional `image` field (path or base64); the target's own chat template and HF processor produce the expanded `input_ids`/`loss_mask` (image region expanded in id space, mask zeros). One image per sample max (v1); text-only samples work in the same run. -- `ServerCaptureLayout.position_ids_feature` → the capture request's - `features["position_ids"]`; the patched server writes the request's mRoPE - positions `(1, L, 3) int64` into Mooncake (`_spec_capture_position_ids` in - the scheduler sink; text requests get the arange broadcast fallback). -- `patches/sglang/v0.5.14/spec-capture.patch`: regenerated with the - `position_ids` artifact (`SpecCaptureSink.put_sample(position_ids=...)`). - Multimodal capture requests ride the stock `input_ids` + `image_data` - `/generate` path with `SGLANG_MM_AVOID_RETOKENIZE=1` (set by the managed - launcher for `input_modality=multimodal`), so the server re-expands - placeholders in id space with zero retokenization drift — and the - passthrough/seq-len checks fail loudly if client and server expansions ever - disagree. -- Training: `OnlineDFlashModel.forward(..., position_ids=None)` gathers 3D - mRoPE positions for context + anchor-offset draft slots - (`(3, B, S + N·bs)`); `DFlashTrainStrategy` passes the collated - `position_ids` tensor through. Text runs are byte-identical to before. -- Recipe: `examples/configs/qwen3.5-4b-vl-dflash-disaggregated.yaml` - (single-node Ascend NPU managed stack). +- `patches/sglang/v0.5.14/spec-capture.patch`: tracks upstream's rewritten + async streaming sink. Multimodal capture requests ride the stock + `input_ids` + `image_data` `/generate` path with + `SGLANG_MM_AVOID_RETOKENIZE=1` (set by the managed launcher for + `input_modality=multimodal`), so the server re-expands placeholders in id + space with zero retokenization drift - and the passthrough/seq-len checks + fail loudly if client and server expansions ever disagree. +- Training: `OnlineDFlashModel._forward_draft_blocks` builds positions with + the unconditional text-path 1D `arange` convention; multimodal batches flow + through the identical forward as text batches. Text runs are + byte-identical to before. +- Recipe: `examples/configs/online/disaggregated/external/qwen3.5-4b-vl-dflash-disaggregated.yaml` + (single-node Ascend NPU managed stack; draft config + `configs/qwen3.5-4b-dflash.json`). ## Not ported (by design) - HF-backend VLM capture (`dflash_target_model.py`, `_build_vlm_reqs`, `mm_token_type_ids`) and the `train_dflash.py --is-vlm` plumbing from the - pre-#678 script stack — superseded by server capture. -- `QwenVLOnlineDFlashModel` wiring — PR #585 referenced this class but never + pre-#678 script stack - superseded by server capture. +- `QwenVLOnlineDFlashModel` wiring - PR #585 referenced this class but never defined it; the unified runtime needs no separate VLM wrapper class. +- mRoPE draft support (`Qwen3InterleavedMultiRotaryEmbedding`, the + `rope_scaling.mrope_interleaved` switch, and the server `position_ids` + capture artifact): the draft consumes visual information only through the + captured target hidden states, so the (3, L) target positions carry no + signal for it. Retired in favor of the single plain-rope convention above. - Two accidental reverts in the original #585 diff (domino projector code, D-PACE CLI args) were dropped during the cherry-pick. - Offline (precomputed hidden states) multimodal capture: the offline path @@ -77,14 +82,13 @@ Mooncake → collator → training forward with 3D position ids**. - **Verified (CPU, this repo)**: registration parity and provider gates, request/payload construction, expansion math, collator, golden - topology/recipe tests — `tests/test_algorithms/test_dflash_multimodal.py` - plus the updated `test_config` suites (80 passed locally, 2 torch tests - skipped pending GPU). + topology/recipe tests - `tests/test_algorithms/test_dflash_multimodal.py` + plus the updated `test_config` suites. - **Verified statically**: the regenerated patch applies cleanly both ways to pristine sglang v0.5.14 (`git apply --check` / `--reverse --check`). - **Not yet verified (needs GPU/NPU)**: a live multimodal capture run - (Qwen3-VL / Qwen3.5 target + ViT + mRoPE positions) and an end-to-end - training run. This is the next step; see the recipe above. + (Qwen3-VL / Qwen3.5 target + ViT) and an end-to-end training run. This is + the next step; see the recipe above. ## Reference results from PR #585 (HF stack, author-validated) diff --git a/examples/configs/README.md b/examples/configs/README.md index 2a092a3c7..08969e0a3 100644 --- a/examples/configs/README.md +++ b/examples/configs/README.md @@ -69,12 +69,12 @@ target hidden states are concatenated and projected. It remains registered as the `eagle3` strategy; EAGLE3.1 is a draft-model configuration variant, not a second runtime or launch path. -Recipes under `online/disaggregated/managed-local` are opt-in, single-node - Multimodal (image+text) training is supported for `training.strategy=dflash` via `model.input_modality: multimodal` (see -`qwen3.5-4b-vl-dflash-disaggregated.yaml`); the remaining catalog is -text-only. +`online/disaggregated/external/qwen3.5-4b-vl-dflash-disaggregated.yaml`); the +remaining catalog is text-only. + +Recipes under `online/disaggregated/managed-local` are opt-in, single-node full-stack examples. Their typed `managed_local` blocks own Mooncake, one or more patched SGLang capture servers, and the trainer GPU allocation; the same `specforge train -c ...` command starts and cleans up the complete stack. @@ -198,6 +198,7 @@ should make their training strategy and topology explicit. | `model.mask_token_id` | `null` | DFlash-family/P-EAGLE mask token override. Otherwise it resolves from the draft config and then the tokenizer. | | `model.tokenizer_pad_token_id` | `null` | Explicit non-negative tokenizer pad ID. Use it for released tokenizers that omit padding metadata. | | `model.sglang_attention_backend` | `flashinfer` | SGLang attention implementation for an in-process or managed capture server. | +| `model.sglang_mm_attention_backend` | `null` | Vision-encoder attention backend for capture servers. On Ascend NPU with a non-text `model.input_modality` it defaults to `ascend_attn` (fused); `sdpa` materializes N² vision scores and can OOM on large images. | | `model.sglang_mem_fraction_static` | `0.4` | SGLang static-memory fraction in `(0, 1]`; inherited by managed capture servers unless they override it. | | `model.sglang_disable_radix_cache` | `true` | Preserve the historical managed-capture behavior. Set `false` for hybrid targets such as Inkling that require the radix tree. Unique per-attempt cache namespaces still force complete capture prefills. | | `model.sglang_context_length` | `null` | Positive explicit context limit. Managed capture requires at least `data.max_length + 7`; omitting it derives that value. | @@ -460,9 +461,10 @@ unless tuning throughput or memory pressure. requires batch size 1. - `training.compact_teacher` is offline text EAGLE3 only. - Multimodal (image+text) training requires `training.strategy=dflash` with - `model.input_modality: multimodal`, a VLM draft config - (`configs/*-dflash-vlm-*.json`), and the v0.5.14 spec-capture patch with the - `position_ids` artifact. Vendor modalities such as `qwen2_5_vl` remain + `model.input_modality: multimodal`, a plain-rope DFlash draft config (e.g. + `configs/qwen3.5-4b-dflash.json`; the draft needs no mRoPE - visual + information arrives through the captured target hidden states), and the + v0.5.14 spec-capture patch. Vendor modalities such as `qwen2_5_vl` remain unsupported. - Online evaluation is not supported. Offline `data.eval_hidden_states_path` and `training.eval_interval` must be configured together. diff --git a/examples/configs/qwen3.5-4b-vl-dflash-disaggregated.yaml b/examples/configs/online/disaggregated/external/qwen3.5-4b-vl-dflash-disaggregated.yaml similarity index 91% rename from examples/configs/qwen3.5-4b-vl-dflash-disaggregated.yaml rename to examples/configs/online/disaggregated/external/qwen3.5-4b-vl-dflash-disaggregated.yaml index 6b846bfba..099ac219a 100644 --- a/examples/configs/qwen3.5-4b-vl-dflash-disaggregated.yaml +++ b/examples/configs/online/disaggregated/external/qwen3.5-4b-vl-dflash-disaggregated.yaml @@ -1,13 +1,13 @@ # Multimodal (image+text) DFlash on Qwen3.5-4B: disaggregated online training # against an external patched SGLang capture server. Requires sglang v0.5.14 -# with patches/sglang/v0.5.14/spec-capture.patch (position_ids artifact) and +# with patches/sglang/v0.5.14/spec-capture.patch and # SGLANG_MM_AVOID_RETOKENIZE=1 on the capture server. # # Data: ShareGPT-style JSONL with an optional top-level "image" field (file # path or base64) per record; text-only records are supported in the same run. model: target_model_path: Qwen/Qwen3.5-4B - draft_model_config: configs/qwen3.5-4b-vl-dflash.json + draft_model_config: configs/qwen3.5-4b-dflash.json target_backend: sglang trust_remote_code: true input_modality: multimodal diff --git a/patches/sglang/v0.5.14/spec-capture.patch b/patches/sglang/v0.5.14/spec-capture.patch index efe9f4a89..5fd29d1e1 100644 --- a/patches/sglang/v0.5.14/spec-capture.patch +++ b/patches/sglang/v0.5.14/spec-capture.patch @@ -1,5 +1,5 @@ diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py -index a99d252677..f5d0b35e79 100644 +index a99d252..bde8a84 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -94,6 +94,9 @@ class LogitsProcessorOutput: @@ -12,7 +12,7 @@ index a99d252677..f5d0b35e79 100644 ## Part 2: This part will be assigned in python/sglang/srt/layers/sampler.py::Sampler # he log probs of output tokens, if SGLANG_RETURN_ORIGINAL_LOGPROB = True, will get the log probs before applying temperature. If False, will get the log probs before applying temperature. -@@ -361,6 +364,16 @@ class LogitsProcessor(nn.Module): +@@ -361,6 +364,30 @@ class LogitsProcessor(nn.Module): sample_indices, logits_metadata, ) @@ -26,10 +26,24 @@ index a99d252677..f5d0b35e79 100644 + ) + else None + ) ++ # muP targets pass LM-head-scaled hidden states into LogitsProcessor, ++ # while SpecForge folds the same multiplier into its frozen target ++ # head. Restore the pre-head-scale representation before capture so ++ # the multiplier is applied exactly once when training recomputes logits. ++ logits_mup_width_multiplier = getattr( ++ self.config, "logits_mup_width_multiplier", None ++ ) ++ if ( ++ last_hidden_states_to_store is not None ++ and logits_mup_width_multiplier ++ ): ++ last_hidden_states_to_store = ( ++ last_hidden_states_to_store * float(logits_mup_width_multiplier) ++ ) del hidden_states if not logits_metadata.extend_return_logprob: -@@ -374,6 +387,7 @@ class LogitsProcessor(nn.Module): +@@ -374,6 +401,7 @@ class LogitsProcessor(nn.Module): return LogitsProcessorOutput( next_token_logits=sampled_logits, hidden_states=hidden_states_to_store, @@ -38,7 +52,7 @@ index a99d252677..f5d0b35e79 100644 # workaround since ForwardBatch is local to forward_batch_generation(). # They should be moved to GenerationBatchResult to keep this class clean. diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py -index b05334deaa..2853c82315 100644 +index b05334d..2853c82 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/python/sglang/srt/managers/detokenizer_manager.py @@ -441,6 +441,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin): @@ -50,20 +64,20 @@ index b05334deaa..2853c82315 100644 indexer_topk=indexer_topk, customized_info=recv_obj.customized_info, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py -index 951f35495e..4de78534a4 100644 +index 951f354..cd6ebf2 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py -@@ -280,6 +280,10 @@ class GenerateReqInput(BaseReq): - image_max_dynamic_patch: Optional[int] = None - video_max_dynamic_patch: Optional[int] = None - +@@ -188,6 +188,10 @@ class GenerateReqInput(BaseReq): + log_metrics: bool = True + # Whether to return hidden states + return_hidden_states: Union[List[bool], bool] = False + # Spec-training capture sink instructions (see spec_capture_sink.py). + # Batch-level: List[Optional[dict]]; per-request after __getitem__. + spec_capture: Optional[Union[List[Optional[Dict]], Dict]] = None + - # Pre-computed delimiter indices for multi-item scoring. - # Batch-level: List[List[int]] (one per request). After __getitem__: List[int]. - multi_item_delimiter_indices: Optional[Union[List[List[int]], List[int]]] = None + # Whether to return captured routed experts + return_routed_experts: bool = False + return_indexer_topk: bool = False @@ -743,6 +747,11 @@ class GenerateReqInput(BaseReq): if self.multi_item_delimiter_indices is not None else None @@ -107,7 +121,7 @@ index 951f35495e..4de78534a4 100644 token_steps: List[List[int]] = None diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py -index f1dc81d174..acefaace89 100755 +index f1dc81d..acefaac 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -707,6 +707,7 @@ class Req(ReqDllmMixin): @@ -134,7 +148,7 @@ index f1dc81d174..acefaace89 100755 # extra key for classifying the request (e.g. cache_salt) if lora_id is not None: diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py -index abba37441d..3018b82473 100644 +index abba374..241da3f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -576,6 +576,19 @@ class Scheduler( @@ -165,11 +179,103 @@ index abba37441d..3018b82473 100644 ) req.tokenizer = self.tokenizer +@@ -3097,6 +3111,13 @@ class Scheduler( + # GenerationBatchResult.extra_keep_alive_refs after forward returns. + self.batch_record_buf[self.batch_record_ct] = [batch, attr_snapshot] + ++ def _should_copy_hidden_states_to_cpu(self, batch: ScheduleBatch) -> bool: ++ """Avoid redundant multi-GiB capture D2H on non-writer TP ranks.""" ++ return batch.return_hidden_states and ( ++ self.ps.attn_tp_rank == 0 ++ or any(req.spec_capture is None for req in batch.reqs) ++ ) ++ + @contextmanager + def _forward_isolation(self, batch: ScheduleBatch, *, overlap: bool): + """Make SB transactional across one forward (overlap and non-overlap). +@@ -3258,7 +3279,9 @@ class Scheduler( + batch_result.copy_done = self.device_module.Event() + batch_result.copy_to_cpu( + return_logprob=batch.return_logprob, +- return_hidden_states=batch.return_hidden_states, ++ return_hidden_states=self._should_copy_hidden_states_to_cpu( ++ batch ++ ), + ) + else: + kwargs = ( +@@ -3369,6 +3392,10 @@ class Scheduler( + batch: ScheduleBatch, + result: Union[GenerationBatchResult, EmbeddingBatchResult], + ): ++ # Complete capture transfers on the scheduler thread before processing ++ # any kind of next result (prefill, decode, or idle). This preserves ++ # the output socket's single-thread ownership under mixed traffic. ++ self.batch_result_processor.drain_spec_captures() + self.publish_load_snapshot(force=batch.forward_mode.is_extend()) + + if batch.forward_mode.is_decode(): +@@ -3448,6 +3475,15 @@ class Scheduler( + + def on_idle(self): + """Idle housekeeping: guard, check, metrics, reset, sleep.""" ++ # A capture response is intentionally delayed until its background ++ # sink transfer completes; finish it on this thread so ownership of ++ # the ZeroMQ output socket remains single-threaded. Do not enter the ++ # idle sleeper while a future is outstanding or the producer waiting ++ # for that response could deadlock. ++ self.batch_result_processor.drain_spec_captures() ++ if self.batch_result_processor.has_pending_spec_captures(): ++ time.sleep(0.001) ++ return + if not self.is_fully_idle(): + return + +@@ -3496,6 +3532,7 @@ class Scheduler( + and (self.last_batch is None or self.last_batch.is_empty()) + and (self.cur_batch is None or self.cur_batch.is_empty()) + and (not self.enable_overlap or len(self.result_queue) == 0) ++ and not self.batch_result_processor.has_pending_spec_captures() + and self._pp_microbatches_drained() + ) + diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py -index a9d5f0c28b..c9f8fdf7d3 100644 +index a9d5f0c..f134354 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py -@@ -257,6 +257,18 @@ class SchedulerBatchResultProcessor: +@@ -1,7 +1,9 @@ + from __future__ import annotations + + import logging +-from dataclasses import dataclass ++import os ++import time ++from dataclasses import dataclass, field + from typing import ( + TYPE_CHECKING, + Callable, +@@ -78,6 +80,12 @@ class SchedulerBatchResultProcessor: + logprob_result_processor: SchedulerLogprobResultProcessor + output_streamer: SchedulerOutputStreamer + abort_request: Callable ++ _spec_capture_batches: List = field( ++ default_factory=list, ++ init=False, ++ repr=False, ++ compare=False, ++ ) + + def process_batch_result_prebuilt(self, batch: ScheduleBatch): + assert self.disaggregation_mode == DisaggregationMode.DECODE +@@ -181,6 +189,7 @@ class SchedulerBatchResultProcessor: + result: Union[GenerationBatchResult, EmbeddingBatchResult], + ): + skip_stream_req = None ++ pending_spec_captures: List = [] + + if self.is_generation: + if result.copy_done is not None: +@@ -257,6 +266,20 @@ class SchedulerBatchResultProcessor: ) if ( @@ -183,12 +289,43 @@ index a9d5f0c28b..c9f8fdf7d3 100644 + hidden_state_offset=hidden_state_offset, + ) + if req.finished(): -+ self._sink_spec_capture(req) ++ pending = self._sink_spec_capture(req) ++ if pending is not None: ++ pending_spec_captures.append(pending) + elif ( req.return_hidden_states and logits_output.hidden_states is not None ): -@@ -462,6 +474,95 @@ class SchedulerBatchResultProcessor: +@@ -329,9 +352,25 @@ class SchedulerBatchResultProcessor: + req.inflight_middle_chunks -= 1 + req.time_stats.set_last_chunked_prefill_finish_time() + +- self.output_streamer.stream_output( +- batch.reqs, batch.return_logprob, skip_stream_req +- ) ++ if pending_spec_captures: ++ self._queue_spec_captures( ++ pending_spec_captures, ++ return_logprob=batch.return_logprob, ++ ) ++ capture_req_ids = {id(item[0]) for item in pending_spec_captures} ++ ready_reqs = [ ++ req ++ for req in batch.reqs ++ if id(req) not in capture_req_ids and req is not skip_stream_req ++ ] ++ if ready_reqs: ++ self.output_streamer.stream_output( ++ ready_reqs, batch.return_logprob ++ ) ++ else: ++ self.output_streamer.stream_output( ++ batch.reqs, batch.return_logprob, skip_stream_req ++ ) + + can_run_cuda_graph = result.can_run_cuda_graph + self.metrics_reporter.report_prefill_stats( +@@ -462,6 +501,183 @@ class SchedulerBatchResultProcessor: f"Placeholder zeros would be appended to output_ids." ) @@ -207,20 +344,43 @@ index a9d5f0c28b..c9f8fdf7d3 100644 + """ + start = hidden_state_offset + end = start + len(req.origin_input_ids) -+ req.spec_capture_aux.append( -+ logits_output.hidden_states[start:end].cpu().clone() ++ # Only attention-TP rank 0 owns the Mooncake sink. Copying these very ++ # large tensors to host on every TP rank creates identical D2H ++ # transfers whose results are immediately discarded on TP > 1. ++ if self.output_streamer.ps.attn_tp_rank != 0: ++ return end ++ # Materialize each scheduler result on host once, rather than issuing ++ # one synchronous D2H transfer per request. In overlap mode the writer ++ # rank has already copied both tensors asynchronously on copy_stream; ++ # ``.cpu()`` is then a no-op. The fallback keeps non-overlap correct. ++ features = dict(req.spec_capture.get("features") or {}) ++ aux_cpu = getattr(logits_output, "_spec_capture_aux_cpu", None) ++ if "aux" in features and aux_cpu is None: ++ aux_cpu = logits_output.hidden_states.cpu() ++ logits_output._spec_capture_aux_cpu = aux_cpu ++ last_hidden_cpu = getattr( ++ logits_output, "_spec_capture_last_hidden_cpu", None + ) -+ if logits_output.last_hidden_states is not None: -+ req.spec_capture_last_hidden.append( -+ logits_output.last_hidden_states[start:end].cpu().clone() -+ ) ++ if ( ++ "last_hidden" in features ++ and logits_output.last_hidden_states is not None ++ and last_hidden_cpu is None ++ ): ++ last_hidden_cpu = logits_output.last_hidden_states.cpu() ++ logits_output._spec_capture_last_hidden_cpu = last_hidden_cpu ++ if "aux" in features and aux_cpu is not None: ++ req.spec_capture_aux.append(aux_cpu[start:end]) ++ if "last_hidden" in features and last_hidden_cpu is not None: ++ req.spec_capture_last_hidden.append(last_hidden_cpu[start:end]) + return end + -+ def _sink_spec_capture(self, req: Req) -> None: -+ """Write a finished capture request's tensors to the Mooncake sink. ++ def _sink_spec_capture(self, req: Req): ++ """Stage a finished capture request for the background Mooncake sink. + -+ Runs on the attention-TP rank that streams output; the per-request result -+ (or an ``{"error": ...}`` marker) is set on ``req.spec_capture_result``, ++ Runs on the attention-TP rank that streams output; returns the pending ++ tuple queued by ``_queue_spec_captures`` (or ``None`` off the writer ++ rank). The per-request result (or an ``{"error": ...}`` marker) is set ++ on ``req.spec_capture_result`` when the batch transfer completes, and + returned to the client via the dedicated ``spec_capture`` output field + (a per-request channel, unlike per-token ``customized_info``). + """ @@ -228,67 +388,147 @@ index a9d5f0c28b..c9f8fdf7d3 100644 + + sink = spec_capture_sink.get_sink() + if sink is None or self.output_streamer.ps.attn_tp_rank != 0: ++ return None ++ timing_enabled = os.environ.get("SGLANG_SPEC_CAPTURE_TIMING", "0") == "1" ++ cat_start = time.perf_counter() ++ # With chunked prefill disabled (the recommended capture-server ++ # configuration), each request owns exactly one contiguous view into ++ # the batch-level D2H buffer. torch.cat([view]) needlessly copied the ++ # whole capture a second time on CPU. Keep that view zero-copy; ++ # concatenate only the genuinely chunked case. ++ aux = self._coalesce_spec_capture_chunks(req.spec_capture_aux) ++ last_hidden = self._coalesce_spec_capture_chunks( ++ req.spec_capture_last_hidden ++ ) ++ cat_ms = (time.perf_counter() - cat_start) * 1000.0 ++ req.spec_capture_aux = [] ++ req.spec_capture_last_hidden = [] ++ return req, req.spec_capture, aux, last_hidden, timing_enabled, cat_ms ++ ++ def _queue_spec_captures(self, pending, *, return_logprob: bool) -> None: ++ """Queue one batch transfer while the scheduler runs the next prefill.""" ++ if not pending: + return -+ aux = torch.cat(req.spec_capture_aux, dim=0) if req.spec_capture_aux else None -+ last_hidden = ( -+ torch.cat(req.spec_capture_last_hidden, dim=0) -+ if req.spec_capture_last_hidden -+ else None ++ from sglang.srt import spec_capture_sink ++ ++ sink = spec_capture_sink.get_sink() ++ samples = [ ++ (spec, aux, last_hidden) ++ for _, spec, aux, last_hidden, _, _ in pending ++ ] ++ self._spec_capture_batches.append( ++ ( ++ pending, ++ sink.submit_samples(samples), ++ return_logprob, ++ time.perf_counter(), ++ ) + ) -+ position_ids = self._spec_capture_position_ids(req) -+ try: -+ req.spec_capture_result = sink.put_sample( -+ req.spec_capture, -+ aux=aux, -+ last_hidden=last_hidden, -+ position_ids=position_ids, ++ # Bound retained D2H buffers. Reaching this point means the target ++ # prefill for batch N+1 already overlapped host transfer N; wait only ++ # if N is still finishing before admitting N+2. ++ max_pending = int( ++ os.environ.get("SGLANG_SPEC_CAPTURE_MAX_PENDING_BATCHES", "2") ++ ) ++ if max_pending < 1: ++ raise ValueError("SGLANG_SPEC_CAPTURE_MAX_PENDING_BATCHES must be >= 1") ++ if len(self._spec_capture_batches) >= max_pending: ++ self.drain_spec_captures(block=True, max_batches=1) ++ ++ def has_pending_spec_captures(self) -> bool: ++ return bool(self._spec_capture_batches) ++ ++ def drain_spec_captures( ++ self, *, block: bool = False, max_batches: Optional[int] = None ++ ) -> int: ++ """Finish ready transfers and stream their responses on this thread.""" ++ completed = 0 ++ while self._spec_capture_batches: ++ if max_batches is not None and completed >= max_batches: ++ break ++ pending, future, return_logprob, queued_at = self._spec_capture_batches[0] ++ if not block and not future.done(): ++ break ++ self._spec_capture_batches.pop(0) ++ self._complete_spec_capture_batch( ++ pending, ++ future, ++ return_logprob=return_logprob, ++ queued_at=queued_at, + ) ++ completed += 1 ++ return completed ++ ++ def _complete_spec_capture_batch( ++ self, pending, future, *, return_logprob: bool, queued_at: float ++ ) -> None: ++ try: ++ results = future.result() ++ if len(results) != len(pending): ++ raise RuntimeError( ++ f"spec-capture sink returned {len(results)} results for " ++ f"{len(pending)} samples" ++ ) + except Exception as e: -+ logger.error("spec-capture sink failed for %s: %s", req.rid, e) -+ req.spec_capture_result = { -+ "sample_id": req.spec_capture.get("sample_id"), -+ "error": str(e), -+ } -+ req.spec_capture_aux = [] -+ req.spec_capture_last_hidden = [] ++ logger.error( ++ "spec-capture batch sink failed for %d requests: %s", len(pending), e ++ ) ++ for req, spec, _, _, _, _ in pending: ++ req.spec_capture_result = { ++ "sample_id": spec.get("sample_id"), ++ "error": str(e), ++ } ++ else: ++ for pending_item, result in zip(pending, results): ++ req, _, _, _, _, _ = pending_item ++ req.spec_capture_result = result ++ ++ timing_enabled = any(item[4] for item in pending) ++ if timing_enabled: ++ logger.info( ++ "[spec-capture-timing] async_batch_complete samples=%d " ++ "queue_to_stream_ms=%.3f cat_ms=%.3f", ++ len(pending), ++ (time.perf_counter() - queued_at) * 1000.0, ++ sum(item[5] for item in pending), ++ ) ++ self.output_streamer.stream_output( ++ [item[0] for item in pending], return_logprob ++ ) + + @staticmethod -+ def _spec_capture_position_ids(req: Req) -> Optional[torch.Tensor]: -+ """Prompt position ids for the sink's ``position_ids`` artifact. -+ -+ Multimodal (mRoPE) requests carry the processor-computed positions on -+ ``req.multimodal_inputs.mrope_positions`` ([3, L]); text requests fall -+ back to the plain arange broadcast the model uses for non-mm prompts. -+ Returns an (L, 3) CPU int64 tensor, or None when the client did not -+ request the artifact. -+ """ -+ features = (req.spec_capture or {}).get("features") or {} -+ if features.get("position_ids") is None: ++ def _coalesce_spec_capture_chunks( ++ chunks: List[torch.Tensor], ++ ) -> Optional[torch.Tensor]: ++ if not chunks: + return None -+ seq_len = len(req.origin_input_ids) -+ mrope_positions = ( -+ req.multimodal_inputs.mrope_positions -+ if req.multimodal_inputs is not None -+ else None -+ ) -+ if mrope_positions is not None: -+ return mrope_positions[:, :seq_len].t().contiguous().cpu() -+ return ( -+ torch.arange(seq_len, dtype=torch.int64) -+ .unsqueeze(0) -+ .expand(3, -1) -+ .t() -+ .contiguous() -+ ) ++ if len(chunks) == 1: ++ return chunks[0] ++ return torch.cat(chunks, dim=0) + def _append_prefill_hidden_states( self, *, diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py -index f95c59f6f2..9e64326a13 100644 +index f95c59f..6c51abc 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py -@@ -273,6 +273,7 @@ class _GenerationStreamAccumulator: +@@ -145,6 +145,14 @@ class SchedulerOutputStreamer: + for req in reqs: + if req is skip_req: + continue ++ if self.server_args.enable_spec_capture and req.spec_capture is not None: ++ # Capture requests are not complete until the background sink ++ # has durably published every feature object. Several ++ # scheduler paths can ask the common streamer to emit a ++ # finished request; centralize the completion barrier here so ++ # none of them can win the race and send a metadata-less 200. ++ if req.finished() and req.spec_capture_result is None: ++ continue + if req.finished() and req.finished_output: + # With the overlap schedule, a request will try to output twice and hit this line twice + # because of the one additional delayed token. This "continue" prevented the dummy output. +@@ -273,6 +281,7 @@ class _GenerationStreamAccumulator: spec_correct_drafts_histogram: list = field(default_factory=list) retraction_counts: list = field(default_factory=list) output_hidden_states: Optional[list] = None @@ -296,7 +536,7 @@ index f95c59f6f2..9e64326a13 100644 routed_experts: Optional[list] = None indexer_topk: Optional[list] = None customized_info: dict = field(default_factory=dict) -@@ -482,6 +483,8 @@ class _GenerationStreamAccumulator: +@@ -482,6 +491,8 @@ class _GenerationStreamAccumulator: self.output_hidden_states.append(hs) else: self.output_hidden_states.append(None) @@ -305,7 +545,7 @@ index f95c59f6f2..9e64326a13 100644 if self.return_routed_experts: self.routed_experts.append( req.routed_experts if req.return_routed_experts else None -@@ -540,6 +543,7 @@ class _GenerationStreamAccumulator: +@@ -540,6 +551,7 @@ class _GenerationStreamAccumulator: output_token_ids_logprobs_idx=self.output_token_ids_logprobs_idx, output_token_entropy_val=None, output_hidden_states=self.output_hidden_states, @@ -314,7 +554,7 @@ index f95c59f6f2..9e64326a13 100644 indexer_topk=self.indexer_topk, customized_info=self.customized_info, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py -index bf932611a5..f0e821ec47 100644 +index bf93261..f0e821e 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1172,6 +1172,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): @@ -336,11 +576,32 @@ index bf932611a5..f0e821ec47 100644 if getattr(recv_obj, "routed_experts", None): val = recv_obj.routed_experts[i] if val is not None: +diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py +index e9ede57..1d624fd 100644 +--- a/python/sglang/srt/managers/utils.py ++++ b/python/sglang/srt/managers/utils.py +@@ -77,6 +77,16 @@ class GenerationBatchResult: + Only the tensors which are needed for processing results are copied, + e.g., next_token_ids, logits outputs + """ ++ if ( ++ return_hidden_states ++ and self.logits_output.last_hidden_states is not None ++ ): ++ # Spec-capture's post-norm last hidden rides the same async D2H ++ # as the aux hidden states (the field exists only when the ++ # spec-capture patch is applied). ++ self.logits_output.last_hidden_states = ( ++ self.logits_output.last_hidden_states.to("cpu", non_blocking=True) ++ ) + if return_logprob: + if self.logits_output.next_token_logprobs is not None: + self.logits_output.next_token_logprobs = ( diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py -index 1cff5c9839..b842aeccb9 100644 +index 1cff5c9..5309965 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py -@@ -515,6 +515,22 @@ class ModelRunner(ModelRunnerKVCacheMixin): +@@ -515,6 +515,42 @@ class ModelRunner(ModelRunnerKVCacheMixin): draft_num_layers=int(draft_num_layers), ) @@ -348,26 +609,46 @@ index 1cff5c9839..b842aeccb9 100644 + # Aux capture without a draft worker, routed to the strategy's own + # capture method (they wire different submodules — e.g. VL models + # only populate layers_to_capture via the dflash path). -+ if getattr(server_args, "spec_capture_method", "eagle3") == "dflash": ++ capture_method = getattr(server_args, "spec_capture_method", "eagle3") ++ if capture_method in ("dflash", "dspark"): ++ # DSpark rides the DFlash aux plumbing: this patch's supported ++ # targets expose dflash capture hooks (targets with a native ++ # dspark hook are served by the kimi-k3 patch variant). + self.dflash_use_aux_hidden_state = True + self.dflash_target_layer_ids = server_args.spec_capture_aux_layer_ids ++ if hasattr(self, "spec_aux_config"): ++ self.spec_aux_config.dflash_use_aux_hidden_state = True ++ self.spec_aux_config.dflash_target_layer_ids = ( ++ server_args.spec_capture_aux_layer_ids ++ ) + self.dflash_family_use_aux_hidden_state = True + self.dflash_family_target_layer_ids = ( + server_args.spec_capture_aux_layer_ids + ) -+ else: ++ elif capture_method == "eagle3": + self.eagle_use_aux_hidden_state = True + self.eagle_aux_hidden_state_layer_ids = ( + server_args.spec_capture_aux_layer_ids + ) ++ if hasattr(self, "spec_aux_config"): ++ self.spec_aux_config.eagle_use_aux_hidden_state = True ++ self.spec_aux_config.eagle_aux_hidden_state_layer_ids = ( ++ server_args.spec_capture_aux_layer_ids ++ ) ++ else: ++ raise ValueError( ++ "--spec-capture-method must be one of: eagle3, dflash, dspark; " ++ f"got {capture_method!r}" ++ ) ++ # Apply the rank zero filter to logger if server_args.show_time_cost: enable_show_time_cost() diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py -index c7162c16d5..d515b3c699 100644 +index c7162c1..cfbce7a 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py -@@ -2127,6 +2127,25 @@ class ServerArgs: +@@ -2127,6 +2127,26 @@ class ServerArgs: bool, "Enable returning hidden states with responses.", ] = False @@ -386,19 +667,20 @@ index c7162c16d5..d515b3c699 100644 + ] = None + spec_capture_method: A[ + str, -+ "Capture method for --enable-spec-capture: 'eagle3' or 'dflash'. Must " -+ "match the draft strategy being trained; they wire capture onto " -+ "different submodules (VL targets only populate the dflash path).", ++ "Capture method for --enable-spec-capture: 'eagle3', 'dflash', or " ++ "'dspark'. Must match the draft strategy being trained; they wire " ++ "capture onto different submodules (VL targets only populate the " ++ "dflash path; dspark rides the dflash aux plumbing).", + ] = "eagle3" enable_return_routed_experts: A[ bool, "Enable returning routed experts of each layer with responses.", diff --git a/python/sglang/srt/spec_capture_sink.py b/python/sglang/srt/spec_capture_sink.py new file mode 100644 -index 0000000000..a678bd85e7 +index 0000000..d54b920 --- /dev/null +++ b/python/sglang/srt/spec_capture_sink.py -@@ -0,0 +1,256 @@ +@@ -0,0 +1,395 @@ +# Copyright 2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. @@ -419,8 +701,7 @@ index 0000000000..a678bd85e7 +Request schema:: + + {"store_id", "sample_id", "gen", "replace", # key namespace / retry policy -+ "features": {"aux": , "last_hidden": , # artifact -> feature -+ "position_ids": }, # (L, 3) prompt positions ++ "features": {"aux": , "last_hidden": }, # artifact -> feature + "passthrough": [{"name", "data", "shape", "dtype"}]} # client tensors verbatim + +Response (``meta_info["spec_capture"]``): ``{"sample_id", "store_id", "gen", @@ -435,7 +716,9 @@ index 0000000000..a678bd85e7 +import logging +import os +import threading -+from typing import Any, Dict, List, Optional ++import time ++from concurrent.futures import Future, ThreadPoolExecutor ++from typing import Any, Dict, List, Optional, Tuple + +import torch + @@ -459,7 +742,6 @@ index 0000000000..a678bd85e7 + +_ARTIFACT_AUX = "aux" +_ARTIFACT_LAST_HIDDEN = "last_hidden" -+_ARTIFACT_POSITION_IDS = "position_ids" + + +class SpecCaptureSink: @@ -470,9 +752,13 @@ index 0000000000..a678bd85e7 + self._store = None + self._put_config = None + self._lock = threading.Lock() -+ # Retried HTTP requests reuse deterministic keys. Striped locks keep -+ # replacement atomic per key without retaining one lock per sample. -+ self._write_locks = [threading.Lock() for _ in range(256)] ++ # One store writer is sufficient: Mooncake already stripes a batched ++ # transfer internally. The executor decouples that host transfer from ++ # the scheduler so the next target prefill can run concurrently. ++ self._executor = ThreadPoolExecutor( ++ max_workers=1, ++ thread_name_prefix="spec-capture-batch-put", ++ ) + + # -- connection --------------------------------------------------------- + def _connect(self): @@ -508,7 +794,13 @@ index 0000000000..a678bd85e7 + # before the trainer consumes it. + cfg = ReplicateConfig() + cfg.replica_num = 1 -+ cfg.with_hard_pin = True ++ # `with_hard_pin` exists only on newer Mooncake builds; older ROCm ++ # sglang images expose only `with_soft_pin`. Map the hard-pin intent ++ # onto whatever the installed build supports. ++ if hasattr(cfg, "with_hard_pin"): ++ cfg.with_hard_pin = True ++ elif hasattr(cfg, "with_soft_pin"): ++ cfg.with_soft_pin = True + self._put_config = cfg + self._store = store + logger.info("spec-capture mooncake sink connected") @@ -519,83 +811,118 @@ index 0000000000..a678bd85e7 + def _tkey(store_id: str, sample_id: str, gen: int, name: str) -> str: + return f"{store_id}/{sample_id}/g{gen}/{name}" + -+ def _put_tensor( -+ self, key: str, t: torch.Tensor, *, replace: bool = False -+ ) -> None: -+ store = self._connect() -+ t = t.detach().to("cpu").contiguous() -+ nbytes = t.element_size() * t.numel() -+ lock = self._write_locks[hash(key) % len(self._write_locks)] -+ with lock: -+ if replace: -+ # Do not probe with is_exist(): Mooncake existence checks can -+ # acquire a read lease that prevents the following removal. -+ self._remove_quiet(key) -+ try: -+ store.register_buffer(t.data_ptr(), nbytes) -+ except Exception: -+ pass # some builds auto-register -+ try: -+ rc = store.put_from(key, t.data_ptr(), nbytes, self._put_config) -+ finally: -+ try: -+ store.unregister_buffer(t.data_ptr()) -+ except Exception: -+ pass -+ if rc is not None and int(rc) < 0: -+ raise RuntimeError(f"spec-capture put_from failed (status {rc}) for {key}") -+ + def _remove_quiet(self, key: str) -> None: + try: + self._connect().remove(key) + except Exception: + pass + -+ # -- the one entry point -------------------------------------------------- -+ def put_sample( ++ def _remove_many_quiet(self, keys: List[str]) -> None: ++ if not keys: ++ return ++ store = self._connect() ++ batch_remove = getattr(store, "batch_remove", None) ++ if batch_remove is not None: ++ try: ++ batch_remove(keys) ++ return ++ except Exception: ++ pass ++ for key in keys: ++ self._remove_quiet(key) ++ ++ # -- the batch entry point ------------------------------------------------ ++ def submit_samples( + self, -+ spec: Dict[str, Any], -+ *, -+ aux: Optional[torch.Tensor], -+ last_hidden: Optional[torch.Tensor], -+ position_ids: Optional[torch.Tensor] = None, -+ ) -> Dict[str, Any]: -+ """Write one sample's artifacts; return the meta_info result dict. ++ samples: List[ ++ Tuple[ ++ Dict[str, Any], Optional[torch.Tensor], Optional[torch.Tensor] ++ ] ++ ], ++ ) -> Future[List[Dict[str, Any]]]: ++ """Queue one scheduler batch without blocking the scheduler thread.""" ++ return self._executor.submit(self.put_samples, samples) + -+ ``aux``/``last_hidden`` are the per-request (L, W) captured tensors, -+ ``position_ids`` the per-request (L, 3) prompt positions (mRoPE rows -+ for multimodal requests, arange broadcast for text), all stored with a -+ leading batch dim of 1. On any failure the keys already written are -+ best-effort removed (no partial sample is consumable). ++ def put_samples( ++ self, ++ samples: List[ ++ Tuple[ ++ Dict[str, Any], Optional[torch.Tensor], Optional[torch.Tensor] ++ ] ++ ], ++ ) -> List[Dict[str, Any]]: ++ """Publish a scheduler batch with one native Mooncake batch RPC. ++ ++ A capture prefill batch normally finishes many samples together. ++ Calling ``put_from`` once per feature object paid dozens of metadata ++ and transport round trips per batch on the writer rank and serialized ++ the response behind them. ++ ``batch_put_from`` preserves the existing per-feature keys and raw ++ tensor layout while amortizing that fixed cost across the whole ++ scheduler batch. The response is still emitted only after every ++ status succeeds, so refs can never point at incomplete samples. + """ -+ store_id = str(spec["store_id"]) -+ sample_id = str(spec["sample_id"]) -+ gen = int(spec.get("gen", 1)) -+ replace = bool(spec.get("replace", False)) -+ features: Dict[str, str] = dict(spec.get("features") or {}) ++ if not samples: ++ return [] + -+ written: List[str] = [] -+ result_feats: Dict[str, Dict[str, Any]] = {} ++ store = self._connect() ++ timing_enabled = os.environ.get("SGLANG_SPEC_CAPTURE_TIMING", "0") == "1" ++ started = time.perf_counter() ++ keys: List[str] = [] ++ tensors: List[torch.Tensor] = [] ++ sizes: List[int] = [] ++ replace_keys: List[str] = [] ++ results: List[Dict[str, Any]] = [] + -+ def _write(name: str, t: torch.Tensor) -> None: ++ def _stage( ++ result_feats: Dict[str, Dict[str, Any]], ++ *, ++ store_id: str, ++ sample_id: str, ++ gen: int, ++ replace: bool, ++ name: str, ++ tensor: torch.Tensor, ++ ) -> None: ++ tensor = tensor.detach().to("cpu").contiguous() + key = self._tkey(store_id, sample_id, gen, name) -+ self._put_tensor(key, t, replace=replace) -+ written.append(key) ++ keys.append(key) ++ tensors.append(tensor) ++ sizes.append(tensor.element_size() * tensor.numel()) ++ if replace: ++ replace_keys.append(key) + result_feats[name] = { -+ "shape": list(t.shape), -+ "dtype": _DTYPE_STR.get(t.dtype, str(t.dtype).replace("torch.", "")), ++ "shape": list(tensor.shape), ++ "dtype": _DTYPE_STR.get( ++ tensor.dtype, str(tensor.dtype).replace("torch.", "") ++ ), + } + -+ try: ++ for spec, aux, last_hidden in samples: ++ store_id = str(spec["store_id"]) ++ sample_id = str(spec["sample_id"]) ++ gen = int(spec.get("gen", 1)) ++ replace = bool(spec.get("replace", False)) ++ features: Dict[str, str] = dict(spec.get("features") or {}) ++ result_feats: Dict[str, Dict[str, Any]] = {} ++ + aux_name = features.get(_ARTIFACT_AUX) + if aux_name is not None: + if aux is None: + raise RuntimeError( + "spec_capture requested 'aux' but no aux hidden states were " -+ "captured — launch the server with --enable-spec-capture " ++ "captured -- launch the server with --enable-spec-capture " + "(and optionally --spec-capture-aux-layer-ids)" + ) -+ _write(aux_name, aux.unsqueeze(0)) ++ _stage( ++ result_feats, ++ store_id=store_id, ++ sample_id=sample_id, ++ gen=gen, ++ replace=replace, ++ name=aux_name, ++ tensor=aux.unsqueeze(0), ++ ) + lh_name = features.get(_ARTIFACT_LAST_HIDDEN) + if lh_name is not None: + if last_hidden is None: @@ -603,15 +930,15 @@ index 0000000000..a678bd85e7 + "spec_capture requested 'last_hidden' but the logits " + "processor did not return it (is aux capture enabled?)" + ) -+ _write(lh_name, last_hidden.unsqueeze(0)) -+ pos_name = features.get(_ARTIFACT_POSITION_IDS) -+ if pos_name is not None: -+ if position_ids is None: -+ raise RuntimeError( -+ "spec_capture requested 'position_ids' but no prompt " -+ "positions were collected on the scheduler" -+ ) -+ _write(pos_name, position_ids.unsqueeze(0)) ++ _stage( ++ result_feats, ++ store_id=store_id, ++ sample_id=sample_id, ++ gen=gen, ++ replace=replace, ++ name=lh_name, ++ tensor=last_hidden.unsqueeze(0), ++ ) + for item in spec.get("passthrough") or []: + dtype = _STR_DTYPE.get(str(item.get("dtype", "int64"))) + if dtype is None: @@ -619,22 +946,116 @@ index 0000000000..a678bd85e7 + f"spec_capture passthrough {item.get('name')!r}: " + f"unsupported dtype {item.get('dtype')!r}" + ) -+ t = torch.tensor(item["data"], dtype=dtype).reshape( ++ tensor = torch.tensor(item["data"], dtype=dtype).reshape( + [int(d) for d in item["shape"]] + ) -+ _write(str(item["name"]), t) ++ _stage( ++ result_feats, ++ store_id=store_id, ++ sample_id=sample_id, ++ gen=gen, ++ replace=replace, ++ name=str(item["name"]), ++ tensor=tensor, ++ ) ++ results.append( ++ { ++ "sample_id": sample_id, ++ "store_id": store_id, ++ "gen": gen, ++ "aux_layer_ids": self.aux_layer_ids, ++ "features": result_feats, ++ } ++ ) ++ ++ materialize_ms = (time.perf_counter() - started) * 1000.0 ++ self._remove_many_quiet(replace_keys) ++ registered: List[torch.Tensor] = [] ++ register_started = time.perf_counter() ++ try: ++ for tensor, nbytes in zip(tensors, sizes): ++ try: ++ store.register_buffer(tensor.data_ptr(), nbytes) ++ registered.append(tensor) ++ except Exception: ++ pass # TCP and some Mooncake builds auto-register ++ register_ms = (time.perf_counter() - register_started) * 1000.0 ++ put_started = time.perf_counter() ++ batch_put = getattr(store, "batch_put_from", None) ++ if batch_put is None: ++ statuses = [ ++ store.put_from(key, tensor.data_ptr(), nbytes, self._put_config) ++ for key, tensor, nbytes in zip(keys, tensors, sizes) ++ ] ++ else: ++ statuses = batch_put( ++ keys, ++ [tensor.data_ptr() for tensor in tensors], ++ sizes, ++ self._put_config, ++ ) ++ put_ms = (time.perf_counter() - put_started) * 1000.0 + except Exception: -+ for key in written: -+ self._remove_quiet(key) ++ self._remove_many_quiet(keys) + raise ++ finally: ++ for tensor in registered: ++ try: ++ store.unregister_buffer(tensor.data_ptr()) ++ except Exception: ++ pass + -+ return { -+ "sample_id": sample_id, -+ "store_id": store_id, -+ "gen": gen, -+ "aux_layer_ids": self.aux_layer_ids, -+ "features": result_feats, -+ } ++ if statuses is None: ++ statuses = [0] * len(keys) ++ if len(statuses) != len(keys): ++ self._remove_many_quiet(keys) ++ raise RuntimeError( ++ "spec-capture batch_put_from returned " ++ f"{len(statuses)} statuses for {len(keys)} keys" ++ ) ++ failed = [ ++ (key, status) ++ for key, status in zip(keys, statuses) ++ if status is not None and int(status) < 0 ++ ] ++ if failed: ++ self._remove_many_quiet(keys) ++ raise RuntimeError( ++ "spec-capture batch_put_from failed for " ++ f"{len(failed)}/{len(keys)} keys; first={failed[0]}" ++ ) ++ ++ if timing_enabled: ++ logger.info( ++ "[spec-capture-timing] batch_sink samples=%d objects=%d " ++ "bytes=%d materialize_ms=%.3f register_ms=%.3f put_ms=%.3f " ++ "total_ms=%.3f", ++ len(samples), ++ len(keys), ++ sum(sizes), ++ materialize_ms, ++ register_ms, ++ put_ms, ++ (time.perf_counter() - started) * 1000.0, ++ ) ++ return results ++ ++ def put_sample( ++ self, ++ spec: Dict[str, Any], ++ *, ++ aux: Optional[torch.Tensor], ++ last_hidden: Optional[torch.Tensor], ++ ) -> Dict[str, Any]: ++ """Write one sample's artifacts; return the meta_info result dict. ++ ++ ``aux``/``last_hidden`` are the per-request (L, W) captured tensors, ++ stored with a leading batch dim of 1. On any failure the keys already ++ written are best-effort removed (no partial sample is consumable). ++ """ ++ # Keep the single-sample entry point for compatibility with tests and ++ # callers outside the scheduler; production uses put_samples(). ++ return self.put_samples([(spec, aux, last_hidden)])[0] + + +_SINK: Optional[SpecCaptureSink] = None diff --git a/scripts/regenerate_train_data.py b/scripts/regenerate_train_data.py index f2f737fd2..252c989b2 100644 --- a/scripts/regenerate_train_data.py +++ b/scripts/regenerate_train_data.py @@ -244,6 +244,29 @@ def build_query_kwargs(args, messages, max_tokens=None): return query_kwargs +def _extract_record_images(data: Dict[str, Any]) -> List[str]: + """Resolve image references of a record (``image``/``image_path`` string or + ``images`` list), in insertion order.""" + refs: List[str] = [] + single = data.get("image") or data.get("image_path") + if isinstance(single, str): + refs.append(single) + images = data.get("images") + if isinstance(images, list): + refs.extend(r for r in images if isinstance(r, str)) + return refs + + +def _image_url_part(path: str) -> Dict[str, Any]: + import base64 + import mimetypes + + mime = mimetypes.guess_type(path)[0] or "image/jpeg" + with open(path, "rb") as f: + b64 = base64.b64encode(f.read()).decode("ascii") + return {"type": "image_url", "image_url": {"url": f"data:{mime};base64,{b64}"}} + + def call_sglang( args, server_address: str, @@ -260,6 +283,8 @@ def call_sglang( messages = data["conversations"] regenerated_messages = [] + record_images = _extract_record_images(data) + image_attached = False # ignore data which starts with an assistant message if messages[0]["role"] == "assistant": @@ -273,6 +298,26 @@ def call_sglang( elif message["role"] == "assistant": continue elif message["role"] == "user": + # Multimodal records: attach the record's images to the first user + # turn that carries the placeholder (OpenAI content parts). + content = message.get("content") + if ( + record_images + and not image_attached + and isinstance(content, str) + and "" in content + ): + try: + parts = [_image_url_part(p) for p in record_images] + except OSError as exc: + data["status"] = "error" + data["error"] = f"unreadable image file: {exc}" + return data + text = content.replace("\n", "").replace("", "") + parts.append({"type": "text", "text": text}) + message = dict(message) + message["content"] = parts + image_attached = True regenerated_messages.append(message) query_kwargs = build_query_kwargs(args, regenerated_messages, max_tokens) diff --git a/specforge/algorithms/common/dflash_family_model.py b/specforge/algorithms/common/dflash_family_model.py index faf5320b6..2fde4a624 100644 --- a/specforge/algorithms/common/dflash_family_model.py +++ b/specforge/algorithms/common/dflash_family_model.py @@ -305,7 +305,6 @@ def _forward_draft_blocks( input_ids: torch.Tensor, hidden_states: torch.Tensor, loss_mask: torch.Tensor, - position_ids: Optional[torch.Tensor] = None, max_valid_anchors: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: bsz, seq_len = input_ids.shape @@ -322,32 +321,11 @@ def _forward_draft_blocks( input_ids, anchor_positions, block_keep_mask ) - if position_ids is None: - context_position_ids = ( - torch.arange(seq_len, device=device).unsqueeze(0).expand(bsz, -1) - ) - draft_position_ids = self._create_position_ids(anchor_positions) - full_position_ids = torch.cat( - [context_position_ids, draft_position_ids], dim=1 - ) - else: - if not getattr(self.draft_model, "use_interleaved_mrope", False): - raise ValueError( - "multimodal capture carries mRoPE position_ids, but the " - "draft config does not enable rope_scaling.mrope_interleaved; " - "use a VLM draft config (e.g. configs/*-dflash-vlm-*.json)" - ) - # Server-produced mRoPE positions, (B, S, 3) -> (3, B, S + N*bs). - offsets = torch.arange(self.block_size, device=device).view(1, 1, -1) - draft_indices = (anchor_positions.unsqueeze(-1) + offsets).view(bsz, -1) - draft_position_ids = torch.gather( - position_ids, - 1, - draft_indices.unsqueeze(-1).expand(-1, -1, 3), - ) - full_position_ids = torch.cat( - [position_ids, draft_position_ids], dim=1 - ).permute(2, 0, 1) + context_position_ids = ( + torch.arange(seq_len, device=device).unsqueeze(0).expand(bsz, -1) + ) + draft_position_ids = self._create_position_ids(anchor_positions) + full_position_ids = torch.cat([context_position_ids, draft_position_ids], dim=1) mask_builder = ( create_dflash_block_mask @@ -465,15 +443,10 @@ def forward( input_ids: torch.Tensor, hidden_states: torch.Tensor, loss_mask: torch.Tensor, - position_ids: Optional[torch.Tensor] = None, max_valid_anchors: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, object]]: """Parallel block-wise training forward pass; returns - (loss, accuracy, metrics) — same shape as Domino's forward. - - ``position_ids`` is the optional server-captured mRoPE position tensor - ``(B, S, 3)`` for multimodal runs; text runs leave it None and use - internally synthesized flat positions. + (loss, accuracy, metrics) - same shape as Domino's forward. """ if self.attention_backend == "flex_attention" and not FLEX_ATTENTION_AVAILABLE: raise ValueError( @@ -486,7 +459,6 @@ def forward( input_ids=input_ids, hidden_states=hidden_states, loss_mask=loss_mask, - position_ids=position_ids, max_valid_anchors=max_valid_anchors, ) diff --git a/specforge/algorithms/common/hidden_states_data.py b/specforge/algorithms/common/hidden_states_data.py index f774d9e61..770ed710c 100644 --- a/specforge/algorithms/common/hidden_states_data.py +++ b/specforge/algorithms/common/hidden_states_data.py @@ -168,29 +168,6 @@ def collate(features): return collate -def build_vlm_collator(): - """Collator for multimodal capture: text tensors + mRoPE position ids. - - ``position_ids`` arrives as ``(1, L, 3)`` int64 per sample (temporal / - height / width mRoPE rows produced by the capture server) and is padded - along the sequence axis like every other per-token feature. - """ - - def collate(features): - return pad_and_concatenate_features( - features, - sequence_axes={ - "input_ids": 1, - "loss_mask": 1, - "hidden_states": 1, - "position_ids": 1, - }, - required_keys=("input_ids", "loss_mask", "hidden_states", "position_ids"), - ) - - return collate - - def build_dspark_collator(): def collate(features): return pad_and_concatenate_features( @@ -304,7 +281,6 @@ def collate(features): "build_mtp_offline_reader", "build_offline_normalizer", "build_offline_reader", - "build_vlm_collator", "normalize_dspark_offline_sample", "normalize_mtp_offline_sample", "normalize_offline_sample", diff --git a/specforge/algorithms/common/providers.py b/specforge/algorithms/common/providers.py index 9d1114a4c..b3ae2facf 100644 --- a/specforge/algorithms/common/providers.py +++ b/specforge/algorithms/common/providers.py @@ -511,25 +511,18 @@ def __post_init__(self) -> None: @dataclass(frozen=True) class ServerCaptureLayout: - """Maps generic server artifacts onto algorithm-ready feature names. - - ``position_ids_feature`` names an optional server-produced position-id - artifact (mRoPE positions for multimodal targets), stored ``(1, L, 3)`` - int64. Text providers leave it unset. - """ + """Maps generic server artifacts onto algorithm-ready feature names.""" aux_feature: str | None last_hidden_feature: str | None passthrough: Tuple[Tuple[str, str, Tuple[int, ...]], ...] attention_mask_feature: str | None = None - position_ids_feature: str | None = None def __post_init__(self) -> None: for field_name in ( "aux_feature", "last_hidden_feature", "attention_mask_feature", - "position_ids_feature", ): value = getattr(self, field_name) if value is not None: @@ -761,7 +754,6 @@ def make_registration( layout.aux_feature, layout.last_hidden_feature, layout.attention_mask_feature, - layout.position_ids_feature, ) if feature is not None ), diff --git a/specforge/algorithms/common/vlm_input.py b/specforge/algorithms/common/vlm_input.py index b81c32593..36250d12a 100644 --- a/specforge/algorithms/common/vlm_input.py +++ b/specforge/algorithms/common/vlm_input.py @@ -16,7 +16,7 @@ - ``prepare_prompts``: ShareGPT-style JSONL (+ optional ``image`` field) -> payload dicts with expanded ``input_ids``/``loss_mask`` (what the trainer and the passthrough capture use) plus ``request_input_ids`` (single - placeholder) and base64 ``image_data`` (what the capture request sends). + placeholder) and ``data:``-URI ``image_data`` (what the capture request sends). - ``build_request_inputs``: batch payloads -> the ``/generate`` model-input fields ``{"input_ids", "image_data"}``. diff --git a/specforge/algorithms/dflash/providers.py b/specforge/algorithms/dflash/providers.py index d12901d03..98832d5c2 100644 --- a/specforge/algorithms/dflash/providers.py +++ b/specforge/algorithms/dflash/providers.py @@ -13,7 +13,6 @@ build_collator, build_offline_normalizer, build_offline_reader, - build_vlm_collator, ) from specforge.algorithms.common.providers import ( AlgorithmProviders, @@ -135,7 +134,6 @@ def needs_input_tools(config, draft_model): def algorithm_spec() -> AlgorithmSpec: ready = {"input_ids", "loss_mask", "hidden_states"} - vlm_ready = {"input_ids", "loss_mask", "hidden_states", "position_ids"} return AlgorithmSpec( name=ALGORITHM_NAME, draft=DraftRequirement( @@ -162,7 +160,7 @@ def algorithm_spec() -> AlgorithmSpec: FeatureContract( mode=FeatureMode.STREAMING, modality="multimodal", - required_tensors=vlm_ready, + required_tensors=ready, ), ), capabilities=AlgorithmCapabilities( @@ -245,9 +243,8 @@ def algorithm_providers() -> AlgorithmProviders: ("input_ids", "input_ids", ()), ("loss_mask", "loss_mask", ()), ), - position_ids_feature="position_ids", ), - build_collator=build_vlm_collator, + build_collator=collator, build_input_adapter=build_vlm_input_adapter, ), ), diff --git a/specforge/config/schema.py b/specforge/config/schema.py index 27870d0ca..eb955e7c5 100644 --- a/specforge/config/schema.py +++ b/specforge/config/schema.py @@ -80,6 +80,11 @@ class ModelConfig(StrictConfigModel): tokenizer_pad_token_id: Optional[int] = Field(default=None, ge=0) #: SGLang target-engine tuning. Ignored by hf/custom backends. sglang_attention_backend: str = "flashinfer" + #: Multimodal (vision encoder) attention backend for capture servers. On + #: Ascend NPU with a non-text input_modality it defaults to ascend_attn + #: (fused vision attention); sdpa materializes N^2 scores and OOMs on + #: large images. + sglang_mm_attention_backend: Optional[str] = None sglang_mem_fraction_static: float = Field(default=0.4, gt=0.0, le=1.0) #: Keep the historical managed-local behavior by default. Hybrid targets #: such as Inkling require the radix tree and can opt back in explicitly. @@ -967,13 +972,24 @@ def apply_overrides(config: Config, overrides: List[str]) -> Config: node = raw keys = path.split(".") for key in keys[:-1]: - if not isinstance(node.get(key), dict): + child = node.get(key) + if child is None: + # Optional sections default to None (e.g. disaggregated + # .managed_local on a server_urls recipe); create them so + # their declared fields can be overridden. Re-validation with + # extra="forbid" still rejects typo'd paths, so this stays safe. + child = {} + node[key] = child + if not isinstance(child, dict): raise ValueError(f"override path {path!r} does not exist") - node = node[key] - if keys[-1] not in node: - raise ValueError(f"override path {path!r} does not exist") - current = node[keys[-1]] - if isinstance(current, (dict, list)) and value.lstrip().startswith(("[", "{")): + node = child + current = node.get(keys[-1]) + # Structured (YAML) values are parsed when the target is a dict/list + # or a yet-unset leaf (None); schema re-validation does the final + # type check either way. + if (current is None or isinstance(current, (dict, list))) and ( + value.lstrip().startswith(("[", "{")) + ): import yaml try: diff --git a/specforge/data/vlm_preprocessing.py b/specforge/data/vlm_preprocessing.py index dc6ae3121..fe9a2fdc9 100644 --- a/specforge/data/vlm_preprocessing.py +++ b/specforge/data/vlm_preprocessing.py @@ -19,13 +19,16 @@ (``SGLANG_MM_AVOID_RETOKENIZE=1``), which is guaranteed to reproduce the client expansion because both sides run the target model's own HF processor on the same image. -- ``image_data``: a base64 image string for the capture request (the image +- ``image_data``: a ``data:``-URI image string for the capture request (the image itself never enters the feature store; only text tokens and captured hidden states do). Limitations (v1): at most one image per sample; the image is attached to the first user turn (matching the Qwen-VL chat layout); text-only samples are -supported in the same run (``image_data=None``). +supported in the same run (``image_data=None``). Image references are read +from ``image`` / ``image_path`` (string) or ``images`` (one-element list); +unreadable images and multi-image samples are fatal ``ImageDataError``s, +never a silent text-only downgrade. """ from __future__ import annotations @@ -52,32 +55,92 @@ def _image_token_count(image_grid_thw, merge_size: int) -> int: return count // (merge_size * merge_size) +class ImageDataError(ValueError): + """Fatal image-contract violation (unreadable image, multi-image sample). + + Raised instead of silently degrading an image-bearing sample to text-only + training; the prompt-preparation loop re-raises it to abort loudly. + """ + + +def _extract_image_field(record: Dict[str, Any], *, source: str) -> Optional[str]: + """Resolve the single image reference of one record, or None. + + Accepts ``image`` / ``image_path`` (string) and ``images`` (a one-element + list, the convention used by common VLM corpora). A record carrying more + than one image is a fatal error in v1 (single-image contract), and so is + an ``images`` list whose element is not a string. + """ + + image_field = record.get("image") or record.get("image_path") + images = record.get("images") + if image_field and images: + raise ImageDataError( + f"{source}: record has both 'image' and 'images' fields; use one" + ) + if images is None: + return image_field + if not isinstance(images, list): + raise ImageDataError( + f"{source}: 'images' must be a list, got {type(images).__name__}" + ) + if len(images) == 0: + return image_field + if len(images) > 1: + raise ImageDataError( + f"{source}: multi-image samples are not supported (got " + f"{len(images)} images); the v1 contract is one image per sample" + ) + first = images[0] + if not isinstance(first, str): + raise ImageDataError( + f"{source}: 'images[0]' must be a path or base64 string, got " + f"{type(first).__name__}" + ) + return first + + +def _as_data_uri(raw: bytes, *, media_type: str = "image/jpeg") -> str: + """Wrap raw image bytes as a data URI. + + Some SGLang builds (e.g. the Ascend 0.5.14 line) only route ``data:`` + URIs through base64 decoding and treat every other string as a file + path; the data-URI form is accepted by all supported server versions. + """ + + return f"data:{media_type};base64," + base64.b64encode(raw).decode("ascii") + + def _load_image(image_field: Any, *, source: str): - """Return (pil_image, base64_str) from a path / base64 / data-URI field.""" + """Return (pil_image, data_uri) from a path / base64 / data-URI field.""" from PIL import Image if not isinstance(image_field, str) or not image_field: - raise ValueError( + raise ImageDataError( f"{source}: image field must be a file path or base64 string, got " f"{type(image_field).__name__}" ) if os.path.isfile(image_field): with open(image_field, "rb") as image_file: raw = image_file.read() - return Image.open(io.BytesIO(raw)).convert("RGB"), base64.b64encode(raw).decode( - "ascii" - ) + return Image.open(io.BytesIO(raw)).convert("RGB"), _as_data_uri(raw) encoded = image_field if encoded.startswith("data:"): + media_type = encoded[5 : encoded.index(";")] if ";" in encoded else "image/jpeg" encoded = encoded.split(",", 1)[-1] + else: + media_type = "image/jpeg" try: raw = base64.b64decode(encoded, validate=True) except Exception as exc: - raise ValueError( + raise ImageDataError( f"{source}: image field is neither an existing file nor valid base64" ) from exc - return Image.open(io.BytesIO(raw)).convert("RGB"), encoded + return ( + Image.open(io.BytesIO(raw)).convert("RGB"), + _as_data_uri(raw, media_type=media_type), + ) def _render_conversation_text( @@ -181,6 +244,7 @@ def build_vlm_prompt_payloads( merge_size = int(getattr(processor.image_processor, "merge_size", 2)) payloads: List[Dict[str, Any]] = [] + skipped = 0 for line_number, record in _iter_records(path): source = f"{path}:{line_number}" try: @@ -195,14 +259,24 @@ def build_vlm_prompt_payloads( max_length=max_length, min_loss_tokens=min_loss_tokens, ) + except ImageDataError: + # Image-contract violations are fatal: never degrade an + # image-bearing sample to text-only training silently. + raise except ValueError as exc: print(f"WARNING: skipping {source}: {exc}") + skipped += 1 continue if prepared is None: + skipped += 1 continue payloads.append(prepared) if max_prompts not in (None, 0) and len(payloads) >= max_prompts: break + print( + f"VLM prompt preparation done: {len(payloads)} prepared, " + f"{skipped} skipped ({path})" + ) return payloads @@ -223,7 +297,7 @@ def _prepare_one_record( conversations = record.get("conversations") if not conversations: raise ValueError("record has no 'conversations' field") - image_field = record.get("image") or record.get("image_path") + image_field = _extract_image_field(record, source=source) text = _render_conversation_text( tokenizer, @@ -284,4 +358,4 @@ def _prepare_one_record( } -__all__ = ["build_vlm_prompt_payloads"] +__all__ = ["ImageDataError", "build_vlm_prompt_payloads"] diff --git a/specforge/inference/adapters/server_capture.py b/specforge/inference/adapters/server_capture.py index 2633436a1..6647c627f 100644 --- a/specforge/inference/adapters/server_capture.py +++ b/specforge/inference/adapters/server_capture.py @@ -48,16 +48,12 @@ class ServerCaptureSchema: ``(feature_name, payload_key, trailing_shape)`` for client tensors stored verbatim (``trailing_shape`` is appended after ``(1, L)``). ``attention_mask_feature`` is synthesized all-ones (PromptTasks are unpadded). - ``position_ids_feature`` names the server-produced position-id artifact - (mRoPE positions for multimodal targets, stored ``(1, L, 3)`` int64); - None = not requested. """ aux_feature: Optional[str] last_hidden_feature: Optional[str] passthrough: Tuple[Tuple[str, str, Tuple[int, ...]], ...] attention_mask_feature: Optional[str] = None - position_ids_feature: Optional[str] = None @dataclass(frozen=True) @@ -212,8 +208,6 @@ def _spec_capture_payload(self, task: PromptTask) -> Dict[str, Any]: features["aux"] = self.schema.aux_feature if self.schema.last_hidden_feature is not None: features["last_hidden"] = self.schema.last_hidden_feature - if self.schema.position_ids_feature is not None: - features["position_ids"] = self.schema.position_ids_feature passthrough: List[Dict[str, Any]] = [] for feature_name, payload_key, trailing in self.schema.passthrough: if payload_key == "input_ids": diff --git a/specforge/launch_plan.py b/specforge/launch_plan.py index 357ed926f..74460483f 100644 --- a/specforge/launch_plan.py +++ b/specforge/launch_plan.py @@ -515,6 +515,16 @@ def _managed_local_services( ): # flashinfer does not exist on Ascend; default to the NPU backend. attention_backend = "ascend" + mm_attention_backend = cfg.model.sglang_mm_attention_backend + if ( + mm_attention_backend is None + and cfg.model.input_modality != "text" + and device_visibility_env == "ASCEND_RT_VISIBLE_DEVICES" + ): + # The sdpa vision backend materializes [heads, N, N] attention + # scores and OOMs on large images; ascend_attn is the fused + # (flash-style) vision attention on NPU. + mm_attention_backend = "ascend_attn" argv.extend( _sglang_argv( cfg.model, @@ -526,6 +536,7 @@ def _managed_local_services( else cfg.model.sglang_mem_fraction_static ), "sglang_attention_backend": attention_backend, + "sglang_mm_attention_backend": mm_attention_backend, }, ) ) diff --git a/specforge/modeling/draft/dflash.py b/specforge/modeling/draft/dflash.py index 91c12b98a..39f86568e 100644 --- a/specforge/modeling/draft/dflash.py +++ b/specforge/modeling/draft/dflash.py @@ -108,64 +108,6 @@ def _prepare_dflash_eager_mask( return additive_mask, valid_queries -def get_rope_scaling_value(config: Qwen3Config, key: str, default=None): - rope_scaling = getattr(config, "rope_scaling", None) - if rope_scaling is None: - return default - if isinstance(rope_scaling, dict): - return rope_scaling.get(key, default) - return getattr(rope_scaling, key, default) - - -class Qwen3InterleavedMultiRotaryEmbedding(Qwen3RotaryEmbedding): - """Interleaved mRoPE for Qwen3-VL style multimodal position ids.""" - - def __init__(self, config: Qwen3Config): - super().__init__(config) - self.mrope_section = get_rope_scaling_value( - config, "mrope_section", [24, 20, 20] - ) - - def _apply_interleaved_mrope(self, freqs: torch.Tensor) -> torch.Tensor: - freqs_t = freqs[0] - for dim_idx, offset in enumerate((1, 2), start=1): - length = self.mrope_section[dim_idx] * 3 - idx_slice = slice(offset, length, 3) - freqs_t[..., idx_slice] = freqs[dim_idx, ..., idx_slice] - return freqs_t - - @torch.no_grad() - def forward( - self, x: torch.Tensor, position_ids: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - if position_ids.ndim == 2: - position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) - - inv_freq_expanded = ( - self.inv_freq[None, None, :, None] - .float() - .expand(3, position_ids.shape[1], -1, 1) - ) - position_ids_expanded = position_ids[:, :, None, :].float() - - device_type = ( - x.device.type - if isinstance(x.device.type, str) and x.device.type != "mps" - else "cpu" - ) - with torch.autocast(device_type=device_type, enabled=False): - freqs = ( - inv_freq_expanded.float() @ position_ids_expanded.float() - ).transpose(2, 3) - interleaved_freqs = self._apply_interleaved_mrope(freqs) - emb = torch.cat((interleaved_freqs, interleaved_freqs), dim=-1) - scaling = getattr(self, "attention_scaling", 1.0) - cos = emb.cos() * scaling - sin = emb.sin() * scaling - - return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) - - class Qwen3DFlashAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" @@ -451,13 +393,7 @@ def __init__( build_target_layer_ids(config.num_target_layers, config.num_hidden_layers), ) self.norm = kernels.make_rms_norm(config.hidden_size, config.rms_norm_eps) - self.use_interleaved_mrope = bool( - get_rope_scaling_value(config, "mrope_interleaved", False) - ) - if self.use_interleaved_mrope: - self.rotary_emb = Qwen3InterleavedMultiRotaryEmbedding(config) - else: - self.rotary_emb = Qwen3RotaryEmbedding(config) + self.rotary_emb = Qwen3RotaryEmbedding(config) self.fc = nn.Linear( len(self.target_layer_ids) * config.hidden_size, config.hidden_size, diff --git a/specforge/training/disaggregated.py b/specforge/training/disaggregated.py index 2f31fb7f2..47d175e80 100644 --- a/specforge/training/disaggregated.py +++ b/specforge/training/disaggregated.py @@ -639,7 +639,6 @@ def _build_online( last_hidden_feature=layout.last_hidden_feature, passthrough=layout.passthrough, attention_mask_feature=layout.attention_mask_feature, - position_ids_feature=layout.position_ids_feature, ) adapters = [ SGLangServerCaptureAdapter( diff --git a/specforge/training/strategies/base.py b/specforge/training/strategies/base.py index d7a748118..69495a295 100644 --- a/specforge/training/strategies/base.py +++ b/specforge/training/strategies/base.py @@ -474,15 +474,11 @@ def forward_loss( t = batch.tensors device = self._device() max_valid_anchors = _cpu_max_valid_anchors(t["loss_mask"]) - # Multimodal capture additionally stores server-produced mRoPE - # position ids (B, S, 3); text runs do not carry the tensor at all. - position_ids = t.get("position_ids") loss, accuracy, model_metrics = self.dflash_model( input_ids=t["input_ids"].to(device, non_blocking=True), hidden_states=t["hidden_states"].to(device, non_blocking=True), loss_mask=t["loss_mask"].to(device, non_blocking=True), max_valid_anchors=max_valid_anchors, - position_ids=None if position_ids is None else position_ids.to(device), ) metrics = {"accuracy": accuracy.detach()} if "accuracy_denom" in model_metrics: diff --git a/tests/test_algorithms/test_builtin_providers.py b/tests/test_algorithms/test_builtin_providers.py index 59b8750a2..cc10fa18f 100644 --- a/tests/test_algorithms/test_builtin_providers.py +++ b/tests/test_algorithms/test_builtin_providers.py @@ -161,6 +161,7 @@ def test_builtin_modalities_are_pinned_per_algorithm(self): "domino": {"text"}, "dspark": {"text"}, "eagle3": {"text"}, + "mtp": {"text"}, "peagle": {"text"}, } for registration in self.registry: diff --git a/tests/test_algorithms/test_dflash_multimodal.py b/tests/test_algorithms/test_dflash_multimodal.py index 813bbffbb..21c30cd85 100644 --- a/tests/test_algorithms/test_dflash_multimodal.py +++ b/tests/test_algorithms/test_dflash_multimodal.py @@ -28,11 +28,10 @@ def test_dflash_registers_the_multimodal_streaming_contract(self): ) self.assertEqual( set(contract.required_tensors), - {"input_ids", "loss_mask", "hidden_states", "position_ids"}, + {"input_ids", "loss_mask", "hidden_states"}, ) provider = registration.providers.server_streaming_for("multimodal") self.assertEqual(provider.capture_method, "dflash") - self.assertEqual(provider.layout.position_ids_feature, "position_ids") self.assertEqual(provider.layout.aux_feature, "hidden_states") def test_other_builtins_have_no_multimodal_contract(self): @@ -78,6 +77,75 @@ def test_expand_image_region_requires_exactly_one_placeholder(self): ) +class ExtractImageFieldTest(unittest.TestCase): + def _extract(self, record): + from specforge.data.vlm_preprocessing import _extract_image_field + + return _extract_image_field(record, source="t") + + def test_image_and_image_path_fields(self): + self.assertEqual(self._extract({"image": "a.jpg"}), "a.jpg") + self.assertEqual(self._extract({"image_path": "b.jpg"}), "b.jpg") + self.assertIsNone(self._extract({})) + + def test_images_list_takes_the_single_element(self): + self.assertEqual(self._extract({"images": ["a.jpg"]}), "a.jpg") + self.assertIsNone(self._extract({"images": []})) + + def test_multi_image_sample_is_fatal(self): + from specforge.data.vlm_preprocessing import ImageDataError + + with self.assertRaises(ImageDataError): + self._extract({"images": ["a.jpg", "b.jpg"]}) + + def test_non_list_images_and_non_string_element_are_fatal(self): + from specforge.data.vlm_preprocessing import ImageDataError + + with self.assertRaises(ImageDataError): + self._extract({"images": "a.jpg"}) + with self.assertRaises(ImageDataError): + self._extract({"images": [123]}) + + def test_conflicting_image_fields_are_fatal(self): + from specforge.data.vlm_preprocessing import ImageDataError + + with self.assertRaises(ImageDataError): + self._extract({"image": "a.jpg", "images": ["b.jpg"]}) + + def test_unreadable_image_is_fatal_not_skipped(self): + from specforge.data.vlm_preprocessing import ImageDataError, _load_image + + with self.assertRaises(ImageDataError): + _load_image("/nonexistent/path/to/image.jpg", source="t") + with self.assertRaises(ImageDataError): + _load_image("not-valid-base64!!!", source="t") + + def test_load_image_returns_data_uri(self): + import base64 as b64mod + + from specforge.data.vlm_preprocessing import _load_image + + # 1x1 white PNG + png_b64 = ( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4" + "z8DwHwAFAAH/q842iQAAAABJRU5ErkJggg==" + ) + _, uri = _load_image(png_b64, source="t") + self.assertTrue(uri.startswith("data:image/jpeg;base64,")) + raw = b64mod.b64decode(uri.split(",", 1)[1]) + self.assertEqual(raw, b64mod.b64decode(png_b64)) + + def test_load_image_preserves_input_media_type(self): + from specforge.data.vlm_preprocessing import _load_image + + png_b64 = ( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4" + "z8DwHwAFAAH/q842iQAAAABJRU5ErkJggg==" + ) + _, uri = _load_image(f"data:image/png;base64,{png_b64}", source="t") + self.assertTrue(uri.startswith("data:image/png;base64,")) + + class VlmRequestInputsTest(unittest.TestCase): def test_build_request_inputs_uses_collapsed_ids_and_image_data(self): adapter = VlmServerInputAdapter(config=SimpleNamespace()) @@ -102,111 +170,56 @@ def test_build_request_inputs_uses_collapsed_ids_and_image_data(self): self.assertEqual(request["image_data"], ["aGVsbG8=", None]) -class ServerCapturePositionIdsTest(unittest.TestCase): - def _adapter(self, position_ids_feature): - from specforge.inference.adapters.server_capture import ( - ServerCaptureSchema, - SGLangServerCaptureAdapter, - ) - - class _FakeStore: - store_id = "store" - - def adopt(self, ref): - pass - - def discard_external_attempts(self, *args, **kwargs): - pass +@unittest.skipUnless(TORCH_AVAILABLE, "requires torch") +class DraftPositionsTest(unittest.TestCase): + """Drafts always train on the plain 1D position convention, identical to + the text-only path; multimodal capture stores no position ids.""" - def track_external_attempt(self, *args, **kwargs): - pass + def _build_model(self): + import torch as t + from torch import nn - schema = ServerCaptureSchema( - aux_feature="hidden_states", - last_hidden_feature=None, - passthrough=( - ("input_ids", "input_ids", ()), - ("loss_mask", "loss_mask", ()), - ), - position_ids_feature=position_ids_feature, - ) - return SGLangServerCaptureAdapter( - "http://localhost:1", - _FakeStore(), - run_id="run", - algorithm="dflash", - schema=schema, - ) + from specforge.algorithms.common.dflash_family_model import OnlineDFlashModel - def _task(self): - return SimpleNamespace( - task_id="t0", - attempt=0, - payload={"input_ids": [5, 6, 7], "loss_mask": [0, 1, 1]}, - metadata={}, + class _StubDraftModel(nn.Module): + def __init__(self): + super().__init__() + self.recorded = {} + self.sliding_window = None + + def forward( + self, + position_ids=None, + noise_embedding=None, + target_hidden=None, + attention_mask=None, + ): + self.recorded["position_ids"] = position_ids + return t.zeros(1) + + return OnlineDFlashModel( + draft_model=_StubDraftModel(), + target_lm_head=nn.Linear(8, 32, bias=False), + target_embed_tokens=nn.Embedding(32, 8), + mask_token_id=31, + block_size=2, + attention_backend="sdpa", + num_anchors=4, ) - def test_payload_requests_position_ids_artifact_when_configured(self): - adapter = self._adapter("position_ids") - payload = adapter._spec_capture_payload(self._task()) - self.assertEqual( - payload["features"], - {"aux": "hidden_states", "position_ids": "position_ids"}, - ) - - def test_payload_omits_position_ids_artifact_when_unset(self): - adapter = self._adapter(None) - payload = adapter._spec_capture_payload(self._task()) - self.assertEqual(payload["features"], {"aux": "hidden_states"}) - - -@unittest.skipUnless(TORCH_AVAILABLE, "requires torch") -class VlmCollatorTest(unittest.TestCase): - def test_collator_pads_position_ids_like_other_features(self): - from specforge.algorithms.common.dflash_family_data import build_vlm_collator - - collate = build_vlm_collator() - features = [ - { - "input_ids": torch.tensor([[1, 2, 3]]), - "loss_mask": torch.tensor([[0, 1, 1]]), - "hidden_states": torch.zeros(1, 3, 8), - "position_ids": torch.arange(9).reshape(1, 3, 3), - }, - { - "input_ids": torch.tensor([[4]]), - "loss_mask": torch.tensor([[1]]), - "hidden_states": torch.zeros(1, 1, 8), - "position_ids": torch.arange(3).reshape(1, 1, 3), - }, - ] - batch = collate(features) - self.assertEqual(tuple(batch["input_ids"].shape), (2, 3)) - self.assertEqual(tuple(batch["position_ids"].shape), (2, 3, 3)) - # Padding is zeros on the sequence axis. - self.assertTrue((batch["position_ids"][1, 1:] == 0).all()) - self.assertEqual(batch["position_ids"][0, 2].tolist(), [6, 7, 8]) - - -@unittest.skipUnless(TORCH_AVAILABLE, "requires torch") -class MropeDraftPositionsTest(unittest.TestCase): - def test_gathered_draft_positions_follow_anchor_offsets(self): + def test_positions_follow_the_1d_convention(self): import torch as t - from specforge.algorithms.common.dflash_family_model import OnlineDFlashModel - - model = OnlineDFlashModel.__new__(OnlineDFlashModel) - model.block_size = 2 - anchors = t.tensor([[1, 3]]) - stored = t.arange(5 * 3).reshape(1, 5, 3) - offsets = t.arange(model.block_size).view(1, 1, -1) - draft_indices = (anchors.unsqueeze(-1) + offsets).view(1, -1) - gathered = t.gather(stored, 1, draft_indices.unsqueeze(-1).expand(-1, -1, 3)) - full = t.cat([stored, gathered], dim=1).permute(2, 0, 1) - self.assertEqual(tuple(full.shape), (3, 1, 5 + 4)) - # Draft slot for anchor=1: positions of indices 1 and 2. - self.assertEqual(full[:, 0, 5].tolist(), [3, 4, 5]) - self.assertEqual(full[:, 0, 6].tolist(), [6, 7, 8]) + model = self._build_model() + b, s = 2, 8 + input_ids = t.randint(0, 31, (b, s)) + hidden_states = t.randn(b, s, 16) + loss_mask = t.ones(b, s) + t.manual_seed(0) + model._forward_draft_blocks(input_ids, hidden_states, loss_mask) + got = model.draft_model.recorded["position_ids"] + self.assertEqual(got.ndim, 2) + self.assertTrue(t.equal(got[:, :s], t.arange(s).unsqueeze(0).expand(b, -1))) if __name__ == "__main__": diff --git a/tests/test_config/test_schema.py b/tests/test_config/test_schema.py index 03edbbd6a..690835fba 100644 --- a/tests/test_config/test_schema.py +++ b/tests/test_config/test_schema.py @@ -648,6 +648,48 @@ def test_override_bad_path_or_form_raises(self): with self.assertRaises(ValueError): apply_overrides(cfg, ["not-an-assignment"]) + def test_override_creates_none_optional_sections(self): + payload = { + "model": { + "target_model_path": "some/target", + "draft_model_config": "draft.json", + }, + "data": {"train_data_path": "train.jsonl"}, + "deployment": { + "mode": "disaggregated", + "trainer": {"nnodes": 1, "nproc_per_node": 2}, + "disaggregated": { + "control_dir": "outputs/x/control", + "backend": "mooncake", + "server_urls": ["http://127.0.0.1:30000"], + }, + }, + } + cfg = Config.model_validate(payload) + out = apply_overrides( + cfg, + [ + # managed_local and server_urls are mutually exclusive. + "deployment.disaggregated.server_urls=[]", + 'deployment.disaggregated.managed_local.trainer_cuda_visible_devices=["8","9"]', + 'deployment.disaggregated.managed_local.capture_servers=[{port: 40000, cuda_visible_devices: ["0"], tp_size: 1}]', + ], + ) + self.assertEqual(out.deployment.disaggregated.server_urls, []) + managed = out.deployment.disaggregated.managed_local + self.assertIsNotNone(managed) + self.assertEqual(managed.trainer_cuda_visible_devices, ["8", "9"]) + self.assertEqual(len(managed.capture_servers), 1) + self.assertEqual(managed.capture_servers[0].port, 40000) + # Defaults fill in untouched managed sub-configs. + self.assertGreater(managed.mooncake.rpc_port, 0) + # Typo'd paths are still rejected by re-validation. + with self.assertRaises(Exception): + apply_overrides( + cfg, + ["deployment.disaggregated.managed_local.no_such_field=1"], + ) + def test_load_config_applies_overrides(self): path = _write(MINIMAL, ".json") self.addCleanup(os.unlink, path) diff --git a/tests/test_config/test_server_only_online.py b/tests/test_config/test_server_only_online.py index 9be95463e..a5bc1443f 100644 --- a/tests/test_config/test_server_only_online.py +++ b/tests/test_config/test_server_only_online.py @@ -73,7 +73,7 @@ def test_multimodal_modality_resolves_for_dflash(self): resolved = resolve_run(Config.model_validate(payload)) self.assertEqual(resolved.algorithm.name, "dflash") streaming = resolved.algorithm.providers.server_streaming_for("multimodal") - self.assertEqual(streaming.layout.position_ids_feature, "position_ids") + self.assertEqual(streaming.layout.aux_feature, "hidden_states") self.assertIsNotNone(streaming.create_input_adapter(resolved.config)) def test_multimodal_modality_is_rejected_for_text_only_algorithms(self):