The O(Actual) Indexer: Validating 128K Context at 500 tok/s on Blackwell

"The Triton indexer holds up remarkably well at 128K context—maintaining 96-98% of the throughput compared to 8K, which means the early-exit optimization is actually working as intended even at 16× the context length."

In the engineering of high-performance inference systems, there is a moment of truth that separates a clever idea from a working solution: the benchmark that proves the optimization actually delivers at scale. Message <msg id=12639> captures exactly such a moment. The assistant, having just deployed a custom Triton indexer kernel with early-exit per page for the DeepSeek-V4-Flash model on 8× RTX PRO 6000 Blackwell GPUs, presents the results of running that kernel at 128K context length — a 16× increase over the previous 8K cap — and shows that throughput holds at 96-98% of the baseline. This is the validation that the O(actual) kernel design works, that the context-length tax has been eliminated, and that the deployment can now serve long-context requests at hundreds of tokens per second.

The Bottleneck That Was

To understand why this message matters, one must understand the problem it solves. The DeepSeek-V4-Flash model uses a paged KV cache with a "DSA indexer" — a component that computes attention scores between the query and all cached key-value positions to determine which pages are relevant for the next token. In the original PyTorch fallback implementation, this indexer computed scores over the entire max context length every single decode step, regardless of how many tokens the actual sequence contained. With --context-length 262144 (the model's maximum), this meant computing scores over 262,144 positions per decode step, even when the actual context was only 512 tokens.

The computational waste was staggering. The indexer was performing a batch matrix multiply on tensors of shape [32, 262208, 64] — that is, 32 heads, 262,208 positions, and 64-dimensional queries — only to have the vast majority of those positions masked to negative infinity. The aten::copy_, mul, clamp_min, sum, and bmm operations on these enormous tensors accounted for roughly 69% of GPU time, making the indexer the single largest bottleneck in the entire decode pipeline.

The initial fix was pragmatic: cap --context-length to 8192, which reduced the indexer's work by a factor of 32 and delivered a dramatic 17.9× throughput improvement (from 29.7 to 531.7 tok/s at C=64). But this was a workaround, not a solution. It meant the model could never serve sequences longer than 8192 tokens, which defeated the purpose of having a model capable of 128K or 256K context.

The proper fix was the Triton indexer kernel with early-exit per page — a kernel that processes only the pages that actually contain valid tokens for each request, while still producing a full-size output tensor (required by the downstream topk operation and CUDA graph capture). The grid is fixed (capture-safe), but each program checks whether its assigned page is within the request's actual sequence length and, if not, takes a fast path that writes negative infinity without doing any expensive gather or dot-product computation. This makes the kernel O(actual sequence length) rather than O(max context length).

The Benchmark Results

Message <msg id=12639> presents the decisive comparison:

| C | torch indexer @ 8K | Triton indexer @ 128K | |---|---|---| | 1 | 58.9 | 58.7 | | 16 | 285 | 279 | | 64 | 532 | 509 |

These numbers tell a remarkable story. At C=1 (single concurrent request), the Triton indexer at 128K context achieves 58.7 tok/s — essentially identical to the 58.9 tok/s of the torch fallback at 8K context. At C=16, the Triton indexer delivers 279 tok/s versus 285 tok/s (97.9% retention). At C=64, the highest concurrency tested, it achieves 509 tok/s versus 532 tok/s (95.7% retention).

The key insight is that these throughput numbers are within 96-98% of the 8K baseline despite operating at 16× the context length. The early-exit optimization is working exactly as designed: the kernel spends compute only on the pages that matter, and the pages beyond the actual sequence length are handled with a trivial write of negative infinity. The ~2-4% remaining gap at higher concurrency is attributed to the downstream topk operation (which must operate over the full 128K logits tensor) and allocation overhead — both of which are O(max_c4) rather than O(actual), but are cheap single operations compared to the per-step indexer compute.

The Profile: Proof in the GPU Timeline

The assistant doesn't stop at throughput numbers. They run a GPU profile at 128K context to confirm the indexer kernel is no longer a bottleneck:

total GPU-kernel us: 487579.1

 27.9%    135.97ms  n=  774  cutlass::device_kernel (MoE)
 19.3%     94.25ms  n=  783  ncclDevKernel_AllReduce_Sum_bf16_RING_LL
 17.6%     85.89ms  n=  672  _mma_sparse_decode_split_kernel
  5.7%     27.66ms  n=   84  _mma_sparse_decode_kernel
  5.0%     24.60ms  n= 2124  _w8a8_block_fp8_matmul
  1.8%      8.93ms  n=  344  ...

This is a healthy profile for a production inference system. The top three kernels — MoE computation (27.9%), NCCL all-reduce communication (19.3%), and the custom MMA sparse-MLA decode attention kernel (17.6%) — are exactly what you would expect from a well-optimized transformer decode pipeline. The indexer kernel, which previously dominated at 69% of GPU time, has been reduced to a negligible fraction (it doesn't even appear in the top entries, having been absorbed into the "other" category).

The profile confirms that the system is now compute-bound and communication-bound in the right places: MoE expert computation, inter-GPU all-reduce, and attention. This is the signature of a pipeline where the major architectural components are the dominant costs, not a software workaround or a fallback kernel.

The Reasoning Process

The assistant's reasoning in this message reveals a disciplined engineering mindset. Having just received the benchmark results, the assistant immediately interprets them in context:

"The remaining ~4% difference at C=64 comes down to the topk operation over the larger logits tensor and allocation overhead, which is negligible and could be optimized further if needed."

This is not speculation — it's a grounded inference based on understanding the system architecture. The assistant knows that the Triton indexer itself is O(actual), so its cost is identical at 8K and 128K for the same actual sequence length. The only operations that scale with max context length are the topk (which operates on the full logits tensor) and the allocation of that tensor. A 4% overhead at 16× context is remarkably small, confirming that these operations are indeed cheap relative to the decode pipeline.

The assistant then immediately pivots to the next step: "I should profile at 128K to confirm the indexer kernel stays small, then commit this work and present the final results." This shows a systematic approach — validate the hypothesis, confirm with profiling, then lock in the change. The profile results arrive and confirm the indexer is invisible in the top kernels.

What This Message Creates

This message produces several forms of output knowledge:

  1. Quantitative validation: The Triton indexer kernel maintains 96-98% throughput at 128K context vs 8K, across concurrency levels from 1 to 64.
  2. Profile evidence: The GPU time breakdown confirms the indexer is no longer a bottleneck, with MoE, NCCL, and attention as the dominant kernels.
  3. Attribution of remaining overhead: The ~2-4% gap is attributed to topk and allocation, not the indexer itself.
  4. A decision point: The assistant decides the results are good enough to commit and present as final, rather than chasing further optimization of the topk operation.

What This Message Assumes

The assistant makes several assumptions that are worth examining:

The Broader Significance

This message represents the culmination of a multi-day optimization campaign that transformed the DeepSeek-V4-Flash deployment from a system struggling at ~30 tok/s to one delivering 500+ tok/s at 128K context. The indexer fix alone accounted for the largest single leap — from 29.7 to 531.7 tok/s at C=64 — but it was the Triton kernel that made this performance context-length-independent, unlocking the ability to serve long-context requests without a per-step penalty.

For the user deploying this system, the practical implication is profound. They can now set --context-length 131072 (or higher) without worrying about a proportional throughput drop. The system will serve short-context requests at the same speed as long-context requests, because the indexer's cost is determined by the actual sequence length, not the configured maximum. This is the difference between a system that works "in theory" and one that works "in production."

Conclusion

Message <msg id=12639> is a validation message — the moment when a carefully designed optimization is put to the test and proven to work. The Triton indexer kernel with early-exit per page, running at 128K context on Blackwell GPUs, delivers 509 tok/s at C=64 — 96% of the 8K baseline. The profile confirms a healthy compute-bound pipeline where the indexer is no longer visible. The O(actual) design has been validated, and the deployment can now serve long-context sequences at production-grade throughput. This is engineering at its most satisfying: a clean design, a decisive benchmark, and a result that speaks for itself.