[zhou9402]

Ports miles' day-0 model support: opt-in mbridge HF→Megatron load, DSA indexer/topk/kpool/sparse-attention Triton kernels, the DeepSeek-V4 plugin (ops, THD packing, CP row-balance, indexer + compressor, attention, converter, bridge), DSA indexer top-k replay (vLLM channel via vllm-indexer-topk.patch + trainer wiring), and the Kimi-K3 plugin (line-for-line with miles).

Source: radixark/miles@ba28991834. Port is additive (new modules + new add_*_arguments() groups); slime-origin files are not restructured so future [Sync] Mirror Slime PRs stay clean. One commit per sub-PR; review each commit separately.

Commits:

• d97bc5f F07-0: converter verifier harness (2-layer logits allclose + round-trip bitwise)
• ab47a77 F07-1: opt-in mbridge HF->Megatron load mode
• cedf399 F07-2a: DSA indexer + topk Triton kernels
• 8389153 F07-2b: DSA kpool + sparse_attention kernels
• 4a17002 F07-3a-i: IndexerReplayManager replay_base (pulled forward from F07-8b)
• 04ede28 F07-3a: DeepSeek-V4 plugin small ops + arguments
• 196f866 F07-3a2: DeepSeek-V4 THD packing + CP row-balance ops
• 4fd6038 F07-3a3: DeepSeek-V4 indexer + compressor
• 7b4c166 F07-3a4: DeepSeek-V4 attention module + get_dsv4_spec
• 7548d5e F07-3b: DeepSeek-V4 converter + mbridge bridge
• 1c14255 F07-8a: DSA indexer top-k vLLM channel + rollout decode
• 3444dd9 F07-8b: trainer-side indexer top-k replay wiring
• 69d5a46 F07-6a: shared linear-attn helper layer for Kimi-K3
• b7e4003 F07-6b: Kimi-K3 model plugin + converter + mbridge bridge

Verification: F07-0 converter harness (2-layer logits allclose + round-trip bitwise). The vLLM patch applies in full Dockerfile order (after vllm.patch); docker build still pending. Follow-ups: DSv4.1, GLM5-next, Qwen3.8-next, ROCm side of the indexer patch.

Merge order of the miles-port series: F10 → F02 → F06 → F11a → F11b → F05 → F09 → F04 → F03 → F07 → F08 → F01. train.py / vime/utils/arguments.py are touched by several branches; later ones are rebased after each merge.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces support for Day-0 models (DeepSeek-V4 and Kimi-K3) by adding custom model plugins, Triton-based DSA (DeepSeek Sparse Attention) kernels, and an opt-in mbridge loading mode for HuggingFace checkpoints. It also implements a rollout indexer top-k replay mechanism to record and replay indexer decisions during training. The review feedback highlights three critical issues: an incorrect offset calculation in the indexer replay preparation when context parallelism is enabled, an erroneous tensor-parallel sharding of replay data under sequence parallelism, and a missing mask for invalid columns in the fallback PyTorch top-k implementation when sglang is unavailable.

[gemini-code-assist]

When context parallelism is enabled (cp_size > 1), local is a sliced version of the sample's top-k indices. Incrementing offset by local.shape[0] (the local sliced size) rather than token_len (the global size of the sample) will cause incorrect indexing for subsequent samples in the packed sequence, as their keys in the all-gathered k are located after the entire previous sample. We should increment offset by token_len instead.

Suggested change

	# Rebase per-sample KV positions onto the packed training sequence.
	local = torch.where(local != -1, local + offset, local)
	offset += local.shape[0]
	sliced.append(local)
	# Rebase per-sample KV positions onto the packed training sequence.
	local = torch.where(local != -1, local + offset, local)
	offset += token_len
	sliced.append(local)

[gemini-code-assist]

The DSA indexer always gathers the inputs to the full sequence length when sequence_parallel is enabled, meaning the attention computation and the indexer run on the full sequence length on each TP rank. Sharding replay_data by TP rank here will cause a shape mismatch and assertion failure during training. This sharding block should be removed.

    return replay_data

[gemini-code-assist]

When sglang is not installed and topk_backend == "canonical", logits contains garbage values in the invalid columns because clean_logits is set to False. Since torch_dsa_topk does not use row_starts and row_ends to mask out invalid columns, it will incorrectly select garbage indices. We should mask out the invalid columns using row_starts and row_ends if they are provided.

Suggested change

	def torch_dsa_topk(logits: torch.Tensor, topk: int, row_starts=None, row_ends=None) -> torch.Tensor:
	# torch.topk errors on k > n; a context shorter than index_topk is legal, so pad with -1
	k = min(topk, logits.shape[-1])
	score, indices = torch.topk(logits, k, dim=-1)
	indices = indices.to(torch.int32)
	indices = indices.masked_fill(score == -torch.inf, -1)
	if k < topk:
	pad = indices.new_full((*indices.shape[:-1], topk - k), -1)
	indices = torch.cat([indices, pad], dim=-1)
	return indices
	def torch_dsa_topk(logits: torch.Tensor, topk: int, row_starts=None, row_ends=None) -> torch.Tensor:
	if row_starts is not None and row_ends is not None:
	columns = torch.arange(logits.shape[-1], device=logits.device)
	in_range = (columns >= row_starts.long().unsqueeze(-1)) & (columns < row_ends.long().unsqueeze(-1))
	logits = logits.masked_fill(~in_range, float("-inf"))
	# torch.topk errors on k > n; a context shorter than index_topk is legal, so pad with -1
	k = min(topk, logits.shape[-1])
	score, indices = torch.topk(logits, k, dim=-1)
	indices = indices.to(torch.int32)
	indices = indices.masked_fill(score == -torch.inf, -1)
	if k < topk:
	pad = indices.new_full((*indices.shape[:-1], topk - k), -1)
	indices = torch.cat([indices, pad], dim=-1)
	return indices

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #35444835 | 📁 Comparing b7e4003 against latest (d7e40fe)

  🔍 Preview build  

62 files changed · + 4 added · ± 58 modified

+ Added

• advanced/day0-models.html
• advanced/dsa-kernels.html
• advanced/mbridge-load-mode.html
• advanced/model-converter-harness.html

± Modified

• genindex.html
• library.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/delta-weight-sync.html
• advanced/external-rollout-engines.html
• advanced/fault-tolerance.html
• advanced/low-precision.html
• advanced/megatron-config.html
• advanced/observability.html
• and 48 more...