[FlyDSL] DeepSeek-V4.1-Flash mono decode layer (WIP)

[Fangzhou-Ai]

Summary

Work in progress: an ATOM-style mono decode layer for DeepSeek-V4.1-Flash on MI355X (gfx950), written in FlyDSL for vLLM to import. One decoder layer runs in two persistent launches:

Launch	Stages
K1	the mHC attention seam, then the attention front: wqkv, q/kv norms, the KV insert, wq_b
K2	the attention back, the attention TP all-reduce (in-kernel, over symmetric peer memory), the FFN seam, then the MoE: router, top-6, shared expert, MXFP4 experts, and its own in-kernel all-reduce

The launches replace roughly 40 kernels per layer on vLLM's ROCm path. Both are drop-ins for DeepseekV4DecoderLayer.forward (same tensors at the layer boundary).

Scope:

• Layers: the 30 backbone layers with a standard seam and no compressor or indexer: 3–7, 9–13, 15–19, 21–23, 25–27, 29–31, 33–35, 37–39. Layers 0, 1 and the compressor/indexer layers keep vLLM's path.
• Steps: decode only, M ≤ 48 rows (8 requests × 6 DSpark tokens), TP2 and TP4.

Opened as a draft for early review of the structure. Same-node baselines are still running; see Status.

Layout

aiter/ops/flydsl/dsv41_mono.py            host runner: DSV41MonoLayer (scratch, epoch, peer memory, launches)
aiter/ops/flydsl/kernels/dsv41_mono/
    layer.py      the two launches (K1, K2) composed from the stages below
    attention/    attention stages (new): front.py (K1), back.py (K2), gemv.py, device.py, plan.py
    stages/       seam.py (mHC seams + attention TP reduce), moe.py, moe_shape.py, moe_plan.py, router.py, gemv.py, dims.py
    common/       mailbox, device ops, build keys, peer memory
    README.md     design, numerics, validation

stages/ and common/ are adapted from ROCm/ATOM (MIT License; commit 97359d6, atom/models/deepseek_v41/mono and atom/mono). Every adapted file names its origin in its header. The changes from ATOM:

• The seams are specialized to this layer.
• The MoE follows vLLM's numerics:
  • router logits stay f32;
  • the shared expert reads vLLM's row-major FP8 weights;
  • mxfp8_e4m3_quantize and silu_and_mul_with_clamp match vLLM.
• Hand-offs carry epoch tags instead of per-step mailbox clears, and peer regions alternate by epoch parity, so no step fence is needed.
• ATOM's engine-side, attention, indexer and debug modules are dropped.

Design notes

• Execution model: 256 CTAs × 512 threads, all resident; each stage is a task table. Hand-offs use device-scope tagged pairs or plain data plus flags.
• No host step hook: the kernels move their own epoch, so they are safe to capture in graphs.
• Weights: read in vLLM's loaded layouts, so there is no second copy. Attention uses row-major MXFP8 with 32×32 E8M0 scales; routed experts use AITER-shuffled MXFP4.
• All-reduces are exact (fp32 sum in rank order, then bf16). vLLM's INT4 quick reduce only applies at 8 MB and above (16 MB at TP4), so vLLM's own decode all-reduces at M ≤ 48 are exact too. An INT4 codec option is planned.

Test results (MI355X, real checkpoint weights)

The reference is vLLM's own modules in a TP process group: DeepseekV4MoE, the mHC seam ops, the ROCm attention chain, and the all-reduce. The harness lives outside this repo for now. Full matrix: https://claude.ai/artifact/Duic63s8CfzydY2zhKRdWp

Kernel-level performance, µs per layer. Each run chains 8 layers of one type back to back in one HIP graph, with cold weights. vLLM's path is captured under its graph_capture, so its all-reduces use the graph-registered custom all-reduce. 96 points: TP2/TP4 × concurrency 1–8 × 1–6 tokens per request × 2 layer types. Every point is faster.

Ratio	TP	Speedup (min–max)	Median	DSpark e2e points c1 · c2 · c4 · c8 (vLLM → mono µs)
1	TP2	1.11–2.04×	1.57×	209→122 · 262→179 · 336→280 · 474→427
2	TP2	1.10–2.05×	1.57×	210→123 · 255→172 · 325→269 · 446→406
1	TP4	1.15–2.18×	1.82×	180→94 · 218→123 · 256→192 · 324→282
2	TP4	1.17–2.18×	1.83×	180→95 · 216→121 · 248→185 · 316→270

