Inference for Interpretable Text-Diffusion Models

Introduction

At Guide Labs, we build models whose representations are transparent, so that humans can audit and understand them. Our recent release, Steerling-8B, is a diffusion language model that generates text in 64-token blocks, filling in masked tokens over several steps rather than one token at a time. Every token’s prediction passes through a layer of human-understandable concepts, so the model can report which concepts drove each output chunk (concept attribution) and users can shift the output by adjusting those concepts (concept steering).

Clarity, our platform, puts these capabilities in front of users. To serve Steerling-8B in Clarity, we sought a serving engine that provides production-grade latency and throughput while treating the interpretability properties of Steerling as first-class citizens.

Below, Steerling writes one 64-token block at a time: each block starts as a row of masks and comes into focus over many steps, its words landing out of order rather than left to right. As each word lands, the concepts behind it build up in the panel beside the text, and once a block is complete, its words and their concepts are fixed for good.

Loading…

Goals for an Interpretable Serving Engine

We set out to build a serving engine that is:

• Diffusion-native: decodes 64-token diffusion blocks iteratively.
• Built for concept attribution: returns concept attributions for every chunk of generated tokens.
• Steerable at request time: accepts per-request steering concepts and schedules without recompiling the model.
• Memory efficient: limits fragmentation in KV cache and minimizes memory allocation in concept bottleneck layer.
• Low-latency: block latency that is comparable to other diffusion serving engines.
• High-throughput: competitive throughput with interpretability enabled.

We found that no existing engines met these goals. vLLM and SGLang assume autoregressive, text-only generation. Steering extensions such as EasySteer inherit those assumptions. Diffusion engines such as Diffulex lack block-causal attention support, compilation, dynamic batching, and any notion of interpretable outputs. So we built the Causal Diffusion Engine (CDE).

Background

Three properties of Steerling shape every design decision in CDE: it decodes block by block, it routes predictions through concepts, and it accepts concept-level interventions.

• Block-causal diffusion decoding. A diffusion block is a set of contiguous tokens processed together. Generation starts with a block of masked tokens and unmasks them over several forward passes. Attention is block-causal: each token attends bidirectionally to every token in its own block, causally to tokens in earlier blocks of the same sequence, and never to later blocks. This provides the block-causal property CDE depends on: a completed block’s representation cannot change once it is complete, so its attention keys and values can be cached, as shown below.

D = 2 tokens per diffusion-block · B = 2 tokens per block-tile
3 diffusion-blocks · Each small square = one query–key token pair

OverviewWithin a blockPrevious blocksFuture blocksKV cache

← query tokens

block 1t1 · t2

block 2t3 · t4

block 3t5 · t6

block 1t1 · t2

block 2t3 · t4

block 3t5 · t6

× × ×

key/value tokens →

BlockMask Tiles Fully included · skip masking ×Fully skipped · skip attention

Purple tiles are fully included. Empty tiles are fully skipped.

• Concept bottleneck. As shown below, Steerling has a layer that translates its hidden representations into human understandable concepts. These hidden states feed two concept heads. Each head outputs weighted concept embeddings, concept indices, and logits. The sum of concept embeddings plus a residual term passes through the language modeling head to produce token logits. Consequenlty, concept attribution is part of the forward pass rather than a separate model.

• Steering. Adding a concept vector to the hidden states at a chosen layer shifts generation toward that concept. In serving, this is a per-request intervention and can change from one decoding step to the next.

For the full model, see Scaling Inherently Interpretable Language Models.

System Overview

We will now discuss the high-level design of CDE.

Core Components

Six core components carry requests through the four stages described below, coordinating request scheduling, model execution, and the delivery of generated text with concept attributions:

• Scheduler: Handles admission control, dynamic batching, and request state transitions.
• Model runner: Executes specialized compiled graphs for prefill and decode.
• KV cache manager & page table: Allocates cache pages and translates attention masks to the physical page layout.
• Steering injector: Constructs per-request concept embeddings and applies their schedules.
• Sampler: Selects which masked tokens to unmask at each step.
• Concept attribution module: Processes concept indices and logits to produce token-level attributions.

Inbound Request

A request to CDE flows through four main stages:

1. Admission: The request, containing a prompt and optional steering concepts, is admitted by the scheduler into a batch.
2. Prefill: Completed blocks from the prompt are processed, and their keys and values are written into the paged KV cache.
3. Block-by-Block Decode: The model unmasks the next 64-token block, one step at a time. When a block has no masked tokens left, a final pass commits it to the cache.
4. Attribution: Concept attributions are computed for each sampled token and returned with the generated text.

