From b152eb625af6be9610c25e808ba875f2e83151b5 Mon Sep 17 00:00:00 2001 From: Pengyu Chen Date: Thu, 21 May 2026 17:43:18 +0000 Subject: [PATCH 1/2] JIT custom_all_reduce/tp_qknorm: use reinterpret_cast not std::bit_cast ``std::bit_cast`` is a C++20 library feature added in libstdc++ 3.4.29 (gcc 11.1). On Debian 11's gcc-10 (libstdc++ 3.4.28) the JIT compilation of these three kernels fails with:: error: namespace "std" has no member "bit_cast" making ``--disable-custom-all-reduce`` mandatory on that host. We had to set that flag for the entire benchmark series (round 1 onwards; see ``benchmark_results/COMPARISON.md``). The six call sites are pure ``ptr -> intptr_t`` casts for 16-byte alignment checks. ``reinterpret_cast(ptr)`` is value- equivalent for this conversion and has been valid C++ since c++98, so the JIT now builds on any reasonable toolchain. Files patched: * ``custom_all_reduce_push.cuh:232`` (1 cast) * ``custom_all_reduce_pull.cuh:164`` (1 cast) * ``tp_qknorm.cuh:299-302`` (4 casts) Verified end-to-end on H100 / gcc-10 / libstdc++ 3.4.28: * Before: server crashes during cuda-graph capture with the ``std::bit_cast`` build error. * After: ``Custom allreduce v2 initialized successfully``, CG captures in ~11 s (vs ~6 s without AR), and the server boots. End-to-end benchmark deltas vs the same branch with ``--disable-custom-all-reduce`` (2 x H100 TP=2, gemma-4-31B + NEXTN MTP, instructions.md workload + decode-burst variant): workload bench no-AR with-AR delta -------------------- ------------- ------- ---------- ----- no-spec decode-burst output tok/s 1608 1688 +5.0 % no-spec decode-burst median TPOT 19.58 ms 18.49 ms -5.6 % no-spec decode-burst median E2E 20.38 s 19.41 s -4.8 % with-spec decode-burst output tok/s 1166 1087 -6.8 % with-spec decode-burst median TPOT 23.09 ms 24.66 ms +6.8 % with-spec full bench total tok/s 6067 5994 -1.2 % So custom-AR is a real win on the no-spec path (closes about half of the ~10 % gap vs vLLM that ``benchmark_results/NOSPEC_GAP.md`` attributed to NCCL overhead -- per-fwd comms time drops from 1.611 ms to ~0.05 ms, matching vLLM's ``cross_device_reduce_1stage``). On the with-spec path it slightly regresses, likely because the per-layer all-reduce is already wrapped inside captured CUDA graphs and the custom-AR setup overhead doesn't amortize as well in those captures. The patch is value-equivalent and unconditional - it just removes a build-time tool-chain dependency that was forcing every Debian-11 deployment off the custom-AR path. Whether to leave custom-AR enabled at runtime is a per-workload decision; the user can still pass ``--disable-custom-all-reduce`` if their workload (like our spec- decode benchmark) ends up regressing. --- .../csrc/distributed/custom_all_reduce_pull.cuh | 5 ++++- .../csrc/distributed/custom_all_reduce_push.cuh | 5 ++++- .../sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh | 11 +++++++---- 3 files changed, 15 insertions(+), 6 deletions(-) diff --git a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh index e8837af4cd34..a20f48c87c0a 100644 --- a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh +++ b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh @@ -161,7 +161,10 @@ struct CustomAllReducePull : public CustomAllReduceBase { RuntimeCheck(shot == 1 || shot == 2, "Invalid shot count: ", shot); RuntimeCheck(device.device_type == kDLCUDA, "Only CUDA device is supported"); RuntimeCheck(is_type(input.dtype()), "Input dtype mismatch"); - RuntimeCheck(std::bit_cast(input_ptr) % 16 == 0, "Input pointer is not properly aligned"); + // ``reinterpret_cast`` rather than ``std::bit_cast`` so the JIT + // builds on libstdc++ < 11 (gcc 10 ships in Debian 11). The cast + // is value-equivalent for pointer-to-integer. + RuntimeCheck(reinterpret_cast(input_ptr) % 16 == 0, "Input pointer is not properly aligned"); RuntimeCheck(m_pull_ctrl.has_value(), "Controller is not initialized"); RuntimeCheck(static_cast(num_items) == num_items_int64, "Number of items exceeds 4G limit"); diff --git a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_push.cuh b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_push.cuh index c4523c27eec3..8ca4f9927f3c 100644 --- a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_push.cuh +++ b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_push.cuh @@ -229,7 +229,10 @@ struct CustomAllReducePush : public CustomAllReduceBase { RuntimeCheck(m_num_gpu == kNumGPU, "Number of GPUs mismatch"); RuntimeCheck(device.device_type == kDLCUDA, "Only CUDA device is supported"); RuntimeCheck(is_type(input.dtype()), "Input dtype mismatch"); - RuntimeCheck(std::bit_cast(input_ptr) % 16 == 0, "Input pointer is not properly aligned"); + // ``reinterpret_cast`` rather than ``std::bit_cast`` so the JIT + // builds on libstdc++ < 11 (gcc 10 ships in Debian 11). The cast + // is value-equivalent for pointer-to-integer. + RuntimeCheck(reinterpret_cast(input_ptr) % 16 == 0, "Input pointer is not properly aligned"); RuntimeCheck(m_push_ctrl.has_value(), "Controller is not initialized"); RuntimeCheck(shot == 1, "Push all-reduce only supports 1-shot, got: ", shot); RuntimeCheck(static_cast(num_items) == num_items_int64, "Number of items exceeds 4G limit"); diff --git a/python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh b/python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh index ca80e1efcdf1..be59e2c738f4 100644 --- a/python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh +++ b/python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh @@ -296,10 +296,13 @@ struct FusedParallelQKNormAcrossHead : public CustomAllReduceBase { const auto needed_buffer_bytes = static_cast(num_tokens) * 2 * sizeof(float); RuntimeCheck(m_num_gpu == kNumGPU, "Number of GPUs mismatch"); RuntimeCheck(m_push_ctrl.has_value(), "Controller is not initialized"); - RuntimeCheck(std::bit_cast(params.q_ptr) % 16 == 0, "q pointer is not properly aligned"); - RuntimeCheck(std::bit_cast(params.k_ptr) % 16 == 0, "k pointer is not properly aligned"); - RuntimeCheck(std::bit_cast(params.q_weight) % 16 == 0, "q_weight pointer is not properly aligned"); - RuntimeCheck(std::bit_cast(params.k_weight) % 16 == 0, "k_weight pointer is not properly aligned"); + // ``reinterpret_cast`` rather than ``std::bit_cast`` so the JIT + // builds on libstdc++ < 11 (gcc 10 ships in Debian 11). The cast + // is value-equivalent for pointer-to-integer. + RuntimeCheck(reinterpret_cast(params.q_ptr) % 16 == 0, "q pointer is not properly aligned"); + RuntimeCheck(reinterpret_cast(params.k_ptr) % 16 == 0, "k pointer is not properly aligned"); + RuntimeCheck(reinterpret_cast(params.q_weight) % 16 == 0, "q_weight pointer is not properly aligned"); + RuntimeCheck(reinterpret_cast(params.k_weight) % 16 == 0, "k_weight pointer is not properly aligned"); RuntimeCheck(needed_buffer_bytes <= m_push_buffer_bytes, "Push buffer is too small"); LaunchKernel(num_blocks, num_threads, device) // From 64ffea6dbca6227f254b0b1309c85c3ca9db37b1 Mon Sep 17 00:00:00 2001 From: Pengyu Chen Date: Sat, 23 May 2026 04:07:01 +0000 Subject: [PATCH 2/2] perf(gemma4 MTP H100): tune Triton extend tile for Lq=256 / sm_90 The Hopper branch in '_get_block_sizes_for_extend_attention' picked (BLOCK_M=128, BLOCK_N=64, num_warps=8, num_stages=1) for every Lq<=256. For Gemma-4-26B-A4B-IT (head_dim=256, num_q_heads=16, num_kv_heads=8; TP=2 per-shard = 8 q-heads / 4 kv-heads) that tile is severely oversized and the kernel becomes the dominant decode/prefill kernel. Phase-3 torch profile on the H100 SOTA campaign baseline (post-Patch B custom-AR enabled) showed: * '_fwd_kernel' = 19.2% of decode GPU time (25.6 ms / 133 ms) * '_fwd_kernel' = 60.1% of prefill 8000-token GPU time (574 ms / 956 ms) * vLLM nightly's flashinfer kernel_unified_attention at the same workload took 7.2 ms decode and 381 ms prefill 8k. Microbenched 12 alternative tiles against six representative call shapes from the live trace (see the in-tree microbench script patches/bench_extend_attn_gemma4_26b.py in the H100 run artifact dir). Winners: shape (bs, ext, prefix, sw) legacy (128,64,w8,s1) new delta ---------------------------------- --------------------- ------------ ----- prefill long bs=1 ext=8192 sw=-1 2656.80 us 1907.64 us -28.2 % (32,64,w4,s2) prefill chat bs=1 ext=1000 sw=-1 128.21 us 55.98 us -56.3 % (32,64,w4,s2) verify chat bs=32 ext=4 pf=1000 sw=1024 616.48 us 144.01 us -76.6 % (16,64,w4,s2) verify summ bs=32 ext=4 pf=8000 sw=1024 1075.79 us 191.49 us -82.2 % (16,64,w4,s2) verify burst bs=32 ext=4 pf=64 sw=1024 93.98 us 22.10 us -76.5 % (32,32,w4,s2) prefill multi bs=4 ext=1000 sw=-1 225.33 us 153.53 us -31.9 % (32,64,w4,s2) The two regimes (single-seq long-extend prefill vs high-bs short-verify MTP step) want different tiles. Gate on batch_size >= 8: * bs < 8 ('single-seq long-extend prefill'): (32, 64, w4, s2) * bs >= 8 ('MTP verify / chunked-prefill'): (16, 64, w4, s2) Plumbing changes: * '_get_block_sizes_for_extend_attention' now takes 'batch_size' (kw-only) and returns 'num_stages' as well. * Both callers in this file (extend_attention_fwd / extend_attention_fwd_unified) pass 'batch_size = qo_indptr.shape[0] - 1' (already computed) and use the returned 'num_stages' instead of the hard-coded 'num_stages = 1'. Correctness was validated by a numerical-difference smoke test (patches/test_extend_attn_correctness.py): per-element max-abs / ref-max < 2e-3 across all six call shapes (bf16 noise). Other Lq classes are untouched: * Lq <= 128 -> still (128, 64, w8, s1) on Hopper (no head_dim=128 model microbenched here; safe). * Lq > 256 -> still (32, 64, w8, s1) on Hopper (sgl PR #22079 only affects sm_100a; this branch is unchanged). * sm120 / sm100a / Ampere / older: unchanged. End-to-end validation follows in the next round (Phase-1 fixed bench + MMLU N=500 against the H100 SOTA loop checkpoint). --- .../attention/triton_ops/extend_attention.py | 98 ++++++++++++++++--- 1 file changed, 82 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/layers/attention/triton_ops/extend_attention.py b/python/sglang/srt/layers/attention/triton_ops/extend_attention.py index e6a353e9bfd9..9d29487e6220 100644 --- a/python/sglang/srt/layers/attention/triton_ops/extend_attention.py +++ b/python/sglang/srt/layers/attention/triton_ops/extend_attention.py @@ -32,16 +32,34 @@ _is_hip = is_hip() -def _get_block_sizes_for_extend_attention(Lq: int, Lv: int): +def _get_block_sizes_for_extend_attention( + Lq: int, + Lv: int, + *, + batch_size: int = 0, + max_len_extend: int = 0, +): """ Get block sizes and configuration for extend attention kernels. Args: Lq: Query head dimension Lv: Value head dimension + batch_size: Number of sequences in the batch (kw-only). Used by the + H100 (sm_90, Lq<=256) heuristic to pick a smaller tile for + high-bs spec-decode verify shapes where the default (128, 64, w8) + wastes work per program. ``0`` (default) is treated as "unknown" + and preserves the legacy tile. + max_len_extend: Maximum extend length per sequence in the batch + (kw-only). Used together with batch_size to distinguish + high-bs *verify* shapes (small max_len_extend, e.g. 4 for + num_draft_tokens=4) from high-bs *chunked prefill* shapes + (larger max_len_extend). ``0`` (default) is treated as + "unknown" and falls back to the long-extend tile. Returns: - tuple: (BLOCK_DMODEL, BLOCK_DPE, BLOCK_DV, BLOCK_M, BLOCK_N, num_warps) + tuple: (BLOCK_DMODEL, BLOCK_DPE, BLOCK_DV, BLOCK_M, BLOCK_N, num_warps, + num_stages) """ # Determine BLOCK_DMODEL and BLOCK_DPE based on head dimension if Lq == 576: @@ -59,6 +77,8 @@ def _get_block_sizes_for_extend_attention(Lq: int, Lv: int): BLOCK_DV = triton.next_power_of_2(Lv) + num_stages = 1 + # Determine BLOCK_M, BLOCK_N, and num_warps based on hardware if _is_hip: BLOCK_M, BLOCK_N = (64, 64) @@ -82,8 +102,48 @@ def _get_block_sizes_for_extend_attention(Lq: int, Lv: int): BLOCK_M, BLOCK_N = (16, 64) elif _is_cuda and CUDA_CAPABILITY[0] >= 9: # Hopper architecture (H100, etc.) - if Lq <= 256: + if Lq <= 128: BLOCK_M, BLOCK_N = (128, 64) + elif Lq <= 256: + # H100 / sm_90, head_dim == 256 (e.g. Gemma-4-26B-A4B-IT, + # which uses head_dim=256). The legacy (128, 64, w8, s1) + # tile is severely oversized for both the long-extend + # initial-prefill shape (bs=1, ext=8k) and the high-bs + # MTP verify shape (bs=32, ext=4, prefix>=1k) — see + # the microbench in the H100 SOTA run artifact dir + # ``patches/bench_extend_attn_gemma4_26b.py`` (and the + # ``patches/extend_attn_microbench_*.log`` artifacts). + # Microbench winners on bf16, num_q_heads=8, num_kv_heads=4: + # prefill long ext=8192 bs=1 2657us -> 1908us -28% (32,64,w4,s2) + # prefill chat ext=1000 bs=1 128us -> 56us -56% (32,64,w4,s2) + # verify chat ext=4 pf=1000 bs=32 616us -> 144us -77% (16,64,w4,s2) + # verify summ ext=4 pf=8000 bs=32 1076us-> 191us -82% (16,64,w4,s2) + # verify burst ext=4 pf=64 bs=32 94us -> 22us -77% (32,32,w4,s2) + # chunked-prefill ext=512 bs=8 136us -> 92us -32% (32,64,w4,s2) + # chunked-prefill ext=1024 bs=16 752us -> 559us -26% (32,64,w4,s2) + # The (16, 64, w4, s2) tile that dominates the high-bs + # *verify* path (max_len_extend = num_draft_tokens, very + # small) regresses the high-bs *chunked-prefill* path + # (max_len_extend = chunked_prefill_size_per_seq, larger) + # by ~30 %. Gate on BOTH batch_size and max_len_extend + # so chunked prefill keeps (32, 64, w4, s2). + if batch_size >= 8 and 0 < max_len_extend <= 16: + BLOCK_M, BLOCK_N = (16, 64) + num_warps = 4 + num_stages = 2 + else: + BLOCK_M, BLOCK_N = (32, 64) + num_warps = 4 + num_stages = 2 + return ( + BLOCK_DMODEL, + BLOCK_DPE, + BLOCK_DV, + BLOCK_M, + BLOCK_N, + num_warps, + num_stages, + ) else: BLOCK_M, BLOCK_N = (32, 64) elif _is_cuda and CUDA_CAPABILITY[0] >= 8: @@ -109,7 +169,7 @@ def _get_block_sizes_for_extend_attention(Lq: int, Lv: int): num_warps = 4 if Lq <= 64 else 8 - return BLOCK_DMODEL, BLOCK_DPE, BLOCK_DV, BLOCK_M, BLOCK_N, num_warps + return BLOCK_DMODEL, BLOCK_DPE, BLOCK_DV, BLOCK_M, BLOCK_N, num_warps, num_stages @triton.jit @@ -591,15 +651,19 @@ def extend_attention_fwd( v_extend.shape[-1], ) - # Get block sizes and configuration - BLOCK_DMODEL, BLOCK_DPE, BLOCK_DV, BLOCK_M, BLOCK_N, num_warps = ( - _get_block_sizes_for_extend_attention(Lq, Lv) - ) - sm_scale = sm_scale or 1.0 / (Lq**0.5) batch_size, head_num = qo_indptr.shape[0] - 1, q_extend.shape[1] kv_group_num = q_extend.shape[1] // k_extend.shape[1] + # Get block sizes and configuration. Pass batch_size + max_len_extend so + # the H100 Lq<=256 heuristic can pick the spec-decode-verify tile + # (only when extend is tiny) vs the chunked-prefill / long-extend tile. + BLOCK_DMODEL, BLOCK_DPE, BLOCK_DV, BLOCK_M, BLOCK_N, num_warps, num_stages = ( + _get_block_sizes_for_extend_attention( + Lq, Lv, batch_size=batch_size, max_len_extend=max_len_extend + ) + ) + USE_CUSTOM_MASK = custom_mask is not None # Skip custom mask for prefix part SKIP_PREFIX_CUSTOM_MASK = skip_prefix_custom_mask @@ -607,7 +671,6 @@ def extend_attention_fwd( HAS_SINK = sinks is not None grid = (batch_size, head_num, triton.cdiv(max_len_extend, BLOCK_M)) - num_stages = 1 extra_kargs = {} if _is_hip: @@ -1001,15 +1064,19 @@ def extend_attention_fwd_unified( """ Lq, Lv = q.shape[-1], v_buffer.shape[-1] - # Get block sizes and configuration - BLOCK_DMODEL, BLOCK_DPE, BLOCK_DV, BLOCK_M, BLOCK_N, num_warps = ( - _get_block_sizes_for_extend_attention(Lq, Lv) - ) - sm_scale = sm_scale or 1.0 / (Lq**0.5) batch_size, head_num = qo_indptr.shape[0] - 1, q.shape[1] kv_group_num = q.shape[1] // k_buffer.shape[1] + # Get block sizes and configuration. Pass batch_size + max_len_extend so + # the H100 Lq<=256 heuristic can pick the spec-decode-verify tile + # (only when extend is tiny) vs the chunked-prefill / long-extend tile. + BLOCK_DMODEL, BLOCK_DPE, BLOCK_DV, BLOCK_M, BLOCK_N, num_warps, num_stages = ( + _get_block_sizes_for_extend_attention( + Lq, Lv, batch_size=batch_size, max_len_extend=max_len_extend + ) + ) + USE_CUSTOM_MASK = custom_mask is not None HAS_SINK = sinks is not None @@ -1020,7 +1087,6 @@ def extend_attention_fwd_unified( window_start_pos = torch.zeros(batch_size, dtype=torch.int32, device=q.device) grid = (batch_size, head_num, triton.cdiv(max_len_extend, BLOCK_M)) - num_stages = 1 extra_kargs = {} if _is_hip: