feat(deepseek_v4): PR1 skeleton — end-to-end inference with triton MoE#650
Draft
feat(deepseek_v4): PR1 skeleton — end-to-end inference with triton MoE#650
Conversation
…arity Adds the foundational scaffolding for DeepSeek-V4-Pro support — a major architecture shift from V3.2 with mHC residuals, hybrid CSA+HCA attention, hash routing, and grouped output LoRA. PR1 ships the eager-mode model code with torch fallback kernels, validated against the official inference implementation at bit-exact parity (max_abs_diff = 0.0). Scope (PR1 only): - New atom/models/deepseek_v4.py: full Compressor / Indexer / Attention / Gate / Expert / MoE / Block / MTPBlock / ParallelHead / Transformer port (~1200 lines). Single-rank only; plain nn.Linear / nn.Embedding for now. - New atom/model_ops/sparse_attn_v4.py: torch fallbacks for sparse_attn and hc_split_sinkhorn (Sinkhorn-Knopp projection on Birkhoff polytope). - New atom/model_ops/quant_v4.py: torch fallbacks for FP8/FP4 inplace QAT round-trip and Walsh-Hadamard transform (replaces fast_hadamard_transform which doesn't build on ROCm). - Register DeepseekV4ForCausalLM in support_model_arch_dict. Out of scope (tracked for PR2-6): - Real HF checkpoint loading (PR2 = FP4 e2m1 loader, PR3 = TP + KV cache). - AITER sparse_attn kernel (PR4; spec at /app/logs_claude/aiter_v4_sparse_attn_spec.md, AITER team kicked off). - MTP integration with EagleProposer (PR5). - @support_torch_compile + CUDAGraph + openai_server (PR6). Verification: /app/logs_claude/v4_pr1_verify.py monkey-patches the reference's TileLang kernel imports with our torch fallbacks, copies the same dummy state_dict into both models, and runs prefill + decode side-by-side. 259 tensors match exactly; max_abs_diff = 0.0 on logits.
DeepSeek-V4-Pro stores routed expert weights as packed FP4 e2m1 (int8 with 2 values per byte, low nibble first) plus per-block ue8m0 scale (block size 32 along input dim). This commit adds `dequant_fp4_e2m1(packed, scale)` in atom/model_ops/quant_v4.py — a pure-torch unpacker that mirrors convert.py exactly but produces BF16 directly instead of repacking into FP8. Validated bit-exactly against an independent reference unpack on a real 22M-element expert tensor from the on-disk checkpoint. Also regression- tested across 5 different shapes/positions (w1/w2/w3 in first/mid/last layer + MTP). All produce values that lie exactly on the FP4 e2m1 grid. Scope: this is the standalone dequant utility. Wiring it into the model loader's safetensors pipeline + tying it to specific param names happens in PR3 alongside TP-aware expert sharding. Test: /app/logs_claude/v4_pr2_dequant_test.py Result: max_abs_diff = 0.0 (bit-exact)
PR3a: replace nn.Linear / nn.Embedding with ATOM tensor-parallel-aware
classes for the BF16 projections in Attention, Indexer, and the model
embedding. Same `weight` parameter naming so dummy state_dicts continue
to load. At TP=1 ATOM's tgemm.mm produces bit-identical output to F.linear,
so PR1's reference parity (max_abs_diff = 0.0) still passes.
Layers refactored (8 total):
- DeepseekV4Model.embed: nn.Embedding -> VocabParallelEmbedding
- DeepseekV4Attention.wq_a: nn.Linear -> ReplicatedLinear
- DeepseekV4Attention.wq_b: nn.Linear -> ColumnParallelLinear
- DeepseekV4Attention.wkv: nn.Linear -> ReplicatedLinear (single shared MQA head)
- DeepseekV4Attention.wo_a: nn.Linear -> ColumnParallelLinear
- DeepseekV4Attention.wo_b: nn.Linear -> RowParallelLinear (with all-reduce)
- Indexer.wq_b: nn.Linear -> ColumnParallelLinear
- Indexer.weights_proj: nn.Linear -> ColumnParallelLinear
Deferred to later PRs (intentional):
- Compressor.wkv/wgate (fp32) -> PR3c with quant_type wiring
- ParallelHead.weight (fp32 LM head) -> PR3c
- Expert.w{1,2,3} -> PR3b (FusedMoE wholesale rewrite)
- MoE.gate.weight (used as raw Parameter, not Linear class) -> kept
Verification: /app/logs_claude/v4_pr1_verify.py (now GPU mode with
init_dist_env) shows max_abs_diff = 0.0 for prefill + decode against
reference at TP=1.
… for real ckpt
PR3c delivers end-to-end real-checkpoint loading for DeepSeek-V4 attention
layers via ATOM's existing FP8/FP4 GEMM infrastructure.
What works after this commit (validated on real /data/DeepSeek-V4-Pro/):
- DeepseekV4ForCausalLM(atom_config) auto-builds a V4QuantConfig that maps
routed-experts -> per_1x32 (FP4) and overrides wo_a / Compressor.wkv /
Compressor.wgate / indexer.weights_proj -> bf16 (no quant). Everything
else inherits the global FP8 (per_1x128) spec from the HF quantization_config.
- load_weights(weights) walks an iterable of (name, tensor) pairs and:
* Remaps ATOM's `weight_scale` -> on-disk `scale` naming.
* Special-cases wo_a: dequantizes FP8+scale -> BF16 on the fly so the
grouped-LoRA einsum (which aiter doesn't support in FP8) works.
* Dispatches to ATOM Linear's weight_loader for FP8 / FP4 / BF16 paths.
* Skips params with shape mismatch (e.g. expert nn.Linear waiting for
PR3b's FusedMoE refactor) without crashing.
- All 23 attention parameters (FP8 q/kv proj + FP4 indexer + BF16 wo_a + fp32
compressor) load successfully on real layer-2 of the V4 checkpoint.
Threading changes:
- DeepseekV4Args gains `quant_config: Optional[Any] = None`.
- DeepseekV4Attention / Indexer / Compressor / Block / MTPBlock / DeepseekV4Model
now accept `prefix: str = ""` and pass `quant_config + prefix` down to each
ATOM Linear constructor so per-layer quant lookup works.
Backward compatibility:
- When `args.quant_config is None` (toy / dummy validation), V4QuantConfig
retains its `QuantType.No` global — Linear layers stay BF16 and the PR1
bit-exact reference parity test (max_abs_diff = 0.0) still passes.
Remaining gaps for end-to-end real-ckpt forward (tracked in design doc):
- PR3b: replace MoE/Expert with FusedMoE so 384 expert FP4 weights load.
- PR3d: refactor V4 attention.forward to accept 2D [num_tokens, dim] input
(ATOM TP linears require 2D — current 3D path raises "GEMM not supported").
PR3d adapts V4 model to ATOM's scheduler convention: model.forward consumes flat 2D `[num_tokens, dim]` tokens (single sequence implicit B=1), matching how ATOM's ModelRunner / scheduler pass tokens. This unblocks ATOM Linear's quantized GEMM kernels (which only accept 2D `[M, K]` input) and enables end-to-end real-checkpoint forward. What changed: - DeepseekV4Attention.forward(x, start_pos): now accepts 2D [num_tokens, dim]. Internally adds a B=1 dim only where needed (RoPE, sparse_attn). The grouped-LoRA einsum string changes from "bsgd,grd->bsgr" to "sgd,grd->sgr". - Compressor.forward / Indexer.forward: accept 2D x; auto-unsqueeze to 3D internally for backward compatibility with the existing logic. - Block.hc_pre / hc_post + ParallelHead.hc_head: refactored to be shape-agnostic in leading dims (use negative indexing on flatten / sum). Both 4D `[B, S, hc, D]` (legacy reference path) and 3D `[num_tokens, hc, D]` (ATOM path) work. - ParallelHead.get_logits: 2D path takes last token via `x[-1:]`; 3D path preserves `x[:, -1]` for legacy [B, S, D] inputs. - MTPBlock.forward: 2D-aware via `e.unsqueeze(-2)` for hc-dim broadcast. - DeepseekV4Model.forward: auto-flattens 2D `[1, S]` input_ids to 1D `[S]` for the new convention; rejects B>1 (proper multi-sequence batching needs attn_metadata, deferred). Validated: - PR1 reference parity (toy 4-layer dummy weights at B=1 S=32): max_abs_diff = 0.0 — still bit-exact after the 2D refactor. - PR3d end-to-end on REAL V4 weights: + Built DeepseekV4ForCausalLM (4 layers, real V4 dims, ~105B params) + load_weights() loaded 36 layer-2 params; 23/23 attn params nonzero + attn(x_2d=[16, 7168], start_pos=0) → output [16, 7168] bf16 + No NaN/Inf; output range [-2.94, 3.08], abs mean 0.42 (sensible) + This is the first successful V4 attention forward on real weights via ATOM Test scripts (under /app/logs_claude/): - v4_pr1_verify.py — toy parity (now uses B=1 + ATOM 2D path) - v4_pr3d_layer_e2e.py — real-weight 2D forward end-to-end - v4_pr3c_layer0_test.py — per-Linear validation against real ckpt Remaining for full model end-to-end: - PR3b: MoE → FusedMoE so 384 expert FP4 weights load (currently shape-skipped) - Multi-sequence support via attn_metadata (currently single-sequence implicit B=1)
PR3b enables ATOM's FusedMoE for V4's 384 routed experts so FP4 expert
weights can load via the existing aiter `gemm_a4w4_quant` kernel and
shard across TP/EP ranks. Also extends `select_experts` in moe.py to
support V4's `sqrtsoftplus` scoring with `e_score_correction_bias`.
Changes in atom/model_ops/moe.py:
- `FusedMoE.select_experts` now handles `scoring_func="sqrtsoftplus"`:
routing_weights = sqrt(softplus(router_logits)) + topk + renormalize.
Mirrors the V4 reference Gate.forward exactly for non-hash layers.
Changes in atom/models/deepseek_v4.py:
- Dual-path MoE: when `quant_config` is set AND ATOM's global atom_config
is initialized, MoE uses ReplicatedLinear gate + FusedMoE experts +
ATOM-Linear shared_experts. Otherwise falls back to the original manual
per-expert nn.Linear path so PR1 toy validation stays bit-exact (the
reference test runs without ATOM's ModelRunner setting the global config).
- Expert class accepts `quant_config + prefix`: when set, w1/w2/w3 become
ColumnParallelLinear/RowParallelLinear (FP8 path); else nn.Linear (toy).
- DeepseekV4ForCausalLM.get_expert_mapping() returns the (param_name,
weight_name, expert_id, shard_id) tuples mapping V4's `w1/w2/w3` ckpt
names to FusedMoE's merged `w13_*`/`w2_*` params.
- load_weights() walks expert_mapping first to dispatch routed expert
tensors via FusedMoE's per-expert weight_loader, then handles the rest:
* ATOM `weight_scale` ↔ on-disk `scale` rename (existing)
* ATOM `gate.e_score_correction_bias` ↔ on-disk `gate.bias` rename (NEW)
* `wo_a` FP8 → BF16 dequant on load (existing)
Validated:
- PR1 toy parity: max_abs_diff = 0.0 (manual MoE path still bit-exact).
- PR3d e2e: real layer-2 attn + 2D forward still works.
- PR3b new: under stub atom_config, FusedMoE path activates correctly.
Layer-3 (non-hash, real V4 dims): gate + e_score_correction_bias +
shared_experts (6/6) loaded; FusedMoE expert mapping returns 1152
entries (384 experts × {w1,w2,w3}).
Known limitations (deferred):
- Hash routing (layers 0/1/2): tid2eid table is loaded but routing logic
still falls through to sqrtsoftplus path → INCORRECT for hash layers.
Proper hash routing requires either a custom path through FusedMoE
or a pre-computed (topk_weights, topk_ids) injection point.
- Multi-sequence batching via attn_metadata (currently single-sequence implicit B=1).
Test: /app/logs_claude/v4_pr3b_fusedmoe_test.py
… prefix Bug: `make_v4_quant_config` matched `"ffn.experts." in layer_name` (with trailing dot). FusedMoE.__init__ asks for the layer's quant_type with prefix `layers.N.ffn.experts` (NO trailing dot — it's the parent module of the per-expert weights, not a per-expert lookup). The check failed, so FusedMoE inherited the global FP8 (per_1x128) spec and allocated the routed expert weights as `float8_e4m3fn` instead of `float4_e2m1fn_x2`. Symptom in PR3b validation output before the fix: FusedMoE experts: 3/5 nonzero (loader couldn't dispatch FP4-shaped on-disk tensors into FP8-typed model params; shape mismatch silently skipped them) After the fix: experts.w13_weight: (385, 6144, 3584) torch.float4_e2m1fn_x2 ✓ experts.w13_weight_scale: (385, 6144, 224) torch.float8_e8m0fnu ✓ experts.w2_weight: (385, 7168, 1536) torch.float4_e2m1fn_x2 ✓ experts.w2_weight_scale: (385, 7168, 96) torch.float8_e8m0fnu ✓ e_score_correction_bias: (384,) torch.float32 ✓ Match condition tightened to `".ffn.experts" in layer_name` so it catches BOTH `layers.N.ffn.experts.M.w1` (per-expert Linear lookups) AND `layers.N.ffn.experts` (FusedMoE parent module lookup). Note: a separate aiter-side issue (HSA_STATUS_ERROR_EXCEPTION on FP4 expert weight_loader, traced to a `direct_copy_kernel` with grid size exceeding HW limits) prevents end-to-end FP4 expert load testing on this box. The dtype/shape correctness above is verified by inspecting the constructed module's params directly. Validated: - PR1 toy parity: max_abs_diff = 0.0 (manual MoE fallback unaffected) - PR3d real-attention forward: still works
PR3b's expert weight loader had three bugs that caused weights to load as
zero or be silently dropped:
1. **Expert mapping pattern mismatch**: `make_expert_params_mapping` returns
`(param_part="experts.w13_", weight_part="experts.0.w1.", ...)` — substring
substitution, not endswith. The old code built `f".experts.{e}.{suffix}"`
which never matched. Switched to longest-prefix substring substitution
matching the standard ATOM loader pattern.
2. **Scale dtype zero-fill**: copying `torch.float8_e8m0fnu` into a `uint8`
destination via `copy_()` silently produces zeros (mismatched dtype, no
reinterpret). FusedMoE allocates `w13_weight_scale` as uint8; force a
`.view(torch.uint8)` on the e8m0 source before passing to the loader.
3. **Param suffix `_scale` vs `.weight_scale`**: after substring sub,
`experts.0.w1.scale` becomes `experts.w13_scale`, but the FusedMoE param is
`experts.w13_weight_scale`. Added `_scale` → `_weight_scale` post-fix.
Plus: gracefully slice on-disk gate.weight / gate.bias when the test caps
n_routed_experts below the checkpoint size (no-op in real serving).
Verified:
- v4_pr3b_fusedmoe_test: 32 params loaded, 5/5 expert + 6/6 shared nonzero
- v4_pr3d_layer_e2e: real attention forward still works
- v4_pr1_verify: bit-exact reference parity preserved (0.0 max diff)
…uting_function V4 uses tid2eid hash lookup (instead of gate-logit topk) for routing in layers where compress_ratio implies hash layer (first 3 layers in standard config). Previously, MoE just declared tid2eid for weight loading but inference fell through to sqrtsoftplus path → wrong routing for those layers. This commit: - Adds an early `custom_routing_function` branch to FusedMoE.select_experts (it was in the signature but never honored — the non-grouped path went straight to scoring_func dispatch). Now any non-None custom fn takes precedence and returns (topk_weights, topk_ids). - Adds DeepseekV4MoE._hash_topk(): topk_ids = tid2eid[input_ids], topk_weights = sqrtsoftplus(router_logits) gathered + renormalized. Stashes input_ids on self before the experts() call so the closure can index tid2eid; clears immediately after. - For hash layers: assigns experts.custom_routing_function = self._hash_topk in MoE.__init__ so FusedMoE picks it up via the moe_forward custom op → forward_impl_graph → quant_method.apply → select_experts plumbing. Verified: - PR3e (new): synthetic tid2eid → _hash_topk produces exact expected ids, renormalized weights match reference math (max_abs_diff = 0.0) - PR3e: FusedMoE.select_experts honors custom_routing_function correctly - PR1 toy parity: still 0.0 max diff (hash path is opt-in via is_hash_layer) - PR3b FusedMoE load: 32 params, all nonzero (no regression) - PR3d real attn forward: still works (non-hash layer)
… real ckpt Three changes converging on the first working V4 layer forward: 1. **weights_mapping**: Add class-level rename dict so the standard ATOM loader (`atom.model_loader.loader.load_model`) can ingest V4 ckpt names without per-model loader.py changes. `.gate.bias` → `.gate.e_score_correction_bias`, `.scale` → `.weight_scale_inv`. Loader's built-in `weight_scale_inv` → `weight_scale` rename then completes the path. Real serving via ModelRunner now works for non-wo_a layers. 2. **process_weights_after_loading hook**: After my custom `model.load_weights` finishes copying tensors, walk all submodules and call `quant_method.process_weights_after_loading(layer)` (or `layer.process_weights_after_loading()` if no quant_method). Without this, FusedMoE's `shuffle_weights` step is skipped and the FP4 ck_moe kernel reads stale weight layout — manifested as HSA_STATUS_ERROR_EXCEPTION mid-forward. Standard loader.py calls this for us; my custom loader had to replicate it. 3. **PR3f end-to-end test** (logs_claude/v4_pr3f_block_e2e.py): - Build 1 dense layer (compress_ratios=[0]) with 8 routed experts - Load real layer-3 weights (32 target params, 33/33 nonzero) - Build mHC residual `[8 tokens, hc_mult=4, dim=7168]` - Call Block.forward(x, start_pos=0, input_ids) - Output: shape preserved, range [-4.1, 4.6], abs mean 0.81, no NaN/Inf This is the first end-to-end forward through V4's full layer: attention (FP8 wq/wkv + BF16 wo grouped LoRA + indexer) + FusedMoE (FP4 experts via aiter ck_moe + sqrtsoftplus routing + bias correction + shared expert) + mHC pre/post Sinkhorn projections. Confirmed no regression on PR1/PR3b/PR3d/PR3e.
…kpts ModelRunner uses atom.model_loader.loader.load_model() — not the model's custom load_weights(). This commit closes that gap so real serving via openai_server works end-to-end: 1. **Expand weights_mapping with prefix renames**: V4 ckpt has bare names (`embed.`, `layers.`, `norm.`, `head.`, `hc_head_`) but our params live under `self.model = ...`. Add prefix substitutions so the loader's `model.get_parameter(name)` lookup hits the right attribute path. 2. **Fix dtype-mismatch silent zero in FusedMoE._load_w13/_load_w2**: PyTorch's `tensor.copy_()` between mismatched float8/uint8 dtypes silently writes zeros. V4's per-1x32 weight scales are stored as `float8_e8m0fnu` on disk but FusedMoE allocates them as `uint8` (raw byte storage). Force a `.view(torch.uint8)` reinterpret on the source so the bytes round-trip correctly. This is a pre-existing bug that was masked because V2/V3 use `float32` scales — V4 is the first ATOM model to use e8m0/e4m3 scales. Verified: - PR3i (new): standard load_model() loads V4 layer-0 from full 805GB ckpt index — 43/43 model params nonzero (100%), 5GB selective load. - PR3g (new): full Model.forward(input_ids) → logits on real ckpt. Output shape (1, 129280), range [-14.2, 15.4], std 3.05, no NaN/Inf. - PR3h (new): hash layer (layers 0/1/2) Block.forward works on real layer-0 ckpt (tid2eid loaded, 773423/775680 nonzero entries, real per-token expert assignments diverge from default sqrtsoftplus path). - All 5 prior tests (PR1/PR3b/PR3d/PR3e/PR3f) still pass — no regression. Net result: V4 inference pipeline is now production-ready for real ckpt loading + forward; remaining gap is multi-layer + multi-batch attn metadata + AITER sparse_attn (parallel work).
…hook PR3i shipped "100% nonzero params" but never ran forward through the standard-loader path. Verifying with PR3j (new) revealed wo_a values were 2768× too large — `torch.copy_(BF16_dst, FP8_src)` does an FP8→BF16 dtype conversion but SKIPS the per-128-block scale multiplication. Result: raw FP8 e4m3 max value (448.0) lands in the BF16 weight buffer instead of the true ~0.04 attention-init magnitude. Fix: stop forcing wo_a to no_spec/BF16 in V4QuantConfig. Let it allocate as FP8 ColumnParallelLinear so the standard FP8 loader fills both `wo_a.weight` (FP8) and `wo_a.weight_scale` (e8m0) correctly. Then DeepseekV4Attention.process_weights_after_loading dequants in place, replacing weight with BF16 + dropping the scale param. Forward continues to use BF16 weight in the grouped LoRA einsum (aiter has no FP8 grouped einsum). Also removes the manual wo_a special-case from custom load_weights() — both load paths (custom + standard) now converge through the same process_weights_after_loading dequant. Verified by PR3j parity test: - Custom path wo_a: abs.mean=0.0214, abs.max=0.4062 - Standard path wo_a: abs.mean=0.0214, abs.max=0.4062 (BIT-EXACT) - Standard-loader Model.forward → logits range [-17.9, 15.8], std 3.04 - Magnitude ratio: 1.00 (was 2768× before fix) - All 9 tests pass — no regression. This was a silent corruption that PR3i's "params nonzero" check missed. The lesson: nonzero != correct. Always verify with forward.
Major changes enabling correct V4 inference (single-prompt verified with 512-token coherent output in both English and Chinese): Model fixes: - WeightsMapper prefix-anchored remapping (fixes 381 silently-skipped params) - wo_a FP8→BF16 dequant with quant_type=No to prevent CK shuffle corruption - Hash routing (first 3 layers) now applies route_scale=2.5 - shared_experts reduce_results=False + unified all_reduce in MoE.forward - KV cache reset on start_pos=0 with score_state=-inf initialization - TP-correct head/group counts for Attention and Indexer MoE routing: - Standard Silu activation (not Swiglu — aiter a16w4+Swiglu has 9× amplitude loss on gfx950). swiglu_limit clamping done in triton post-kernel. - ATOM_USE_TRITON_MOE=1: triton matmul_ogs path with swiglu_limit clamp - ATOM_V4_TORCH_MOE=1: per-expert torch fallback with FP4 dequant (slow) - GFX950MXScaleLayout→CDNA4MXScaleLayout fix in fused_moe_triton.py Loader improvements: - WeightsMapper auto-read from model class attribute - Post-load WARNING listing all unloaded params - Shape-mismatch raises RuntimeError instead of silent skip Config: - deepseek_v4→deepseek_v3 registry mapping with V4 field re-injection - Robust from_hf_config with getattr defaults Known limitations: - Single-sequence only (kv_cache[:1,...] hardcoded); batch>1 needs PR3 - Multi-request KV isolation pending scheduler integration - TPOT ~213ms with --enforce-eager (no CUDAGraph)
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| # populated with extra V4 attrs (some fields may live only in the raw | ||
| # config_dict, not on the config object — `transformers` strips unknown | ||
| # kwargs unless they're in the schema). | ||
| g = lambda k, default=None: getattr(hf_config, k, default) |
Contributor
4 tasks
Oseltamivir
added a commit
to SemiAnalysisAI/InferenceX
that referenced
this pull request
Apr 26, 2026
dsv4-fp4-mi355x-atom (ROCm/ATOM#650 PR1, single-sequence at TP=8 with torch-fallback hc_pre because aiter mhc_pre crashes on this image) runs at ~5 min per request in steady state. With 1k1k at 12 prompts plus 8k1k at the same shape, the full sweep can exceed the 300-min cap that #1148 set for the SGLang-DSv4 path. Bump both the SLURM allocation in runners/launch_mi355x-amds.sh and the GitHub Actions timeout-minutes in benchmark-tmpl.yml together — either expiring first kills the job, so they need to stay aligned. Note: this is a global bump that affects every MI355X benchmark and every job that uses the shared workflow template, not just the dsv4 ATOM one. Drop back to 300 once the slow paths are gone (PR4 CUDAGraph + a working aiter MHC).
This was referenced Apr 26, 2026
Merged
…202) Upstream ref (deepseek-ai/DeepSeek-V4-Pro@a1fd202) changed shared_experts from no swiglu_limit to swiglu_limit=args.swiglu_limit, making it consistent with routed experts.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
DeepSeek-V4 PR1: model skeleton with end-to-end inference verified on real checkpoint (
/data/DeepSeek-V4-Pro, TP=8).ATOM_USE_TRITON_MOE=1) withswiglu_limit=10.0clampingload_modelwithWeightsMapper)Reproduce
Sample output (Chinese, 512 tokens, temperature=0)
Sample output (English, 512 tokens, temperature=0)
Bugs fixed in this PR
weights_mappingsubstring collision (381/2519 params silently skipped)WeightsMapperprefix-anchored remappingwo_aFP8 shuffle after BF16 dequant (attn output cos=-0.002)quant_type=Noto skip CK shuffleroute_scale(FFN output 5.2x too small)topk_weights *= routed_scaling_factorActivationType.Swiglucauses 9x amplitude loss on gfx950Silu+ triton post-kernel clampshared_experts.w2reduce_resultsmismatch with FusedMoEreduce_results=False+ unifiedall_reducestart_pos=0MoE paths
ATOM_USE_TRITON_MOE=1ATOM_V4_TORCH_MOE=1Known limitations
kv_cache[:1,...]hardcoded; batch>1 / multi-request needs PR3--enforce-eagerrequired: No CUDAGraph support yet (PR4)e_score_correction_bias(expected) + 43x MTP params (PR5)Files changed
atom/config.pyatom/models/deepseek_v4.pyatom/model_loader/loader.pyWeightsMapperauto-read + post-load unloaded-params WARNINGatom/model_ops/moe.pyATOM_USE_TRITON_MOEgate +swiglu_limitpassthroughatom/model_ops/fused_moe_triton.pyCDNA4MXScaleLayoutfix +swiglu_limitclampatom/model_ops/sparse_attn_v4.pydevice=for multi-GPUatom/examples/simple_inference.py--max-tokensTest plan
lm_evalGSM8K accuracy (blocked on multi-request KV isolation)