A Learning Method with Gap-Aware Generation for Heterogeneous DAG Scheduling

Ruisong Zhou; Haijun Zou; Li Zhou; Chumin Sun; Zaiwen Wen

cs.LGMar 24, 2026

A Learning Method with Gap-Aware Generation for Heterogeneous DAG Scheduling

Authors: Ruisong Zhou, Haijun Zou, Li Zhou, Chumin Sun, Zaiwen Wen

Organizations: School of Mathematical Science, Peking University, Beijing 100871, China · State Key Laboratory of Mathematical Sciences, Institute of Computational Mathematics and Scientific/Engineering Computing, Academy of Mathematics and Systems Science, Chinese Academy of Sciences, Beijing 100190, China · Theory Lab, Central Research Institute, 2012 Labs, Huawei Technologies Co., Ltd, China · Beijing International Center for Mathematical Research, Peking University, Beijing 100871, China

Abstract

Efficient scheduling of directed acyclic graphs (DAGs) is a core problem in large-scale data-intensive computing systems, where query plans, data-processing workloads, and computation graphs consist of dependent tasks competing for limited heterogeneous resource pools. In practice, achieving high-performance execution requires schedulers to adapt across environments with varying resource pools and task types, while generating schedules under tight runtime budgets. We propose WeCAN, an end-to-end reinforcement learning framework for heterogeneous DAG scheduling that addresses task-pool compatibility coefficients and generation-induced optimality gaps. It adopts a two-stage single-pass design: a single forward pass produces task-pool scores and global parameters, followed by a generation map that constructs schedules without repeated network calls. Its weighted cross-attention encoder models task-pool interactions gated by compatibility coefficients, and is size-agnostic to environment fluctuations. Moreover, widely used list-scheduling maps can incur generation-induced optimality gaps from restricted reachability. We introduce an order-space analysis that characterizes the reachable set of generation maps via feasible schedule orders, explains the mechanism behind generation-induced gaps, and yields sufficient conditions for gap elimination. Guided by these conditions, we design a skip-extended realization with an analytically parameterized decreasing skip rule, which enlarges the reachable order set while preserving single-pass efficiency. Experiments on real-world TPC-H query DAGs, resource-intensive workload datasets, and ML-compiler computation graphs demonstrate improved makespan over strong baselines, with inference time comparable to classical heuristics and faster than multi-round neural schedulers.

Figures & tables

Fig. 1: The illustration for spaces and maps. (a) The schedule space, divided into disjoint fibers for feasible schedule orders. (b) The serial schedule generation scheme $S_{\text{SGS}}$ embeds the feasible schedule order space $\mathcal{B}_{f}$ as a subset of the schedule space $\mathcal{X}$ , which is “orthogonal” to each fiber $T^{-1}(w)$ as only intersect at one point $S_{\text{SGS}}(w)$ . $S_{\text{SGS}}$ maps each schedule order to a best schedule (with minimal objective value) in the corresponding fiber. This relationship fix the optimality gaps. (c) The list scheduling map $S_{\text{list}}$ maps each schedule order to a best schedule in some fiber, but the fiber may not be the corresponding fiber. Therefore, the fiber containing the optimal solution may be missed since the optimal schedule order may be mapped to the schedule with other schedule orders. This mismatch leads to an optimality gap. (d) The skip action map $S_{d}$ is defined on a larger space $\overline{\Omega}_{\text{action}}\supset\Omega_{\text{action}}$ , where the extended part $S_{d}(\overline{\Omega}_{\text{action}}\backslash\Omega_{\text{action}})$ covers the missed part $S_{\text{SGS}}(\mathcal{B}_{f})\backslash S_{\text{list}}(\Omega_{\text{action}})$ , making $TS_{d}$ a surjection to $\mathcal{B}_{f}$ . The extended part closes the optimality gap.

