ZIANG :: LLM_PERF_ENG

Phase 1 — GPU + CUDA Foundations

GPU Mental Model: SMs, Warps, and Tensor Cores

A modern NVIDIA GPU is a throughput-optimized machine, not a latency-optimized one. The key abstraction is the Streaming Multiprocessor (SM) — the unit that executes groups of threads. An H100 has 132 SMs, an A100 has 108 SMs, and a B200 (Blackwell) has roughly 160 SMs across two reticle-limited dies connected by a high-speed NVLink bridge.

Threads execute in groups of 32 called warps. All 32 threads in a warp run the same instruction on different data (SIMT). Divergent branches within a warp serialize both paths.

Tensor Cores are the secret sauce. An H100 tensor core can do one 16×16×16 matrix multiply per clock per sub-core.

• A100 introduced third-gen tensor cores with TF32 and sparsity. HBM2e at 2 TB/s.
• H100 added Hopper: 3× FP16 tensor-core bump, FP8 support, Transformer Engine, DPX instructions, and asynchronous WGMMA. HBM3 at 3.35 TB/s.
• B200/GB200 packs two dies with 208 B transistors. FP8 tensor reaches 4.5 PFLOPS, new FP4 at 9 PFLOPS. HBM3e at 8 TB/s. Second-gen Transformer Engine, NVLink 5 at 1.8 TB/s.

The roofline bottleneck for all three: dense matrix ops are compute-bound on tensor cores; memory-bound ops (element-wise, softmax) are HBM-limited. Arithmetic intensity (FLOP/byte) must increase to keep tensor cores fed.

CUDA Memory Hierarchy

    Register File     ~64K registers/SM      ~0.5 cycles
       ↓
Shared Memory (SMEM)    48-228 KB/SM         ~20 cycles
       ↓
      L2 Cache          40-60 MB/chip        ~200 cycles
       ↓
  Global Memory (HBM)   40-80 GB/chip        ~600-800 cycles

Coalescing rule: 32 threads in a warp accessing contiguous, aligned bytes → one 128-byte transaction. Strided or random → many transactions → bandwidth collapse.

Bank conflicts: SMEM is 32 banks (4-byte wide on H100). Multiple threads → same bank → serialized. Stride-2 causes 2-way conflicts; power-of-two strides are dangerous.

PMPP Essentials (Chapters 1-5)

• Reduction: Tree-based, warp-shuffle for last log2(32) steps with __shfl_down_sync.
• Scan: Hillis-Steele O(N log N) vs Blelloch O(N) work-efficient. GPU wants block-level scan + single-pass block sums.
• Tiled matmul: Partition C into BM×BN tiles in SMEM. Tile 128×128 with BK=64 → 64 KB SMEM. Coalesced 128-byte global loads per warp.
• Histogram: Per-block private bins in SMEM + atomic to global. Pad bins to power-of-two + 1 to dodge conflicts.

Roofline Model + Arithmetic Intensity

Attention math (head, seq N, dim d):

• QK^T: 2N²d FLOPs, reads 8Nd bytes (FP16). Intensity = N/4 FLOP/byte.
• Softmax: 5N² FLOP, 2N² bytes. Intensity = 2.5 FLOP/byte — memory-bound.
• PV: Same as QK^T.

At N=4096, d=128: QK^T intensity = 1024 FLOP/byte. H100 balance point (989 TFLOPS / 3.35 TB/s) = 295 FLOP/byte. Compute-bound. At N=512: 128 FLOP/byte → memory-bound. This is why FlashAttention matters most at medium-long sequences.

Nsight Compute — Key Metrics

• Achieved Occupancy = active warps / max (64 on H100). Tensor-core kernels max out at ~50% occupancy; the tensor pipe is the ceiling, not schedulers.
• Mem Busy % — < 90% in memory-bound kernel = coalescing issue.
• Warp stalls:
• long_scoreboard — waiting on global load (memory-bound).
• short_scoreboard — waiting on SMEM/register.
• not_selected — other warps ready; low occupancy + this = need more warps.
• throttle — dependency on previous instruction (tight reduction).
• barrier — __syncthreads() overuse.
• math_pipe — compute-bound.

Flow: ncu --set full -o profile python train.py → Speed-of-Light tab gives single-page diagnosis.