At M=48 both paths are bound by expert-weight reads: about 200 distinct experts × 8.8 MB per layer at TP2.

Kernel-level accuracy. Layer 0 of every point, both paths from the same inputs and caches:

• Residual cosine ≥ 0.99999.
• MoE output: 1,234 of 1,260 tokens at cosine ≥ 0.999. Each of the remaining 26 is a top-6 routing near-tie: the experts differ between the two paths' FFN inputs, and the reference's own 6th/7th expert score gap is ≤ 3.2e-3, against medians near 1e-2.
• On the same FFN input, the mono MoE matches vLLM's at cosine 0.99999.

Model eval (vLLM integration, separate PR; TP2, the InferenceX#3690 recipe with DSpark block rejection). InferenceX's GSM8K task, 5-shot, full 1,319, temperature 0, concurrency 8:

	strict-match	flexible-extract
mono	0.9750 ± 0.0043	0.9742 ± 0.0044
vLLM	0.9742 ± 0.0044	0.9735 ± 0.0044

DSpark mean acceptance length over that eval: 3.911 with mono vs 3.949 baseline. Single-stream greedy decode, 1,024 tokens: 273.5 vs 208.9 tok/s (1.31×).

AgentX e2e, TP2 (InferenceX#3690 recipe, golden acceptance 3.51, 3600 s per point). The baseline is InferenceX's own run of the same image and recipe (run 37510490434):

conc	p90 interactivity, tok/s/user	total throughput, tok/s
c1	302.6 → 429.1 (1.42×)	22,331 → 25,039 (1.12×)
c2	285.4 → 381.7 (1.34×)	22,800 → 25,695 (1.13×)
c4	267.8 → 332.4 (1.24×)	33,017 → 35,885 (1.09×)
c8	203.9 → 245.3 (1.20×)	58,385 → 60,501 (1.04×)

AgentX e2e, TP4 (same baseline run; ¹ scored from aiperf's complete summary, see below):

conc	p90 interactivity, tok/s/user	total throughput, tok/s
c1 ¹	343.7 → 523.7 (1.52×)	23,691 → 26,796 (1.13×)
c2 ¹	327.4 → 471.1 (1.44×)	23,579 → 27,217 (1.15×)
c4	307.2 → 445.1 (1.45×)	34,831 → 36,791 (1.06×)
c8	265.8 → 358.6 (1.35×)	60,266 → 62,621 (1.04×)

¹ The shared home volume briefly hit its quota mid-run, which left holes in these two points' per-request exports. aiperf's summary covers every request, and on the intact points it reproduces InferenceX's p90 interactivity (1 / p90 ITL) and total throughput ((input + output tokens) / duration) exactly.

TTFT p90 was higher in our runs (0.47–0.81 → 1.08–1.28 s). The mono path does not touch prefill. We ran the four points at once on one node, which shares CPU tokenization of the ~300K-token prompts and the host-resident TP2 Engram tables; InferenceX runs each point alone. A same-node baseline will follow.

Fixes found by these tests

• The router keeps f32 logits, as vLLM's gate does. ATOM's bf16 rounding flipped top-6 picks near ties.
• At TP4 the routed intermediate (576 columns) is padded to 640 by vLLM's loader. The down stage's last 128-wide MFMA step read 64 never-written columns. ATOM zeroes its scratch every step; the epoch scheme does not. Those columns are now masked in registers.

Status / TODO

• Whole-layer kernel benchmark and accuracy sweep (above)
• INT4 quick-reduce codec option for the in-kernel all-reduces
• Self-contained op test (op_tests/)
• vLLM integration (separate PR, opt-in VLLM_ROCM_DSV41_MONO=1) and GSM8K
• AgentX e2e TP2 c1/c2/c4/c8 (above)
• AgentX e2e TP4 c1/c2/c4/c8 (above)
• Same-node mono-off baselines, TP2 and TP4 (running)

AI assistance

This PR was written with substantial AI assistance (Claude Code). I am reviewing it line by line before it leaves draft. Kernel-level results above come from the harness described; AgentX e2e results will be added before review.

🤖 Generated with Claude Code

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 6196 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

[Fangzhou-Ai]

Superseded by vllm-project/vllm#60397