AI-Assisted Software Engineering

Looking back, how did we do additions with computers? A while ago we always wrote assembly, either by hand or using a C compiler, which is one instruction. Today we still do some of that, but we’re also using Python a lot, where operands likely sit on the heap, and the interpreter will run much more than one instruction for the addition itself. Apparently, addition has gotten so much inefficient!

However, the reason we use Python is that it is much easier to write code in those languages than in assembly, especially for complex things due to the higher-level abstractions they provide. Due to that, when a human looks at the code, it’s also much easier to understand and debug. So, the cost of doing the simplest thing (e.g., addition) can be much higher in those languages, but the net cost of doing complex things can be much lower (especially considering the time to ship new things).

Here, a very important property that enables this higher-level abstraction to be useful is it being not leaky. When we write code in Python, we generally don’t worry about how the interpreter is going to boil down our class hierarchy and method resolution into machine code; we only reason at the higher-level API. With the level of abstraction raised, the world of software engineering was able to produce more complex systems and generally do more cool stuff.

The way I view it, AI-assisted software engineering is a similar lift. In 2020, if you wanted to create a new static website for your project, you would have to set up a web framework, write a bunch of boilerplate code, and write the actual content, and still no one could figure out how to center a div. Or, if you find yourself repeating a similar command in your shell, you would briefly consider automating it with a shell script or a CLI tool, until you realize it’s unclear whether the amount of time you’ll spend on writing that thing is less than the amount of time you’ll save by using it.¹ AI-assisted software engineering is a lift in that, the time it takes to build many things that used to take a lot of human time and effort is now significantly reduced; you can ask for a website with a paragraph of natural langauge and get it working the next moment, or you can paste in part of your ~/.zsh_history and get a working CLI tool in minutes.

As of today, however, these AI-assisted tools are not completely reliable. In other words, they are still leaky; especially for tools that matter, or those with potentially large blast radius, you still have to review the code it generates. But in the limit of AI-assisted software engineering, it will probably approach non-leakiness.

Systems Today are Polymorphic

Now let’s switch to system building. Why do we build general systems in the first place? I believe that it is in large part because writing code has been expensive; a human had to write and maintain every line. Much of software engineering practice, like libraries, frameworks, code sharing, and dynamic dispatch, exists so that we write less code; to have concise representations of our intent in code in higher-level abstractions.

ML systems is no exception. One codebase runs every model, at every precision, on every accelerator, so that nobody has to write a system per workload. Identical PyTorch code running on both NVIDIA GPUs and AMD GPUs is a beautiful example of this. Things like PyTorch’s dispatcher allows you write one operator call (e.g., torch.mm), and the right kernel among the many registered ones is dispatched at runtime. This leads to not only runtime overhead² but also code complexity. Different models, precisions, and GPU types route to different backends through layers of dynamic dispatch, and reading the code tells you the list of what could run, and figuring out what actually runs in your environment takes effort.

This isn’t merely an issue of PyTorch of course. vLLM’s feature compatibility matrix shows that it’s trying to cram in all these features into a single codebase and some combinations are not supported, which is sometimes not a fundamental incompatibility but rather due to the difficulty of maintaining a single codebase that supports all combinations. vLLM’s hybrid KV cache manager reads each layer type’s attention spec (e.g., full attention, mamba, MLA) on startup, computes how many bytes its state needs per token, and adjusts block sizes until every layer type has the same bytes-per-page. This is complex runtime resolution logic; it is hard to understand, maintain, keep bug-free, generalize to new model architectures that come out every week, and keep it working for all existing models.

I would call the current systems we’re building polymorphized systems: a single system codebase that is polymorphic over many workloads, and the polymorphism is resolved at runtime.

Monomorphized Systems

I think most of this generalization is a compromise forced on us by the cost of producing code. However, if AI-assisted software engineering makes writing code fast, we ought to reconsider this compromise. We can couple the systems layer tightly with the application layer and build each system monomorphized for one workload, borrowing the compiler term for stamping out a specialized copy of generic code for each concrete type it is used with.