// Tiled GEMM using shared memory (from PMPP)
__global__ void tiled_gemm(float *A, float *B, float *C,
                            int M, int N, int K) {
    __shared__ float As[128][16];
    __shared__ float Bs[16][16];
    int row = blockIdx.y * 128 + threadIdx.y;
    int col = blockIdx.x * 16  + threadIdx.x;
    float sum = 0.0f;
    for (int t = 0; t < K / 16; ++t) {
        As[threadIdx.y][threadIdx.x] = A[row * K + t * 16 + threadIdx.x];
        Bs[threadIdx.y][threadIdx.x] = B[(t * 16 + threadIdx.y) * N + col];
        __syncthreads();
        for (int k = 0; k < 16; ++k)
            sum += As[threadIdx.y][k] * Bs[k][threadIdx.x];
        __syncthreads();
    }
    C[row * N + col] = sum;
}

Phase 2 — Attention Mechanics + LLM Inference

FlashAttention: Three Generations

Standard attention materializes N×N to HBM — O(N²) memory. FlashAttention eliminates it via tiling + online softmax.

v1 (Dao 2022): Split Q, K, V into SMEM tiles. Per Q tile, load K/V tiles sequentially. Maintain running m_i (max) and ℓ_i (sum): - m_new = max(m_old, rowmax(S_ij)) - ℓ_new = e^(m_old - m_new) × ℓ_old + sum(e^(S_ij - m_new)) - Rescale output continuously. No O(N²) HBM writes.

v2 (Dao 2023): Parallelize across seq dim, not just heads. Block-level softmax removes __shfl_sync overhead. ~2× faster than v1.

v3 (Shah 2024, Hopper): Uses WGMMA async warp-group MMA + cp.async shared memory copies:

async_load(K_tile_0, SMEM)
async_load(V_tile_0, SMEM)
sync_wait(read_all)
WGMMA(Q_0, K_0 → S_0)     // async
wgmma_fence()
async_load(K_tile_1, SMEM) // overlap load with compute

FP8 halves memory traffic; TE handles casting.

def flash_attention_v1(Q, K, V, Bc=128, Br=64):
    N, d = Q.shape
    O = np.zeros_like(Q)
    for i in range(0, N, Br):
        Qi = Q[i:i+Br]
        m_i, ell_i, Oi = -inf*ones(Br), zeros(Br), zeros((Br, d))
        for j in range(0, N, Bc):
            Kj, Vj = K[j:j+Bc], V[j:j+Bc]
            Sij = Qi @ Kj.T
            m_new = max(m_i, rowmax(Sij))
            Pij = exp(Sij - m_new[:, None])
            alpha = exp(m_i - m_new)
            ell_i = alpha * ell_i + rowsum(Pij)
            Oi = alpha * Oi + Pij @ Vj
            m_i = m_new
        O[i:i+Br] = Oi / ell_i[:, None]
    return O

Paged KV Cache (vLLM)

Problem: 70B GQA (8 KV heads), FP16, seq 4096 → ~16 MB per request. 100 requests × 2048 tok → 6.4 GB. Contiguous tensors fragment as requests finish at different times.

vLLM's solution: OS-style page table. Cache stored in fixed blocks (16 tokens). A block table maps logical request → physical blocks:

Logical:  [tok0..tok15] [tok16..tok31]
              ↓              ↓
Blocks:   phys_7         phys_42

Internal fragmentation ≤ 1 block per request (16 tokens = ~1%). Beats 30-50% waste of contiguous allocation.

Continuous Batching + Prefill/Decode Split

Prefill: All KV cache for prompt in one shot. Big GEMM, M = batch × prompt_len. Compute-bound.

Decode: One token per request per step. Loads full KV cache, one vector @ weight matrix. Intensity ≈ 1 FLOP/byte. Memory-bound, GPU utilization <5%.

Never mix prefill+decode in one iteration — prefill's big tiles starve decode's quick ops. Split them across GPUs or fractions of the batch.

Speculative Decoding + MTP

Speculative (Leviathan/Chen 2023): Cheap draft model generates γ tokens; target verifies in ONE forward pass. Accept token if target prob ≥ draft prob (or via speculative sampling ratio). Expected tokens/pass = (1 − α^(γ+1))/(1 − α). At α=0.8, γ=4: 3.0 tokens/pass = 3× wall-clock.

