Skip to content

Commit 25c1447

Browse files
zyongyeclaude
andauthored
[Bugfix][Decode] Bound per-plan decode memory: shared Q8KV4 kernel workspace, caller-owned indexer score buffer (#45)
Callers that keep one decode plan or indexer wrapper per batch size (vLLM does, for every eager batch size above its CUDA graph range) grew GPU memory by up to 288 MiB per new size at a 1M-token context and ran out of memory under load. - Every Q8KV4 decode plan allocated its own 32 MiB CUTLASS workspace, which the kernels never read (get_workspace_size is 0). Plans on a device now share one buffer. - The decode indexer wrappers take an optional caller-owned score_buffer (fp32, score_buffer_size(batch, max_pages, num_heads) elements) for the page-score scratch, batch x 8 x num_heads x max_pages fp32 per wrapper, so wrappers that never run concurrently can share one. Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
1 parent 57d639f commit 25c1447

4 files changed

Lines changed: 106 additions & 4 deletions

File tree

‎python/fmha_sm100/cute/q8_indexer_interface.py‎

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -330,7 +330,9 @@ class _BatchDecodeIndexerBase:
330330
that position divided by 128. The eight tokens of all ``num_heads`` heads
331331
are scored together, reading each K page once. A wrapper instance owns its
332332
workspace and buffers and must not be shared by concurrently running
333-
streams.
333+
streams. ``score_buffer`` (fp32, at least ``batch * 8 * num_heads *
334+
max_pages`` elements) replaces the per-wrapper page-score scratch, so
335+
wrappers of different batch sizes that never run concurrently can share it.
334336
335337
Wrappers with ``_supports_query_len`` also plan requests of fewer queries
336338
(``query_len``): ``q`` and the output then hold only those, which take the
@@ -345,7 +347,12 @@ def __init__(
345347
use_cuda_graph: bool = False,
346348
block_table_buffer: torch.Tensor | None = None,
347349
seq_lens_buffer: torch.Tensor | None = None,
350+
score_buffer: torch.Tensor | None = None,
348351
) -> None:
352+
if score_buffer is not None:
353+
_check_cuda_tensor(score_buffer, name="score_buffer", dtype=torch.float32)
354+
if score_buffer.ndim != 1:
355+
raise ValueError("score_buffer must be one-dimensional")
349356
if workspace_buffer is not None:
350357
_check_cuda_tensor(workspace_buffer, name="workspace_buffer", dtype=torch.uint8)
351358
if workspace_buffer.ndim != 1:
@@ -368,6 +375,7 @@ def __init__(
368375
self._seq_lens_buffer = seq_lens_buffer
369376
self._block_table: torch.Tensor | None = None
370377
self._seq_lens: torch.Tensor | None = None
378+
self._score_buffer = score_buffer
371379
self._scores: torch.Tensor | None = None
372380
self._num_valid_pages: torch.Tensor | None = None
373381
self._topk_indices: torch.Tensor | None = None
@@ -434,6 +442,13 @@ def workspace_size(cls, batch_size: int, *, num_heads: int = 1) -> int:
434442
batch_size, _check_num_heads(num_heads, cls._supported_num_heads)
435443
)
436444

445+
@classmethod
446+
def score_buffer_size(cls, batch_size: int, max_pages: int, *, num_heads: int = 1) -> int:
447+
"""Return the fp32 elements ``score_buffer`` needs for a batch."""
448+
449+
num_heads = _check_num_heads(num_heads, cls._supported_num_heads)
450+
return batch_size * _DECODE_QUERY_LENGTH * num_heads * max_pages
451+
437452
@staticmethod
438453
def _check_metadata(block_table: torch.Tensor, seq_lens: torch.Tensor) -> None:
439454
_check_cuda_tensor(seq_lens, name="seq_lens", dtype=torch.int32)
@@ -499,13 +514,24 @@ def plan(
499514

500515
tokens = batch_size * query_len
501516
score_shape = (batch_size, _DECODE_QUERY_LENGTH * self._num_heads, max_pages)
502-
if (
517+
if self._score_buffer is not None:
518+
score_elems = self.score_buffer_size(batch_size, max_pages, num_heads=self._num_heads)
519+
if self._score_buffer.device != device or self._score_buffer.numel() < score_elems:
520+
raise ValueError(
521+
f"score_buffer must be on {device} with at least {score_elems} elements"
522+
)
523+
self._scores = self._score_buffer[:score_elems].view(score_shape)
524+
elif (
503525
self._scores is None
504526
or self._scores.shape != score_shape
505527
or self._scores.device != device
506-
or self._topk_indices.shape[0] != tokens
507528
):
508529
self._scores = torch.empty(score_shape, dtype=torch.float32, device=device)
530+
if (
531+
self._topk_indices is None
532+
or self._topk_indices.shape[0] != tokens
533+
or self._topk_indices.device != device
534+
):
509535
self._num_valid_pages = torch.empty(
510536
(tokens * self._num_heads,), dtype=torch.int32, device=device
511537
)

‎python/fmha_sm100/cute/test_q8_indexer.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -552,6 +552,53 @@ def test_decode_concurrent_streams_do_not_alias(fmt, num_heads):
552552
assert torch.equal(out, expected)
553553

554554

555+
@pytest.mark.parametrize("fmt,num_heads", DECODE_CASES)
556+
def test_decode_wrappers_share_a_score_buffer(fmt, num_heads):
557+
"""Graph-mode wrappers of two batch sizes share one score buffer: replaying one after the
558+
other ran matches its own result, and a buffer too small for the batch is rejected."""
559+
560+
max_pages = 6
561+
inputs = [
562+
_make_decode(fmt, batch, max_pages, seed, num_heads=num_heads)
563+
for batch, seed in ((64, 51), (96, 52))
564+
]
565+
shared = torch.empty(
566+
DECODE_WRAPPERS[fmt].score_buffer_size(96, max_pages, num_heads=num_heads),
567+
dtype=torch.float32,
568+
device="cuda",
569+
)
570+
wrappers, outs, expected = [], [], []
571+
for q, k, block_table, seq_lens in inputs:
572+
own = DECODE_WRAPPERS[fmt](num_heads=num_heads)
573+
own.plan(block_table, seq_lens)
574+
expected.append(own.run(q, k.cache).clone())
575+
wrapper = DECODE_WRAPPERS[fmt](
576+
num_heads=num_heads,
577+
use_cuda_graph=True,
578+
block_table_buffer=torch.empty_like(block_table),
579+
seq_lens_buffer=torch.empty_like(seq_lens),
580+
score_buffer=shared,
581+
)
582+
wrapper.plan(block_table, seq_lens)
583+
wrappers.append(wrapper)
584+
outs.append(torch.empty_like(expected[-1]))
585+
(q0, k0, _, _), (q1, k1, block_table, seq_lens) = inputs
586+
wrappers[0].run(q0, k0.cache, out=outs[0])
587+
graph = torch.cuda.CUDAGraph()
588+
with torch.cuda.graph(graph):
589+
wrappers[0].run(q0, k0.cache, out=outs[0])
590+
wrappers[1].run(q1, k1.cache, out=outs[1])
591+
outs[0].fill_(-777)
592+
graph.replay()
593+
torch.cuda.synchronize()
594+
assert torch.equal(outs[0], expected[0])
595+
assert torch.equal(outs[1], expected[1])
596+
597+
small = DECODE_WRAPPERS[fmt](num_heads=num_heads, score_buffer=shared[: shared.numel() - 1])
598+
with pytest.raises(ValueError, match="score_buffer"):
599+
small.plan(block_table, seq_lens)
600+
601+
555602
@pytest.mark.parametrize("num_heads", DECODE_HEADS["q8kv4"])
556603
@pytest.mark.parametrize("query_len", (1, 2, 4, 7))
557604
@pytest.mark.parametrize("batch,max_pages", [(1, 1), (5, 6), (129, 4)])

‎python/fmha_sm100/decode_q8kv4/csrc/api/decode_attention_api.cpp‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,21 @@ static int get_runtime_sm_count(int device) {
273273
static std::unordered_map<int, int> g_num_cta;
274274
static std::mutex g_num_cta_mutex;
275275

276+
// The CUTLASS kernels need no workspace (get_workspace_size is 0), so every plan on a device
277+
// passes the same buffer instead of allocating its own. Never freed: it outlives the plans.
278+
static auto *g_cute_workspace = new std::unordered_map<int, at::Tensor>();
279+
static std::mutex g_cute_workspace_mutex;
280+
281+
static at::Tensor shared_cute_workspace(int device) {
282+
std::lock_guard<std::mutex> lock(g_cute_workspace_mutex);
283+
auto it = g_cute_workspace->find(device);
284+
if (it == g_cute_workspace->end())
285+
it = g_cute_workspace
286+
->emplace(device, allocate_plan_buffer(32 * 1024 * 1024, device, torch::kUInt8))
287+
.first;
288+
return it->second;
289+
}
290+
276291
static int get_num_cta(int device) {
277292
std::lock_guard<std::mutex> lock(g_num_cta_mutex);
278293
auto [it, inserted] = g_num_cta.try_emplace(device, -1);
@@ -406,7 +421,7 @@ PlanInfo _make_decode_plan_impl(at::Tensor qo_segment_lens, at::Tensor kv_segmen
406421

407422
PlanInfo info;
408423

409-
auto cute_workspace = allocate_plan_buffer(32 * 1024 * 1024, device, torch::kUInt8);
424+
auto cute_workspace = shared_cute_workspace(device);
410425

411426
int num_ctas = get_num_cta(device);
412427
if (usable_SM_count > 0)

‎tests/q8kv4/test_contract.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,3 +121,17 @@ def plan(batch, q_len, ctas=128):
121121
plan(65536, 1, ctas=127) # split 4-item last wave
122122
with pytest.raises(RuntimeError, match="at most 65536 sequences"):
123123
plan(65537, 1, ctas=127)
124+
125+
126+
def test_plans_share_one_kernel_workspace(device):
127+
"""A plan holds only its schedule; the kernel workspace is one buffer per device."""
128+
def plan(batch):
129+
return plan_decode(batch_size=batch, q_len_per_req=8, topk=16, device=device, num_q_heads=16,
130+
num_kv_heads=1)
131+
132+
plans = [plan(8)]
133+
torch.cuda.synchronize(device)
134+
before = torch.cuda.memory_allocated(device)
135+
plans += [plan(batch) for batch in (129, 256, 512, 1024)]
136+
torch.cuda.synchronize(device)
137+
assert torch.cuda.memory_allocated(device) - before < 32 * 2**20 # was 32 MiB per plan

0 commit comments

Comments
 (0)