In such a monomorphized system, every choice that a general system defers to runtime dispatch or resolution is instead resolved and baked in at build time. Places that used to be polymorphic will be replaced by an interface (e.g., prefill attention) and a list of functionally interchangeable implementations (e.g., FlashAttention, FlashInfer, CuDNN), and exactly one of those implementations will be chosen and baked into the system (i.e., the interface will be monomorphized) at build time. Similarly, configurations are hard-coded, statically known shapes and precisions are baked in, and code paths that do not apply to the deployment are never generated in the first place.

If this is doable, there are many immediate benefits. Smaller binary size. Minimal dependencies (and surface for supply chain attacks for at least that specific monomorphization). No runtime dispatch overhead or resolution/warmup time on startup or first-run. Trivial to read and step through code for understanding, debugging, and profiling. No worries about code changes of one part of the system unintentionally breaking some completely unrelated part of the system due to a hidden coupling, because that unrelated part of the system doesn’t even exist in your monomorphized system.

Continuously Monomorphized Systems

But the world does not stay still; production systems sit on top of components that change at high velocity. Key dependencies like kernel libraries release constantly (e.g., FlashAttention is the fastest long context prefill attention kernel for GB300 one day, and FlashInfer may take over tomorrow), and hardware architectures, drivers, and CUDA versions keep shifting underneath as well. On top of that, the target workload the system was built for can itself change: the model, the precision, the request pattern, prefix cache hit rate, etc.

This is why monomorphization must be continuous, which I call continuously monomorphized systems. Each monomorphized system declares its system component dependencies and a workload signature. A centralized benchmarking authority (e.g., one per organization) continuously evaluates the performance of system components whenever a trigger occurs, such as a new release of a kernel library or a change in a workload pattern that corresponds to a workload signature. When the authority detects that, for a certain workload signature, the performance ranking among interchangeable components compatible with an interface has changed, it will signal the downstream monomorphized system to trigger re-monomorphization. Centralization is important here; this kind of continuous benchmarking is costly, but it only needs to happen once, rather than being repeated by every downstream system.

Why? If you’re aiming to get a PhD admission from a US university, interning at a research lab in a US university concretely boosts your chances. You will get to know professors there, and if you do a good job, they may try to hire you as a PhD student themselves. At the same time, you will get a strong letter of recommendation from them, which is crucial for PhD admissions. You see, nearly all letters say that the applicant is very good. What matters is who said that, and whether the reader of the letter knows the letter writer well enough to weigh in on their opinion. A professor at a US university is more likely to be known by the letter readers at other US universities.

I never quite imagined that interning abroad was possible when I was a prospective PhD student in SNU. However, I am now a fifth year PhD student at the University of Michigan CSE, and I can tell you for sure that there are many cases where international students intern abroad at labs in US universities. I especially see Chinese students (particularly those from elite classes like Yao Class, Turing class, ACM class, etc.) doing this. We’re always short on good interns; there are tons of engineering work (research and open-source) to be done, and also more project ideas than good people to work on them.

So, it’s worth it, and it’s possible. Now, is it easy? Not really. You need to clearly demonstrate to the professor that you are worth the trouble of sponsoring a J-1 visiting scholar visa for you. Being from a university that is well-known in the US (e.g., SNU, KAIST) definitely helps, and having a CV & GitHub profile that demonstrate your research potential, engineering skill, and craftsmanship is crucial. Establishing communications with PhD students in the lab beforehand should also help.

You will be traveling to the US for a few months, likely alone in a foreign country with unfamiliar language, culture, and work environment. You will have to keep yourself alive there and handle unfamiliar paperwork on your own, while also scrambling to do a good job at work. It’s a great challenge, but also a great opportunity.