MTP (Gloeckle 2024): Train target with γ output heads, no external draft. Each head predicts one future position. 2-3× speedup on code, no draft overhead.

Real: γ=4 + 1/10 draft on 70B target → 2.5× throughput on H100. MTP on 7B (γ=3) → 2× with zero draft overhead.

Phase 3 — SGLang

Architecture: Router → Scheduler → TP Workers

SGLang's serving process is a three-tier pipeline connected by ZMQ IPC:

• Router — public HTTP endpoint. Terminates OpenAI-compatible requests, tokenizes, forwards to Scheduler over ZMQ PUSH/PULL with backpressure. Zero-copy via send_multipart on the token buffer.
• Scheduler — Python-native event loop. Owns the radix tree, decides which requests enter the next forward pass (prefill vs decode mix), tracks KV budget, handles preemption.
• TP Workers — one process per GPU rank. Executes the forward pass via PyTorch + FlashInfer kernels. Communicates activations across TP ranks via NCCL.

The Python-native scheduler is the tradeoff: easier to hack (RadixAttention landed in weeks) but higher per-request overhead than vLLM's C++ scheduler. At low QPS this doesn't matter; at 10K+ QPS the Python GIL becomes the bottleneck.

RadixAttention

Prefixes in production LLMs repeat: system prompts, few-shot examples, chat history. RadixAttention stores KV cache in a radix tree keyed by token IDs. Each edge = a run of tokens; each node = the KV blocks for that prefix.

Match: On new request, walk the tree from root, matching tokens edge-by-edge until divergence. All matched blocks are cache hits — skip prefill for them.

Insert: After forward pass, extend the matched leaf with the new tokens' KV blocks.

Evict: LRU on leaf nodes only (internal nodes have children still in use). Eviction returns blocks to the free pool.

Vs vLLM paged blocks: vLLM allocates blocks per-request with a hash lookup for shared prefixes. RadixAttention uses the tree structure itself as the index — cheaper matching, natural hierarchical eviction. RadixAttention wins on chat/agent workloads (long shared prefixes); vLLM's simpler block allocator is more predictable under uniform random workloads.

Structured Output: xgrammar + Outlines

Constrained decoding masks logits to only valid tokens per an FSM.

• Outlines compiles a regex or JSON schema to an FSM once. At each step, look up allowed tokens for current state, mask logits with -inf elsewhere.
• xgrammar advances the FSM on GPU in parallel with sampling. The token mask is applied inside the sampling kernel, not as a separate CPU→GPU copy per step. This cuts structured decoding overhead from ~30% to <5%.

SGLang wires xgrammar directly into its sampler. The grammar object is pre-compiled at request time; the FSM state advances alongside the accepted token.

SGLang vs vLLM

Rule of thumb: agents, chat, RAG → SGLang; batch inference, uniform prompts → vLLM. Both wrap FlashInfer for the actual attention kernels.

Phase 4 — FlashInfer

API Surface

FlashInfer separates plan (metadata setup) from forward (kernel launch). Two main wrappers:

from flashinfer import BatchPrefillWithPagedKVCacheWrapper, BatchDecodeWithPagedKVCacheWrapper

# Prefill
prefill = BatchPrefillWithPagedKVCacheWrapper(workspace_buffer, kv_layout="NHD")
prefill.plan(
    qo_indptr, paged_kv_indptr, paged_kv_indices, paged_kv_last_page_len,
    num_qo_heads, num_kv_heads, head_dim, page_size,
    causal=True, sm_scale=1/math.sqrt(head_dim), q_data_type=torch.float16
)
out = prefill.run(q, paged_kv_cache)

# Decode
decode = BatchDecodeWithPagedKVCacheWrapper(workspace_buffer, kv_layout="NHD")
decode.plan(paged_kv_indptr, paged_kv_indices, paged_kv_last_page_len,
            num_qo_heads, num_kv_heads, head_dim, page_size, q_data_type=torch.float16)
out = decode.run(q, paged_kv_cache)

The plan step is called once per batch shape; run is called per forward pass. This amortizes CUDA graph capture and heuristic dispatch.

Cascade Attention + Shared Prefix