Fig. 2: A counterexample showing list scheduling excludes the optimal solution for the DAG scheduling problem with heavy tasks. (a) Problem instance $\mathcal{P}_{0}$ with only one pool offering one type of resource with capacity 3. The resource demand and processing time is shown and the compatibility coefficients $K\equiv 1$ . Each arrow $v\to w$ corresponds to an edge $(v,w)\in E$ . (b) The only solution $x_{0}$ in the image of list scheduling $S_{\text{list}}(\Omega_{\text{action}})=\{x_{0}\}$ . The horizontal axis represents the time while the vertical axis represents the resource usage. (c) Optimal solution $x^{*}$ of the problem. (d) Worse solution for unrestricted waiting, leading to a significantly large objective value (makespan).

Fig. 3: The two-stage framework of our architecture. The first stage involves once network processing by our efficient and adaptive weighted cross-attention architecture. It embeds the problem attribution and provides scores and coefficients for the second stage. The second stage employs a generation map extended from list scheduling. This extension restores surjectivity onto the schedule order space and ensures the inclusion of an optimal solution in the image of the generation map, without significantly slowing down training convergence.

	TPC-H-30, 3 pools			TPC-H-50, 3 pools			TPC-H-100, 3 pools
	Makespan	Improvement	Time	Makespan	Improvement	Time	Makespan	Improvement	Time
SFT	27404	-28.12%	0.23	49172	-32.31%	0.78	84986	-24.27%	3.08
MOPNR	25052	-17.13%	0.30	43545	-17.17%	0.99	77362	-13.12%	3.34
CP	23869	-11.59%	0.29	41597	-11.93%	0.90	74364	-8.74%	3.35
HEFT	23177	-8.36%	0.18	39315	-5.79%	0.54	70137	-2.56%	1.86
Tetris	23170	-8.33%	0.21	38654	-4.01%	0.62	71296	-4.25%	2.13

TABLE I: Experimental results on TPC-H datasets.

	Erdős-Rényi			Layer Graphs			Stochastic Block
	Makespan	Improvement	Time	Makespan	Improvement	Time	Makespan	Improvement	Time
SFT	13317	-19.99%	0.81	16158	-30.04%	0.34	14408	-27.96%	0.53
MOPNR	12771	-15.07%	1.07	14714	-18.42%	0.38	13148	-16.77%	0.68
Tetris	13084	-17.90%	0.52	14271	-14.86%	0.44	13666	-21.37%	0.64
CP	12457	-12.25%	1.08	14797	-19.09%	0.40	13388	-18.90%	0.74
PEFT	11311	-1.92%	0.45	12425	0.00%	0.71	11438	-1.58%	0.56

TABLE II: Experimental results on computation graph datasets with 500 tasks.

	RI-W-30			RI-W-50
	Makespan	Improvement	Time	Makespan	Improvement	Time
Tetris	29324	-14.11%	0.20	44296	-16.46%	0.56
CP	28144	-9.52%	0.26	41378	-8.79%	0.75
HEFT	26856	-4.51%	0.18	39825	-4.71%	0.52
PEFT	28260	-9.97%	0.13	39928	-4.98%	0.48
IPPTS	25697	0.00%	0.21	38034	0.00%	0.51

TABLE III: Full results on RI-W datasets.

Fig. 4: Ablation results on the skip design. (a-b) Results of different skip action score formulations on RI-W-30/50 datasets. “NoSkip”, “ConstSkip” and “DecrSkip” correspond to the variants without skip action, with constant-score skip action and with our decreasing-score skip action, respectively. Error bars denote standard deviation. (c) Relative improvement from enabling skip under different heavy-task ratios.

Fig. 5: For the counterexample instance $P_{0}$ , a constant-score skip extension still misses the optimum and cannot consistently eliminate the optimality gap.

	TPC-H-30		TPC-H-50
Architecture	Makespan	Improvement	Makespan	Improvement
IPPTS	21389	0.0%	37164	0.0%
WeCA + LDDGNN	19908	6.92%	34260	7.81%
WeCA-inside + LDDGNN	20729	3.09%	34980	5.88%
WeCA-alt + LDDGNN	20234	5.40%	34815	6.32%
WeCA-inside-alt + LDDGNN	21981	-2.77%	36984	0.48%

TABLE IV: Ablation study results on TPC-H datasets with different architectures.