I treat AI-assisted programming or writing as an instance of speculative decoding. For those unfamiliar with the term, speculative decoding is a technique that speeds up LLM text generation by using a smaller, faster model (the draft model) to generate multiple tokens, and then having the original, larger model (the verifier model) accept or reject those tokens through rejection sampling. The crucial property of speculative decoding is that once we go through rejection sampling, the resulting accepted tokens exactly follow the distribution of the original larger model; there is no quality loss.

When I use an AI coding assistant (e.g., Copilot, Cursor) to produce code, for instance, the AI model is the draft model, and I (the human programmer) am the verifier model that accepts or rejects AI-generated code based on whether they work correctly and fit my intent. And I make sure that the code that ultimately gets committed follows the distribution of code that I would have written myself.

I put on this mindset not only when I write text, but also when I review others’ text. I assume that whatever text has been put forward for review follows the distribution of text the author would have written themselves, and I don’t care whether they used AI assistance or not; it’s an implementation detail that should be hidden behind the abstraction of the author. That’s why I get annoyed when the answer to my review comment is “Oh, I used AI to generate that part.” It shows that the author doesn’t have a proper explicit intent for that part of the text, signaling that they didn’t verify thoroughly and is shifting accountability to the AI model. This makes authorship a leaky abstraction and breaks the assumption that the text follows the author’s distribution.

This is why I think whatever has been put forward by a human author should be attributed only to the author, and whether AI assistance was used or not is irrelevant. If the text is great, the author should get credit for it; if it’s bad, the author should take accountability for it. I believe this mindset will help us reap the productivity benefits of AI assistance (with increasing draft acceptance rate as AI models get better) while maintaining quality and accountability.

I’ll go over a couple tricks with Matplotlib that I’ve found useful, especially for research.

Making plots compatible with PowerPoint

PNG files exported with Matplotlib don’t look very nice in PowerPoint slides, especially when you enlarge it. Also, you can’t, say, change the color of your lines or remove unnecessary components, because they’re all flattened into an image file.

Instead, export your plot in SVG as well. All you need to do is call fig.savefig once more with a file path that ends with .svg. Then, drag the SVG file onto PowerPoint, right click, and click Convert to Shape. Now, every component (e.g., line, text) will be converted to PowerPoint-native objects that you can tweak freely. See also Microsoft’s documentation on this.

Before you actually call savefig, make sure you have the following configuration:

import matplotlib as mpl
mpl.rcParams["svg.fonttype"] = "none"

Without this, every character in your plot will be saved as a separate text box. That’s beyond annoying.

Making plot files Git-friendly

Let’s say you’re exporting your plot in PDF or SVG and committing it into a Git repository for paper-writing. Normally, if you re-export the exact same plot (e.g., because you had to restart and re-run every Jupyter Notebook cell), git will recognize them as modified. This is because by default, the contents of the PDF/SVG file will change slightly. This is particularly bad for PDF files, because they are binaries – git will have to store the whole PDF file because it doesn’t know how to do line-level diff on binary files. However, it is possible to make Matplotlib generate the exact same file each time deterministically.

PDF files

fig.savefig("plot.pdf", metadata={"CreationDate": None})

Notice that we’re removing the file creation date from the PDF metadata.

SVG files

fig.savefig("plot.svg", metadata={"Date": None})

The savefig call is basically the same for SVG, except that it needs a slightly different metadata key for the creation date. However, we do need one more thing before we call savefig:

import matplotlib as mpl
mpl.rcParams["svg.hashsalt"] = "42"

Without fixing the hash salt, SVG generation is non-deterministic and git will (rightfully) detect the file as modified.

Avoiding Type 3 fonts

This may not apply to all publishers, but at least USENIX and ACM papers require final PDFs to not have any Type 3 fonts. They just want Type 1 fonts. So, when we export our plots into PDF files and embed them into our LaTeX document, it’s better to make sure they don’t contain any Type 3 fonts in the first place.

