The Kernel Reconnaissance: Mapping the Blackwell SM120 Optimization Frontier in SGLang
Introduction
In the high-stakes world of large language model inference on cutting-edge hardware, the difference between a functioning deployment and a performant one often comes down to the kernels—those hand-tuned GPU programs that execute the mathematical primitives underpinning every transformer layer. When deploying DeepSeek-V4-Flash on NVIDIA's RTX PRO 6000 Blackwell GPUs (compute capability sm_120), a team found themselves stuck at roughly 28 tokens per second—a fraction of the hardware's roofline potential of 300–600 tok/s. The culprit was identified as "sm_120 fallback kernel bottlenecks," but the precise nature of those bottlenecks remained unclear.
This article examines a pivotal subagent session within that optimization campaign—a systematic, read-only exploration of the SGLang inference engine's kernel landscape. Over the course of 18 messages, an AI assistant was tasked with a focused mission: "Map MXFP4 MoE + FP8 GEMM kernels (agent: explore)." The result was a comprehensive optimization specification that identified not just what was slow, but why it was slow at the level of individual GPU instructions, memory access patterns, and tensor-core utilization. This is the story of that exploration—a case study in how systematic codebase reconnaissance can transform an amorphous performance problem into a set of well-defined engineering tasks.
The Mission: From Performance Crisis to Kernel Map
The broader context of this subagent session is a multi-week optimization campaign on a machine equipped with eight NVIDIA RTX PRO 6000 Blackwell GPUs. The root session (segment 67) had already navigated numerous challenges: installing NVIDIA drivers and CUDA Toolkit 13.1, setting up a Python virtual environment with PyTorch, resolving flash-attention build issues by reducing parallel compilation jobs from 128 to 20, and stabilizing the software stack with compatible versions of PyTorch 2.9.1, flash-attn 2.8.3, and vLLM 0.15.1.
The deployment of DeepSeek-V4-Flash had achieved prefill-decode disaggregation with NIXL KV transfer across two NUMA domains, but performance fell far short of expectations. The user's target was approximately 1000 tok/s; the system delivered only ~10 tok/s at batch size 1 and ~25 tok/s at concurrency 16. Every configuration lever had been exhausted—NCCL LL+Ring tuning, CUDA graphs (already enabled), tilelang indexer fusion (which failed JIT compilation on sm_120), non-marlin MoE backends (invalid for FP4 experts), and expert parallelism (worse due to PCIe all-to-all overhead). None moved the needle.
A definitive GPU profile traced 63% of decode time to a single kernel: _tiled_sparse_decode_kernel, the sm_120 Triton fallback for sparse MLA attention. This kernel launched only 64 blocks (1 batch × 64 heads) on approximately 170 SMs, serially iterating all 512 top-k tokens. It was the same low-occupancy pathology that had plagued earlier work on the K2.6 verify kernel.
It was at this point that the subagent was spawned with a specific mandate: not to fix the bottleneck, but to map it—to understand the MXFP4 MoE decode kernel and the FP8 block GEMM infrastructure well enough that a subsequent optimization effort could be precisely targeted.
The Investigation Arc: Systematic Codebase Reconnaissance
The subagent's exploration unfolded across 18 messages, each building on the previous one. What makes this investigation remarkable is its methodical, hypothesis-driven approach. Rather than randomly reading files, the assistant formulated specific questions, selected targeted tools to answer them, and explicitly documented its reasoning at each step.
Phase 1: Reading the Kernel Implementation
The investigation began with the most directly relevant file: the MXFP4 MoE Triton kernel (mxfp4_moe_sm120_triton.py). This file contained the _mxfp4_slot_gemv_kernel—the kernel that actually runs on the GPU during decode. The assistant discovered that this kernel uses a "slot GEMV" approach, where each (token, expert) pair is processed as an independent matrix-vector multiplication. The term "slot" refers to one token-expert pair; every slot independently recomputes a full GEMV against its expert's entire weight matrix.
The critical finding was on lines 156-157 of the kernel:
acc += tl.sum(a_even[None, :] * val_lo, axis=1)
acc += tl.sum(a_odd[None, :] * val_hi, axis=1)
This is a CUDA-core FMA (fused multiply-add) reduction using tl.sum, not a tensor-core operation using tl.dot. The SM120a GPU has native FP4 tensor cores that could perform this computation much more efficiently, but the current kernel does not use them. This single observation would become the central theme of the entire investigation.
Phase 2: Tracing the Integration Layer
Having understood the kernel implementation, the assistant next read the quantization method files that wire the kernel into the broader inference pipeline. The mxfp4_marlin_moe.py file contained the Mxfp4MarlinMoEMethod class, which is selected when the model has FP4 experts. Crucially, the assistant discovered that on SM120 hardware, the _dsv4_mxfp4_backend is set to "sm120_triton", which bypasses the Marlin runner entirely. The class name is misleading—"Marlin" appears in the name, but on Blackwell hardware, the execution path diverges completely to a Triton-based implementation.
The mxfp4.py file contained the SM120-specific tuning patch (_patch_sm120_mxfp4_min_warps), which forces a minimum of 4 warps and uses StridedLayout with num_stages=1. This confirmed that the project had already established SM120 tuning patterns, but they applied to a different code path than the slot kernel under analysis.
Phase 3: Discovering the FP8 Block GEMM Infrastructure
The assistant then pivoted to the FP8 block GEMM, which handles the dense (non-MoE) linear layers. Reading fp8_kernel.py revealed the config-lookup mechanism: the kernel searches for a JSON file matching the exact problem dimensions (N, K), device name, data type, and block shape. The device name is obtained from torch.cuda.get_device_name(0) with spaces replaced by underscores.
A bash command enumerated the quantization configs directory, revealing 636 JSON files—but critically, none for "NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition." Configs existed for B200, H100, H200, A100, and other devices, but the RTX PRO 6000 Blackwell was absent. This meant the FP8 block GEMM always fell back to default parameters: BLOCK_M=64, BLOCK_N=128, BLOCK_K=128, GROUP_M=32, warps=4, stages=3.
The assistant also located the tuning script at benchmark/kernels/quantization/tuning_block_wise_kernel.py, which could generate the missing config by running on the actual hardware.
Phase 4: The Discovery of the Grouped FP4 GEMM
Perhaps the most significant discovery came when the assistant found an existing in-tree FP4 grouped GEMM implementation called cutlass_fp4_group_mm_sm100a_sm120a, exposed through Python bindings in nvfp4.py. This was a CUTLASS-based kernel specifically designed for MoE experts, supporting both SM100 (Hopper) and SM120 (Blackwell) architectures.
The grouped GEMM promised to solve the fundamental inefficiency of the slot GEMV approach. Instead of processing each (token, expert) pair independently—launching many small, low-utilization kernels—the grouped GEMM could fuse multiple expert computations into larger, more efficient matrix-matrix operations that fully utilize Blackwell's FP4 tensor cores.
However, the assistant quickly identified a critical blocker: both in-tree CUTLASS paths were designed for NVFP4 quantization, not MXFP4. The differences were fundamental:
- NVFP4: group size = 16, scale type =
float_ue4m3_t(E4M3, 4 exponent bits, 3 mantissa bits) - MXFP4: group size = 32, scale type =
float_ue8m0_t(E8M0, 8 exponent bits, 0 mantissa bits—essentially a power-of-two scale) The DeepSeek-V4-Flash checkpoint used MXFP4: E2M1 values with group-32, UE8M0 scales. These could not be fed directly to the NVFP4 kernels without conversion.
Phase 5: Confirming the CUDA Graph Constraint
A critical architectural question remained: did the MoE operations run inside or outside the decode CUDA graph? The answer would determine whether any replacement kernel needed to be graph-compatible—a constraint that rules out approaches relying on dynamic expert routing, Python-level loops over experts, or host-device synchronization.
The assistant located the CUDA graph runner files (decode_cuda_graph_runner.py, base_cuda_graph_runner.py) and traced the capture mechanism. The capture_one_shape function invokes model_runner.model.forward, capturing the entire forward pass—attention, MoE, and all—as a single CUDA graph. The mxfp4_moe_sm120_triton.py file's comments made this explicit, noting that the design avoids operations like .item(), .nonzero(), and .unique() specifically because the MoE runs inside the captured graph.
This confirmed that any replacement kernel must be CUDA-graph-compatible: no host synchronizations, no data-dependent shapes, and fixed grid dimensions per captured batch size.
Phase 6: Verifying the Hardware Detection Mechanism
The final phase of investigation focused on the is_sm120_supported() function, which gates SM120-specific code paths throughout the codebase. The assistant had seen references to this function in fp8_utils.py and mxfp4_marlin_moe.py, but had not yet located its definition.
A grep search initially failed because the function is defined using a decorator pattern (is_sm120_supported = lru_cache(maxsize=1)(...) rather than a traditional def statement. A refined search found it at line 277 of common.py. Reading the surrounding code revealed the pattern: the function uses _check_cuda_device_version with device_capability_majors=[12] and cuda_version=(12,8), cached via lru_cache for single-computation efficiency.
This confirmed that the RTX PRO 6000 Blackwell Server Edition would be correctly identified as SM120 hardware, enabling the SM120-specific kernel paths.
The Synthesis: A Comprehensive Optimization Specification
The culmination of this investigation was message 17, a comprehensive specification that synthesized all findings into a structured document. The spec covered four major areas:
1. The MXFP4 MoE Decode Path
The spec documented the exact weight tensor shapes and data types: w13_weight as [256, 4096, 2048] int8 (packed FP4), w2_weight as [256, 4096, 1024] int8, and scale tensors with group size 32 using float8_e8m0fnu format. It corrected a misconception from the task prompt—the group size is 32, not 128.
The root-cause analysis identified five reasons for the slot kernel's slowness:
- No tensor-core utilization: The kernel uses
tl.sum(CUDA-core FMA) instead oftl.dot(tensor-core matrix multiply) - Memory-bound with no weight reuse: Each slot streams its expert's entire packed weight (8 MB for w13, 4 MB for w2) with no amortization across tokens
- Per-element dequantization overhead: Every weight byte requires
exp2and selection operations to decode FP4 values - Eager-mode epilogue: SiLU, clamp, multiply, and reduction operations launch several extra kernels
- Tiny tiles with no batch dimension: The kernel uses
BLOCK_N ≤ 128with M=1 per slot, leading to low occupancy
2. What a Tuned SM120 MXFP4 MoE Kernel Needs
The spec identified two existing tensor-core paths in the codebase—the dense FP4 GEMM (nvfp4_scaled_mm_sm120.cuh) and the grouped FP4 MoE GEMM (nvfp4_blockwise_moe.cuh)—but noted that both target NVFP4 format, not MXFP4.
Three optimization paths were recommended in priority order:
- Best: Grouped FP4 tensor-core GEMM using
cutlass_fp4_group_mm, either by re-quantizing experts from MXFP4 to NVFP4 at load time or by adding an MXFP4 variant of the CUTLASS kernel - Medium: A Triton grouped GEMM that sorts slots by expert and runs
tl.dotblock GEMM withBLOCK_M > 1per expert segment - Cheapest: Better autotuning of the existing slot kernel with larger
BLOCK_K,num_stages=1, and more warps
3. FP8 Block GEMM Tuning
The spec documented that no SM120-specific FP8 block config exists, and the dispatch logic in fp8_utils.py forces the Triton backend on SM120. The tuning script at benchmark/kernels/quantization/tuning_block_wise_kernel.py could generate the missing config with the invocation:
python benchmark/kernels/quantization/tuning_block_wise_kernel.py \
--N 4096 --K 512 --input-type fp8 --out-dtype bfloat16 \
--block-n 128 --block-k 128
4. Integration Points and CUDA Graph Status
The spec traced the complete dispatch chain from method selection to kernel invocation, identifying the exact swap-in point in Mxfp4MarlinMoEMethod.apply. It documented the CUDA graph constraints and noted that per-expert token offsets for a grouped GEMM must be produced with pure tensor operations (histogram via scatter_add, argsort on routing) to remain capture-safe.
The NVFP4 vs MXFP4 Distinction: The Central Challenge
Throughout the investigation, one theme recurred with increasing clarity: the tension between NVFP4 and MXFP4 quantization formats. This distinction became the central challenge documented in the spec.
NVFP4 (NVIDIA's FP4 format) uses a group size of 16 with float8_e4m3fn scales—an 8-bit floating-point format with 4 exponent bits and 3 mantissa bits. The in-tree CUTLASS kernels (cutlass_fp4_group_mm, nvfp4_scaled_mm_sm120.cuh) are designed for this format, leveraging Blackwell's native FP4 tensor-core operations.
MXFP4 (Microscaling FP4, an industry standard format) uses a group size of 32 with E8M0 scales—an 8-bit unsigned integer format with no mantissa, essentially a pure exponent. The DeepSeek-V4-Flash checkpoint uses MXFP4, and the current Triton slot kernel handles this format through software dequantization.
The group size difference (16 vs 32) determines how many elements share a single scale factor, affecting both precision and memory access patterns. The scale format difference (E4M3 vs E8M0) determines whether native tensor-core operations can be used. The CUTLASS kernels expect NVFP4's E4M3 scales; the MXFP4 weights use E8M0 scales.
The spec noted that adapting the CUTLASS grouped GEMM for MXFP4 is feasible because "CUTLASS block-scaled supports UE8M0; only the ElementSFType/group_size/scale-layout need changing." This observation opened a concrete path forward: either re-quantize the experts at load time (changing numerics slightly) or add an MXFP4 variant of the kernel.
The Thinking Process: A Model for Systematic Investigation
What makes this subagent session particularly instructive is the transparency of the assistant's reasoning. At each step, the assistant articulated its hypotheses, selected tools to test them, and documented what it learned—including when searches failed.
The investigation followed a clear pattern:
- Formulate a hypothesis: "The grouped FP4 MoE GEMM might replace the Triton slot-GEMV kernel."
- Identify evidence needed: Check if the kernel handles MXFP4 format, if it's CUDA graph compatible, and how it's invoked.
- Execute targeted reads: Read the specific kernel file and the dispatch logic.
- Connect to prior knowledge: Build on previous discoveries to construct an increasingly complete mental model. This pattern recurred across every phase of the investigation. When a grep for
def is_sm120_supportedreturned nothing (because the function uses a decorator pattern rather than a traditionaldefstatement), the assistant didn't give up—it refined the search and found the definition. When a file read was truncated, the assistant read the surrounding context and inferred the missing information through pattern recognition. The assistant also demonstrated a sophisticated understanding of the tool environment. By dispatching independent reads and searches in parallel, it minimized the number of rounds needed to gather information. This parallel strategy was particularly effective in the early phases, where the assistant read the MXFP4 kernel, the quantization method files, and the FP8 kernel files simultaneously.
The Broader Significance: From Bottleneck to Blueprint
The specification produced in message 17 transformed an amorphous performance problem into a set of well-defined engineering tasks with clear success criteria. For the engineering team, the path forward was now clear:
- Run the FP8 tuning script to generate SM120-specific configs for dense layers
- Implement MXFP4 support in the CUTLASS grouped GEMM for MoE layers
- Or implement a Triton grouped GEMM with expert sorting as a fallback
- Or autotune the existing slot kernel as a quick improvement Each task had a defined scope (specific files to modify), known constraints (CUDA graph compatibility), and measurable outcomes (throughput improvement). More broadly, this investigation reveals a fundamental truth about GPU kernel optimization on emerging hardware architectures: the fast, optimized paths are often arch-gated to specific compute capabilities. The SGLang codebase had been primarily developed and tuned for datacenter Blackwell GPUs (B200/B100 with SM100 architecture), and the workstation-class RTX PRO 6000 GPUs (SM120) were left to fallback implementations. The subagent's exploration revealed the precise contours of this gap: which paths had SM120 support (MXFP4 MoE Triton backend), which paths lacked it (FP8 block GEMM configs), and how the system navigated between them.
Conclusion
The subagent session examined in this article is a masterclass in systematic codebase analysis for performance optimization. Over 18 messages, the assistant transformed a vague performance bottleneck into a precisely documented specification with prioritized action items. It discovered an in-tree grouped GEMM that could replace the slow slot GEMV kernel, identified the NVFP4/MXFP4 format mismatch as the central integration challenge, confirmed the CUDA graph constraint that shapes all kernel design decisions, and documented the exact tuning infrastructure needed to generate missing FP8 configurations.
The investigation's greatest strength is its ability to connect multiple levels of abstraction: from the arithmetic of FP4 dequantization (exp2((exp_bits - 1).to(tl.float32))) to the system architecture of CUDA graph capture (capture_one_shape in decode_cuda_graph_runner.py), from the data layout of weight tensors ([256, 4096, 2048] int8) to the strategic tradeoffs of format conversion (MXFP4→NVFP4 at load time vs. native MXFP4 kernel).
For anyone studying how to approach complex systems optimization, this session offers a template: understand the current implementation before proposing changes, verify assumptions against source code, document constraints explicitly, and present findings in a structured, actionable format. These principles apply far beyond GPU kernel engineering, making this investigation instructive for anyone who works with complex software systems.