Skip to content

triton attention: mixed-batch split dispatch (decode rows -> decode kernels) (reorg 6/9) - #6

Open
mike-char wants to merge 3 commits into
cai-pr5-triton-attn-tuningfrom
cai-pr6-mixed-split-dispatch
Open

triton attention: mixed-batch split dispatch (decode rows -> decode kernels) (reorg 6/9)#6
mike-char wants to merge 3 commits into
cai-pr5-triton-attn-tuningfrom
cai-pr6-mixed-split-dispatch

Conversation

@mike-char

@mike-char mike-char commented Aug 26, 2026

Copy link
Copy Markdown

Part 6 of 9 — last OSS-eligible layer; the private stack (sglang-internal PRs) forks off this tip.

--enable-mixed-chunk folds running decodes into the extend batch as 1-token varlen rows; the extend kernel re-scans each row's whole KV prefix without GQA amortization or split-KV (~4x per-decode-token on gemma-4; capacity collapse under trickle arrivals). Serve a MIXED batch as two contiguous sub-calls instead: forward_extend over the prefill prefix (sliced indptr views, whole-batch KV save), then forward_decode over the 1-token suffix with real decode metadata. Upstream's forward_mixed hook exists but is NPU-gated; this generalizes it to a supports_forward_mixed capability flag. Kill switch SGLANG_TRITON_MIXED_SPLIT_DISPATCH=0; auto-fallback for MLA/DCP/spec/deterministic. Plus boot-time preallocation of the MIXED decode mid-buffers (~0.5 GB allocator churn per mixed step otherwise; only allocated when mixed-chunk is on).

Measured (prod-shaped replay, MI325X, 26b fp8): +25% served over interval-2 baseline at saturation, errors halved, TTFT p95 -35%. CPU unit tests + serving test included.

Upstream candidate: strongest in the stack — every upstream CUDA/ROCm mixed-chunk user currently pays the fold-into-extend tax.

🤖 Generated with Claude Code
mike-char and others added 3 commits August 26, 2026 17:24
@mike-char @claude
triton-attn: mixed-batch split dispatch (decode rows -> decode kernels)
6a23cb9 
--enable-mixed-chunk folds running decodes into the extend batch as 1-token
varlen rows; the Triton extend kernel then re-scans each row's whole KV
prefix without GQA amortization or split-KV (~4x per-decode-token on
Gemma-4, the mixed-chunk capacity collapse). Serve a MIXED batch as two
sub-calls with each row type's native kernel instead:

- mix_with_running appends the running decodes after the prefill rows, so
  the split is two contiguous slices found from the CPU extend lens
  (sync-free, once per step): forward_extend over the prefill prefix via
  sliced indptr views of the step metadata; forward_decode over the 1-token
  suffix via decode-mode metadata (full seq_lens, windowed SWA indices,
  split-KV policy, fresh mid-buffers) built in init_forward_metadata.
- The extend sub-call's whole-batch KV save runs first, so the decode
  kernel reads the suffix rows' current-token KV like a pure decode step.
- Dispatch via a new supports_forward_mixed backend capability replacing
  the NPU-only forward_mixed gate; each sub-call keeps its existing g4-HIP
  / tuned-Triton dispatch and tuning-table shape keys.
- SGLANG_TRITON_MIXED_SPLIT_DISPATCH=0 restores fold-into-extend (A/B kill
  switch); spec/MLA/DCP/deterministic and gpu-only metadata fall back the
  same way. First engagement is boot-logged for activation proof.

Unit test covers the partition/slicing/orchestration logic; the new
registered Gemma4 serving test drives MIXED steps under concurrency.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
(cherry picked from commit 83bb8eed41327138637ead4d3ae2f6d34b0725f7)
@mike-char @claude
test: drop maybe_stub_sgl_kernel from mixed-split unit test
6fde09a 
The stub helper's import chain (test_utils -> benchmark.serving ->
quantization) pulls in aiter and GPU probes, which the test does not need;
the sibling backend unit tests import triton_backend directly without it.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
(cherry picked from commit aec8c5ade96c1d5a1c7dc0af05650ccfbccaef32)
@mike-char @claude
triton-attn: preallocate MIXED decode mid-buffers (fix HBM exhaustion)
5660452 
Under sustained mixed-chunk load the server aborts with
HSA_STATUS_ERROR_OUT_OF_RESOURCES while the KV pool is only ~half full:
every mixed step allocated ~0.5 GB of freshly-shaped decode mid-buffers
(attn_logits/swa_attn_logits/attn_lse sized to that step's decode row
count, which grows with load), churning the caching allocator on top of
the 16k-chunk activation transients until the transient region hit zero.

Allocate the buffers once at max batch size and row-slice per step
(views keep the strides of a same-shape contiguous tensor, so kernel
contracts are unchanged). Allocation happens at boot, and only when
--enable-mixed-chunk is set, so plain Triton deployments pay nothing and
the ~0.5 GB is claimed while headroom is guaranteed rather than
mid-traffic; a lazy fallback covers any future MIXED producer. The small
indptr/split tensors move to persistent buffers too (zeros-init keeps
the cumsum base valid across reuse); kv_indices stays per-step
(variable-length, same scale as the extend path's own per-step indices).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
(cherry picked from commit cd88875fdbd41e6694f49680f97a05a3a4389345)
@mike-char
mike-char force-pushed the cai-pr5-triton-attn-tuning branch from 364e93e to 25b01af Compare August 27, 2026 00:27
@mike-char
mike-char force-pushed the cai-pr6-mixed-split-dispatch branch from d66dbbb to 5660452 Compare August 27, 2026 00:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

1 participant

@mike-char