import matplotlib as mpl
mpl.rcParams["pdf.fonttype"] = "42"
mpl.rcParams["ps.fonttype"] = "42"

Unlike svg.hashsalt where 42 was just a random number, 42 for the font type actually means TrueType fonts. The default is 3, which means Type 3 fonts.

An mplstyle file

You will definitely forget all of these next time. Instead of trying to remember, put everything in a style file like this (call it something like paper.mplstyle):

# Fonts
font.size : 9

# SVG export now doesn't render font as path, but saves fonts as text objects.
# This allows for easier integration with MS Office.
svg.fonttype : none
# Make SVG generation deterministic
svg.hashsalt : 42

# Avoid type 3 font usage
pdf.fonttype : 42
ps.fonttype : 42

Then, import the file with a one-liner:

import matplotlib.pyplot as plt
plt.style.use("./paper.mplstyle")

I think it’s because it gives me a sense of scoring extra points in a game. Alas, when I played the video I have decided to spend some time chilling instead of working, but if I actually make progress in work during that bulk of time that was anyway going to be wasted, I feel like going beyond a set goal. And that feels good.

However, the YouTube trick is not always useful because I can’t really accomplish challenging tasks that require full focus with a YouTube video on. That’s how I came to think of Structured Slacking, designed to exploit this fake sense of extra accomplishment.

It works like this. Sometimes, you don’t want to do anything today. Then think to yourself, “Okay, then let’s actually not do anything today. After all, I can use some rest from time to time. But, I’ll just do this very small thing today. It’s like an extra point I can score on a day off to feel good about myself.” Then suddenly you’re motivated to do that small thing. And once you get that started, you’ll probably continue working. That said, it’s also important to set up that tiny task to be the tip of the iceberg of a bigger set of tasks.

As many would guess from the name, Structured Slacking was also inspired by Structured Procrastination. The two are similar in that both aim to enhance the productivity of an individual and operate on top of self-deception; Deep inside your heart you know Structured Slacking is bullshit, but you manage to trick yourself to almost believe that you will take today off, and you will actually stop working after accomplishing that small task. The difference is that Structured Procrastination haunts you with a pending important task that you’re procrastinating, which is honestly not the best feeling. On the other hand, with Structured Slacking, you’re pretty chill. Of course you are – you’re taking the day off!

That’s why type stubs exist. They have .pyi extensions and are like C headers, which only declare class, function, or method names and their parameter types without implementation. These type stubs do not have to be coupled with the actual library. So virtually anyone can just create type stubs for an existing library and ship it separately.

Microsoft’s pylance ships with type stubs bundled for popular libraries without native type annotations. But especially for large libraries type stubs cannot be perfect, so they have a repository called python-type-stubs to collaboratively work on creating type stubs together with the community, and these stubs are bundled together with pylance.

However, pylance is closed source, and is only available inside VS Code. As a Neovim person, I instead have to use the open-source version of pylance, which is pyright. However, by default, pyright doesn’t ship with pylance’s type stubs.

So the question is, how do I use python-type-stubs with pyright? It’s actually simple enough, but at the time of writing, it seems like nowhere on the Internet just has a straightforward guide on this.

Say you have a Python project proj managed with git.

Add python-type-stubs as a git submodule under the directory stubs:

$ cd proj
# Assuming you have GitHub SSH authentication set up.
$ git submodule add git@github.com:microsoft/python-type-stubs stubs

Then, point pyright to the stubs inside the submodule.

If you’re using pyproject.toml:

[tool.pyright]
stubPath = "./stubs/stubs"

If you’re not using pyproject.toml, you need to have pyrightconfig.json in the root of your workspace:

{
    "stubPath": "./stubs/stubs"
}

When you see glitches in the type stubs provided by python-type-stubs, just post a PR fixing the issue. When the PR gets merged, update the submodule (e.g., git submodule update).

GPT is Stateless

