Coding with AI Agents to Supercharge Your Spark UDFs with GPU-Acceleration

Summary

NVIDIA Project Aether uses AI agents with four specialized skills—test generation, GPU conversion, benchmarking, and profiling—to convert Spark UDFs into GPU-compatible implementations using cuDF or CUDA primitives.
Customer results show a 22x reduction in operation time and a 3x wall-clock speedup, and the agent-based approach achieves a 100% success rate on UDF conversion versus 50% with direct prompting.
Agents validate each GPU conversion bite-for-bite against CPU behavior tests before benchmarking, and use NVIDIA Nsight in an optimization loop to prevent reward hacking.

Coding with AI Agents to Supercharge Your Spark UDFs with GPU-Acceleration

NVIDIA cuDF enables GPU-accelerated query execution for Apache Spark workloads. However, a common performance bottleneck for vectorized query engines are User-Defined Functions (UDFs). Since UDFs contain opaque custom code, the engine must fall back to row-at-a-time execution. We present a new set of agents skills in Project Aether to convert Spark UDFs into GPU-compatible implementations: either native SQL expressions or GPU UDFs invoking cuDF or CUDA primitives. The agent reasons through tasks in an end-to-end workflow: 1) generating CPU-behavior tests, 2) producing a GPU implementation, 3) running benchmarks, and 4) profiling and optimizing performance, with correctness verified through iterative feedback loops. By transforming UDFs to run on GPUs, agents enable full query acceleration for the cuDF plugin, RAPIDS Accelerator for Apache Spark, unlocking significant 20x+ UDF speedups. In this video we will go into the practical learnings from building an AI agent for Spark.
Talk By: Felix Cheung, Product, NVIDIA ; Rishi Chandra, Systems Software Engineer, NVIDIA ;

Chapters

FAQs

What is NVIDIA Project Aether?

Project Aether is an open-source framework from NVIDIA that uses AI agents to convert Apache Spark UDFs into GPU-accelerated implementations using cuDF Java APIs or CUDA kernels. It follows a four-step workflow: test generation, GPU conversion, benchmarking, and profiling with NVIDIA Nsight.

Why do Spark UDFs prevent GPU acceleration?

Spark UDFs contain opaque custom code, so vectorized query engines like the RAPIDS Accelerator for Apache Spark must fall back to row-at-a-time execution. This bottleneck prevents full GPU acceleration of the query even when the rest of the pipeline runs on GPU.

What speedups are possible when converting Spark UDFs to GPU?

Customer results show a 22x reduction in operation time and a 3x wall-clock speedup when UDFs are converted to GPU-compatible implementations using Project Aether. The RAPIDS Accelerator for Apache Spark enables 20x-plus speedups on converted UDFs more broadly.

How does the agent achieve better conversion rates than direct prompting?

Project Aether achieves 100% UDF conversion success versus 50% with direct LLM prompting by using structured agent skills with clearly defined rewards and a restricted action space. Each conversion is validated against CPU behavior tests before acceptance, preventing incorrect GPU implementations.

Full transcript

[00:11] So here's the session. Um, so I'm Felix. Today we're going to talk about AI agent rewriting a Spark job and then you get benefit of GPU acceleration. We're going to talk about that. So here's a breakdown. Um, at a very high level,
[00:24] we're going to talk a little bit more about this, you know, GPU spark thing, right? And then we'll go into how UDF could accelerate it and why you want to do that. And then the more important part is actually go through more of the
[00:36] details of uh how do you actually get an AI agent to work with your code? What is the best practices? What are the learnings we have? So I think hopefully that will be a good takeaway for you.
[00:50] Well, before we go in there, this is Nvidia QDF. Uh I'm Felix. I'm from Nvidia. What we're doing here is we're accelerating uh structured data processing with GPU. So before I get started with this kind
[01:03] of premise of it, can I ask for a quick show of hand? How many of you know about GPU accelerator spark in here? Not a whole lot, maybe 20%. That's great. Uh how many of you actually use
[01:18] agent with your spark job today or maybe writing one? Okay, a little bit you know about the same but yeah I think that's the area we want to focus on and talk to you about
[01:30] actually today. So um problem with data processing uh hopefully this is not news to you. I think people have been saying that for a long time. Data is growing at a enormous pace every single year,
[01:42] right? I think every single year will come in and say and tell you data is growing at enormous pace. Okay, it is the case. But what we're seeing more with agents is that agents going to be
[01:54] asking more questions and they demands faster respond. like they won't even like I'm gonna wait until you know five minutes you know get this query get processed and then I'm going to do more
[02:06] stuff with it and that's just not the pace that is acceptable with a gentic sort of workflow. So what we're really seeing here is that with agents we expect analytics being a bigger much
[02:19] bigger problem very soon. That's kind of what we're saying here on the u your left hand side or your right hand side sorry. So what are we doing at Nvidia? Um I'm happy to say that we've been working
[02:32] through this with many many partners and many many platform out there. We're accelerating pretty much everything you can think of. Uh we have the foundation like I said the Nvidia QDF which is cool
[02:45] structure data processing uh data frame I guess it's called. uh it integrates with all these different open source library whether is spark polers presto ray po uh duct db pandas I see you over
[03:00] there uh and all these different uh platform partners uh are integrating this as well and specifically today we're going to talk about spark as you know spark is
[03:13] very popular I'm sure all of you uses spark every single day uh what we're trying to do here with GPU is that we're accelerating your Spark job requiring no code change, right? And then we're able
[03:26] to deliver not only faster result, but add a better TCO. So what that means is that if your job run faster, it takes less resource to run, right? You save
[03:38] your money, you get the same result, and then you get a result faster. So that's kind of the gist of it. If there's one takeaway if you want to remember GPU GPU spark, what this does is faster to
[03:52] respond and better TCO. And so that's kind of the gist of what this GPU spark is about. And you can pretty much get it anywhere you want and then it would just work and then you know you're welcome to try it out. We
[04:04] actually have a bunch of things uh to support you there. But the reason why I think this is the time for it right now for all of you to think about is that we
[04:16] are really seeing this turning to mainstream like not only GPU is everywhere their partners talk about how they adopting GPU for data analytics services and platform right we also have
[04:28] use case like uh story like snap a few months ago they talk about how they were able to use this technology and apply buy it you know now in production actually getting the benefit
[04:41] of it you know the cost saving and so on so forth so if you're interested into it um unfortunately I think the captions override but if you do a search on the snap engineing blog you're welcome to look into it there's a lot more details
[04:54] in terms of the learnings they have how it apply and so on and so forth um but the quick recap here is that they were able to save 76% of the cost with this
[05:06] technology Right. So, I think that is like a phenomenal sort of result. Um, like I said, the GTC talk that they gave a few months ago has a recording. You're welcome to click through it. There's a
[05:18] deck blog that you can actually go through with more details as well. So, that's where we are. Uh, without further ado, let me introduce Rishi, my colleague here, and he was going to talk more about
[05:35] Okay. Yeah. So I'm going to start talking about UDFs and how to accelerate them on GPUs and you know why we want to do that in the first place. So I think everyone here is pretty familiar but just to cover our bases you
[05:48] know UDFs these userdefined functions are these custom routines to extend built-in spark functionality right usually to do some sort of business logic or some bespoke transformations on your data and the challenge is that
[06:01] firstly you know they're used everywhere in modern pipelines and problematically they are essentially a black box to optimizers and accelerators and as we'll see in more detail in a moment this can
[06:13] make them a really common performance bottleneck, especially when you're trying to vectorize your job. So, as a quick sneak peek of why we want to GPU accelerate UDFs, this is an
[06:26] example from a customer that I'll get into more later. It's a large consumer technology company that was able to GPU accelerate UDF that was affecting thousands of jobs and get a 22x savings
[06:38] in op time for this particular UDF. and you know in one case saw 3x wall clock speed up for their job. >>
[06:50] know, a simple example, I have a UDF that does some sort of custom business logic. It's not, you know, the details are not really important, but it's some combination of some conditional logic to bucket users into rates, do some
[07:02] arithmetic to figure out discounts basically. Um, and by default, if you run this UDF with the Spark QDF plugin that Felix was talking about, or really any vectorized engine for that matter,
[07:14] you end up incurring overhead that looks something like this, right? Most of your job, whether it's SQL or dataf frame, is vectorzed and then you have these islands of fallback around your UDF
[07:26] where you pay this roundtrip data conversion cost to go from, you know, columner format to rows and then roundtrip back. Um, and not to mention in this case you also have the data moving costs to go across the PCI
[07:38] Express bus from GPU to CPU and back to run your UDF. Um,
[07:50] instance, you know, the conditional logic can be performed through, you know, vectorz mask and a select to figure out your rates. From there, we just do a sequence of maps to get to your result. Um, and this UDF, it's a
[08:02] pretty simple example, but many much more complex EDFs can be performantly implemented on the GPU with the right primitives. And QDF and the Spark QDF plug-in have done a lot of the hard work
[08:14] to expose those performant uh GPU primitives in a standard recognizable dataf frame API. So ideally if we're able to write a UDF using some GPU accelerated API surface
[08:28] which you know I'll get into the options in a moment we can avoid that fallback entirely keep the full pipeline on GPU so we both avoid both that data overhead and get better compute performance for
[08:40] the function itself. Okay. So what does rewriting the UDF look like? So I'll talk about these three ways starting from you know more abstraction to more fine grain control.
[08:55] Okay. So the most straightforward and probably most recognizable option is just rewriting a UDF using some built-in SQL operators. Right? So for instance the s simple example can be rewritten using your standard case statements and
[09:08] arithmetic in SQL. And under the hood, this will go through Spark's logical plan and the QDF plugin will intercept at the physical plan layer and turn each of these SQL operations into a GPU
[09:20] equivalent that's powered by QDF. And so this will work for most simple functions that do some basic relational logic. Um the main drawbacks are that for one sometimes finer grain logic like you
[09:33] know string parsing or dealing with nested types is possible but really uh far less performant than it could be if you had more control as a programmer. Um and another drawback is that this typically requires a change to your
[09:46] application code you know which makes it a little more invasive especially if the UDF is being reused in a lot of places.
[10:00] UDF using QDF APIs. So the spark QDF plugin provides an API interface with a single method that you can override called evaluate columner. So instead of
[10:12] taking raw Java objects, we're taking and returning GPU column vectors. So you know in this method body we're essentially implementing a columner or vectorized analog of the row-based CPU
[10:25] logic using the QDF Java APIs. And in terms of the plan once again the plugin will intercept things at the physical plan layer and it will recognize that the UDF extends our
[10:38] interface and hand it columns directly on the GPU without any data transfer. Um and so the benefits is that you know with this programming uh paradigm we get a lot more control than pure SQL. You
[10:50] know we have access to a much richer set of lower level dataf frame APIs. Um and we also preserve the encapsulation without necessarily having to touch the application code.
[11:05] Now in cases where you need to go even lower than that um the third conversion target is rewriting the UDF using custom cruer kernels. So just like before we would override the same interface but instead of using KUDF Java APIs we can
[11:18] delegate some or all of that logic to a native CUDA kernel using the Java native interface. So in the second box I have a CUDA C++ kernel that's doing you know
[11:30] most of the work business logic. Um and at the physical plan level this looks the exact same as before. our GPU udf is getting invoked but internally it's dispatching to native code. So for our
[11:42] very simple UDF, this is obviously complete overkill. Um, and you know, it does require more complexity by having to dip into C++. But there are cases where you have, you know, much more
[11:54] complex processing in a single function, maybe doing a ton of compute. And in such cases, fusing a lot of that logic into one specialized kernel can give you the maximum performance.
[12:10] Okay, so to recap, we have these three approaches that range from sort of higher uh to lower levels of abstraction. And generally the decision of what you want to use will make sense based on how complex your UDF is starting from, you know, simple relational operators to increasingly
[12:24] complex logic and types. Okay, so we want to make this conversion process easy and reliable. Um but you know the main challenges of course is
[12:36] first of all GPU programming um is non-trivial. Um and the second thing is in data engineering you know as was mentioned throughout the keynotes you know not only is the code critical but data correctness is critical right
[12:48] changes need to be carefully audited and so fortunately for us um most frontier coding models you know your claude or GPT are actually quite good at
[13:00] GPU programming and with the right set of workflows we can approach pretty rigorous levels of verification and drastically reduce um developer hours for this conversion process.
[13:16] And so that's what we're introducing today. Um it's a set of agent skills that help with the full workflow for converting a UDF for GPU acceleration. So you know that doesn't just include the implementation, but also triaging,
[13:29] testing, validation, proving out speedups in benchmarks. So it's the full endto-end development process to get this from your CPU implementation to a validated GPU implementation that you
[13:41] can actually start to integrate in a production pipeline.
[13:53] around a four-phase workflow. So you know we start with generating tests, converting the UDF to GPU, benchmarking the performance and finally optimizing the implementation. Um and each phase
[14:07] itself is a feedback loop with its own verification. And I'll get into these one by one. Okay, so the first skill is generating unit tests. Um you know the main purpose is to essentially capture a
[14:19] specification of the CPU EDF. And I'll briefly walk through these steps. Um, essentially we start by analyzing the input domain of the EDF. You know, the types, but more specifically the actual
[14:32] domain of valid and invalid inputs that the function handles. We'll generate a bunch of test cases for those input domains along with edge cases as a Spark data frame.
[14:44] We'll put that in a unit test with assertions and run that against the UDF fixing any errors. And then finally we'll generate a code coverage report to ensure that we're capturing all the relevant branches and lines and again feeding back if the coverage is too low.
[15:00] So once we have that spec defined we go to our second skill which is essentially you know a test-driven development to convert the UDF to GPU using the unit test as the ground truth.
[15:13] So the agent will start by searching the API docs, you know, in this case looking for the right KUDF APIs to use, for instance. Um, it'll implement the KUDF version of the UDF using the QDF Java
[15:25] APIs. Um, we'll run that against the unit test and we'll do a bite for bite comparison of the CPU and GPU outputs feeding back if there are any errors. And then finally we have an independent
[15:39] judge sub agent check for variance of cheating to make sure the implementation is general and legitimate. Okay. So the third phase once we have a proper GPU implementation is
[15:52] benchmarking to prove out speedups. So these steps look like um
[16:04] oops we'll start by generating a large synthetic data set using spark. We'll take a sample of that and run it through both UDFs as a sanity check to fix any schema or data errors. Um and then we'll run the full benchmark using
[16:17] both CPU and GPU in a minimal spark job. And finally, we'll compare both the end-to-end wall clock run times as well as finer grain metrics like op times
[16:29] through the event logs to assess how much the change will move the needle in a production job. Okay. And the fourth workflow phase which is optional if we want to further
[16:41] improve the performance is to optimize the implementation. So what this looks like is we point an NSIS profiler at our benchmark which will give us things like the time spent in different GPU kernels,
[16:53] memory allocations, QDF APIs. From there, the agent can interpret and make a diff to address the hotspots. And for correctness, we'll rerun the
[17:05] unit test to check for any regressions. And if there are regressions, we'll revert that diff to the last working snapshot. And then finally, we'll rerun the benchmarks to assess whether performance
[17:17] actually improved and reverting the changes if it didn't. And we'll repeat this process in an optimization loop. Okay, so those are the four phases of the workflow encapsulated in these four
[17:30] skills that we provide. Okay, so next I want to talk a little bit about some of the design decisions and learnings that came out of building
[17:42] these skills. So I've summarized some of the ways we thought about designing the skills into these three points on the screen. Um the first is defining clear rewards. The
[17:54] second is restricting the action space. And the third is preventing reward hacking. So I'll explain what I mean in a moment. Um these are not meant to be, you know, scientifically rigorous um design decisions. I think you'll find
[18:06] that the ideas are really, you know, pretty obvious and pretty intuitive. But uh I think it may be informative to those of you who are using agents building skills to share how we thought about these in the context of the UDF
[18:18] problem. So the first on defining clear rewards. So the goal is essentially to steer the agent in this open-ended objective. Um and agents are stochastic. You know
[18:31] their outputs have high variance and we want to leverage that variance when it's beneficial. Um so I see that as most helpful when the task is essentially to search a solution space that is large and input dependent right so some
[18:44] examples are you know writing test cases for a specific UDF designing a correct GP algorithm and finding optimizations um and the problem is you know even though we have the CPU UDF as our ground
[18:56] truth if we don't define some sort of proxy feedback for these tasks there are still these clear failure modes because the objective is underspecified So you know some of the concrete failure
[19:09] modes that we've observed are you know models will generate these very plausible looking test cases and edge cases but will still miss concrete branches of the UDF. Um and secondly
[19:22] even if the model generates a good set of test cases the way it enforces that spec is weak. like it'll have these vague assertions that just check for the contents of a string.
[19:48] coverage as a useful metric to encourage, you know, better test cases as well as having this adversarial reviewer that is explicitly looking for uncovered inputs and delivering a pass
[20:00] or fail rating. And for weak assertions, you know, we provide a test harness that enforces this bite forbite comparison between the CPU and the GPU outputs. So essentially we still try to give the
[20:12] agent flexibility to come up with a solution but we just try to improve the reward signals for the specific task.
[20:26] all about trying to limit variance in the areas where it's not helpful. Um so those are areas where the solution space to me is sort of small or input independent and the error consequences are non-trivial. you know, essentially
[20:39] we really care about consistency. So, some examples are build configurations, versioning, um, you know, benchmark configurations, how are you doing your speed up measurement? Um, and to give
[20:51] some concrete failure modes, you know, some things we've seen are, you know, first of all, misconfigured builds where the agent just spends, you know, half of its trajectory just aligning dependencies, fixing CUDA errors. Um,
[21:03] and another is variable, uh, benchmark configs. So you know one example is it turns out the choice of GPU memory allocator you use can vary benchmark runtimes by an order of magnitude and
[21:16] agents just you know aren't aware of that domain knowledge. So essentially small deviations in these areas just make the workflow painfully inconsistent. You can get different speedups across different iterations
[21:28] etc. And to resolve this we basically try to give a frozen environment to the agent right. So we give it a pre-provided palm file with a validated configuration and
[21:40] a benchmark harness with how to do the measurements, all the best practices and consistent configs.
[21:52] Um the last principle is on uh preventing reward hacking. So sometimes there are just competing objectives. So 100% compliance with the CPU is sometimes either impossible or just not worth the performance on GPU. Um and
[22:07] overprescribing these kinds of strict objectives are where we've seen the most cheating. Um so for example there are just you know fundamental CPU and GPU differences in floatingpoint math. Sometimes there are library calls that
[22:19] depend on the runtime. Um sometimes there are pathological edge cases that just might be out of scope for your data pipeline. >>
[22:32] is in the GPU implementation, the model will cheat by silently copying rows one by one back to the CPU, invoking the original function and then copying it back to GPU. Um, and we've also seen the
[22:44] model, you know, write these switch statements that essentially hardcode handling of different test cases. And to combat this, you know, firstly, as I mentioned, we have this independent
[22:57] judge sub agent that explicitly checks for these kinds of cheating. And we also just try to make it easier for the model to admit deviations than to cheat on the implementation. You know, so we allow the agent to bypass certain edge cases
[23:11] if it essentially leaves an audit trail for a user to decide whether a certain edge case or input is in scope for the specific application.
[23:25] So with all that said, this is what we end up bundling with our skills. You know, it's a lot. We essentially provide a full project template to work with with all the harnesses for testing, benchmarking setup as shown on the left
[23:37] along with references, encoding, you know, the best practices and examples of conversion that are shown on the right.
[23:49] So if we peek at the unit test for example, you'll see what this enables is we give the agent flexibility in things like test design with these open to-dos and we also provide a lot of the scaffolding to do strict comparisons using you know helper methods that we
[24:02] predefined to do bite forbite comparisons ensure that the UDF actually runs on GPU and things like that.
[24:17] Okay. So finally I want to talk a little bit about what the skills look like in practice. Okay. Starting with some evaluation. So you know a natural question is does using the skills versus just directly
[24:30] prompting an LLM actually help? Right. So these are some evaluation results we did with a simple AB test with and without the skills on a set of six real customer UDFs. um using claude and
[24:44] you'll see that without the skills first first of all we see a lot more cheating you know 50% of the UDF silently fall back to the CPU we see a lot more variance in the build setup you know different versions across iterations and
[24:56] there are just a lot more iterations uh figuring out the environment version conflicts and just a lot more total errors in general and
[25:08] and as we'd expect by giving the agent this frozen environment. There's one consistent build. There's far fewer iterations spent on these auxiliary errors and we end up with better
[25:20] speedups across the board. Okay, so here's an anecdote from a customer that I briefly shared at the beginning of the talk. Um, so they had
[25:32] this CDF that was being used in a ton of workloads that they converted to KUDF. So on this on the screen is an example of one of these big workloads that was processing 16 terabytes of data. And you
[25:44] can see the op time of that particular UDF decreased from 275 task summed hours to just 13 hours. And so the net effect was unblocking
[25:57] that fallback for over 3,000 of their jobs. um improving you know the op time of that UDF by over 20x and by removing the fallback you know speeding up the wall clock of this overall job by around
[26:10] 3x. Okay, so to finish um with what you should take away from this talk, you know, first of all, UDFs, they can be a
[26:22] performance challenge for vectorzed engines and moving them to GPU can unlock, you know, big speedups and total cost optimization. Um, and coding agents along with skills can make this process
[26:36] easier and more reliable. Um, and skills are easy to install and use. They can be pulled directly from our open-source repo. Um, and we see agents as a big enabler for developers to work with
[26:49] GPUs, reach for higher performance, and we see these skills as a step in that direction in the data engineering space.
[27:01] So with that, I want to thank you all for listening and direct you to some links where you can find the skills and generally learn more about data processing with KUDA.

Learn more about the Databricks Data and AI platform.

The information provided herein is for general informational purposes only and may not reflect the most current product capabilities or configurations.