Fig. 6: Training results with different generation maps and their explanation by toy illustration. (a) Training loss on the TPC-H-30 dataset, serial schedule generation scheme $S_{\text{SGS}}$ exhibits significantly slower convergence and inferior performance compared to our generation map $S_{d}$ . (b) The toy illustration for landscape by $S_{\text{SGS}}$ . The flat basins inside black lines represent the image of list scheduling. The extension introduces extra regions containing the optimal solution and the numerous extra mountains corresponding to worse solutions. This irregular mixed structure slows down the convergence. (c) The toy illustration for landscape by $S_{d}$ . The extended spaces brings worse solutions but the worst solutions are mostly far away in the axis of $\alpha$ while the optimal solutions are usually near the boarder line. This regular structure stabilizes and accelerates the convergence.

Fig. 7: Evaluations of models on TPC-H with different environment fluctuations under fixed training conditions. The percent of improvement over best heuristics are labeled.

	3 pools						12 pools
	TPC-H-150			TPC-H-200			TPC-H-150			TPC-H-200
	Makespan	Improvement	Time	Makespan	Improvement	Time	Makespan	Improvement	Time	Makespan	Improvement	Time
Tetris	109079	-2.54%	3.65	140438	-2.52%	6.09	30552	-16.54%	4.80	37737	-11.18%	8.51
HEFT	108285	-1.80%	3.26	139094	-1.53%	4.98	27415	-4.58%	3.11	34834	-2.63%	5.21
PEFT	108490	-1.99%	3.05	140154	-2.31%	5.04	27943	-6.59%	3.60	35539	-4.71%	4.98
IPPTS	106372	0.00%	2.74	136992	0.00%	4.80	26215	0.00%	5.75	33942	0.00%	4.88

TABLE V: Full generalization results on large-scale TPC-H datasets

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

Setting	Formulation	Variables	Constraints
Homogeneous	Problem ( 9 )	$n$	$\|E\|+O(n)$
Homogeneous	MILP ( 10 )	$3n^{2}$	$6n^{2}$
Heterogeneous	Problem ( 9 )	$2n$	$\|E\|+O(n)$
Heterogeneous	MILP ( 11 )	$(3+n_{c})n^{2}$	$(6+2n_{c})n^{2}$

Appendix

TABLE VI: Complexity comparison of problem formulations

Fig. 8: Local search with projection. (a) The optimal schedule is excluded by the list scheduling map $S_{\text{list}}$ due to a different schedule order. (b) Local search $LS$ on the schedule order $T(S_{\text{list}}(\omega))$ , followed by the serial schedule generation scheme $S_{\text{SGS}}$ , recovers the optimal schedule $S_{n}(LS(T(S_{\text{list}}(\omega))))$ .

	TPC-H-30		TPC-H-50		TPC-H-100
	Makespan	Time	Makespan	Time	Makespan	Time
WeCAN-Greedy(NAR)	19578	0.15	33428	0.50	62587	1.72
WeCAN-S(256)(NAR)	18964 $\pm$ 10	2.43	32814 $\pm$ 47	2.86	61373 $\pm$ 28	10.43
WeCAN-AR-Greedy	19386	1.96	33786	3.39	(out of GPU-memory)
WeCAN-AR-S(256)	18910 $\pm$ 26	3.57	33201 $\pm$ 75	7.05	(out of GPU-memory)

Appendix

TABLE VII: Performance comparison on TPC-H dataset.

	Action set	Variables	Constraints	Nonzeros
Homogeneous	$n$	$3n^{2}$	$6n^{2}$	$19n^{2}$
TPC-H-150, 1 pool	1500	$6.8\times 10^{6}$	$1.4\times 10^{7}$	$4.3\times 10^{7}$
TPC-H-300, 1 pool	3000	$2.7\times 10^{7}$	$5.4\times 10^{7}$	$1.7\times 10^{8}$
Heterogeneous	$nn_{c}$	$(3+n_{c})n^{2}$	$(6+2n_{c})n^{2}$	$(18+8n_{c}+rn_{c})n^{2}$
Computation-Graph, 3 pool	1500	$1.5\times 10^{6}$	$3\times 10^{6}$	$1.2\times 10^{7}$
TPC-H-100, 3 pool	3000	$6.0\times 10^{6}$	$1.2\times 10^{7}$	$4.8\times 10^{7}$