GPT will remember the details of the conversation in the same thread, but that’s only because they cram in the entire conversation in their context window. Outside that window, they are basically stateless; when you start a new conversation with GPT, it’ll have an empty context and won’t remember anything about your previous conversations.

Advisors are often also stateless. Not that they’re stateless intentionally or by design, but due to the sheer amount of things going on around them, it’s more convenient for students to assume them to be stateless and forget everything. It’s just like how you don’t ask ChatGPT, “Hey I think I asked you something about NP-Hardness in another thread last month, do you remember that?”

Context is Important

That’s why initializing GPT’s context right is important. You would have experienced some conversations with ChatGPT where you screwed up the initial description of your problem, and it takes more words to correct ChatGPT’s understanding than what would have taken if you had described it right in the first place. In such cases you can just mutter “Crap.” and click ‘New Chat’ (because ChatGPT is stateless). However, unfortunately, that’s not so easy if you were talking to your advisor. Therefore, I try to make sure my advisor’s context is initialized with a concise and accurate picture of where my research is.

I think this is especially observable when I sometimes hear contradicting advice from my advisor. Not contradicting with my opinion, but with his own advice in the past. That’s probably because the contexts I gave to my advisor that led to those contradictory advices were inconsistent in some way. Therefore, usually my next action is to prompt my advisor further to find out if there are any misunderstandings, either in the previous meeting or this one.

Fine-Tuning GPT

I’ve been saying all along that advisors are stateless, but we all know that they’re not completely stateless all the time. So, I like to think that they take one fine-tuning step at the end of every meeting, and their learning rate depends on how interesting the meeting was (and also on other things that I can’t control). If I excite my advisor with some interesting observation or good result, they’re more likely to remember. Otherwise, they probably won’t remember what happened during the last meeting.

In that sense, I think it’s an effective strategy to present my advisor with a concise summary at the end of the meeting. That way they don’t have to summarize the entire meeting on their own for fine-tuning, but rather just directly use the takeaway messages I present. For that, I also try to set forth a couple TODO bullets that are rooted on the core takeaways of this meeting and roughly represent what my advisor can expect for the next meeting.

As a mentor, my task was to collaboratively design solutions with my mentee, answer technical questions, provide feedback on design decisions and code, and figure out what pieces of knowledge the student was missing and either provide study material or come up with Google search keywords. While I wasn’t doing any actual work, in that I wasn’t studying relevant technology or writing and testing code myself, the sequence of tasks was very difficult to perform satisfactorily because of the very fact that I wasn’t doing any actual work. That is, without a complete picture of what’s going on, I was supposed to have better foresight than my mentee about what would happen if we were to proceed in a certain direction, or my mentee risked hitting a hard wall. Eventually, when we officially merged the feature into Zeus, I felt very happy and proud that our collaboration worked out.

While the mentorship itself did not specifically push my ongoing research project forward, it indirectly helped with doing research, because I learned two important things from this experience.

First, I came to respect the weight of the advice my advisor gives me and understand the mental pressure of providing advice. My mentee and I were simply developing a moderately-sized software feature, and I could reasonably predict what would happen along the way and what things would look like in the end. Still, sometimes I wasn’t entirely sure about the advice I was giving to my mentee, but I anyway needed to at least look confident in order to give faith and motivation to my mentee. Moreover, it was obvious that if I ended up pointing my mentee to a wrong path, my mentee would face frustration.

What makes it more difficult for my advisor is that we are doing research, which is inherently uncertain; you never know if the hole you’re digging is your grave. Yet, PhD students expect their advisors to still provide advice that is roughly in the right direction. I have come to understand that not even the best professors can know with confidence that the way the student is headed is a good direction, and how incredibly good my advisor is in that he usually provides advice that ends up being correct.