Decode Step

Almost all of a request’s time is spent in this loop:

for block in blocks_to_generate:
    while block.has_masked_tokens():
        logits, concepts = model_runner.decode(batch)      # compiled: steering, transformer,
                                                           #   concept heads, LM head
        tokens = sampler.unmask_one(logits, block)         # outside the graph
        attribution.record(concepts, tokens)               # outside the graph
        steering_injector.update(block.unmasked_positions)
    model_runner.commit(block)                             # write the block's KV to the cache

The boundary between compiled and uncompiled code is the most important line in this loop. Everything from steering injection through the transformer, concept heads, and language modeling head runs inside a single torch.compile graph. Sampling and attribution depend on which token was just chosen, so they run outside it. Every step therefore passes from GPU to CPU and back, a cost that grows with batch size (see Dynamic Batching).

Serving Concept Attribution

Serving engines are built to return tokens, and sometimes per-token hidden states, but Clarity also needs a non-token output: the concepts that produced each chunk of tokens. Serving these attributions raises two problems. It must not slow down generation, and the concept heads must fit in GPU memory alongside the model and KV cache. CDE solves the first by keeping almost all of the work inside the compiled graph, and the second by never materializing the largest concept head, which cuts its memory from 16 GB to 1 GB per sequence.

Keeping attribution fast. Producing concept attributions from Steerling’s forward pass involves three steps:

1. Concept heads. Hidden states feed the concept heads. Each outputs weighted concept embeddings, concept indices, and logits.
2. Language modeling head. The concept embeddings are summed with a residual term and passed to the language modeling head to produce token logits.
3. Attribution. After a token is sampled, the attribution module turns the raw concept indices and logits into concept attributions.

Steps 1 and 2 run inside the torch.compile graph, so producing concept embeddings adds no separate pass. Only step 3 runs outside it, because it depends on the sampled token.

Fitting the concept heads in memory. The largest memory cost of step 1 comes from the larger of the two heads: the discovery head, with over 101K concepts. A naive multiplication against the discovery head took 16 GB for a single sequence. The discovery head is stored as two factorized matrices with a small inner dimension r. Multiplying through the r-dimensional space, instead of forming the full 101K-wide result, reduces computation and brings memory to 1 GB, a 16× reduction.

Steering

Steering is difficult to serve for three reasons. Every request in a batch can steer toward different concepts, or not at all. A request’s steering changes from one decoding step to the next. Lastly, turning steering on for a request must not trigger a recompile.

A request can name the concepts it wants to steer toward. The steering injector turns them into a concept embedding and adds it to the hidden states at layer 7, roughly a quarter of the way through Steerling’s 32 layers. Each request in a batch carries its own vector; for requests without steering, the vector is zero.

Each vector follows a schedule. At every step, the injector zeroes the vector at positions that are already unmasked, so steering acts only on tokens still being decided. Once a threshold number of tokens has been unmasked, the injector zeroes the whole vector, ending steering for that request.

Because the concept embedding is always part of the compiled forward pass, CDE never recompiles to turn steering on or off. The cost of this always-on addition is small: 1–2% in our measurements.

Making CDE Fast

Now that we have covered how the CDE serves concept attribution and steering, we will turn to the rest of the components that work together to deliver a fast and performant engine. First, recall the decode loop from the System Overview: the CDE samples one token per forward pass of an 8B model, so a 256-token response takes hundreds of passes. Three things decide how long that takes, and the next three sections take them in turn:

• How much overhead surrounds each pass. Compilation cuts the overhead around each pass, as long as serving never triggers a recompile (Compilation Strategy).
• How much each pass computes. Block-causal attention skips work the attention pattern rules out, and the KV cache means completed blocks are never recomputed (Block-Causal Paged Attention).
• How many requests share each pass. Batching spreads the cost of the 8B weights across several requests (Dynamic Batching).

Together, these three costs set CDE’s latency and throughput. The next three sections tackle them in order, and the GSM8K comparison at the end shows where that leaves CDE against other diffusion serving engines.

Compilation Strategy

In serving, we encounter different configurations and input shapes: prefill vs. decode, varying batch sizes, and the presence or absence of steering and attribution. Our implementation often triggered recompilations or graph breaks for all these differences. We address this by making several key decisions, the most impactful being the ones we describe here:

• Specialized graphs for decode, dynamic shapes for prefill. Prefill uses a dynamic graph, eliminating the need for bucketing. Profiling showed that using specialized, fixed-shape graphs for decode significantly reduced call latency compared to dynamic-shape graphs. Therefore, CDE uses specialized graphs for decode, and dynamic shapes for prefill.

• Steering and concept heads inside the graph; attribution outside: Concept embeddings are always present in the compiled graph, so enabling steering for a request does not trigger a recompilation (see Steering). Concept attribution runs after the main graph executes, keeping the compiled graph’s execution path unchanged (see Serving Concept Attribution).

Performance Comparison

Setup	Time to first block (s)	Inter-block latency (s)	Time to last block (s)	Throughput (tok/s)
Not compiled	1.86	2.26	4.68	15.77
Compiled	0.69	0.78	1.72	44.38

Compiling with torch.compile reduces Steerling’s time-to-last-block from 4.68s to 1.72s.

Compilation also makes debugging harder: out-of-memory errors inside a compiled region report an allocation in a private memory pool rather than identifying the precise operation that caused it (see Serving Concept Attribution). We found that having the option to profile without torch.compile is invaluable.

Block-Causal Paged Attention

CDE features a custom block-causal attention mechanism, implemented as a fused FlexAttention kernel where the BlockMask size is perfectly aligned with the diffusion block size and KV cache page size, so the FlexAttention BlockMask can be perfectly sparse. These optimizations are not possible in mainstream serving engines, prompting us to build our own.

FlexAttention with Aligned Tiles

FlexAttention generates a fused attention kernel using a BlockMask, which represents the attention pattern in B×B tiles. Each tile is either fully included, fully excluded, or partially included. By setting the tile size B equal to the diffusion block size D = 64, boundaries always align, ensuring every tile either participates fully in attention or is excluded entirely. This optimization lets fully included tiles bypass masking computations, while fully excluded tiles skip attention computation altogether.

KV Cache

Without caching, every unmasking step would need to recompute attention over the entire sequence, causing compute requirements to grow quadratically with sequence length. However, since completed blocks are never resampled (Background), CDE can cache their keys and values to accelerate future sampling steps. The process works as follows:

1. Prefill: Completed blocks from the prompt are processed and written to the cache; any prompt tokens that don’t fill a complete block remain fixed within the first block of decoding.
2. Decode: The next block is processed using the cache; tokens within that block are unmasked one per step, in no fixed order, until the block is complete.
3. Commit: After a block is fully unmasked, a final pass is run (without unmasking) to commit it to the cache.

Because prefill and decode operate on different attention shapes, each requires its own BlockMask and thus its own compiled graph (torch.compile).

A KV cache for block-causal attention

D = 4 tokens per diffusion-block 10 prompt tokens max prefill length = 16 tokens

Process max prefill length

Prefill pass

Continue to decode B2 →

Prefill BlockMask

16 query tokens
16 key/value tokens

4 × 4 block tiles (B = D = 4).

Purple: attended. Gray: not attended.

Hover or focus a square to inspect a query–key pair.

KV cache after this step

Completed blocks in cacheB0 · B1

Cached tokens8

Next active diffusion-blockB2

Filled sequence positions10 / 16

p = fixed prompt token ? = masked token g = decoded token cached KV

Choose any step to explore the figure. Focus a control and use ← → to change steps.

← Back Next →

Process max prefill length. 10 of 16 positions filled. B2, 2 masked positions remain.

Paged Attention

To efficiently manage memory for many sequences of variable length, CDE uses paged attention: each request’s KV cache is allocated in fixed-size pages that don’t need to be contiguous in memory. The page table remaps logical BlockMasks onto physical memory pages, so memory is only allocated for the sequence length actually in use. This approach, adapted from PyTorch’s example and Flex-nano-vLLM (with stricter error checking and other customizations), removes the need to copy the KV cache repeatedly, allows concurrent requests with different sequence lengths without padding, and enables memory-intensive features like concept attribution (Serving Concept Attribution).

Dynamic Batching

So far, we have explained how to serve one request at a time (batch size 1). A natural way to increase throughput is to serve several requests together. At batch size 1, each decode step performs one forward pass through the 8B weights to process a single 64-token block, so the GPU is underused. Batching shares each forward pass across several requests in a batch, using the graph compiled for that batch size. The batch also advances to the next block only when every request has finished its current block. Batching is still dynamic: at each block boundary, finished requests leave the batch and new requests can join it. This design keeps the implementation simple, but ties inter-block latency to the slowest request. Batching raises throughput to 61.72 tok/s at batch size 4.