Appendix

TABLE VIII: Problem size of Benchmarks

	TPC-H-30, 3 pools		TPC-H-50, 3 pools		TPC-H-100, 3 pools
	Makespan	Improvement	Makespan	Improvement	Makespan	Improvement
SFT	25556	-26.15%	43289	-19.12%	77494	-18.64%
MOPNR	23080	-13.93%	40724	-12.06%	71166	-8.95%
Tetris	21889	-8.05%	37682	-3.69%	67923	-3.99%
CP	22431	-10.73%	39041	-7.43%	68278	-4.53%
HEFT	22622	-11.67%	38643	-6.34%	67872	-3.91%

Appendix

TABLE IX: Experimental results on TPC-H datasets on fluctuating environment (a), which corresponds to the “more resource pool (A)” in Figure 7.

	TPC-H-30, 3 pools		TPC-H-50, 3 pools		TPC-H-100, 3 pools
	Makespan	Improvement	Makespan	Improvement	Makespan	Improvement
SFT	23741	-29.89%	41476	-33.97%	74061	-27.30%
MOPNR	21869	-19.65%	38129	-23.16%	67597	-16.19%
Tetris	20000	-9.42%	32237	-4.13%	58796	-1.06%
CP	21089	-15.38%	34598	-11.75%	62874	-8.07%
HEFT	18624	-1.89%	31866	-2.93%	58531	-0.60%

Appendix

TABLE X: Experimental results on TPC-H datasets on fluctuating environment (b), which corresponds to the “more resource pool (B)” in Figure 7.

	TPC-H-30, 3 pools		TPC-H-50, 3 pools		TPC-H-100, 3 pools
	Makespan	Improvement	Makespan	Improvement	Makespan	Improvement
SFT	26626	-28.43%	44708	-25.14%	80460	-17.13%
MOPNR	24792	-19.58%	41589	-16.41%	75168	-9.43%
Tetris	22829	-10.12%	36691	-2.70%	68863	-0.25%
CP	22519	-8.58%	38745	-8.45%	72774	-5.94%
HEFT	20732	0.00%	35727	0.00%	68996	-0.44%

Appendix

TABLE XI: Experimental results on TPC-H datasets on fluctuating environment (c), which corresponds to the “more pool types (A)” in Figure 7.

	TPC-H-30, 3 pools		TPC-H-50, 3 pools		TPC-H-100, 3 pools
	Makespan	Improvement	Makespan	Improvement	Makespan	Improvement
SFT	27303	-27.61%	47400	-27.93%	84195	-24.71%
MOPNR	24964	-16.68%	43010	-16.08%	77059	-14.14%
Tetris	23203	-8.45%	38885	-4.95%	71034	-5.21%
CP	23129	-8.10%	41465	-11.91%	73481	-8.84%
HEFT	22416	-4.77%	38648	-4.31%	69443	-2.86%

Appendix

TABLE XII: Experimental results on TPC-H datasets on fluctuating environment (d), which corresponds to the “more pool types (B)” in Figure 7.

	TPC-H-30, 3 pools		TPC-H-50, 3 pools		TPC-H-100, 3 pools
	Makespan	Improvement	Makespan	Improvement	Makespan	Improvement
SFT	27254	-26.86%	49040	-27.47%	83922	-24.88%
MOPNR	24809	-15.47%	43135	-12.12%	77623	-15.51%
Tetris	23629	-9.98%	38484	-0.03%	71358	-6.19%
CP	23488	-9.33%	41564	-8.04%	72609	-8.05%
HEFT	22768	-5.97%	40136	-4.33%	68243	-1.55%

Appendix

TABLE XIII: Experimental results on TPC-H datasets on fluctuating environment (e), which corresponds to the “more task types (A)” in Figure 7.

	TPC-H-30, 3 pools		TPC-H-50, 3 pools		TPC-H-100, 3 pools
	Makespan	Improvement	Makespan	Improvement	Makespan	Improvement