Second, mentoring helped me ask better questions to my advisor. Recalling my mentoring experience, when my mentee asked for clarification about my suggestion, I sometimes failed to give good answers because I myself was operating based on logical guesses and inexplicable instincts. I figured this would be the case for my advisor, too. Thus, instead of interrogating my advisor with a bunch of why-questions, I started to instead ask questions such as “Is it because of A and B that you asked me to do X?”, essentially putting forth my guess of the reason behind his advice. Then, my advisor would answer whether or not he agrees with what I said, or sometimes even make better suggestions when the guess I provided sparked something, and both cases led to highly productive conversations.

All in all, before this mentorship experience, my perception of mentoring was that it’s only good for me if it ends with some nice tangible outcome in a reasonable amount of time. However, now I feel like the process of mentorship taught me a lot about the relationship between me and my advisor, and allowed me to improve the productivity of the meetings with my advisor. Moreover, mentoring itself was an extremely rewarding process, where I could interact with enthusiastic junior researchers.

After a good three-week break, I came back to research in late May, exploring potential future directions. I suggested some ideas and directions during meetings with my advisor, and basically everything was killed. They were not just killed, but killed with very good reasons that seemed obvious in retrospect. Quickly in several weeks, I started to feel very unconfident to the point where I was afraid to talk with my advisor, although I knew logically that killing bad ideas in early stages is only saving me time.

Then suddenly in mid-July, we were notified that Zeus was accepted to appear at NSDI. My mood swung to the other end of confidence, especially because this was our first submission of the paper to any conference. I began to prepare the camera-ready version, together with Zeus’s open source repository and homepage. Finding new research directions was paused for a week or two, since we wanted Zeus to be posted on arXiv as soon as possible and have it collect citations from early on.

While polishing Zeus, I was thinking about what could be a follow-up work of it. Then I came up with another idea, an idea I liked no less than all the dead ideas in June, and presented that to my advisor. And he said it was very good. Now I’m working on that idea, and just found last week that my hypothesis was true at least for a limited set of workloads, and the potential gains can be quite large.

Looking back starting from April, when I submitted Zeus to NSDI, me as a researcher did not change that much. Most of the time I was on vacation, or was pouring time into polishing Zeus. However, my level of confidence in terms of research capability fluctuated greatly, which didn’t make any sense. If me in April and me today are similar researchers, there is no reason to be sometimes confident and sometimes not. Rather, my confidence level should be determined by my best times, which is currently mid-July, when Zeus was accepted. External factors, for example whether my ideas are well accepted by people, do not change my potential to do great work.

The pizza served at Mani’s is literally the best in the world. They’re so good.

What is the most appropriate answer?

1. Hmm, how do you quantify the goodness of a restaurant? Is Google Maps stars a sufficient metric?
2. No, “good” does not mean “the best”. You should always think about the exact meaning of words when you speak.
3. You can’t just make such a statement. Do you have an argument to back that?
4. Lol yeah.

After spending half a year as a PhD student, I start to understand when people say:

Getting a PhD is not about acquiring a set of technical skills, but rather a specific mindset.

I suppose there are many elements that consist such a mindset, including but not limited to:

• Excavating meaningful problems to solve and navigating the uncertain process of solving it (#1 above).
• Communicating facts and arguments in precise language (#2 above).
• Maintaining a critical view of relevant matter and accepting arguments after rigorous reasoning and observation (#3 above).

I believe many will agree that none of these are neither easy nor quick to acquire. One must imbue oneself with principles and constantly self-reflect and self-correct. However, I think one should go one step further. One should make a conscious effort so that no existing mindset is completely replaced by the PhD mindset; the PhD mindset must only append to the list of existing mindsets. Then, one must distinguish situations when it is more appropriate to apply the PhD mindset and when it is not in a fine-grained manner.

Life, at least partly, can be viewed as a multi-task learning problem. While acquiring new capabilities is important, one must make sure not to catastrophically forget other important things in the process, which may not always be easy especially when the new capability requires an immense amount of concentrated effort to learn. However, I believe that such an effort is meaningful in advancing one’s maturity as a person.