Batch size	Time to first block (s)	Inter-block latency (s)	Time to last block (s)	Throughput (tok/s)
1	0.69	0.78	1.72	44.38
2	1.51	1.14	2.87	56.09
3	2.38	1.59	4.31	56.80
4	3.03	1.95	5.49	61.72

From batch size 1 to 4, time-to-last-block increases 3.2× while throughput increases only 1.4×. The following table examines why throughput scales so poorly.

Batch size	Time per forward pass (ms)	Avg. forward passes per decode call	GPU waiting on CPU (ms)	CPU waiting on GPU (ms)
1	20.37	37.2	3.34	16.79
4	33.71	61.4	5.49	28.15

From batch size 1 to 4, each forward pass becomes 1.7× slower, and each decode call needs 1.7× more forward passes, because the batch waits for its slowest request to finish the block. The GPU sits idle as it waits for the CPU to compute concept attributions and sample (which is implemented as a loop over requests), growing 1.6× in step with the forward pass.

Letting each request finish its blocks independently of other requests should give clear improvements.

Results on GSM8K

We compare CDE with Diffulex and Fast-dLLM on GSM8K, first with a single request and then with four concurrent requests. For the other engines, we report both single-token sampling (ST), which matches CDE’s one token per step, and each engine’s native multi-token sampling (N). Fast-dLLM always runs a single batch, so its four-request results are its single-request results.

Single request. CDE has the lowest time to first block (0.69 s) and time to last block (1.72 s) of any configuration, and the second-lowest inter-block latency, behind only Fast-dLLM N, while always serving concept attributions and steering. Its single-request throughput (44.38 tok/s) is close to Diffulex ST (48.05 tok/s). The native configurations are faster, Diffulex N at 79.47 tok/s and Fast-dLLM N at 64.49 tok/s, because they sample several tokens per step (Limitations).

Four requests. With four concurrent requests, CDE falls behind. Its throughput (61.72 tok/s) is the lowest of the configurations shown, compared with 102.06 tok/s for Diffulex ST, and its time to first block (3.03 s) is the highest. Its inter-block latency (1.95 s) is second only to Fast-dLLM N, which runs a single batch. This is the batch-scaling cost described in Dynamic Batching: the whole batch waits for its slowest request at every block.

Limitations

Multi-token sampling. Most diffusion serving engines can sample several tokens per step, which reduces the number of forward passes at the cost of accuracy on most benchmarks. To maximize accuracy, CDE currently samples one token per forward pass.

Batch scaling. Because a batch moves to the next block only when every request has finished the current one, throughput grows slowly with batch size (Dynamic Batching). Letting each request advance through its blocks independently is the clearest next improvement.

Conclusion

In this post, we described the Causal Diffusion Engine (CDE), which serves Steerling-8B on our Clarity platform with single-request latency and throughput competitive with other diffusion serving engines. Unlike those engines, CDE returns concept attributions with every generated token and accepts steering on every request, so that the interpretability built into Steerling remains available in production.

Several lessons from building CDE apply beyond Steerling:

• Interpretability outputs can be served at little cost when they are part of the forward pass. CDE runs the concept heads and steering inside the compiled graph, so attributions require no extra pass and steering adds only 1–2% overhead; only the step that depends on the sampled token runs outside the graph. Memory required the same care: by multiplying through the discovery head’s factorized form, we reduced its memory from 16 GB to 1 GB per sequence.
• Aligning the attention tile, the diffusion block, and the KV cache page at 64 tokens makes the attention mask perfectly sparse and lets every completed block be cached as is.
• torch.compile pays off only when the engine is designed to avoid recompiles. With fixed-shape graphs and a single dynamic graph, compilation reduced time to last block from 4.68s to 1.72s and raised throughput 2.8×.

Finally, batching requests block by block keeps the implementation simple, but ties every batch to its slowest request. At batch size 4, throughput rises only 1.4× to 61.72 tok/s, in part because sampling and attribution run as a loop over requests on the CPU. Allowing each request to advance through its blocks independently is the natural next step. More broadly, CDE shows that interpretability does not have to come at the cost of serving performance: a model can be easy to audit, align, and understand, and still be fast enough to use.