SFT	29738	-18.8%	50118	-27.54%	94002	-21.59%
MOPNR	28402	-13.46%	45473	-15.72%	86306	-11.63%
Tetris	25983	-3.80%	40322	-2.61%	78031	-0.93%
CP	26427	-5.57%	43760	-11.36%	85981	-11.21%
HEFT	26707	-6.69%	41042	-4.44%	83114	-7.50%

Appendix

TABLE XIV: Experimental results on TPC-H datasets on fluctuating environment (f), which corresponds to the “more task types (B)” in Figure 7.

	TPC-H-30, 3 pools		TPC-H-50, 3 pools		TPC-H-100, 3 pools
	Makespan	Improvement	Makespan	Improvement	Makespan	Improvement
SFT	20842	-29.98%	36076	-29.92%	63677	-21.21%
MOPNR	18417	-17.58%	32614	-17.45%	57660	-9.75%
Tetris	18417	-14.86%	30438	-9.62%	54054	-2.89%
CP	17790	-10.95%	29502	-6.24%	53009	-0.90%
HEFT	17360	-8.26%	29010	-4.47%	52935	-0.76%

Appendix

TABLE XV: Experimental results on TPC-H datasets on fluctuating environment (g), which corresponds to the “Larger resource-pool capacity” in Figure 7.

	TPC-H-30, 3 pools		TPC-H-50, 3 pools		TPC-H-100, 3 pools
	Makespan	Improvement	Makespan	Improvement	Makespan	Improvement
SFT	33138	-28.00%	54852	-18.54%	100653	-17.54%
MOPNR	29829	-15.22%	50931	-10.07%	92989	-8.59%
Tetris	29179	-12.71%	48891	-5.66%	91940	-7.36%
CP	27760	-7.23%	46985	-1.54%	86498	-1.01%
HEFT	28119	-8.61%	47602	-2.87%	88638	-3.50%

Appendix

TABLE XVI: Experimental results on TPC-H datasets on fluctuating environment (h), which corresponds to the “Smaller resource-pool capacity” in Figure 7.

	TPC-H-50		TPC-H-100
	Makespan	Improvement	Makespan	Improvement
SFT	49172	-27.2%	84986	-21.2%
MOPNR	43545	-12.7%	74364	-10.3%
CP	41597	-7.62%	74364	-6.03%
Tetris	38654	0.00%	71269	-1.65%
HEFT	36326	-1.71%	70137	0.00%

Appendix

TABLE XVII: Experimental results on TPC-H datasets. OneShot-Train30 and WeCAN-Train30 refer to models trained in TPC-H-30 while OneShot and WeCAN refer to model trained in problems with corresponding size (TPC-H-50 or TPC-H-100)

IPPTS	25697	0.00%	0.21	38034	0.00%	0.51
	RI-W-30			RI-W-50
	Makespan	Improvement	Time	Makespan	Improvement	Time
WeCAN-NoSkip-Greedy	27138	-5.61%	0.17	39457	-3.74%	0.47
WeCAN-NoSkip-S(256)	26159	-1.80%	2.38	38485	-1.19%	5.37
WeCAN-NoSkip-S(256)+LocalSearch(50)	24946	2.92%	92.80	37726	0.81%	275.57
WeCAN-NoSkip-S(256)+LocalSearch(100)	24793	3.52%	192.56	37519	1.35%	526.22

Appendix

TABLE XVIII: Experimental results on RI-W datasets.

	GPU-inference			CPU-inference
Experiment	Total	Generation map	Network	Total	Generation map	Network
TPC-H-30	0.154	0.140	0.014	0.201	0.116	0.085
TPC-H-50	0.480	0.461	0.019	0.587	0.442	0.145
TPC-H-100	1.72	1.71	0.017	1.92	1.53	0.391
Erdős-Rényi graph	0.533	0.512	0.020	0.548	0.393	0.155
Layer Graphs	0.256	0.238	0.018	0.372	0.214	0.158

Appendix

TABLE XIX: Experimental results about inference time (s).

Experiment	Total time	Time per batch
TPC-H-30	6.6h	0.5min
TPC-H-50	14.5h	1.1min
TPC-H-100	26h	2min
Erdős-Rényi graph	19h	1.4 min
Layer Graphs	15.5h	1.2 min
Stochastic Block	13.3h	1 min

Appendix

TABLE XX: Experimental results about training time of WeCAN-Greedy.

Explore similar work

MoSAIC: Learning DAG Scheduling Heuristics from Optimal Schedules through Motif-Aware Recovery and LLM-Guided Reasoning
Oct 10, 2026Shiva Ahir, Hashmath Shaik, Alex Doboli +2Automated Heuristic Design Learning to Optimize
Department of Electrical & Computer Engineering, Stony Brook University, New York, United States of America · Barcelona Supercomputing Center, Barcelona, Catalonia, Spain · Department of Computer Science, Hofstra University, New York, United States of America
Cloud Workflow Scheduling Based on Graph Attention-Driven Hierarchical Reinforcement Learning
Sep 14, 2026Zongjin Li, Shaohan Feng, Chunxi Yang +1Real-Time Scheduling Hierarchical RL
Faculty of Mechanical and Electrical Engineering, Kunming University of Science and Technology, Kunming 650500, China · School of Information and Electronic Engineering (Sussex Artificial Intelligence Institute), Zhejiang Gongshang University, Hangzhou 310018, China
PPO-STGNN: A Proximal Policy Optimization Approach with Spatio-Temporal Graph Neural Networks for DAG Task Scheduling in Cloud-Edge-End Computing
Sep 3, 2026Yangshuo Qi, Chenwei Wang, Zihan Shen +1Edge-Cloud Computing
Beijing University of Posts and Telecommunications, Beijing, China · Google, Mountain View, USA · NanJing University, NanJing, China

Oct 10, 2026J/K move · Enter open · S save

MoSAIC: Learning DAG Scheduling Heuristics from Optimal Schedules through Motif-Aware Recovery and LLM-Guided Reasoning

Shiva Ahir, Hashmath Shaik, Alex Doboli, Nadeem Yaseen+1

Department of Electrical & Computer Engineering, Stony Brook University, New York, United States of America · Barcelona Supercomputing Center, Barcelona, Catalonia, Spain · Department of Computer Science, Hofstra University, New York, United States of America

Automated Heuristic Design Learning to Optimize Heterogeneous Computing

Open paper AskPDF arXiv

Efficient task scheduling in heterogeneous computing systems remains a fundamental challenge in embedded and high-level synthesis flows. This paper presents MoSAIC (Motif-aware Scheduling with AI-driven Cognition), a learning-based framework that bridges the gap between exact and heuristic scheduling approaches for DAG-based applications. The system uses the reasoning capabilities of large language models (LLMs) to learn the relationships between DAG structures and scheduling priority functions. The framework first leverages the OR-Tools CP-SAT solver [4] to generate optimal schedules, which are then used to supervise the learning of compact priority-based scheduling functions through Bayesian Optimization. These learned heuristics capture relationships between scheduling decisions and structural graph features such as rank, depth, and fanout, enabling the construction of a reusable knowledge base across diverse DAG instances. Experimental results across synthetic graph families and real-world benchmarks, including PolyBench and MachSuite, show that the learned heuristics consistently improve over classical approaches such as HEFT [2], while maintaining low runtime overhead. The results indicate that combining optimal supervision with learning-based generalization provides a practical approach for scalable and adaptive task scheduling.

Abstract

Explore similar work

MoSAIC: Learning DAG Scheduling Heuristics from Optimal Schedules through Motif-Aware Recovery and LLM-Guided Reasoning

Cloud Workflow Scheduling Based on Graph Attention-Driven Hierarchical Reinforcement Learning

PPO-STGNN: A Proximal Policy Optimization Approach with Spatio-Temporal Graph Neural Networks for DAG Task Scheduling in Cloud-Edge-End Computing