From d2434c7f2f3ce6fc5bc7dc5e50acf890319f516b Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Mon, 12 Oct 2026 12:42:40 +0000 Subject: [PATCH 1/5] Accelerate cached FP8 attention with context partitions --- sdm/_kernels/triton/fp8_attention.py | 160 +++++++++++++++++++++++---- test/nn/test_compile.py | 2 +- test/nn/test_fp8_attention.py | 14 ++- 3 files changed, 151 insertions(+), 25 deletions(-) diff --git a/sdm/_kernels/triton/fp8_attention.py b/sdm/_kernels/triton/fp8_attention.py index 8a6686bdd..5f3ce9e10 100644 --- a/sdm/_kernels/triton/fp8_attention.py +++ b/sdm/_kernels/triton/fp8_attention.py @@ -20,6 +20,16 @@ def _attention_kernel( k_scale_ptr, v_scale_ptr, out_ptr, + lse_ptr, + n_splits: tl.constexpr, + stride_kb: tl.constexpr, + stride_kh: tl.constexpr, + stride_kr: tl.constexpr, + stride_kc: tl.constexpr, + stride_vb: tl.constexpr, + stride_vh: tl.constexpr, + stride_vc: tl.constexpr, + stride_vr: tl.constexpr, n_q_heads: tl.constexpr, n_kv_heads: tl.constexpr, n_queries: tl.constexpr, @@ -34,9 +44,8 @@ def _attention_kernel( ) -> None: """Write attention outputs for one query tile and head per program. - Each program keeps a ``[block_q, head_dim]`` query tile and loops over - context tiles of ``block_kv`` rows. Both matrix multiplications use - the full head width. + Each program scans its context partition in ``block_kv`` tiles for a + ``[block_q, head_dim]`` query tile. Both matmuls use the full head width. """ query_block, batch_query_head = tl.program_id(0), tl.program_id(1) @@ -46,6 +55,24 @@ def _attention_kernel( batch_kv_head = batch * n_kv_heads + query_head // ( n_q_heads // n_kv_heads ) + kv_head = batch_kv_head % n_kv_heads + split = tl.program_id(2) if n_splits > 1 else 0 + n_tiles = tl.cdiv(n_context, block_kv) # ty: ignore[invalid-argument-type] + first_tile = split * n_tiles // n_splits + last_tile = (split + 1) * n_tiles // n_splits + k_head_offset = batch * stride_kb + kv_head * stride_kh + v_head_offset = batch * stride_vb + kv_head * stride_vh + # Flatten packed heads directly; honor batch/head strides after slicing. + if ( + stride_kb == n_kv_heads * n_context * head_dim + and stride_kh == n_context * head_dim + ): + k_head_offset = batch_kv_head * n_context * head_dim + if ( + stride_vb == n_kv_heads * n_context * head_dim + and stride_vh == n_context * head_dim + ): + v_head_offset = batch_kv_head * n_context * head_dim rows = query_block * block_q + tl.arange(0, block_q) cols = tl.arange(0, block_kv) dims = tl.arange(0, head_dim) @@ -71,14 +98,14 @@ def _attention_kernel( if acc_flush_tiles > 0: outer_acc = tl.full((block_q, head_dim), 0, tl.float32) outer_max = tl.full((block_q,), -float("inf"), tl.float32) - for start in range(tl.cdiv(n_context, block_kv)): # ty: ignore[invalid-argument-type] + for start in range(first_tile, last_tile): # 1. Tiled QK^T: reduce over channels to get [block_q, block_kv]. indices = start * block_kv + cols k = tl.load( k_ptr - + batch_kv_head * n_context * head_dim - + indices[None, :] * head_dim - + dims[:, None], + + k_head_offset + + indices[None, :] * stride_kr + + dims[:, None] * stride_kc, mask=indices[None, :] < n_context, other=0.0, ) @@ -108,9 +135,13 @@ def _attention_kernel( acc *= correction[:, None] # 3. Tiled weights-V: reduce over context to get [block_q, head_dim]. - v_offsets = indices[:, None] + dims[None, :] * n_context + v_offsets = ( + v_head_offset + + indices[:, None] * stride_vr + + dims[None, :] * stride_vc + ) v = tl.load( - v_ptr + batch_kv_head * n_context * head_dim + v_offsets, + v_ptr + v_offsets, mask=indices[:, None] < n_context, other=0.0, ) @@ -124,8 +155,8 @@ def _attention_kernel( ) # Periodically flush the dot accumulator to FP32 for long L4 contexts. if acc_flush_tiles > 0: # noqa: SIM102 - constexpr guard avoids modulo zero. - if ((start + 1) % acc_flush_tiles == 0) | ( # ty: ignore[unsupported-operator] - start + 1 == tl.cdiv(n_context, block_kv) # ty: ignore[invalid-argument-type] + if ((start + 1 - first_tile) % acc_flush_tiles == 0) | ( + start + 1 == last_tile ): outer_acc = ( outer_acc * tl.exp2(outer_max - new_max)[:, None] + acc @@ -139,9 +170,20 @@ def _attention_kernel( out = acc * value_scale / weight_sum[:, None] if not scale_weights_in_exp: out *= 1.0 / 256.0 + output_head = batch_query_head * n_splits + split + if n_splits > 1: + # Preserve each partition's softmax mass for the final reduction. + logsum = row_max + tl.log2(weight_sum) + if scale_weights_in_exp: + logsum -= 8.0 + tl.store( + lse_ptr + output_head * n_queries + rows, + logsum, + mask=rows < n_queries, + ) tl.store( out_ptr - + batch_query_head * n_queries * head_dim + + output_head * n_queries * head_dim + rows[:, None] * head_dim + dims[None, :], out, @@ -149,6 +191,34 @@ def _attention_kernel( ) +@triton.jit +def _combine_attention_kernel( + partial_ptr, + lse_ptr, + out_ptr, + n_queries: tl.constexpr, + head_dim: tl.constexpr, + n_splits: tl.constexpr, + block_splits: tl.constexpr, +) -> None: + """Combine partition outputs, weighted by their softmax masses.""" + row, head = tl.program_id(0), tl.program_id(1) + splits = tl.arange(0, block_splits) + dims = tl.arange(0, head_dim) + offset = (head * n_splits + splits) * n_queries + row + logsum = tl.load( + lse_ptr + offset, mask=splits < n_splits, other=-float("inf") + ) + weights = tl.exp2(logsum - tl.max(logsum, 0)) # ty: ignore[invalid-argument-type] + partial = tl.load( + partial_ptr + offset[:, None] * head_dim + dims[None, :], + mask=splits[:, None] < n_splits, + other=0.0, + ) + out = tl.sum(partial * weights[:, None], 0) / tl.sum(weights, 0) # ty: ignore[invalid-argument-type] + tl.store(out_ptr + (head * n_queries + row) * head_dim + dims, out) + + @triton.jit def _quantize_kernel( x_ptr, @@ -234,13 +304,15 @@ def quantized_attention( accumulation_chunk: int = 0, scale_weights_in_exp: bool = False, fuse_score_scale: bool = False, + context_splits: int = 1, ) -> Tensor: """Run tiled attention on per-head scaled FP8 inputs. Args: q: Contiguous FP8 queries shaped ``[batch, query_heads, queries, D]``. - k: Contiguous FP8 keys shaped ``[batch, kv_heads, context, D]``. - v: Contiguous FP8 values shaped ``[batch, kv_heads, D, context]``. + k: FP8 keys shaped ``[batch, kv_heads, context, D]``. + v: FP8 values shaped ``[batch, kv_heads, D, context]``. + K/V may be strided, including views of selected cache heads. qs: Query scales shaped ``[batch, query_heads, 1, 1]``. ks: Key scales shaped ``[batch, kv_heads, 1, 1]``. vs: Value scales shaped ``[batch, kv_heads, 1, 1]``. @@ -254,6 +326,9 @@ def quantized_attention( scale_weights_in_exp: Apply the weight factor 256 inside exp2 instead of before the weighted-value multiplication. fuse_score_scale: Permit fused score scaling and maximum subtraction. + context_splits: Maximum number of context partitions per query tile. + One scans the entire context; larger counts combine partial + outputs using FP32 softmax masses. Returns: Attention output with the query shape and requested dtype. @@ -268,10 +343,18 @@ def quantized_attention( raise ValueError(f"{name} must be a positive power of two") if k.size(1) == 0 or h % k.size(1): raise ValueError("The number of K/V heads must divide the query heads") + if context_splits < 1: + raise ValueError("context_splits must be positive") n = k.size(-2) + splits = min(context_splits, triton.cdiv(n, context_tile)) # ty: ignore[invalid-argument-type] # Q/K are [batch, heads, sequence, channels]; V is transposed. output = q.new_empty(q.shape, dtype=dtype) - grid = (triton.cdiv(m, query_tile), b * h) # ty: ignore[invalid-argument-type] + partial = output + lse = None + if splits > 1: + partial = q.new_empty((b, h, splits, m, d), dtype=torch.float32) + lse = q.new_empty((b, h, splits, m), dtype=torch.float32) + grid = (triton.cdiv(m, query_tile), b * h, splits) # ty: ignore[invalid-argument-type] with torch.cuda.device(q.device): cast(Any, _attention_kernel)[grid]( q, @@ -280,7 +363,17 @@ def quantized_attention( qs, ks, vs, - output, + partial, + lse, + n_splits=splits, + stride_kb=k.stride(0), + stride_kh=k.stride(1), + stride_kr=k.stride(2), + stride_kc=k.stride(3), + stride_vb=v.stride(0), + stride_vh=v.stride(1), + stride_vc=v.stride(2), + stride_vr=v.stride(3), n_q_heads=h, n_kv_heads=k.size(1), n_queries=m, @@ -295,6 +388,16 @@ def quantized_attention( num_warps=warps, num_stages=stages, ) + if splits > 1: + cast(Any, _combine_attention_kernel)[(m, b * h)]( + partial, + lse, + output, + n_queries=m, + head_dim=d, + n_splits=splits, + block_splits=triton.next_power_of_2(splits), # ty: ignore[invalid-argument-type] + ) return output @@ -333,9 +436,10 @@ def rows(x: Tensor) -> Tensor: qs = qs.clamp_min(1e-12) / 448.0 k, ks = quantize(heads(key)) v, vs = quantize(heads(value)) + v = v.transpose(-1, -2).contiguous() else: - k = heads(cache.key).contiguous() - v = heads(cache.value).contiguous() + k = heads(cache.key) + v = heads(cache.value).transpose(-1, -2) ks, vs, qs = ( heads(t).contiguous() for t in (cache.key_scale, cache.value_scale, cache.query_scale) @@ -347,15 +451,27 @@ def rows(x: Tensor) -> Tensor: blackwell = head_dim == 64 and torch.cuda.get_device_capability( query.device ) == (12, 0) + query_tile = 64 if head_dim == 256 else 128 + context_splits = 1 + if cache is not None and k.size(-2) >= 32768: + programs = triton.cdiv(q.size(-2), query_tile) * q.size(0) * q.size(1) # ty: ignore[invalid-argument-type] + sm_count = torch.cuda.get_device_properties( + query.device + ).multi_processor_count + # Target four programs per SM, with at most 16 scratch partitions. + context_splits = min( + 16, + triton.next_power_of_2(triton.cdiv(4 * sm_count, programs)), # ty: ignore[invalid-argument-type] + ) result = quantized_attention( q, k, - v.transpose(-1, -2).contiguous(), + v, qs, ks, vs, dtype=query.dtype, - query_tile=64 if head_dim == 256 else 128, + query_tile=query_tile, scale=scale, context_tile=64 if wide or blackwell or k.size(-2) < 32768 else 128, warps=8 if wide else 4, @@ -363,11 +479,13 @@ def rows(x: Tensor) -> Tensor: accumulation_chunk=0 if blackwell else 16, scale_weights_in_exp=blackwell, fuse_score_scale=blackwell, + context_splits=context_splits, ) if cache is None: cache = QuantizedKVCacheEntry( key=rows(k), - value=rows(v), + # Keep the logical cache shape while retaining the kernel layout. + value=rows(v.transpose(-1, -2)), key_scale=rows(ks), value_scale=rows(vs), query_scale=rows(qs), diff --git a/test/nn/test_compile.py b/test/nn/test_compile.py index fa626d2b8..e01734c71 100644 --- a/test/nn/test_compile.py +++ b/test/nn/test_compile.py @@ -246,7 +246,7 @@ def test_attention_compile_fp8_cache() -> None: ).eval() module._supports_quantized_attention = True torch.nn.init.normal_(module.out_lin.weight, std=0.05) - context = torch.randn(1, 8193, 128, device="cuda", dtype=dtype) + context = torch.randn(1, 32769, 128, device="cuda", dtype=dtype) compiled = torch.compile(module, fullgraph=True, dynamic=False) with torch.inference_mode(), optimize(attention="fp8"): expected, eager_cache = module(context, return_key_value=True) diff --git a/test/nn/test_fp8_attention.py b/test/nn/test_fp8_attention.py index 198c89eea..1abdf0fa2 100644 --- a/test/nn/test_fp8_attention.py +++ b/test/nn/test_fp8_attention.py @@ -23,9 +23,12 @@ @pytest.mark.parametrize( - ("batch", "queries", "context"), [(2, 129, 513), (1, 129, 8192)] + ("batch", "queries", "context", "context_splits"), + [(2, 129, 513, 1), (1, 129, 8192, 1), (2, 17, 513, 16)], ) -def test_fp8_partial_tiles(batch: int, queries: int, context: int) -> None: +def test_fp8_partial_tiles( + batch: int, queries: int, context: int, context_splits: int +) -> None: with torch.inference_mode(): q = torch.randn( batch, 8, queries, 64, device="cuda", dtype=torch.float16 @@ -37,15 +40,19 @@ def test_fp8_partial_tiles(batch: int, queries: int, context: int) -> None: q8, qs = quantize(q) k8, ks = quantize(k) v8, vs = quantize(v) + if context_splits > 1: + k8 = k8.transpose(1, 2).contiguous().transpose(1, 2) + v8 = v8.transpose(1, 2).contiguous().transpose(1, 2) actual = quantized_attention( q8, k8, - v8.transpose(-1, -2).contiguous(), + v8.transpose(-1, -2), qs, ks, vs, q.dtype, 128, + context_splits=context_splits, ) expected = torch.nn.functional.scaled_dot_product_attention(q, k, v) assert actual.isfinite().all() @@ -84,6 +91,7 @@ def test_grouped_query_attention_and_model_scale( 128, scale, context_tile=context_tile, + context_splits=4, ) expected = torch.nn.functional.scaled_dot_product_attention( q, k, v, enable_gqa=True, scale=scale From a17ecb4a492d9a1e4babecd988e00f025c667026 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Mon, 12 Oct 2026 15:23:02 +0000 Subject: [PATCH 2/5] Use measured context thresholds for cached attention --- sdm/_kernels/triton/fp8_attention.py | 6 +++++- test/nn/test_compile.py | 2 +- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/sdm/_kernels/triton/fp8_attention.py b/sdm/_kernels/triton/fp8_attention.py index 5f3ce9e10..2c950823d 100644 --- a/sdm/_kernels/triton/fp8_attention.py +++ b/sdm/_kernels/triton/fp8_attention.py @@ -453,7 +453,11 @@ def rows(x: Tensor) -> Tensor: ) == (12, 0) query_tile = 64 if head_dim == 256 else 128 context_splits = 1 - if cache is not None and k.size(-2) >= 32768: + # Aligned V loads keep shorter predictions limited by launch overhead. + min_context = ( + 60 * 1024 if v.stride(-1) == 1 and v.stride(-2) % 8 == 0 else 32 * 1024 + ) + if cache is not None and k.size(-2) >= min_context: programs = triton.cdiv(q.size(-2), query_tile) * q.size(0) * q.size(1) # ty: ignore[invalid-argument-type] sm_count = torch.cuda.get_device_properties( query.device diff --git a/test/nn/test_compile.py b/test/nn/test_compile.py index e01734c71..4707032a4 100644 --- a/test/nn/test_compile.py +++ b/test/nn/test_compile.py @@ -246,7 +246,7 @@ def test_attention_compile_fp8_cache() -> None: ).eval() module._supports_quantized_attention = True torch.nn.init.normal_(module.out_lin.weight, std=0.05) - context = torch.randn(1, 32769, 128, device="cuda", dtype=dtype) + context = torch.randn(1, 60 * 1024, 128, device="cuda", dtype=dtype) compiled = torch.compile(module, fullgraph=True, dynamic=False) with torch.inference_mode(), optimize(attention="fp8"): expected, eager_cache = module(context, return_key_value=True) From 6a9e92cea2800d1c3584560c7d1ce578e5930048 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Tue, 13 Oct 2026 01:40:07 +0000 Subject: [PATCH 3/5] Use a 64K cutoff for aligned cached attention --- sdm/_kernels/triton/fp8_attention.py | 2 +- test/nn/test_compile.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/sdm/_kernels/triton/fp8_attention.py b/sdm/_kernels/triton/fp8_attention.py index 2c950823d..270e0aa7d 100644 --- a/sdm/_kernels/triton/fp8_attention.py +++ b/sdm/_kernels/triton/fp8_attention.py @@ -455,7 +455,7 @@ def rows(x: Tensor) -> Tensor: context_splits = 1 # Aligned V loads keep shorter predictions limited by launch overhead. min_context = ( - 60 * 1024 if v.stride(-1) == 1 and v.stride(-2) % 8 == 0 else 32 * 1024 + 64 * 1024 if v.stride(-1) == 1 and v.stride(-2) % 8 == 0 else 32 * 1024 ) if cache is not None and k.size(-2) >= min_context: programs = triton.cdiv(q.size(-2), query_tile) * q.size(0) * q.size(1) # ty: ignore[invalid-argument-type] diff --git a/test/nn/test_compile.py b/test/nn/test_compile.py index 4707032a4..8e606bfdd 100644 --- a/test/nn/test_compile.py +++ b/test/nn/test_compile.py @@ -246,7 +246,7 @@ def test_attention_compile_fp8_cache() -> None: ).eval() module._supports_quantized_attention = True torch.nn.init.normal_(module.out_lin.weight, std=0.05) - context = torch.randn(1, 60 * 1024, 128, device="cuda", dtype=dtype) + context = torch.randn(1, 64 * 1024, 128, device="cuda", dtype=dtype) compiled = torch.compile(module, fullgraph=True, dynamic=False) with torch.inference_mode(), optimize(attention="fp8"): expected, eager_cache = module(context, return_key_value=True) From 9fb927c0636151ecaebeacac0980054bc6197cd7 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Tue, 13 Oct 2026 02:19:36 +0000 Subject: [PATCH 4/5] Use a shared measured cutoff for cached FP8 attention --- sdm/_kernels/triton/fp8_attention.py | 6 +----- test/nn/test_compile.py | 5 +++-- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/sdm/_kernels/triton/fp8_attention.py b/sdm/_kernels/triton/fp8_attention.py index 270e0aa7d..836878b1c 100644 --- a/sdm/_kernels/triton/fp8_attention.py +++ b/sdm/_kernels/triton/fp8_attention.py @@ -453,11 +453,7 @@ def rows(x: Tensor) -> Tensor: ) == (12, 0) query_tile = 64 if head_dim == 256 else 128 context_splits = 1 - # Aligned V loads keep shorter predictions limited by launch overhead. - min_context = ( - 64 * 1024 if v.stride(-1) == 1 and v.stride(-2) % 8 == 0 else 32 * 1024 - ) - if cache is not None and k.size(-2) >= min_context: + if cache is not None and k.size(-2) >= 64 * 1024: programs = triton.cdiv(q.size(-2), query_tile) * q.size(0) * q.size(1) # ty: ignore[invalid-argument-type] sm_count = torch.cuda.get_device_properties( query.device diff --git a/test/nn/test_compile.py b/test/nn/test_compile.py index 8e606bfdd..36bf4acb2 100644 --- a/test/nn/test_compile.py +++ b/test/nn/test_compile.py @@ -233,7 +233,8 @@ def test_attention_compile_key_value_cache(device: torch.device) -> None: @onlyCUDA @pytest.mark.usefixtures("fp8_rng") -def test_attention_compile_fp8_cache() -> None: +@pytest.mark.parametrize("context_rows", [64 * 1024, 64 * 1024 + 1]) +def test_attention_compile_fp8_cache(context_rows: int) -> None: dtype = torch.float16 if torch.cuda.get_device_capability() not in {(8, 9), (9, 0), (12, 0)}: pytest.skip("FP8 integration supports Ada, Hopper, and RTX Blackwell") @@ -246,7 +247,7 @@ def test_attention_compile_fp8_cache() -> None: ).eval() module._supports_quantized_attention = True torch.nn.init.normal_(module.out_lin.weight, std=0.05) - context = torch.randn(1, 64 * 1024, 128, device="cuda", dtype=dtype) + context = torch.randn(1, context_rows, 128, device="cuda", dtype=dtype) compiled = torch.compile(module, fullgraph=True, dynamic=False) with torch.inference_mode(), optimize(attention="fp8"): expected, eager_cache = module(context, return_key_value=True) From 12048bc1c924f17acb8f599c9143798d1b18bf33 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Tue, 13 Oct 2026 02:51:04 +0000 Subject: [PATCH 5/5] Retain cached attention speedups below 64K context --- sdm/_kernels/triton/fp8_attention.py | 8 +++++++- test/nn/test_compile.py | 6 ++++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/sdm/_kernels/triton/fp8_attention.py b/sdm/_kernels/triton/fp8_attention.py index 836878b1c..9342e4f55 100644 --- a/sdm/_kernels/triton/fp8_attention.py +++ b/sdm/_kernels/triton/fp8_attention.py @@ -453,7 +453,13 @@ def rows(x: Tensor) -> Tensor: ) == (12, 0) query_tile = 64 if head_dim == 256 else 128 context_splits = 1 - if cache is not None and k.size(-2) >= 64 * 1024: + # Measured cutoffs depend on V alignment and the query batch size. + min_context = 64 * 1024 + if v.stride(-1) == 1 and v.stride(-2) % 4: + min_context = 48 * 1024 + if v.stride(-2) % 2 == 0 or q.size(-2) >= 1024: + min_context = 32 * 1024 + if cache is not None and k.size(-2) >= min_context: programs = triton.cdiv(q.size(-2), query_tile) * q.size(0) * q.size(1) # ty: ignore[invalid-argument-type] sm_count = torch.cuda.get_device_properties( query.device diff --git a/test/nn/test_compile.py b/test/nn/test_compile.py index 36bf4acb2..367d7bab6 100644 --- a/test/nn/test_compile.py +++ b/test/nn/test_compile.py @@ -233,7 +233,9 @@ def test_attention_compile_key_value_cache(device: torch.device) -> None: @onlyCUDA @pytest.mark.usefixtures("fp8_rng") -@pytest.mark.parametrize("context_rows", [64 * 1024, 64 * 1024 + 1]) +@pytest.mark.parametrize( + "context_rows", [32 * 1024 + 1, 32 * 1024 + 2, 64 * 1024] +) def test_attention_compile_fp8_cache(context_rows: int) -> None: dtype = torch.float16 if torch.cuda.get_device_capability() not in {(8, 9), (9, 0), (12, 0)}: @@ -254,7 +256,7 @@ def test_attention_compile_fp8_cache(context_rows: int) -> None: actual, cache = compiled(context, return_key_value=True) assert isinstance(cache, QuantizedKVCacheEntry) torch.testing.assert_close(actual, expected) - for rows in (17, 1): + for rows in (1024, 512, 17, 1): query = torch.randn(1, rows, 128, device="cuda", dtype=dtype) expected = module(query, eager_cache) torch.testing.assert_close(compiled(query, cache), expected)