Two-level attention when many requests share a prefix:

1. Compute attention over the shared prefix once → produces partial output + log-sum-exp state.
2. Compute attention over each request's unique suffix.
3. Merge the two via flashinfer.merge_state — log-sum-exp combine of the two partial outputs.

Math: if S1 has (o1, lse1) and S2 has (o2, lse2), then merged output = o1 × sigmoid(lse1 - lse2) + o2 × sigmoid(lse2 - lse1) and merged lse = logaddexp(lse1, lse2). Correct because softmax is associative under log-sum-exp merge.

Win: 100 requests sharing a 4K-token system prompt → shared prefix compute is done once, not 100×. Effective throughput scales linearly with batch on the suffix, near-constant on the prefix.

Reading a FlashInfer Triton Decode Kernel

Skeleton of the decode kernel (one head per program):

@triton.jit
def decode_kernel(Q, K_cache, V_cache, block_table, seq_lens, O,
                  BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr):
    batch_idx = tl.program_id(0)
    head_idx = tl.program_id(1)

    # 1. Load Q for this batch/head
    q = tl.load(Q + batch_idx * num_heads * HEAD_DIM + head_idx * HEAD_DIM
                + tl.arange(0, HEAD_DIM))

    seq_len = tl.load(seq_lens + batch_idx)
    m = -float('inf')
    l = 0.0
    o = tl.zeros([HEAD_DIM], dtype=tl.float32)

    # 2. Iterate blocks via block_table indirection
    for blk_idx in range(0, seq_len, BLOCK_SIZE):
        phys_blk = tl.load(block_table + batch_idx * max_blocks + blk_idx // BLOCK_SIZE)
        k_ptr = K_cache + phys_blk * BLOCK_SIZE * HEAD_DIM
        v_ptr = V_cache + phys_blk * BLOCK_SIZE * HEAD_DIM

        # Load K, V tile
        k = tl.load(k_ptr + tl.arange(0, BLOCK_SIZE)[:, None] * HEAD_DIM
                    + tl.arange(0, HEAD_DIM)[None, :])
        v = tl.load(v_ptr + tl.arange(0, BLOCK_SIZE)[:, None] * HEAD_DIM
                    + tl.arange(0, HEAD_DIM)[None, :])

        # 3. Scores + online softmax
        s = tl.sum(q[None, :] * k, axis=1) * sm_scale
        m_new = tl.maximum(m, tl.max(s))
        p = tl.exp(s - m_new)
        alpha = tl.exp(m - m_new)
        l = alpha * l + tl.sum(p)
        o = alpha * o + tl.sum(p[:, None] * v, axis=0)
        m = m_new

    # 4. Final normalize + write out
    o = o / l
    tl.store(O + batch_idx * num_heads * HEAD_DIM + head_idx * HEAD_DIM
             + tl.arange(0, HEAD_DIM), o)

Reading tip: the four zones are (1) Q load, (2) K/V iteration via page table indirection, (3) online softmax + accumulate, (4) normalize + write.

Phase 5 — TransformerEngine + Low Precision

FP8 Formats

E4M3 has tighter range but more mantissa bits → higher precision within range. E5M2 matches BF16's range but with only 2 mantissa bits — accepts wider dynamic range for gradients where large magnitudes matter more than precision.

Delayed vs Current Scaling

FP8 needs per-tensor scale factors to fit values into the narrow range. Two strategies:

• Current scaling: Compute amax (absolute max) of the tensor, derive scale, cast. Requires a full pass over the tensor before the FP8 op. Cost: extra kernel launch + memory traffic.
• Delayed scaling: Maintain a history buffer of past amax values (typically last 1024 steps). Use the max of history to set the next scale. Updates amax_history_buffer during the same forward pass. Overlaps amax collection with the matmul, near-zero extra latency.

Delayed scaling is the default in TE for training. Current scaling used for inference where distributions are stable.

TE Fused Kernels

The wins come from eliminating intermediate HBM writes. A standard MLP: LN write → up-proj read/write → GeLU read/write → down-proj read/write. Fused: one read, one write. On memory-bound decode, this is 2-3× throughput.

FP4 + Block Scaling on Blackwell

FP4 has only 4 bits total — insufficient dynamic range for a whole tensor. Block scaling groups elements and gives each block its own scale:

• MXFP4 (OCP standard) — 32-element blocks, each with an E8M0 (8-bit exponent-only) scale.
• NVFP4 (NVIDIA) — 16-element blocks, each with an E4M3 (FP8) scale. Finer granularity, higher accuracy.

Blackwell's second-gen Transformer Engine hardware-accelerates block scaling; no software overhead.

TE Recipe Differences

import transformer_engine.pytorch as te
from transformer_engine.common.recipe import Format, DelayedScaling

# Delayed scaling recipe (training)
fp8_recipe = DelayedScaling(
    margin=0,
    fp8_format=Format.HYBRID,  # E4M3 fwd, E5M2 bwd
    amax_history_len=1024,
    amax_compute_algo="max",
)

model = te.LayerNormMLP(hidden_size=4096, ffn_hidden_size=16384)

with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
    output = model(input_tensor)
    loss = criterion(output, target)
loss.backward()  # gradients auto-cast to E5M2

Phase 6 — cuDNN Frontend

Graph API

Legacy cuDNN was op-based: one call per operation, opaque descriptors. The graph API (v9+) is declarative:

1. Build — describe the compute as a DAG of tensor ops.
2. Validate — cuDNN checks shapes, dtypes.
3. Heuristics — Mode A (fast, may miss best config) or Mode B (thorough).
4. Engine — cuDNN picks a fused implementation.
5. Plan — compiled, ready to execute.

Same graph can be replayed on new tensors without rebuild. Enables cuDNN's kernel fusion engine to see the whole pattern (matmul + bias + softmax + matmul) and emit one fused kernel.

Fused SDPA Graph

cuDNN wins over FlashInfer when you need arbitrary attention masks (causal + sliding window, block-diagonal, custom bias):

#include <cudnn_frontend.h>
namespace fe = cudnn_frontend;

auto graph = std::make_shared<fe::graph::Graph>();
graph->set_io_data_type(fe::DataType_t::HALF)
     .set_compute_data_type(fe::DataType_t::FLOAT);

auto Q = graph->tensor(fe::graph::Tensor_attributes().set_name("Q")
    .set_dim({B, H, S, D}).set_stride({H*S*D, S*D, D, 1}));
auto K = graph->tensor(fe::graph::Tensor_attributes().set_name("K")
    .set_dim({B, H, S, D}).set_stride({H*S*D, S*D, D, 1}));
auto V = graph->tensor(fe::graph::Tensor_attributes().set_name("V")
    .set_dim({B, H, S, D}).set_stride({H*S*D, S*D, D, 1}));

auto sdpa_attrs = fe::graph::SDPA_attributes()
    .set_is_inference(false)
    .set_causal_mask(true)
    .set_sliding_window_length(4096)   // combine causal + sliding
    .set_attn_scale(1.0f / std::sqrt(D));

auto [O, stats] = graph->sdpa(Q, K, V, sdpa_attrs);
O->set_output(true).set_dim({B, H, S, D});

graph->validate();
graph->build_operation_graph(handle);
graph->create_execution_plans({fe::HeurMode_t::A});
graph->build_plans(handle);

FlashInfer wins on straightforward causal attention with paged KV. cuDNN wins when you need mask combinations FlashInfer doesn't ship, or when you want kernel fusion across attention + LayerNorm boundaries.

Phase 7 — RL Post-Training + Parallels

7.1 Objective Cheat-Sheet

PPO (Schulman 2017) — the workhorse. Clipped surrogate + value + entropy:

L_PPO = E[ min(r_t(θ) A_t, clip(r_t(θ), 1-ε, 1+ε) A_t) ]
        - c1 · E[(V_θ(s_t) - V_target)²]
        + c2 · E[H(π_θ(·|s_t))]
        - β · KL(π_θ || π_ref)
where r_t(θ) = π_θ(a_t|s_t) / π_old(a_t|s_t)

Needs a critic network (value model) — usually a copy of the LM with a scalar head. Doubles training memory.

GRPO (DeepSeek-Math 2024) — group-relative, no value model:

A_i = (r_i - mean(r_1..r_G)) / std(r_1..r_G)
L_GRPO = E[ min(π_θ/π_old · A_i, clip(...) · A_i) ] - β · KL(π_θ || π_ref)

For each prompt, sample G rollouts (typically 8-16). Advantage = z-score of the group's rewards. No critic. Cuts training memory ~40% and removes a whole training loop.

DPO (Rafailov 2023) — offline, no rollouts, no reward model:

L_DPO = -E[(x,yw,yl)~D] log σ( β · log(π_θ(yw|x)/π_ref(yw|x))
                                - β · log(π_θ(yl|x)/π_ref(yl|x)) )

Uses paired (winner, loser) preference data directly. Implicit reward = log-ratio. Cheap but ceiling is lower than online RL.

RLOO (Ahmadian 2024) — leave-one-out baseline:

A_i = r_i - (1/(G-1)) · sum(r_j for j != i)

Like GRPO but subtracts the mean of the other rollouts. Unbiased. Simpler than GRPO (no std normalization), competitive on many tasks.

Why GRPO won for reasoning: No critic to train/pretrain/tune. Group mean as baseline is high-variance-reducing when rollouts are diverse. Cheaper wall-clock than PPO. DeepSeek-R1 proved it scales to frontier reasoning.

7.2 Deep Dive: slime (THUDM)

slime = SGLang rollout engine + Megatron trainer, disaggregated.

Why disaggregation matters: - Rollout = autoregressive decode = memory-bound. Wants many small GPUs or high-bandwidth chips. - Training = big matmul on gradients = compute-bound. Wants fewest, fattest GPUs. - Co-locating wastes both. Disaggregating lets each side scale independently.

Placement: e.g. 32 GPUs for rollout (SGLang serving), 64 GPUs for training (Megatron TP+PP). They talk over NCCL.

Weight sync: After a training step, the new weights are broadcast rollout-side via NCCL broadcast. slime uses tensor-parallel-aware sharding: each rollout worker gets its own shard directly, no gather-then-scatter. Sync cost ~5-15s per iteration for a 7B model on 100G Ethernet.

Rollout loop: 1. Sampler on train side picks prompts. 2. Send to SGLang rollout cluster. 3. SGLang generates G rollouts per prompt (with RadixAttention for shared prefix if any). 4. Rollouts + logprobs return to trainer. 5. Reward model scores; GRPO/PPO loss computed. 6. Megatron backward + optimizer step. 7. Broadcast new weights to rollout.

7.3 Deep Dive: Miles (radixark)

Miles is a slime fork with enterprise deltas:

• Schedulers — native K8s operator + Slurm submitter. slime is bare torchrun; Miles adds resource requests, node selectors, gang scheduling.
• Checkpointing — async sharded. slime writes checkpoints synchronously (blocks training for minutes on 70B). Miles snapshots to CPU RAM, then background thread flushes to object storage. Recovery restore in seconds instead of minutes.
• Observability — Prometheus metrics: rollout tok/s, KL divergence per step, reward mean/std/histogram, weight sync latency, GPU util per role. Grafana dashboards shipped.
• Fault tolerance — rollout worker crash doesn't kill training; trainer keeps going with reduced group size, spawns replacement.

Miles ≠ new algorithm; it's slime made production-safe.

7.4 The Four RL Camps

All four eventually call SGLang or vLLM for rollout. The real differentiator is the trainer stack + orchestration layer.

HybridFlow (verl): Single controller programs multi-controller execution. You write rollout(), reward(), train() as normal Python functions; verl schedules them across arbitrary GPU placements (co-located, disaggregated, mixed). Most flexible but hardest to debug.

7.5 Quick-Start Decision Tree

• Prototyping ≤ 7B on 1 node → TRL
• Scaling 70B+ on Ray infra → OpenRLHF
• Need flexible placement, experimenting → verl
• Production reasoning models, disaggregated → slime → Miles

Phase 8 — Capstone

Fine-tune Qwen2.5-7B on NuminaMath-CoT with Miles + GRPO on 1× H100, serve with SGLang + FlashInfer + FP8 TE, profile with Nsight Compute. Budget: ~$11.

8.1 Provision H100 (RunPod)

# Via RunPod CLI (see runpod skill)
runpodctl create pod \
  --name miles-grpo \
  --gpuType "NVIDIA H100 80GB HBM3" \
  --imageName nvcr.io/nvidia/pytorch:24.10-py3 \
  --containerDiskInGb 100 \
  --volumeInGb 200 \
  --ports "22/tcp,8000/http"

# ssh in

Spot pricing ~$2.50-3/hr. Expect ~3 hr training = ~$9.

8.2 Install Stack

git clone https://github.com/radixark/miles && cd miles
pip install -e ".[all]"
pip install "sglang[all]>=0.4" flashinfer-python
pip install transformer-engine[pytorch]

Verify:

python -c "import sglang, flashinfer, transformer_engine.pytorch as te; print('ok')"

8.3 Prep Dataset

python -c "
from datasets import load_dataset
ds = load_dataset('AI-MO/NuminaMath-CoT', split='train').select(range(20000))
ds.to_parquet('numina_20k.parquet')
"

20K prompts is enough for GRPO convergence on math in 3 hours.

8.4 Train (GRPO)

torchrun --nproc_per_node=1 miles/train_grpo.py \
  --model Qwen/Qwen2.5-7B \
  --dataset numina_20k.parquet \
  --rollout_engine sglang \
  --rollouts_per_prompt 8 \
  --learning_rate 1e-6 \
  --kl_coef 0.001 \
  --max_new_tokens 1024 \
  --batch_size 4 \
  --gradient_accumulation 8 \
  --max_steps 500 \
  --save_dir ./ckpt/qwen-grpo-math \
  --reward_fn math_verify \
  --log_prometheus_port 9090

Expected timeline (1× H100): - Rollout: ~40s/step (8 rollouts × 4 prompts × 1024 tok) - Train: ~4s/step - ~45s/step × 240 steps ≈ 3 hr

Watch reward mean climb from ~0.15 → ~0.55 on math_verify.

8.5 Quantize + Serve

# TE FP8 calibration
python miles/quantize_fp8.py \
  --model ./ckpt/qwen-grpo-math \
  --calib_dataset numina_20k.parquet \
  --calib_samples 512 \
  --out ./ckpt/qwen-grpo-math-fp8

# Launch SGLang w/ FlashInfer + FP8
python -m sglang.launch_server \
  --model-path ./ckpt/qwen-grpo-math-fp8 \
  --quantization fp8 \
  --attention-backend flashinfer \
  --port 8000 \
  --enable-radix-cache

8.6 Benchmark

python -m sglang.bench_serving \
  --backend sglang \
  --host 127.0.0.1 --port 8000 \
  --num-prompts 500 --request-rate 8 \
  --dataset-name random --random-input 512 --random-output 512

Expected on H100: - batch=1 decode: ~140 tok/s - batch=16 decode: ~1200 tok/s - Prefill (512 tok): ~15 ms - E2E (in=512, out=512) @ batch=16: ~450 ms/req

8.7 Profile

# Single-request profile
ncu --set full --target-processes all \
    -o qwen_decode.ncu-rep \
    python -c "
import sglang
llm = sglang.Runtime(model_path='./ckpt/qwen-grpo-math-fp8',
                     quantization='fp8', attention_backend='flashinfer')
print(llm.generate('Solve: 2x+3=7', {'max_new_tokens': 64}))
"

# Open in Nsight Compute GUI

Check: - Attention kernel: Achieved Occupancy > 40%, Mem Busy > 85% (decode is memory-bound → this is expected). - MLP FP8 kernel: Tensor Core Utilization > 70%. - Warp stalls dominated by long_scoreboard on decode (waiting on KV load from HBM) — normal.

If tensor cores underutilized on prefill, check FP8 recipe applied. If memory bandwidth < 70% on decode, KV cache layout is wrong (want contiguous heads).

8.8 Deliverables

• Trained checkpoint: ./ckpt/qwen-grpo-math-fp8
• Training curves: http://<pod-ip>:9090/metrics scraped to Grafana
• Serving benchmark: bench.json
• Profile: qwen_decode.ncu-rep
• Total cost: ~$11 (3 hr train + 1 hr eval on spot H100)

Repos

• sgl-project/sglang
• flashinfer-ai/flashinfer
• NVIDIA/TransformerEngine
• radixark/miles
• THUDM/slime (upstream)
• volcengine/verl
• OpenRLHF/OpenRLHF
• huggingface/trl