"""Step-5 Triton kernels for DSA attention, FP8 experts and fused decoder operations. Requires the Triton 3.x runtime specified in requirements.txt. Importing does not initialize CUDA; executing kernels requires CUDA tensors. """ # FP8 quantizer/GEMM portions: Copyright (c) 2023 DeepSeek, MIT License. from __future__ import annotations import math from dataclasses import dataclass import torch import triton import triton.language as tl import triton.language.extra.libdevice as _libdevice # DSA indexer: normalize, compress regions, score and select. Summaries use # logical tensor layout shared by the producers and consumers below. # Step-5 compresses each region of 8 consecutive tokens into one summary vector. REGION_BLOCK_SIZE = 8 PROXY_DIM = 256 def shared_table_stride(block_table: torch.Tensor) -> int: """Row stride to pass for a block table, zero when the table is shared. A single-row table means "every request uses this mapping". Expressing that as a zero row stride rather than a branch inside the kernel keeps the two cases on one code path, and -- more importantly -- makes the shared case safe: the alternative, indexing a one-row table by a nonzero request id, reads past the tensor and returns plausible garbage instead of failing. """ if block_table.ndim != 2: raise ValueError( f"block_table must be [num_reqs, pages], got {tuple(block_table.shape)}" ) return 0 if block_table.shape[0] == 1 else block_table.stride(0) @triton.jit def _affine_normalize( raw, mean, rstd, weight_ptr, bias_ptr, offsets, mask, weight_bias, HAS_BIAS: tl.constexpr, ): """Apply ``(x - mean) * rstd * (w + weight_bias) [+ b]`` to one dim slice. A device function rather than a closure because Triton cannot compile nested ``def`` inside a kernel, and the normalization has to be applied to three different dim slices of the same head: the tail, and each half of the rotated span. """ scale = tl.load(weight_ptr + offsets, mask=mask, other=0.0).to(tl.float32) result = (raw - mean) * rstd * (scale + weight_bias) if HAS_BIAS: result += tl.load(bias_ptr + offsets, mask=mask, other=0.0).to(tl.float32) return result @triton.jit def _indexer_norm_rope_kernel( x_ptr, out_ptr, weight_ptr, bias_ptr, cos_ptr, sin_ptr, positions_ptr, stride_x_token, stride_out_token, stride_cos, eps, weight_bias, head_dim: tl.constexpr, rotary_dim: tl.constexpr, SUBTRACT_MEAN: tl.constexpr, HAS_BIAS: tl.constexpr, BLOCK_D: tl.constexpr, BLOCK_R: tl.constexpr, ) -> None: """Normalize one head of one token, then rotate its leading ``2 * rotary_dim`` dims. ``SUBTRACT_MEAN`` selects LayerNorm over RMSNorm. The indexer needs both: q is RMSNorm with a zero-centered weight, k is LayerNorm with a weight and a bias. They differ by one reduction and are trivial to confuse, which would shift every key without changing any shape. The rotated span is written from a separate load rather than by patching the normalized vector in place, because Triton cannot index a value tensor by another tensor -- and overlapping stores from one program would depend on store ordering the compiler does not promise. """ token = tl.program_id(0) head = tl.program_id(1) base = token * stride_x_token + head * head_dim out_base = token * stride_out_token + head * head_dim dims = tl.arange(0, BLOCK_D) dim_ok = dims < head_dim values = tl.load(x_ptr + base + dims, mask=dim_ok, other=0.0).to(tl.float32) sum_sq = tl.sum(values * values) if SUBTRACT_MEAN: # Match the deployed LayerNorm reduction and fast-rsqrt path. Computing # ``mean((x - mean) ** 2)`` is algebraically equivalent, but changes the # FP32 reduction/rounding before every sparse-indexer key. mean = tl.sum(values) / head_dim variance = sum_sq / head_dim - mean * mean else: mean = 0.0 variance = sum_sq / head_dim rstd = tl.rsqrt(variance + eps) # Dims past the rotated span keep the plain normalized value. tl.store( out_ptr + out_base + dims, _affine_normalize( values, mean, rstd, weight_ptr, bias_ptr, dims, dim_ok, weight_bias, HAS_BIAS=HAS_BIAS, ).to(out_ptr.dtype.element_ty), mask=dim_ok & (dims >= 2 * rotary_dim), ) half = tl.arange(0, BLOCK_R) half_ok = half < rotary_dim position = tl.load(positions_ptr + token).to(tl.int32) cos = tl.load(cos_ptr + position * stride_cos + half, mask=half_ok, other=0.0).to(tl.float32) sin = tl.load(sin_ptr + position * stride_cos + half, mask=half_ok, other=0.0).to(tl.float32) real = _affine_normalize( tl.load(x_ptr + base + half, mask=half_ok, other=0.0).to(tl.float32), mean, rstd, weight_ptr, bias_ptr, half, half_ok, weight_bias, HAS_BIAS=HAS_BIAS, ).to(out_ptr.dtype.element_ty).to(tl.float32) imaginary = _affine_normalize( tl.load(x_ptr + base + rotary_dim + half, mask=half_ok, other=0.0).to(tl.float32), mean, rstd, weight_ptr, bias_ptr, rotary_dim + half, half_ok, weight_bias, HAS_BIAS=HAS_BIAS, ).to(out_ptr.dtype.element_ty).to(tl.float32) tl.store( out_ptr + out_base + half, (real * cos - imaginary * sin).to(out_ptr.dtype.element_ty), mask=half_ok, ) tl.store( out_ptr + out_base + rotary_dim + half, (real * sin + imaginary * cos).to(out_ptr.dtype.element_ty), mask=half_ok, ) def _norm_rope( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None, cos: torch.Tensor, sin: torch.Tensor, positions: torch.Tensor, *, num_heads: int, head_dim: int, rotary_dim: int, eps: float, weight_bias: float, subtract_mean: bool, ) -> torch.Tensor: out = torch.empty_like(x) _indexer_norm_rope_kernel[(x.shape[0], num_heads)]( x, out, weight, bias if bias is not None else weight, cos, sin, positions, x.stride(0), out.stride(0), cos.stride(0), eps, weight_bias, head_dim=head_dim, rotary_dim=rotary_dim, SUBTRACT_MEAN=subtract_mean, HAS_BIAS=bias is not None, BLOCK_D=triton.next_power_of_2(head_dim), BLOCK_R=triton.next_power_of_2(rotary_dim), ) return out def indexer_norm_rope( index_q: torch.Tensor, index_k: torch.Tensor, index_z: torch.Tensor, q_norm_weight: torch.Tensor, k_norm_weight: torch.Tensor, k_norm_bias: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, positions: torch.Tensor, *, head_dim: int, num_q_heads: int, num_k_heads: int, rotary_dim: int, eps: float = 1e-6, q_norm_weight_bias: float = 1.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Prepare the indexer's q/k/z from their fused projection. Args: index_q: ``[tokens, num_q_heads * head_dim]``. index_k: ``[tokens, num_k_heads * head_dim]``. index_z: ``[tokens, num_k_heads * head_dim]`` -- returned untouched. cos, sin: ``[seq, rotary_dim]``, half the width of the rotated span. q_norm_weight_bias: ``1.0`` for a zero-centered RMSNorm weight (``x * (1 + w)``), ``0.0`` otherwise. This follows the checkpoint's norm class, so it must be threaded through from config rather than hardcoded. """ return ( _norm_rope( index_q, q_norm_weight, None, cos, sin, positions, num_heads=num_q_heads, head_dim=head_dim, rotary_dim=rotary_dim, eps=eps, weight_bias=q_norm_weight_bias, subtract_mean=False, ), _norm_rope( index_k, k_norm_weight, k_norm_bias, cos, sin, positions, num_heads=num_k_heads, head_dim=head_dim, rotary_dim=rotary_dim, eps=eps, weight_bias=0.0, subtract_mean=True, ), index_z.contiguous(), ) @triton.jit def _csa_compress_regions_kernel( index_k_ptr, index_z_ptr, token_start_ptr, token_count_ptr, summary_ptr, mean_ptr, stride_token, stride_head, stride_summary_region, stride_summary_head, num_heads: tl.constexpr, proxy_dim: tl.constexpr, BLOCK_T: tl.constexpr, BLOCK_D: tl.constexpr, ) -> None: """One program per (region, head): compress a region's tokens into one vector. The weight is a softmax over the token axis of ``index_z`` taken **per dimension**, so this is an element-wise weighted mean, not a convex combination of whole key vectors: s_d = sum_t exp(z_td - m) * k_td / sum_t exp(z_td - m) ``m`` is one scalar maximum over both axes. It keeps ``exp`` in range; sharing the shift across dimensions also fixes the floating-point rounding used by Step-5's region-summary quantization. """ region = tl.program_id(0) head = tl.program_id(1) start = tl.load(token_start_ptr + region).to(tl.int64) count = tl.load(token_count_ptr + region).to(tl.int64) tokens = tl.arange(0, BLOCK_T) dims = tl.arange(0, BLOCK_D) token_ok = tokens < count dim_ok = dims < proxy_dim mask = token_ok[:, None] & dim_ok[None, :] offsets = ( (start + tokens.to(tl.int64))[:, None] * stride_token + head * stride_head + dims[None, :] ) logits = tl.load(index_z_ptr + offsets, mask=mask, other=float("-inf")).to(tl.float32) values = tl.load(index_k_ptr + offsets, mask=mask, other=0.0).to(tl.float32) shift = tl.max(tl.where(mask, logits, float("-inf"))) weights = tl.where(mask, tl.exp(logits - shift), 0.0) denominator = tl.sum(weights, axis=0) numerator = tl.sum(weights * values, axis=0) summary = tl.where( denominator > 0.0, numerator / tl.maximum(denominator, 1e-20), 0.0 ) out = region * stride_summary_region + head * stride_summary_head + dims tl.store(summary_ptr + out, summary, mask=dim_ok) # Step-5 rounds through BF16 before E4M3. BF16 has 7 stored fraction bits # and E4M3 has 3; the intermediate rounding can move a value onto an E4M3 # midpoint. Direct FP32-to-E4M3 conversion can therefore change the bytes # checked by the quantization gate. tl.store(mean_ptr + out, summary.to(tl.bfloat16).to(tl.float8e4nv), mask=dim_ok) def csa_region_layout( seq_lens: torch.Tensor, *, regions_per_seq: int, region_size: int = REGION_BLOCK_SIZE ) -> tuple[torch.Tensor, torch.Tensor]: """Token span of every region, for sequences packed back to back. ``regions_per_seq`` is the stride of the region id space, not the number of regions a sequence actually fills -- a sequence's regions have to land at a fixed offset so the summary store can be indexed without a second lookup. Regions past a sequence's end get a count of zero and are skipped rather than written, which keeps a stale summary from a previous request out of the selection. """ device = seq_lens.device lengths = seq_lens.to(torch.int64) token_starts = torch.cumsum(lengths, dim=0) - lengths regions = torch.arange(regions_per_seq, device=device, dtype=torch.int64) offset = regions.view(1, -1) * region_size count = (lengths.view(-1, 1) - offset).clamp(0, region_size) start = token_starts.view(-1, 1) + offset return start.reshape(-1).to(torch.int32), count.reshape(-1).to(torch.int32) def csa_compress_regions( index_k: torch.Tensor, index_z: torch.Tensor, token_start: torch.Tensor, token_count: torch.Tensor, *, region_size: int = REGION_BLOCK_SIZE, ) -> tuple[torch.Tensor, torch.Tensor]: """Compress packed tokens into per-region summaries. Args: index_k: ``[tokens, heads, proxy_dim]`` proxy keys. index_z: ``[tokens, heads, proxy_dim]`` per-dimension softmax logits. token_start: ``[regions]`` int32, first token of each region. token_count: ``[regions]`` int32, tokens present in each region; 0 = skip. Returns: ``(summary_fp32, summary_fp8)``, both ``[regions, heads, proxy_dim]``. The fp8 tensor is what selection and sparse attention consume; the fp32 one exists because it is the only place the unrounded value is observable, which is what lets a parity failure be attributed to the reduction rather than the rounding. """ if index_k.shape != index_z.shape: raise ValueError( f"index_k and index_z must match, got {tuple(index_k.shape)} and " f"{tuple(index_z.shape)}" ) if index_k.ndim != 3: raise ValueError(f"index_k must be [tokens, heads, proxy_dim], got {tuple(index_k.shape)}") if index_k.stride(2) != 1: raise ValueError("index_k must be contiguous in the proxy dimension") regions = int(token_start.numel()) _, heads, proxy_dim = index_k.shape summary = torch.zeros((regions, heads, proxy_dim), device=index_k.device, dtype=torch.float32) mean = torch.zeros((regions, heads, proxy_dim), device=index_k.device, dtype=torch.float8_e4m3fn) _csa_compress_regions_kernel[(regions, heads)]( index_k, index_z, token_start, token_count, summary, mean, index_k.stride(0), index_k.stride(1), summary.stride(0), summary.stride(1), num_heads=heads, proxy_dim=proxy_dim, BLOCK_T=triton.next_power_of_2(region_size), BLOCK_D=triton.next_power_of_2(proxy_dim), ) return summary, mean @triton.jit def _indexer_logits_kernel( index_q_ptr, index_k_ptr, weights_ptr, out_ptr, seq_q, seq_k, stride_q_token, stride_q_group, stride_q_head, stride_k_token, stride_w_token, stride_w_group, stride_out_row, heads_per_group: tl.constexpr, proxy_dim: tl.constexpr, BLOCK_Q: tl.constexpr, BLOCK_K: tl.constexpr, ) -> None: """Weighted-ReLU indexer scores for one (query tile, key tile, provider group). score[q, r] = sum_h relu(dot(q_qh, k_r)) * w_qh The ReLU is per head and *inside* the sum, so the heads cannot be folded into a single matmul -- hence the loop. The per-head weights are signed, so applying ReLU to each dot does **not** make the final weighted sum non-negative. Exact zeros are still common because a score is zero whenever every head dot is negative; for Step-5's 4 heads per group that event is about 6.25%. Those zeros are what the top-k tie-break then has to resolve, so they are load-bearing, not noise. """ q_block = tl.program_id(0) k_block = tl.program_id(1) group = tl.program_id(2) queries = q_block * BLOCK_Q + tl.arange(0, BLOCK_Q) keys = k_block * BLOCK_K + tl.arange(0, BLOCK_K) dims = tl.arange(0, proxy_dim) query_ok = queries < seq_q key_ok = keys < seq_k key_tile = tl.load( index_k_ptr + dims[:, None] + keys[None, :] * stride_k_token, mask=key_ok[None, :], other=0.0, ) accumulator = tl.zeros((BLOCK_Q, BLOCK_K), dtype=tl.float32) for head in tl.static_range(heads_per_group): query_tile = tl.load( index_q_ptr + queries[:, None] * stride_q_token + group * stride_q_group + head * stride_q_head + dims[None, :], mask=query_ok[:, None], other=0.0, ) head_weight = tl.load( weights_ptr + queries * stride_w_token + group * stride_w_group + head, mask=query_ok, other=0.0, ).to(tl.float32) scores = tl.dot(query_tile, key_tile, out_dtype=tl.float32) accumulator += tl.maximum(scores, 0.0) * head_weight[:, None] rows = group * seq_q + queries tl.store( out_ptr + rows[:, None] * stride_out_row + keys[None, :], accumulator, mask=query_ok[:, None] & key_ok[None, :], ) def round_activations_e4m3(tensor: torch.Tensor) -> torch.Tensor: """Apply the indexer's e4m3 activation rounding, then widen back to bfloat16. The deployed kernel quantizes indexer activations to e4m3 before the score matmul, even though the checkpoint stores these weights in bf16. Skipping the rounding changes which regions get selected, so it is part of the semantics and not an optimization we can decline. Widening back to bf16 instead of feeding e4m3 into the matmul is deliberate. It is lossless -- e4m3 keeps 3 mantissa bits and bf16 keeps 8 -- and it keeps the kernel off Triton's fp8 MMA path, which is where portability across Triton versions and non-NVIDIA backends breaks down. The cost is peak throughput on Hopper, not accuracy. """ return tensor.to(torch.float8_e4m3fn).to(torch.bfloat16) def indexer_logits( index_q: torch.Tensor, weights: torch.Tensor, index_k: torch.Tensor, *, block_q: int = 64, block_k: int = 128, ) -> torch.Tensor: """Weighted-ReLU indexer scores over proxy keys. Args: index_q: ``[seq_q, groups, heads_per_group, proxy_dim]``. weights: ``[seq_q, groups, heads_per_group]``, already carrying the ``heads_per_group ** -0.5`` prescale. index_k: ``[seq_k, 1, proxy_dim]`` -- one shared key head (the indexer is MQA even though the main attention is GQA). Returns: ``[groups * seq_q, seq_k]`` float32, group-major then query. That row order is what the region selector consumes, so it is a contract rather than a convenience. """ if index_k.shape[1] != 1: raise ValueError(f"indexer keys are MQA; expected one head, got {index_k.shape[1]}") seq_q, groups, heads_per_group, proxy_dim = index_q.shape seq_k = index_k.shape[0] if index_k.shape[2] != proxy_dim: raise ValueError(f"proxy_dim mismatch: q={proxy_dim} k={index_k.shape[2]}") quantized_q = round_activations_e4m3(index_q) quantized_k = round_activations_e4m3(index_k) out = torch.empty((groups * seq_q, seq_k), device=index_q.device, dtype=torch.float32) _indexer_logits_kernel[ (triton.cdiv(seq_q, block_q), triton.cdiv(seq_k, block_k), groups) ]( quantized_q, quantized_k, weights, out, seq_q, seq_k, quantized_q.stride(0), quantized_q.stride(1), quantized_q.stride(2), quantized_k.stride(0), weights.stride(0), weights.stride(1), out.stride(0), heads_per_group=heads_per_group, proxy_dim=proxy_dim, BLOCK_Q=block_q, BLOCK_K=block_k, ) return out # Prefill selection packs one entry as ``phys_region | (valid_tokens << 24)``. # These are ``tl.constexpr`` because Triton kernels cannot read plain module globals. PREFILL_REGION_VALID_SHIFT = tl.constexpr(24) RADIX_BITS = tl.constexpr(8) RADIX_BUCKETS = tl.constexpr(256) RADIX_PASSES = tl.constexpr(4) @triton.jit def _ordered_float_key(scores): """Map fp32 values to monotonically ordered unsigned-32 keys held in int64. Negative IEEE-754 values reverse their bit ordering, while non-negative values become ordered after setting the sign bit. Keeping the resulting 32-bit pattern in a non-negative int64 lets ordinary Triton integer comparisons implement the required unsigned ordering. Branching on ``scores < 0`` deliberately maps ``-0.0`` and ``+0.0`` to the same key, matching numeric equality. Raw ordered-float keys place NaNs above ``+inf``. Production indexer scores are expected to be finite; this helper does not add a costly full-tensor finiteness scan to the hot path. """ bits = scores.to(tl.int32, bitcast=True).to(tl.int64) & 0xFFFFFFFF return tl.where( scores < 0.0, (~bits) & 0xFFFFFFFF, bits | 0x80000000, ) @triton.jit def _radix_topk_threshold( logits_ptr, row, stride_logits_row, visible, budget, seq_regions, BLOCK_R: tl.constexpr, ): """Find the exact score threshold separating the top ``budget`` entries of one row. Returns ``(threshold, remaining)``: the fp32 bit pattern of the ``budget``-th largest score, and how many entries *at* that exact key still belong in the selection. A threshold plus a tie count is what lets the caller take a **prefix** of a tied group rather than all or none of it, without ever materializing a sort. Scores may be signed. :func:`_ordered_float_key` converts their IEEE-754 bit patterns into unsigned-monotone keys before the radix passes. """ threshold = budget.to(tl.int64) * 0 remaining = budget bucket_ids = tl.arange(0, RADIX_BUCKETS) # Narrow the threshold one byte at a time, most significant first. for radix_pass in tl.static_range(RADIX_PASSES): shift = 32 - RADIX_BITS * (radix_pass + 1) histogram = tl.zeros((RADIX_BUCKETS,), dtype=tl.int32) for base in tl.range(0, seq_regions, BLOCK_R): regions = base + tl.arange(0, BLOCK_R) live = regions < visible scores = tl.load( logits_ptr + row * stride_logits_row + regions, mask=live, other=0.0 ) keys = _ordered_float_key(scores) if radix_pass == 0: # No high bits are fixed yet. Testing them would need a shift by 32, # which is undefined. matches = live else: high = shift + RADIX_BITS matches = live & ((keys >> high) == (threshold >> high)) # Park non-matching lanes in bucket 0 and subtract them back out. Relying on # tl.histogram to drop an out-of-range sentinel would corrupt the bucket # totals if it clamps instead -- and the totals are what pick the cut, so the # failure would look like a wrong threshold rather than a bad histogram. buckets = tl.where( matches, ((keys >> shift) & (RADIX_BUCKETS - 1)).to(tl.int32), 0, ) counts = tl.histogram(buckets, RADIX_BUCKETS) parked = tl.sum((~matches).to(tl.int32)) histogram += counts - tl.where(bucket_ids == 0, parked, 0) total = tl.sum(histogram) above = total - tl.cumsum(histogram, axis=0) # ``above`` is non-increasing, so the buckets satisfying above < remaining form a # suffix; the smallest of them is the one holding the cut. chosen = tl.min(tl.where(above < remaining, bucket_ids, RADIX_BUCKETS)) remaining -= tl.sum(tl.where(bucket_ids > chosen, histogram, 0)) # Cast before shifting the MSB byte. An int32 bucket >= 128 would otherwise # become negative first and sign-extend into the int64 threshold. threshold |= chosen.to(tl.int64) << shift return threshold, remaining @triton.jit def _region_topk_pack_kernel( logits_ptr, lengths_ptr, q_positions_ptr, block_table_ptr, request_ptr, out_ptr, seq_regions, stride_logits_row, stride_table_row, stride_out_row, topk: tl.constexpr, region_size: tl.constexpr, regions_per_page: tl.constexpr, BLOCK_R: tl.constexpr, BLOCK_TOPK: tl.constexpr, ) -> None: """Select the top-``topk`` visible regions of one row and pack them. Selection is ``torch.topk`` with ties broken by ascending region id: order the visible scores by ``(-score, id)`` and keep the first ``min(topk, visible)``. Ties at the cut take a *prefix* of the tied group, they do not drop it -- and that case is common rather than exotic, because post-ReLU indexer scores are exactly zero about 6.25% of the time, so the group sitting at the cut is usually the zero group. """ row = tl.program_id(0) visible = tl.load(lengths_ptr + row).to(tl.int32) visible = tl.minimum(tl.maximum(visible, 0), seq_regions) q_position = tl.load(q_positions_ptr + row).to(tl.int32) request = tl.load(request_ptr + row).to(tl.int32) slots = tl.arange(0, BLOCK_TOPK) tl.store(out_ptr + row * stride_out_row + slots, -1, mask=slots < topk) if visible <= 0: return budget = tl.minimum(topk, visible) threshold, remaining = _radix_topk_threshold( logits_ptr, row, stride_logits_row, visible, budget, seq_regions, BLOCK_R=BLOCK_R ) # Emit in ascending region order. The selection contract specifies a set; # ascending order also matches the standalone metadata converter. cursor = 0 ties_taken = 0 for base in tl.range(0, seq_regions, BLOCK_R): regions = base + tl.arange(0, BLOCK_R) live = regions < visible scores = tl.load( logits_ptr + row * stride_logits_row + regions, mask=live, other=0.0 ) keys = _ordered_float_key(scores) greater = live & (keys > threshold) tied = live & (keys == threshold) tied_rank = ties_taken + tl.cumsum(tied.to(tl.int32), axis=0) - 1 selected = greater | (tied & (tied_rank < remaining)) positions = cursor + tl.cumsum(selected.to(tl.int32), axis=0) - 1 safe_regions = tl.where(selected, regions, 0) pages = tl.load( block_table_ptr + request * stride_table_row + safe_regions // regions_per_page, mask=selected, other=0, ) physical = pages * regions_per_page + safe_regions % regions_per_page valid_tokens = tl.minimum( tl.maximum(q_position + 1 - safe_regions * region_size, 0), region_size ) packed = physical | (valid_tokens << PREFILL_REGION_VALID_SHIFT) # An unmapped page is a legal block-table state, not an error. It has to become a # ``-1`` slot rather than an address computed from a negative page, which would # read from before the cache. packed = tl.where(pages >= 0, packed, -1) tl.store( out_ptr + row * stride_out_row + positions, packed, mask=selected & (positions < topk), ) cursor += tl.sum(selected.to(tl.int32)) ties_taken += tl.sum(tied.to(tl.int32)) @triton.jit def _region_topk_ids_kernel( logits_ptr, lengths_ptr, out_ptr, seq_regions, stride_logits_row, stride_out_row, topk: tl.constexpr, BLOCK_R: tl.constexpr, BLOCK_TOPK: tl.constexpr, ) -> None: """Select the top-``topk`` visible regions and emit **logical ids**, ascending. The same selection as :func:`_region_topk_pack_kernel` with the packing left off. Prefill and decode pack the identical selection into two different words, and the decode word needs the logical region id, which the prefill word does not carry -- so decode metadata cannot be derived from prefill metadata, only from the raw ids. """ row = tl.program_id(0) visible = tl.load(lengths_ptr + row).to(tl.int32) visible = tl.minimum(tl.maximum(visible, 0), seq_regions) slots = tl.arange(0, BLOCK_TOPK) tl.store(out_ptr + row * stride_out_row + slots, -1, mask=slots < topk) if visible <= 0: return budget = tl.minimum(topk, visible) threshold, remaining = _radix_topk_threshold( logits_ptr, row, stride_logits_row, visible, budget, seq_regions, BLOCK_R=BLOCK_R ) cursor = 0 ties_taken = 0 for base in tl.range(0, seq_regions, BLOCK_R): regions = base + tl.arange(0, BLOCK_R) live = regions < visible scores = tl.load( logits_ptr + row * stride_logits_row + regions, mask=live, other=0.0 ) keys = _ordered_float_key(scores) greater = live & (keys > threshold) tied = live & (keys == threshold) tied_rank = ties_taken + tl.cumsum(tied.to(tl.int32), axis=0) - 1 selected = greater | (tied & (tied_rank < remaining)) positions = cursor + tl.cumsum(selected.to(tl.int32), axis=0) - 1 tl.store( out_ptr + row * stride_out_row + positions, regions, mask=selected & (positions < topk), ) cursor += tl.sum(selected.to(tl.int32)) ties_taken += tl.sum(tied.to(tl.int32)) def region_topk_ids( logits: torch.Tensor, lengths: torch.Tensor, *, topk: int, block_r: int = 1024, ) -> torch.Tensor: """Top-k region selection, returning logical region ids in ascending order. Args: logits: ``[rows, seq_regions]`` float32 indexer scores, which may be signed. lengths: ``[rows]`` int32, candidate region count per row. For DSA this is the **history** length, ``q_position // region_size``, which excludes the region containing the query itself -- that one is appended by the metadata builder rather than made to compete. Returns: ``[rows, topk]`` int32 logical region ids, ``-1`` padding the tail. """ if logits.dtype != torch.float32: raise ValueError(f"selector expects float32 scores, got {logits.dtype}") rows, seq_regions = logits.shape out = torch.empty((rows, int(topk)), device=logits.device, dtype=torch.int32) _region_topk_ids_kernel[(rows,)]( logits, lengths, out, seq_regions, logits.stride(0), out.stride(0), topk=int(topk), BLOCK_R=min(block_r, triton.next_power_of_2(seq_regions)), BLOCK_TOPK=triton.next_power_of_2(int(topk)), ) return out def region_topk_pack( logits: torch.Tensor, lengths: torch.Tensor, q_positions: torch.Tensor, block_table: torch.Tensor, request_indices: torch.Tensor | None = None, *, topk: int, region_size: int = REGION_BLOCK_SIZE, regions_per_page: int, block_r: int = 1024, ) -> torch.Tensor: """Top-k region selection fused with prefill metadata packing. Args: logits: ``[rows, seq_regions]`` float32 indexer scores, which may be signed, where ``rows = groups * seq_q`` in group-major order. lengths: ``[rows]`` int32, causally visible region count per row. q_positions: ``[rows]`` int32, absolute token position of each row's query. block_table: ``[num_reqs, pages]`` int32, logical page -> physical page, one row per request. ``-1`` marks an unmapped page. A table with a single row is treated as shared by every request, regardless of ``request_indices``. request_indices: ``[rows]`` int32, which block-table row each score row belongs to. Defaults to all-zero. Returns: ``[rows, topk]`` int32, each valid slot holding ``phys_region | (valid_tokens << 24)``, with ``-1`` padding the tail. A row fills exactly ``min(topk, visible)`` slots. """ if logits.dtype != torch.float32: raise ValueError(f"selector expects float32 scores, got {logits.dtype}") if block_table.ndim != 2: raise ValueError( f"block_table must be [num_reqs, pages], got {tuple(block_table.shape)}" ) rows, seq_regions = logits.shape if request_indices is None: request_indices = torch.zeros((rows,), device=logits.device, dtype=torch.int32) out = torch.empty((rows, int(topk)), device=logits.device, dtype=torch.int32) _region_topk_pack_kernel[(rows,)]( logits, lengths, q_positions, block_table, request_indices, out, seq_regions, logits.stride(0), shared_table_stride(block_table), out.stride(0), topk=int(topk), region_size=int(region_size), regions_per_page=int(regions_per_page), BLOCK_R=min(block_r, triton.next_power_of_2(seq_regions)), BLOCK_TOPK=triton.next_power_of_2(int(topk)), ) return out # Sparse metadata. The query's current region is appended after selecting # past regions. Prefill packs phys | valid_tokens << 24; decode packs # start_token | phys << 32 and uses region_counts to exclude padding. DECODE_PHYS_SHIFT = tl.constexpr(32) # Larger than any real packed word, so invalid slots sort to the tail. DECODE_SORT_SENTINEL = tl.constexpr(1 << 62) @triton.jit def _append_current_region_kernel( history_ptr, q_positions_ptr, block_table_ptr, request_ptr, out_ptr, counts_ptr, stride_history_row, stride_table_row, stride_out_row, topk: tl.constexpr, region_size: tl.constexpr, regions_per_page: tl.constexpr, BLOCK_K: tl.constexpr, ) -> None: """Widen prefill history metadata to ``topk + 1`` slots and append the current region. The history is copied verbatim -- including its ``-1`` slots, which stay *in place* rather than being compacted, because prefill's ``region_counts`` is positional (``min(current_region, topk) + 1``) and not a count of live entries. The appended slot lands at exactly the index the selector stopped filling. The two stores are kept **disjoint** (``slots != history_count``) rather than letting the second overwrite the first. Writing the whole row and then patching one slot is a cross-thread race: the row store is a vector store spread over lanes, the patch is a scalar store from one lane, and nothing orders them. It loses roughly one row in a thousand, so correctness requires disjoint writes even when small tests pass. """ row = tl.program_id(0) q_position = tl.load(q_positions_ptr + row).to(tl.int32) request = tl.load(request_ptr + row).to(tl.int32) current_region = tl.maximum(q_position // region_size, 0) history_count = tl.minimum(current_region, topk) slots = tl.arange(0, BLOCK_K) history = tl.load( history_ptr + row * stride_history_row + slots, mask=slots < topk, other=-1 ) tl.store( out_ptr + row * stride_out_row + slots, history, mask=(slots < topk + 1) & (slots != history_count), ) page = tl.load( block_table_ptr + request * stride_table_row + current_region // regions_per_page ) physical = page * regions_per_page + current_region % regions_per_page valid_tokens = tl.minimum( tl.maximum(q_position + 1 - current_region * region_size, 0), region_size ) packed = physical | (valid_tokens << PREFILL_REGION_VALID_SHIFT) tl.store( out_ptr + row * stride_out_row + history_count, tl.where(page >= 0, packed, -1), ) tl.store(counts_ptr + row, history_count + 1) @triton.jit def _decode_pack_kernel( ids_ptr, seqlens_ptr, block_table_ptr, request_ptr, out_ptr, counts_ptr, stride_ids_row, stride_table_row, stride_out_row, topk: tl.constexpr, region_size: tl.constexpr, regions_per_page: tl.constexpr, BLOCK_K: tl.constexpr, ) -> None: """Pack selected logical regions into the int64 decode word. One row fits in one tile, so the compaction is a sort rather than a scatter: invalid slots are mapped to a sentinel above every real word, and sorting by the packed value both drops them to the tail and produces physical-major order. The current region is written at the surviving count after sorting, so it stays last. """ row = tl.program_id(0) request = tl.load(request_ptr + row).to(tl.int32) kv_seqlen = tl.load(seqlens_ptr + request).to(tl.int32) valid_regions = (kv_seqlen + region_size - 1) // region_size slots = tl.arange(0, BLOCK_K) ids = tl.load(ids_ptr + row * stride_ids_row + slots, mask=slots < topk, other=-1) live = (slots < topk) & (ids >= 0) & (ids < valid_regions) safe_ids = tl.where(live, ids, 0) pages = tl.load( block_table_ptr + request * stride_table_row + safe_ids // regions_per_page, mask=live, other=-1, ) live = live & (pages >= 0) physical = pages * regions_per_page + safe_ids % regions_per_page packed = (safe_ids * region_size).to(tl.int64) | ( physical.to(tl.int64) << DECODE_PHYS_SHIFT ) packed = tl.where(live, packed, DECODE_SORT_SENTINEL) packed = tl.sort(packed, descending=False) count = tl.sum(live.to(tl.int32)) tl.store( out_ptr + row * stride_out_row + slots, tl.where(slots < count, packed, 0), # Slot ``count`` is left to the scalar store below. Overlapping the two would be a # cross-thread race between a vector store and a scalar one -- see # :func:`_append_current_region_kernel`. mask=(slots < topk + 1) & (slots != count), ) current_region = tl.maximum((kv_seqlen - 1) // region_size, 0) current_page = tl.load( block_table_ptr + request * stride_table_row + current_region // regions_per_page ) current_live = (current_region < valid_regions) & (current_page >= 0) current_physical = ( current_page * regions_per_page + current_region % regions_per_page ) current_packed = (current_region * region_size).to(tl.int64) | ( current_physical.to(tl.int64) << DECODE_PHYS_SHIFT ) # Written unconditionally so the slot is defined even when the current region is # unmapped, in which case it becomes ordinary zero padding. tl.store( out_ptr + row * stride_out_row + count, tl.where(current_live, current_packed, 0), ) tl.store(counts_ptr + row, count + current_live.to(tl.int32)) def prefill_sparse_meta( logits: torch.Tensor, q_positions: torch.Tensor, block_table: torch.Tensor, request_indices: torch.Tensor | None = None, *, topk: int, region_size: int = REGION_BLOCK_SIZE, regions_per_page: int, ) -> tuple[torch.Tensor, torch.Tensor]: """Select regions and build prefill sparse metadata. Args: logits: ``[rows, seq_regions]`` float32 indexer scores, which may be signed, group-major (``rows = num_kv_groups * total_q``). The scores are weighted sums of per-head ReLU terms and unconstrained signed indexer weights. q_positions: ``[rows]`` int32, absolute token position of each row's query. block_table: ``[num_reqs, pages]`` int32, logical page -> physical page, one row per request. ``-1`` marks an unmapped page, which yields a ``-1`` slot rather than an error. request_indices: ``[rows]`` int32, which block-table row each score row uses. Defaults to all-zero, i.e. every row shares one table. Returns: ``(packed [rows, topk + 1] int32, region_counts [rows] int32)``, ready for :func:`inference.sparse_attn.sparse_attention_prefill`. The history length is derived here as ``q_position // region_size`` rather than taken from the caller, because getting it wrong by one is the difference between the current region being force-included and it competing for a slot. """ rows = int(logits.shape[0]) if request_indices is None: request_indices = torch.zeros((rows,), device=logits.device, dtype=torch.int32) history_lengths = torch.div( q_positions.to(torch.int32), region_size, rounding_mode="floor" ).contiguous() history = region_topk_pack( logits, history_lengths, q_positions, block_table, request_indices, topk=topk, region_size=region_size, regions_per_page=regions_per_page, ) packed = torch.empty((rows, int(topk) + 1), device=logits.device, dtype=torch.int32) counts = torch.empty((rows,), device=logits.device, dtype=torch.int32) _append_current_region_kernel[(rows,)]( history, q_positions, block_table, request_indices, packed, counts, history.stride(0), shared_table_stride(block_table), packed.stride(0), topk=int(topk), region_size=int(region_size), regions_per_page=int(regions_per_page), BLOCK_K=triton.next_power_of_2(int(topk) + 1), ) return packed, counts def decode_sparse_meta( logits: torch.Tensor, kv_seqlens: torch.Tensor, block_table: torch.Tensor, request_indices: torch.Tensor | None = None, *, topk: int, region_size: int = REGION_BLOCK_SIZE, regions_per_page: int, ) -> tuple[torch.Tensor, torch.Tensor]: """Select regions and build decode sparse metadata. Args: logits: ``[num_kv_groups * num_reqs, seq_regions]`` float32 indexer scores, which may be signed, group-major. The selector orders them numerically rather than by the magnitude encoded in their IEEE-754 bit pattern. kv_seqlens: ``[num_reqs]`` int32, context length *including* the current token. block_table: ``[num_reqs, pages]`` int32, logical page -> physical page. request_indices: ``[rows]`` int32. Defaults to the group-major layout implied by the row count, ``arange(num_reqs)`` tiled once per KV group. Returns: ``(packed [rows, topk + 1] int64, region_counts [rows] int32)``, ready for :func:`inference.sparse_attn.sparse_attention_decode`. ``region_counts`` here is a count of *surviving* entries, so it varies with how much of the block table is mapped. The tail beyond it is zero, and a zero word is not self-identifying -- the count is the only bound. """ rows, _ = logits.shape num_reqs = int(kv_seqlens.shape[0]) if rows % num_reqs: raise ValueError(f"{rows} score rows do not divide into {num_reqs} requests") if request_indices is None: request_indices = ( torch.arange(num_reqs, device=logits.device, dtype=torch.int32) .repeat(rows // num_reqs) .contiguous() ) # Decode's query position is the last token, so the history is everything strictly # before the current region. history_lengths = torch.div( (kv_seqlens.to(torch.int32) - 1).clamp_min(0), region_size, rounding_mode="floor" )[request_indices.to(torch.long)].contiguous() ids = region_topk_ids(logits, history_lengths, topk=topk) packed = torch.empty((rows, int(topk) + 1), device=logits.device, dtype=torch.int64) counts = torch.empty((rows,), device=logits.device, dtype=torch.int32) _decode_pack_kernel[(rows,)]( ids, kv_seqlens, block_table, request_indices, packed, counts, ids.stride(0), shared_table_stride(block_table), packed.stride(0), topk=int(topk), region_size=int(region_size), regions_per_page=int(regions_per_page), BLOCK_K=triton.next_power_of_2(int(topk) + 1), ) return packed, counts # Sparse attention over each query's selected regions. Causality is encoded # in metadata; softmax covers only gathered tokens. REGION_VALID_SHIFT = tl.constexpr(24) REGION_ID_MASK = tl.constexpr((1 << 24) - 1) DECODE_LOW_MASK = tl.constexpr((1 << 32) - 1) NEGATIVE_INFINITY = tl.constexpr(float("-inf")) @triton.jit def _sparse_attn_prefill_kernel( q_ptr, k_ptr, v_ptr, packed_ptr, counts_ptr, out_ptr, lse_ptr, stride_q_token, stride_q_head, stride_k_token, stride_k_group, stride_v_token, stride_v_group, stride_packed_row, stride_out_token, stride_out_head, stride_lse_token, softmax_scale, total_q, topk: tl.constexpr, heads_per_group: tl.constexpr, head_dim: tl.constexpr, region_size: tl.constexpr, BLOCK_H: tl.constexpr, BLOCK_D: tl.constexpr, BLOCK_R: tl.constexpr, ) -> None: """Attention for one query against one KV group's selected regions. Online softmax over region tiles. The accumulator is per (head, dim) and the running max is per head, exactly as in dense flash attention -- the only difference is that the K/V rows come from a gather driven by the packed region list instead of a contiguous range. """ query = tl.program_id(0) group = tl.program_id(1) row = group * total_q + query count = tl.load(counts_ptr + row).to(tl.int32) heads = tl.arange(0, BLOCK_H) dims = tl.arange(0, BLOCK_D) head_ok = heads < heads_per_group dim_ok = dims < head_dim tile_ok = head_ok[:, None] & dim_ok[None, :] queries = tl.load( q_ptr + query * stride_q_token + (group * heads_per_group + heads)[:, None] * stride_q_head + dims[None, :], mask=tile_ok, other=0.0, ) running_max = tl.full((BLOCK_H,), NEGATIVE_INFINITY, dtype=tl.float32) running_sum = tl.zeros((BLOCK_H,), dtype=tl.float32) accumulator = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32) offsets_in_region = tl.arange(0, region_size) for base in tl.range(0, topk, BLOCK_R): slots = base + tl.arange(0, BLOCK_R) slot_ok = slots < count meta = tl.load( packed_ptr + row * stride_packed_row + slots, mask=slot_ok, other=-1 ) physical = meta & REGION_ID_MASK valid = meta >> REGION_VALID_SHIFT # A slot is live only if it was within the count, not a -1 sentinel, and actually # exposes tokens. The third condition is not redundant: a region whose first token # is past the query position packs a valid id with zero visible tokens. live = slot_ok & (meta >= 0) & (valid > 0) tokens = tl.reshape( physical[:, None] * region_size + offsets_in_region[None, :], (BLOCK_R * region_size,), ) token_ok = tl.reshape( live[:, None] & (offsets_in_region[None, :] < valid[:, None]), (BLOCK_R * region_size,), ) safe_tokens = tl.where(token_ok, tokens, 0) keys = tl.load( k_ptr + safe_tokens[:, None] * stride_k_token + group * stride_k_group + dims[None, :], mask=token_ok[:, None] & dim_ok[None, :], other=0.0, ) values = tl.load( v_ptr + safe_tokens[:, None] * stride_v_token + group * stride_v_group + dims[None, :], mask=token_ok[:, None] & dim_ok[None, :], other=0.0, ) scores = tl.dot(queries, tl.trans(keys), out_dtype=tl.float32) * softmax_scale scores = tl.where(token_ok[None, :] & head_ok[:, None], scores, NEGATIVE_INFINITY) tile_max = tl.maximum(running_max, tl.max(scores, axis=1)) # A tile can be entirely masked -- every slot exhausted, or a row with no # selection at all. Then both maxima are -inf and ``m - m`` is NaN, which would # poison the accumulator for the rest of the loop. Substituting a finite shift # keeps the arithmetic well defined and leaves the running state untouched. empty = tile_max == NEGATIVE_INFINITY shift = tl.where(empty, 0.0, tile_max) rescale = tl.where(empty, 1.0, tl.exp(running_max - shift)) probabilities = tl.exp(scores - shift[:, None]) running_sum = running_sum * rescale + tl.sum(probabilities, axis=1) accumulator = accumulator * rescale[:, None] + tl.dot( probabilities.to(values.dtype), values, out_dtype=tl.float32 ) running_max = tile_max selected = running_sum > 0 output = tl.where(selected[:, None], accumulator / running_sum[:, None], 0.0) tl.store( out_ptr + query * stride_out_token + (group * heads_per_group + heads)[:, None] * stride_out_head + dims[None, :], output.to(out_ptr.dtype.element_ty), mask=tile_ok, ) tl.store( lse_ptr + query * stride_lse_token + group * heads_per_group + heads, tl.where(selected, running_max + tl.log(running_sum), NEGATIVE_INFINITY), mask=head_ok, ) def sparse_attention_prefill( query: torch.Tensor, key_cache: torch.Tensor, value_cache: torch.Tensor, packed_regions: torch.Tensor, region_counts: torch.Tensor, *, num_kv_groups: int, region_size: int = 8, softmax_scale: float | None = None, block_regions: int = 16, ) -> tuple[torch.Tensor, torch.Tensor]: """Sparse attention over per-query selected regions. Args: query: ``[total_q, num_heads, head_dim]``. key_cache, value_cache: ``[tokens, num_kv_groups, head_dim]``, i.e. the paged cache already flattened across pages. Flattening is safe because a region's tokens are contiguous in the physical token space by construction. packed_regions: ``[num_kv_groups * total_q, topk]`` int32 from the selector, each slot ``phys_region | (valid_tokens << 24)``, ``-1`` padding the tail. region_counts: ``[num_kv_groups * total_q]`` int32, how many slots to read. This truncates *positionally* -- it does not re-select by score. Returns: ``(out, lse)`` shaped ``[total_q, num_heads, head_dim]`` and ``[total_q, num_heads]``. Rows with no selected region get a zero output and ``lse = -inf``. The caller must guarantee the selected regions are distinct. Our selector does (the stage 3 gate checks it), but a union-with-sliding-window path would not, and duplicate regions would silently double-weight those tokens in the softmax rather than fail. """ total_q, num_heads, head_dim = query.shape if num_heads % num_kv_groups: raise ValueError(f"{num_heads} heads do not divide into {num_kv_groups} groups") if packed_regions.shape[0] != num_kv_groups * total_q: raise ValueError( f"expected {num_kv_groups * total_q} metadata rows, got {packed_regions.shape[0]}" ) if key_cache.shape != value_cache.shape: raise ValueError( f"key/value cache shape mismatch: {tuple(key_cache.shape)} vs {tuple(value_cache.shape)}" ) heads_per_group = num_heads // num_kv_groups topk = int(packed_regions.shape[1]) scale = softmax_scale if softmax_scale is not None else head_dim**-0.5 out = torch.empty_like(query) lse = torch.empty((total_q, num_heads), device=query.device, dtype=torch.float32) _sparse_attn_prefill_kernel[(total_q, num_kv_groups)]( query, key_cache, value_cache, packed_regions, region_counts, out, lse, query.stride(0), query.stride(1), key_cache.stride(0), key_cache.stride(1), value_cache.stride(0), value_cache.stride(1), packed_regions.stride(0), out.stride(0), out.stride(1), lse.stride(0), scale, total_q, topk=topk, heads_per_group=heads_per_group, head_dim=head_dim, region_size=region_size, BLOCK_H=max(16, triton.next_power_of_2(heads_per_group)), BLOCK_D=triton.next_power_of_2(head_dim), BLOCK_R=min(block_regions, triton.next_power_of_2(topk)), ) return out, lse @triton.jit def _sparse_attn_decode_kernel( q_ptr, k_ptr, v_ptr, packed_ptr, counts_ptr, seqlens_ptr, out_ptr, lse_ptr, stride_q_token, stride_q_head, stride_k_token, stride_k_group, stride_v_token, stride_v_group, stride_packed_row, stride_out_split, stride_out_token, stride_out_head, stride_lse_split, stride_lse_token, softmax_scale, num_reqs, topk: tl.constexpr, heads_per_group: tl.constexpr, head_dim: tl.constexpr, region_size: tl.constexpr, slots_per_split: tl.constexpr, BLOCK_H: tl.constexpr, BLOCK_D: tl.constexpr, BLOCK_R: tl.constexpr, ) -> None: """One decode query against one KV group's selected regions, over one slot window. Structurally the same online softmax as prefill, but the metadata contract differs in the one way that matters: decode packs ``start_token | (phys_region << 32)`` in int64 and does **not** carry ``valid_tokens``. The visible length is recomputed here as ``clamp(kv_seqlen - start_token, 0, region_size)``, which is why the logical start token has to be in the word at all -- the physical region alone cannot say how much of the last region exists yet. That also disposes of both padding sentinels without a special case. ``-1`` and ``0x7FFF_FFFF_FFFF_FFFF`` both leave ``0xFFFF_FFFF`` in the low word, so ``kv_seqlen - start_token`` is hugely negative and the slot clamps to zero tokens. """ request = tl.program_id(0) group = tl.program_id(1) split = tl.program_id(2) row = group * num_reqs + request count = tl.load(counts_ptr + row).to(tl.int32) kv_seqlen = tl.load(seqlens_ptr + request).to(tl.int64) heads = tl.arange(0, BLOCK_H) dims = tl.arange(0, BLOCK_D) head_ok = heads < heads_per_group dim_ok = dims < head_dim tile_ok = head_ok[:, None] & dim_ok[None, :] queries = tl.load( q_ptr + request * stride_q_token + (group * heads_per_group + heads)[:, None] * stride_q_head + dims[None, :], mask=tile_ok, other=0.0, ) running_max = tl.full((BLOCK_H,), NEGATIVE_INFINITY, dtype=tl.float32) running_sum = tl.zeros((BLOCK_H,), dtype=tl.float32) accumulator = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32) offsets_in_region = tl.arange(0, region_size) window_begin = split * slots_per_split for base in tl.range(window_begin, window_begin + slots_per_split, BLOCK_R): slots = base + tl.arange(0, BLOCK_R) slot_ok = (slots < count) & (slots < topk) meta = tl.load( packed_ptr + row * stride_packed_row + slots, mask=slot_ok, other=-1 ) physical = (meta >> DECODE_PHYS_SHIFT) & DECODE_LOW_MASK start_token = meta & DECODE_LOW_MASK valid = tl.minimum(tl.maximum(kv_seqlen - start_token, 0), region_size) live = slot_ok & (valid > 0) # ``physical`` is a sentinel-sized integer on dead slots, so it must be # neutralised before it becomes an address -- masking the load is not enough, # since the pointer arithmetic itself would overflow. safe_physical = tl.where(live, physical, 0) tokens = tl.reshape( safe_physical[:, None] * region_size + offsets_in_region[None, :], (BLOCK_R * region_size,), ) token_ok = tl.reshape( live[:, None] & (offsets_in_region[None, :] < valid[:, None]), (BLOCK_R * region_size,), ) keys = tl.load( k_ptr + tokens[:, None] * stride_k_token + group * stride_k_group + dims[None, :], mask=token_ok[:, None] & dim_ok[None, :], other=0.0, ) values = tl.load( v_ptr + tokens[:, None] * stride_v_token + group * stride_v_group + dims[None, :], mask=token_ok[:, None] & dim_ok[None, :], other=0.0, ) scores = tl.dot(queries, tl.trans(keys), out_dtype=tl.float32) * softmax_scale scores = tl.where(token_ok[None, :] & head_ok[:, None], scores, NEGATIVE_INFINITY) tile_max = tl.maximum(running_max, tl.max(scores, axis=1)) empty = tile_max == NEGATIVE_INFINITY shift = tl.where(empty, 0.0, tile_max) rescale = tl.where(empty, 1.0, tl.exp(running_max - shift)) probabilities = tl.exp(scores - shift[:, None]) running_sum = running_sum * rescale + tl.sum(probabilities, axis=1) accumulator = accumulator * rescale[:, None] + tl.dot( probabilities.to(values.dtype), values, out_dtype=tl.float32 ) running_max = tile_max selected = running_sum > 0 output = tl.where(selected[:, None], accumulator / running_sum[:, None], 0.0) tl.store( out_ptr + split * stride_out_split + request * stride_out_token + (group * heads_per_group + heads)[:, None] * stride_out_head + dims[None, :], output.to(out_ptr.dtype.element_ty), mask=tile_ok, ) tl.store( lse_ptr + split * stride_lse_split + request * stride_lse_token + group * heads_per_group + heads, tl.where(selected, running_max + tl.log(running_sum), NEGATIVE_INFINITY), mask=head_ok, ) @triton.jit def _merge_split_states_kernel( partial_out_ptr, partial_lse_ptr, out_ptr, lse_ptr, stride_po_split, stride_po_token, stride_po_head, stride_pl_split, stride_pl_token, stride_out_token, stride_out_head, stride_lse_token, num_splits, head_dim: tl.constexpr, BLOCK_S: tl.constexpr, BLOCK_D: tl.constexpr, ) -> None: """Combine per-split attention states into one, weighting by the partial lse. Each split produced an output already divided by its own softmax denominator, so the combination is a softmax over the partial log-sum-exps rather than a plain sum: ``out = sum_s softmax(lse)_s * out_s`` and ``lse = logsumexp_s(lse_s)``. A split that covered no live slot reports ``lse = -inf`` and drops out with weight zero; if *every* split does, the row selected nothing and must come back as a zero output with ``lse = -inf`` rather than ``0/0``. """ token = tl.program_id(0) head = tl.program_id(1) splits = tl.arange(0, BLOCK_S) dims = tl.arange(0, BLOCK_D) split_ok = splits < num_splits dim_ok = dims < head_dim partial_lse = tl.load( partial_lse_ptr + splits * stride_pl_split + token * stride_pl_token + head, mask=split_ok, other=NEGATIVE_INFINITY, ) peak = tl.max(partial_lse) empty = peak == NEGATIVE_INFINITY shift = tl.where(empty, 0.0, peak) weights = tl.where(split_ok, tl.exp(partial_lse - shift), 0.0) total = tl.sum(weights) states = tl.load( partial_out_ptr + splits[:, None] * stride_po_split + token * stride_po_token + head * stride_po_head + dims[None, :], mask=split_ok[:, None] & dim_ok[None, :], other=0.0, ).to(tl.float32) merged = tl.sum(states * weights[:, None], axis=0) / tl.maximum(total, 1.0e-20) tl.store( out_ptr + token * stride_out_token + head * stride_out_head + dims, merged.to(out_ptr.dtype.element_ty), mask=dim_ok, ) tl.store( lse_ptr + token * stride_lse_token + head, tl.where(empty, NEGATIVE_INFINITY, shift + tl.log(total)), ) def merge_split_states( partial_out: torch.Tensor, partial_lse: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: """Reduce ``[splits, tokens, heads, dim]`` partial states to ``[tokens, heads, dim]``. Each request uses the same split count. Partial outputs are combined using their natural-log softmax normalization values. """ num_splits, tokens, num_heads, head_dim = partial_out.shape out = torch.empty((tokens, num_heads, head_dim), device=partial_out.device, dtype=partial_out.dtype) lse = torch.empty((tokens, num_heads), device=partial_out.device, dtype=torch.float32) _merge_split_states_kernel[(tokens, num_heads)]( partial_out, partial_lse, out, lse, partial_out.stride(0), partial_out.stride(1), partial_out.stride(2), partial_lse.stride(0), partial_lse.stride(1), out.stride(0), out.stride(1), lse.stride(0), num_splits, head_dim=head_dim, BLOCK_S=triton.next_power_of_2(num_splits), BLOCK_D=triton.next_power_of_2(head_dim), ) return out, lse def sparse_attention_decode( query: torch.Tensor, key_cache: torch.Tensor, value_cache: torch.Tensor, packed_regions: torch.Tensor, region_counts: torch.Tensor, kv_seqlens: torch.Tensor, *, num_kv_groups: int, region_size: int = 8, softmax_scale: float | None = None, num_splits: int = 1, block_regions: int = 16, ) -> tuple[torch.Tensor, torch.Tensor]: """Sparse attention for one query token per request. Args: query: ``[num_reqs, num_heads, head_dim]``. key_cache, value_cache: ``[tokens, num_kv_groups, head_dim]``, pages flattened. packed_regions: ``[num_kv_groups * num_reqs, topk]`` **int64**, each slot ``start_token | (phys_region << 32)``. Note this is a different packing from prefill's int32 ``phys_region | (valid_tokens << 24)``: decode carries the logical start token and derives the visible length from ``kv_seqlens``. region_counts: ``[num_kv_groups * num_reqs]`` int32, slots to read per row. kv_seqlens: ``[num_reqs]`` int32, context length including the current token. num_splits: split the slot window this many ways and return partial states. With ``num_splits == 1`` the returned tensors are already the final result; otherwise pass them to :func:`merge_split_states`. Returns: ``num_splits == 1``: ``(out [num_reqs, num_heads, head_dim], lse [num_reqs, num_heads])``. Otherwise the same with a leading ``num_splits`` axis. As in prefill, the selected regions must be distinct. Sparse decode is stricter about this than the prefill one: it concatenates the gathered spans with no dedup at all, so a duplicated region would be counted twice in the softmax. """ num_reqs, num_heads, head_dim = query.shape if num_heads % num_kv_groups: raise ValueError(f"{num_heads} heads do not divide into {num_kv_groups} groups") if packed_regions.shape[0] != num_kv_groups * num_reqs: raise ValueError( f"expected {num_kv_groups * num_reqs} metadata rows, got {packed_regions.shape[0]}" ) if packed_regions.dtype != torch.int64: raise ValueError(f"decode metadata must be int64, got {packed_regions.dtype}") if kv_seqlens.shape[0] != num_reqs: raise ValueError(f"expected {num_reqs} kv_seqlens, got {kv_seqlens.shape[0]}") heads_per_group = num_heads // num_kv_groups topk = int(packed_regions.shape[1]) scale = softmax_scale if softmax_scale is not None else head_dim**-0.5 # Every split walks the same number of slots so the window bound is a constexpr; the # tail split simply finds every slot masked out. slots_per_split = -(-topk // num_splits) out = torch.empty((num_splits, num_reqs, num_heads, head_dim), device=query.device, dtype=query.dtype) lse = torch.empty((num_splits, num_reqs, num_heads), device=query.device, dtype=torch.float32) _sparse_attn_decode_kernel[(num_reqs, num_kv_groups, num_splits)]( query, key_cache, value_cache, packed_regions, region_counts, kv_seqlens, out, lse, query.stride(0), query.stride(1), key_cache.stride(0), key_cache.stride(1), value_cache.stride(0), value_cache.stride(1), packed_regions.stride(0), out.stride(0), out.stride(1), out.stride(2), lse.stride(0), lse.stride(1), scale, num_reqs, topk=topk, heads_per_group=heads_per_group, head_dim=head_dim, region_size=region_size, slots_per_split=slots_per_split, BLOCK_H=max(16, triton.next_power_of_2(heads_per_group)), BLOCK_D=triton.next_power_of_2(head_dim), BLOCK_R=min(block_regions, triton.next_power_of_2(slots_per_split)), ) if num_splits == 1: return out[0], lse[0] return out, lse # DSA state: main K/V, summaries of completed regions and pending region tokens. # Block tables map request-local pages to physical cache storage. def build_rope_cache( *, rotary_span: int, theta: float, max_position: int, device: torch.device, dtype: torch.dtype = torch.float32, scaling: dict | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Cos/sin tables for a NEOX-style partial rotation. Args: rotary_span: how many of the head's dimensions rotate. The returned tables are ``[max_position, rotary_span // 2]`` -- *half* the span, because each frequency drives the pair ``(d, d + rotary_span // 2)``. The kernels take that halved width as their ``rotary_dim``, which is easy to misread as the span itself. scaling: the ``rope_scaling`` block, or ``None``. Only ``llama3`` is implemented, because it is the only type step5's config asks for; an unrecognised type raises rather than silently falling back to unscaled, which would look like a long-context quality regression rather than a bug. """ if rotary_span % 2: raise ValueError(f"rotary span must be even, got {rotary_span}") half = rotary_span // 2 # Compute frequencies and trigonometry in FP32, then round to the requested # cache dtype. This precision boundary determines the BF16 RoPE table. inv_freq = 1.0 / ( theta ** (torch.arange(0, rotary_span, 2, device=device, dtype=torch.float32) / rotary_span) ) if scaling: inv_freq = _apply_llama3_scaling(inv_freq, scaling) positions = torch.arange(max_position, device=device, dtype=torch.float32) freqs = torch.outer(positions, inv_freq) assert freqs.shape[1] == half return freqs.cos().to(dtype).contiguous(), freqs.sin().to(dtype).contiguous() def _apply_llama3_scaling(inv_freq: torch.Tensor, scaling: dict) -> torch.Tensor: rope_type = str(scaling.get("rope_type", scaling.get("type", ""))).lower() if rope_type != "llama3": raise NotImplementedError(f"unsupported rope_scaling type {rope_type!r}") factor = float(scaling["factor"]) low_freq_factor = float(scaling["low_freq_factor"]) high_freq_factor = float(scaling["high_freq_factor"]) original_max = float(scaling["original_max_position_embeddings"]) low_wavelen = original_max / low_freq_factor high_wavelen = original_max / high_freq_factor wavelen = 2 * math.pi / inv_freq # Long wavelengths are stretched by the full factor, short ones are left alone, and the # band between is interpolated. Applying the stretch everywhere is the classic mistake: # it destroys the high-frequency positional detail the short-range heads rely on. scaled = torch.where(wavelen > low_wavelen, inv_freq / factor, inv_freq) smooth = (original_max / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor) smoothed = (1 - smooth) * scaled / factor + smooth * scaled in_band = (wavelen >= high_wavelen) & (wavelen <= low_wavelen) return torch.where(in_band, smoothed, scaled) @dataclass(frozen=True) class DSAGeometry: """The region/page arithmetic, in one place. ``regions_per_page`` is derived rather than configured: it is the ratio the metadata kernels use to split a logical region id into a page lookup and an offset, so a configured value that disagreed with ``page_size // region_size`` would address the wrong tokens. """ proxy_dim: int topk: int region_size: int = REGION_BLOCK_SIZE page_size: int = 16 @property def regions_per_page(self) -> int: if self.page_size % self.region_size: raise ValueError( f"page size {self.page_size} must be a multiple of region size {self.region_size}" ) return self.page_size // self.region_size class DSALayerCache: """Per-layer DSA state for one batch of in-flight sequences. Allocated once for a whole generation, so ``max_tokens`` is a capacity rather than a length; ``lengths`` tracks how much of it is live. Every tensor is addressed by the conventions the module docstring lists, and mixing them up is the failure this class exists to prevent. ``num_kv_groups`` is the *local* count: it is one for the production TP=8 layout, with the corresponding global KV head selected while pre-sharding. """ def __init__( self, *, batch: int, max_tokens: int, num_kv_groups: int, head_dim: int, geometry: DSAGeometry, device: torch.device, dtype: torch.dtype, ) -> None: pages = -(-max_tokens // geometry.page_size) self.geometry = geometry self.batch = batch # Rounded up to a whole page so a request's token span never straddles the boundary # between two requests' page ranges. self.max_tokens = pages * geometry.page_size self.max_regions = self.max_tokens // geometry.region_size self.key = torch.zeros( (batch * self.max_tokens, num_kv_groups, head_dim), device=device, dtype=dtype ) self.value = torch.zeros_like(self.key) self.summary = torch.zeros( (batch * self.max_regions, 1, geometry.proxy_dim), device=device, dtype=torch.float8_e4m3fn, ) self.pending_key = torch.zeros( (batch, geometry.region_size, 1, geometry.proxy_dim), device=device, dtype=torch.bfloat16 ) self.pending_z = torch.zeros_like(self.pending_key) self.lengths = torch.zeros((batch,), device=device, dtype=torch.int32) # Request b owns pages [b * pages, (b + 1) * pages), so the mapping is an offset # identity. Materialising it as a real table rather than special-casing the # contiguous layout keeps the one code path the metadata gates actually cover. self.block_table = ( torch.arange(pages, device=device, dtype=torch.int32).view(1, -1) + torch.arange(batch, device=device, dtype=torch.int32).view(-1, 1) * pages ).contiguous() def token_slice(self, request: int, start: int, count: int) -> slice: base = request * self.max_tokens + start return slice(base, base + count) def write_kv(self, request: int, start: int, key: torch.Tensor, value: torch.Tensor) -> None: """Store ``key``/``value`` for one request's tokens at logical position ``start``.""" span = self.token_slice(request, start, key.shape[0]) self.key[span] = key self.value[span] = value class Step5SparseIndexer: """The indexer's projections and scoring, borrowing its parameters from the attention. Not an ``nn.Module``: the checkpoint stores the indexer weights directly under ``self_attn`` (``self_attn.sparse_indexer_q.weight``), so wrapping them in a submodule would rename every key and require a conversion step to load. """ def __init__(self, owner, *, geometry: DSAGeometry, num_kv_groups: int) -> None: self.owner = owner self.geometry = geometry # Provider groups and attention KV groups are the same partition. Under TP they # are both local counts; for TP=8, one provider group (four indexer heads) drives # one local KV group (eight attention Q heads). self.num_kv_groups = num_kv_groups self.num_heads = int(owner.sparse_indexer_w.weight.shape[0]) q_out = int(owner.sparse_indexer_q.weight.shape[0]) if q_out != self.num_heads * geometry.proxy_dim: raise ValueError( "local sparse_indexer_q/w shapes disagree: " f"q rows={q_out}, w heads={self.num_heads}, proxy_dim={geometry.proxy_dim}. " "This commonly means stale TP shards from the old replicated-provider layout." ) if self.num_heads % num_kv_groups: raise ValueError( f"{self.num_heads} indexer heads do not divide into {num_kv_groups} groups" ) self.heads_per_group = self.num_heads // num_kv_groups # The indexer's groups *are* the attention's KV groups, so this ratio is fixed under # every tensor-parallel degree and the prescale below is a TP invariant. self.weight_prescale = float(self.heads_per_group) ** -0.5 def project( self, hidden_states: torch.Tensor, positions: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Project and normalise the indexer's q/k/z, and scale its per-head weights. Returns: ``(index_q [tokens, groups, heads_per_group, proxy], index_k [tokens, 1, proxy], index_z [tokens, 1, proxy], weights [tokens, groups, heads_per_group] fp32)``. """ owner = self.owner proxy_dim = self.geometry.proxy_dim tokens = hidden_states.shape[0] index_q = owner.sparse_indexer_q(hidden_states) index_k = owner.sparse_indexer_k(hidden_states) index_z = owner.sparse_indexer_z(hidden_states) index_q, index_k, index_z = indexer_norm_rope( index_q.contiguous(), index_k.contiguous(), index_z.contiguous(), # Keep registered norm parameters in FP32 for checkpoint loading; # round them to the activation dtype before the fused operation. owner.sparse_indexer_q_norm.weight.to(dtype=index_q.dtype), owner.sparse_indexer_k_norm.weight.to(dtype=index_k.dtype), owner.sparse_indexer_k_norm.bias.to(dtype=index_k.dtype), owner.sparse_indexer_rope_cos, owner.sparse_indexer_rope_sin, positions, head_dim=proxy_dim, num_q_heads=self.num_heads, num_k_heads=1, rotary_dim=owner.sparse_indexer_rope_cos.shape[1], eps=owner.sparse_indexer_q_norm.variance_epsilon, q_norm_weight_bias=1.0, ) # Store the checkpoint weight in FP32, round it to the activation dtype # for the GEMM, then widen the result before applying the prescale. weights = torch.nn.functional.linear( hidden_states, owner.sparse_indexer_w.weight.to(dtype=hidden_states.dtype), ).float() weights = weights.view(tokens, self.num_kv_groups, self.heads_per_group) weights = weights * self.weight_prescale return ( index_q.view(tokens, self.num_kv_groups, self.heads_per_group, proxy_dim), index_k.view(tokens, 1, proxy_dim), index_z.view(tokens, 1, proxy_dim), weights.contiguous(), ) def update_summaries_prefill( cache: DSALayerCache, index_k: torch.Tensor, index_z: torch.Tensor, seq_lens: torch.Tensor, ) -> None: """Compress a packed prefill into per-region summaries, buffering the ragged tail. Only regions that are *complete* get a summary. The tail tokens are held instead, and the region they belong to is summarised later, once it fills -- see the module docstring for why a partial summary must never be stored. """ geometry = cache.geometry region_size = geometry.region_size starts, counts = csa_region_layout( seq_lens, regions_per_seq=cache.max_regions, region_size=region_size ) # A count below ``region_size`` marks the ragged tail; zero makes the kernel skip it. counts = torch.where(counts == region_size, counts, torch.zeros_like(counts)) _, summary_fp8 = csa_compress_regions( index_k, index_z, starts, counts, region_size=region_size ) live = counts > 0 cache.summary[live] = summary_fp8[live] token_starts = torch.cumsum(seq_lens.to(torch.int64), dim=0) - seq_lens.to(torch.int64) for request in range(cache.batch): length = int(seq_lens[request]) tail = length % region_size cache.pending_key[request, :tail] = index_k[ int(token_starts[request]) + length - tail : int(token_starts[request]) + length ].to(cache.pending_key.dtype) cache.pending_z[request, :tail] = index_z[ int(token_starts[request]) + length - tail : int(token_starts[request]) + length ].to(cache.pending_z.dtype) def update_summaries_decode( cache: DSALayerCache, index_k: torch.Tensor, index_z: torch.Tensor ) -> None: """Append one token per request and summarise any region that just filled. ``cache.lengths`` must already have been advanced past the new token, so the pending slot is derived from it rather than tracked separately -- one counter that can be wrong instead of two that can disagree. """ geometry = cache.geometry region_size = geometry.region_size lengths = cache.lengths.tolist() for request in range(cache.batch): length = int(lengths[request]) slot = (length - 1) % region_size cache.pending_key[request, slot] = index_k[request].to(cache.pending_key.dtype) cache.pending_z[request, slot] = index_z[request].to(cache.pending_z.dtype) if slot != region_size - 1: continue region = (length - 1) // region_size starts = torch.zeros((1,), device=index_k.device, dtype=torch.int32) counts = torch.full((1,), region_size, device=index_k.device, dtype=torch.int32) _, summary_fp8 = csa_compress_regions( cache.pending_key[request].contiguous(), cache.pending_z[request].contiguous(), starts, counts, region_size=region_size, ) cache.summary[request * cache.max_regions + region] = summary_fp8[0] def score_regions( indexer: Step5SparseIndexer, cache: DSALayerCache, index_q: torch.Tensor, weights: torch.Tensor, seq_lens: torch.Tensor, past_lens: torch.Tensor, ) -> torch.Tensor: """Indexer scores for every query against its own request's region summaries. Returns ``[groups * total_q, width]`` float32 in the group-major row order the selector requires, where ``width`` is the largest number of regions any row can see. Rows whose own history is shorter are left at zero past their bound; that is safe only because the selector's candidate range is bounded per row by its own history length and never reads further. ``width`` is computed rather than taken as ``cache.max_regions``, and the difference is not cosmetic. ``max_regions`` is the cache's *capacity*: at a 512k context it is 65536, so a 4k-token prefill would allocate ``4 groups x 4096 queries x 65536`` fp32 -- 4.3 GB, for a tensor whose live corner is a thousandth of that. Sizing to the live width makes the allocation track the context rather than the reservation. The loop is over requests, not a batched GEMM, because each request scores against a different key set. At decode that is one tiny matmul per request; at prefill it is dwarfed by the attention it feeds. """ groups = indexer.num_kv_groups total_q = int(index_q.shape[0]) region_size = cache.geometry.region_size # A row can see every region completed before its query, so the widest row is the one with # the largest final position. last_positions = [ int(past_lens[request]) + int(seq_lens[request]) - 1 for request in range(cache.batch) if int(seq_lens[request]) > 0 ] width = max((position // region_size for position in last_positions), default=0) logits = torch.zeros( (groups * total_q, max(width, 1)), device=index_q.device, dtype=torch.float32 ) offset = 0 for request in range(cache.batch): length = int(seq_lens[request]) if length == 0: continue regions = (int(past_lens[request]) + length - 1) // region_size if regions == 0: offset += length continue base = request * cache.max_regions block = indexer_logits( index_q[offset : offset + length].contiguous(), weights[offset : offset + length].contiguous(), cache.summary[base : base + regions].to(index_q.dtype), ) for group in range(groups): rows = slice(group * total_q + offset, group * total_q + offset + length) logits[rows, :regions] = block[group * length : (group + 1) * length] offset += length return logits def prefill_metadata( indexer: Step5SparseIndexer, cache: DSALayerCache, index_q: torch.Tensor, weights: torch.Tensor, seq_lens: torch.Tensor, past_lens: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Score, select, and pack prefill metadata for a packed batch.""" groups = indexer.num_kv_groups device = index_q.device logits = score_regions(indexer, cache, index_q, weights, seq_lens, past_lens) positions, requests = _packed_positions(seq_lens, past_lens, device) return prefill_sparse_meta( logits, positions.repeat(groups).contiguous(), cache.block_table, requests.repeat(groups).contiguous(), topk=indexer.geometry.topk, region_size=indexer.geometry.region_size, regions_per_page=indexer.geometry.regions_per_page, ) def decode_metadata( indexer: Step5SparseIndexer, cache: DSALayerCache, index_q: torch.Tensor, weights: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Score, select, and pack decode metadata for one token per request.""" ones = torch.ones((cache.batch,), device=index_q.device, dtype=torch.int32) logits = score_regions( indexer, cache, index_q, weights, ones, (cache.lengths - 1).clamp_min(0) ) return decode_sparse_meta( logits, cache.lengths, cache.block_table, topk=indexer.geometry.topk, region_size=indexer.geometry.region_size, regions_per_page=indexer.geometry.regions_per_page, ) def _packed_positions( seq_lens: torch.Tensor, past_lens: torch.Tensor, device: torch.device ) -> tuple[torch.Tensor, torch.Tensor]: """Absolute position and owning request of every token in a packed batch.""" positions = [] requests = [] for request in range(seq_lens.shape[0]): length = int(seq_lens[request]) past = int(past_lens[request]) positions.append(torch.arange(past, past + length, device=device, dtype=torch.int32)) requests.append(torch.full((length,), request, device=device, dtype=torch.int32)) return torch.cat(positions), torch.cat(requests) # FP8 expert GEMM. Activations use 128-wide K blocks; weights use 128x128 # blocks. Both scales multiply the FP8 values when reconstructing real values. # A missing weight scale selects the BF16 path. # fp8_e4m3fn has 4 exponent bits and 3 mantissa bits, with range [-448, 448]. # The activation quantizer maps each block's range into that envelope by # dividing by ``max(amax, 1e-10) / 448``. The 1e-10 floor preserves # low-amplitude blocks in Step-5's FP8 activation path. FP8_E4M3_MAX: float = 448.0 FP8_AMAX_FLOOR: float = 1e-10 # K-block size shared by the activation quantizer and the weight scale grid. # The kernel asserts BLOCK_K == GROUP_K so one K-iteration consumes exactly # one act scale entry and one weight scale entry per row. BLOCK_K: int = 128 # Weight N-block. The kernel asserts BLOCK_N == GROUP_N so one output tile # sits inside exactly one weight N-block, keeping the weight scale a scalar # per (tile, K-iter) instead of a vector. BLOCK_N: int = 128 @triton.jit def _act_quant_kernel( x_ptr, y_ptr, scale_ptr, M, K, stride_xm, stride_xk, stride_ym, stride_yk, stride_sm, stride_sk, BLOCK_M: tl.constexpr, BLOCK_K: tl.constexpr, ) -> None: """Per-row 1-D block quantization to fp8_e4m3fn along K. One program handles a ``[BLOCK_M, BLOCK_K]`` tile -- exactly one quantization group wide. The amax is computed over the *live* elements of each row, so the kernel handles a final ragged K-block where the tile extends past K: the scale covers the live elements only, and the OOB quantized slots are not written (the GEMM's masked load will read them as zero, so the round trip stays numerically faithful). """ pid_m = tl.program_id(0) pid_k = tl.program_id(1) rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) cols = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) row_mask = rows < M col_mask = cols < K tile_mask = row_mask[:, None] & col_mask[None, :] x = tl.load( x_ptr + rows[:, None] * stride_xm + cols[None, :] * stride_xk, mask=tile_mask, other=0.0, ).to(tl.float32) # OOB elements are zero via ``other=0.0``, so they do not inflate the amax. amax = tl.max(tl.abs(x), axis=1) amax = tl.maximum(amax, 1e-10) scale = amax * (1.0 / 448.0) inv_scale = 1.0 / scale # Clamp to [-448, 448] before the round-to-fp8: a value just above 448 would # otherwise wrap to NaN instead of saturating. q = tl.clamp(x * inv_scale[:, None], -448.0, 448.0) tl.store( y_ptr + rows[:, None] * stride_ym + cols[None, :] * stride_yk, q.to(y_ptr.dtype.element_ty), mask=tile_mask, ) tl.store( scale_ptr + rows * stride_sm + pid_k * stride_sk, scale, mask=row_mask, ) def act_quant( x: torch.Tensor, block_size: int = 128 ) -> tuple[torch.Tensor, torch.Tensor]: """Block-wise FP8 quantization along the last dim of ``x``. Args: x: input tensor of shape ``[..., K]``. The last dim is the reduction dim K. Any leading dims are flattened; the result is reshaped back. block_size: K-block size, 128 by default. Must be a power of two so Triton accepts it as a constexpr bound on ``tl.arange``. Returns: ``(y, scale)`` where ``y`` has ``x``'s shape and dtype ``torch.float8_e4m3fn``, and ``scale`` has shape ``[..., ceil(K/block_size)]`` and dtype ``torch.float32``. The recovery formula is ``real_x = y.float() * scale.unsqueeze(-1)``. K need not be divisible by block_size. For a partial final block, amax includes only live elements. """ if block_size & (block_size - 1): raise ValueError(f"block_size must be a power of two, got {block_size}") leading = x.shape[:-1] K = x.shape[-1] x2 = x.contiguous() flat = x2.view(-1, K) M = flat.shape[0] y = torch.empty_like(flat, dtype=torch.float8_e4m3fn) scale = torch.empty( M, (K + block_size - 1) // block_size, device=x.device, dtype=torch.float32 ) BLOCK_M = 16 if M < 32 else 32 grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(K, block_size)) _act_quant_kernel[grid]( flat, y, scale, M, K, flat.stride(0), flat.stride(1), y.stride(0), y.stride(1), scale.stride(0), scale.stride(1), BLOCK_M=BLOCK_M, BLOCK_K=block_size, ) return y.view(*leading, K), scale.view(*leading, -1) @triton.jit def _fp8_gemm_kernel( A_ptr, B_ptr, C_ptr, A_scale_ptr, W_scale_ptr, M, N, K, stride_am, stride_ak, stride_bn, stride_bk, stride_cm, stride_cn, stride_asm, stride_ask, stride_wsn, stride_wsk, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, GROUP_N: tl.constexpr, ) -> None: """fp8 GEMM with per-128 block scaling on both operands. One program computes a ``[BLOCK_M, BLOCK_N]`` output tile. ``BLOCK_N`` must equal ``GROUP_N`` (128) so the tile sits inside exactly one weight N-block, and ``BLOCK_K`` must equal the activation block size (128) so one K-iter consumes exactly one act scale entry and one weight scale entry per row. Loads B transposed (as ``[BLOCK_K, BLOCK_N]``) so ``tl.dot`` sees ``A @ B`` in the natural orientation. The strides are the *original* N/K strides of B; the transpose is by index, not by stride. The two scales are folded into the fp32 accumulator after the fp8 dot, not into the fp8 operands. Folding into the operands would re-quantize and lose the precision fp8 GEMM is supposed to keep; the post-dot multiply is the whole reason the accumulator is fp32. """ pid_m = tl.program_id(0) pid_n = tl.program_id(1) offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) offs_k = tl.arange(0, BLOCK_K) m_mask = offs_m < M n_mask = offs_n < N a_ptrs = A_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak b_ptrs = B_ptr + offs_n[None, :] * stride_bn + offs_k[:, None] * stride_bk acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) # Weight scale N index: tile (pid_n) sits in weight N-block ``pid_n`` when # ``BLOCK_N == GROUP_N``. Tiles never span weight N-block boundaries. n_block = pid_n * BLOCK_N // GROUP_N k_iters = (K + BLOCK_K - 1) // BLOCK_K for k_idx in tl.range(0, k_iters): k_offset = k_idx * BLOCK_K k_mask = (k_offset + offs_k) < K a = tl.load( a_ptrs + k_offset * stride_ak, mask=m_mask[:, None] & k_mask[None, :], other=0.0, ) b = tl.load( b_ptrs + k_offset * stride_bk, mask=n_mask[None, :] & k_mask[:, None], other=0.0, ) # fp8 x fp8 -> fp32 on Hopper; this is the WGMMA path that makes the # scheme worth the precision cost. The accumulator is fp32 so the post- # dot scale multiply keeps full precision. partial = tl.dot(a, b, out_dtype=tl.float32) a_s = tl.load( A_scale_ptr + offs_m * stride_asm + k_idx * stride_ask, mask=m_mask, other=0.0, ) # Scalar weight scale for this (n_block, k_idx) -- one fp32 per tile. w_s = tl.load(W_scale_ptr + n_block * stride_wsn + k_idx * stride_wsk) acc += partial * a_s[:, None] * w_s c_ptrs = C_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn tl.store( c_ptrs, acc.to(C_ptr.dtype.element_ty), mask=m_mask[:, None] & n_mask[None, :], ) def fp8_gemm( x_fp8: torch.Tensor, x_scale: torch.Tensor, w_fp8: torch.Tensor, w_scale_inv: torch.Tensor, ) -> torch.Tensor: """Block-scaled fp8 matmul ``y = x @ w.T``. Args: x_fp8: ``[M, K]`` ``torch.float8_e4m3fn`` activation, already quantized by :func:`act_quant`. x_scale: ``[M, ceil(K/128)]`` ``torch.float32`` activation scales. w_fp8: ``[N, K]`` ``torch.float8_e4m3fn`` weight (one expert's slice). w_scale_inv: ``[ceil(N/128), ceil(K/128)]`` ``torch.float32`` weight scales stored in the *inverse* convention (``real_weight = w_fp8 * w_scale_inv[block]``). Returns: ``y`` of shape ``[M, N]`` in ``torch.bfloat16``. """ if x_fp8.dtype != torch.float8_e4m3fn or w_fp8.dtype != torch.float8_e4m3fn: raise TypeError("operands must be float8_e4m3fn") if x_scale.dtype != torch.float32 or w_scale_inv.dtype != torch.float32: raise TypeError("scales must be float32") if x_fp8.dim() != 2 or w_fp8.dim() != 2: raise ValueError("operands must be 2-D; group by expert at the call site") M, K = x_fp8.shape N, K_w = w_fp8.shape if K != K_w: raise ValueError(f"K mismatch: x has {K}, weight has {K_w}") k_blocks = (K + BLOCK_K - 1) // BLOCK_K if x_scale.shape != (M, k_blocks): raise ValueError( f"x_scale shape {tuple(x_scale.shape)} != ({M}, {k_blocks})" ) if w_scale_inv.shape != ((N + BLOCK_N - 1) // BLOCK_N, k_blocks): raise ValueError( f"w_scale_inv shape {tuple(w_scale_inv.shape)} != " f"({(N + BLOCK_N - 1) // BLOCK_N}, {k_blocks})" ) x_fp8 = x_fp8.contiguous() w_fp8 = w_fp8.contiguous() x_scale = x_scale.contiguous() w_scale_inv = w_scale_inv.contiguous() y = torch.empty((M, N), device=x_fp8.device, dtype=torch.bfloat16) BLOCK_M = 64 if M >= 64 else 32 grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N)) _fp8_gemm_kernel[grid]( x_fp8, w_fp8, y, x_scale, w_scale_inv, M, N, K, x_fp8.stride(0), x_fp8.stride(1), w_fp8.stride(0), w_fp8.stride(1), y.stride(0), y.stride(1), x_scale.stride(0), x_scale.stride(1), w_scale_inv.stride(0), w_scale_inv.stride(1), BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, GROUP_N=BLOCK_N, ) return y def linear_fp8_or_bf16( x: torch.Tensor, weight: torch.Tensor, weight_scale_inv: torch.Tensor | None, ) -> torch.Tensor: """``y = x @ weight.T`` with fp8 path when a scale is present, BF16 fallback. The dispatcher is the only place that decides fp8 vs BF16: layers 88/89/90 ship BF16 experts with no scale, and the criterion is *presence of the scale tensor*, never the layer index. A layer can graduate off the not-convert list without code changes here. Args: x: ``[..., K]`` activation. Cast to BF16 internally on the fallback path (the fp8 path quantizes through BF16 anyway). weight: ``[N, K]`` weight; ``torch.float8_e4m3fn`` when ``weight_scale_inv`` is present, otherwise BF16. weight_scale_inv: ``[ceil(N/128), ceil(K/128)]`` fp32 inverse scales, or ``None`` to select the BF16 fallback path. Returns: ``[..., N]`` BF16 activation. """ if weight_scale_inv is None: if weight.dtype != torch.bfloat16: weight = weight.to(torch.bfloat16) x_bf = x.to(torch.bfloat16) if x.dtype != torch.bfloat16 else x return torch.nn.functional.linear(x_bf, weight) leading = x.shape[:-1] K = x.shape[-1] x_bf = x.to(torch.bfloat16) if x.dtype != torch.bfloat16 else x x_flat = x_bf.reshape(-1, K).contiguous() x_fp8, x_s = act_quant(x_flat) y = fp8_gemm(x_fp8, x_s, weight, weight_scale_inv) return y.view(*leading, -1) # MoE primitives. SwiGLU rounds after FP32 clamps and multiplication; # weighted gather accumulates expert contributions in top-k slot order. _BLOCK_SIZE = 1024 @triton.jit def _clamped_swiglu_kernel( gate_ptr, up_ptr, output_ptr, num_elements, LIMIT: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): """FP32 SWIGLUstep over two independent, contiguous input tensors.""" block = tl.program_id(0).to(tl.int64) offsets = block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = offsets < num_elements gate = tl.load(gate_ptr + offsets, mask=mask, other=0.0).to(tl.float32) up = tl.load(up_ptr + offsets, mask=mask, other=0.0).to(tl.float32) gate_silu = gate * tl.sigmoid(gate) gate_clamped = tl.minimum(gate_silu, LIMIT) up_clamped = tl.minimum(tl.maximum(up, -LIMIT), LIMIT) output = gate_clamped * up_clamped tl.store( output_ptr + offsets, output.to(output_ptr.dtype.element_ty), mask=mask, ) @triton.jit def _weighted_topk_gather_kernel( contributions_ptr, weights_ptr, output_ptr, hidden_size, TOP_K: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): """Sum one ``[top_k, hidden]`` slice in logical slot order.""" token = tl.program_id(0).to(tl.int64) hidden_block = tl.program_id(1).to(tl.int64) hidden_offsets = hidden_block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) hidden_mask = hidden_offsets < hidden_size accumulator = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) # TOP_K is constexpr, just as ``topk_num`` is in Optimus ep_gather. # Keeping the update in this loop produces the required slot-ordered # FP32 accumulation instead of a tree reduction over the K dimension. for slot in range(0, TOP_K): contribution_offset = (token * TOP_K + slot) * hidden_size + hidden_offsets contribution = tl.load( contributions_ptr + contribution_offset, mask=hidden_mask, other=0.0, ).to(tl.float32) weight = tl.load(weights_ptr + token * TOP_K + slot).to(tl.float32) accumulator += contribution * weight output_offset = token * hidden_size + hidden_offsets tl.store( output_ptr + output_offset, accumulator.to(output_ptr.dtype.element_ty), mask=hidden_mask, ) def _require_cuda(tensors: tuple[torch.Tensor, ...]) -> None: if any(tensor.device.type != "cuda" for tensor in tensors): devices = ", ".join(str(tensor.device) for tensor in tensors) raise RuntimeError( f"Step-5 MoE Triton kernels require CUDA tensors, got [{devices}]" ) def clamped_swiglu( gate: torch.Tensor, up: torch.Tensor, limit: float = 7.0, ) -> torch.Tensor: """Apply Step-5's clamped SwiGLU to independent BF16 projections. Computes ``silu(gate.float()).clamp(max=limit) * up.float().clamp(-limit, limit)`` in FP32 and casts the product to the input dtype once. ``gate`` and ``up`` must have the same shape and be contiguous BF16 CUDA tensors. """ if gate.shape != up.shape: raise ValueError( f"gate/up shape mismatch: {tuple(gate.shape)} != {tuple(up.shape)}" ) if gate.dtype != torch.bfloat16 or up.dtype != torch.bfloat16: raise TypeError( f"gate and up must both be torch.bfloat16, got {gate.dtype} and {up.dtype}" ) if gate.device != up.device: raise ValueError(f"gate/up device mismatch: {gate.device} != {up.device}") if not gate.is_contiguous() or not up.is_contiguous(): raise ValueError("gate and up must both be contiguous") try: limit_value = float(limit) except (TypeError, ValueError) as exc: raise TypeError(f"limit must be a real scalar, got {limit!r}") from exc if not math.isfinite(limit_value) or limit_value < 0.0: raise ValueError(f"limit must be finite and non-negative, got {limit_value}") _require_cuda((gate, up)) output = torch.empty_like(gate) if gate.numel() == 0: return output grid = (triton.cdiv(gate.numel(), _BLOCK_SIZE),) _clamped_swiglu_kernel[grid]( gate, up, output, gate.numel(), LIMIT=limit_value, BLOCK_SIZE=_BLOCK_SIZE, num_warps=4, ) return output def weighted_topk_gather( contributions: torch.Tensor, weights: torch.Tensor, ) -> torch.Tensor: """Weight and sum top-k contributions in deterministic slot order. Args: contributions: contiguous BF16 tensor of shape ``[T, K, H]``. weights: contiguous FP32 tensor of shape ``[T, K]``. Returns: Contiguous BF16 tensor of shape ``[T, H]``. For every token and hidden element the kernel starts from FP32 zero and performs ``acc += contributions[:, slot].float() * weights[:, slot]`` for slots ``0, 1, ..., K - 1``, then casts the accumulator to BF16 once. """ if contributions.ndim != 3: raise ValueError( f"contributions must have shape [T, K, H], got {tuple(contributions.shape)}" ) if weights.ndim != 2: raise ValueError(f"weights must have shape [T, K], got {tuple(weights.shape)}") tokens, top_k, hidden_size = contributions.shape if weights.shape != (tokens, top_k): raise ValueError(f"weights shape {tuple(weights.shape)} != ({tokens}, {top_k})") if top_k <= 0: raise ValueError(f"top-k dimension must be positive, got {top_k}") if hidden_size <= 0: raise ValueError(f"hidden dimension must be positive, got {hidden_size}") if contributions.dtype != torch.bfloat16: raise TypeError( f"contributions must be torch.bfloat16, got {contributions.dtype}" ) if weights.dtype != torch.float32: raise TypeError(f"weights must be torch.float32, got {weights.dtype}") if contributions.device != weights.device: raise ValueError( "contributions/weights device mismatch: " f"{contributions.device} != {weights.device}" ) if not contributions.is_contiguous() or not weights.is_contiguous(): raise ValueError("contributions and weights must both be contiguous") _require_cuda((contributions, weights)) output = torch.empty( (tokens, hidden_size), dtype=torch.bfloat16, device=contributions.device, ) if tokens == 0: return output grid = (tokens, triton.cdiv(hidden_size, _BLOCK_SIZE)) _weighted_topk_gather_kernel[grid]( contributions, weights, output, hidden_size, TOP_K=top_k, BLOCK_SIZE=_BLOCK_SIZE, num_warps=4, ) return output # Fused Q/K normalization and NeoX RoPE. The BF16 materialization between # normalization and rotary arithmetic, and the reduction order, are deliberate. HEAD_DIM = 192 _SUPPORTED_ROTARY_PAIRS = frozenset((32, 96)) @triton.jit def _lane_square_sum_192( qkv_ptr, head_base, LANE: tl.constexpr, ): """Step-5's 32-pair layout: three four-value vectors for one of 16 lanes.""" acc = 0.0 # The order is part of the numerical contract: # block 0 dims [0, 64), block 1 [64, 128), block 2 [128, 192), # with four adjacent values owned by each logical lane in every block. for block in range(3): for element in range(4): dim = block * 64 + LANE * 4 + element value = tl.load(qkv_ptr + head_base + dim).to(tl.float32) acc = _libdevice.fma_rn(value, value, acc) return acc @triton.jit def _wide_lane_square_sum_192( qkv_ptr, head_base, LANE: tl.constexpr, ): """32-lane normalization layout for the 96-pair Step-5 path. Each lane covers eight adjacent values, padding the 192-value row to 256 values. Lanes 24..31 contribute eight explicit zeros apiece. """ acc = 0.0 for element in range(8): dim = LANE * 8 + element value = tl.load( qkv_ptr + head_base + dim, mask=dim < 192, other=0.0, ).to(tl.float32) acc = _libdevice.fma_rn(value, value, acc) return acc @triton.jit def _qknorm_rope_192_kernel( qkv_ptr, q_out_ptr, k_out_ptr, q_weight_ptr, k_weight_ptr, cos_ptr, sin_ptr, positions_ptr, stride_qkv_token, stride_q_out_token, stride_k_out_token, stride_cos_token, stride_sin_token, NUM_Q_HEADS: tl.constexpr, NUM_KV_HEADS: tl.constexpr, ROTARY_PAIRS: tl.constexpr, EPS: tl.constexpr, NORM_WEIGHT_BIAS: tl.constexpr, BLOCK_D: tl.constexpr, BLOCK_R: tl.constexpr, ): """One program per token and Q/K head. The physical Triton layout is deliberately not used to define the reduction. Explicit logical lane sums followed by the exact ascending XOR tree make the result independent of Triton's evolving warp/layout lowering. Step-5's 32-pair path uses 16 lanes; the 96-pair path uses 32 so all rotary pairs are covered. """ token = tl.program_id(0) packed_head = tl.program_id(1) is_q = packed_head < NUM_Q_HEADS head_base = token * stride_qkv_token + packed_head * 192 if ROTARY_PAIRS == 96: # Use 32 lanes for the 96-pair path. Normalization covers eight # contiguous values per lane, padded to 256. Preserve the local # FMA chains and XOR(1,2,4,8,16) reduction order. p0 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=0) p1 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=1) p2 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=2) p3 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=3) p4 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=4) p5 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=5) p6 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=6) p7 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=7) p8 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=8) p9 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=9) p10 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=10) p11 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=11) p12 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=12) p13 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=13) p14 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=14) p15 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=15) p16 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=16) p17 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=17) p18 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=18) p19 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=19) p20 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=20) p21 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=21) p22 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=22) p23 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=23) p24 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=24) p25 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=25) p26 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=26) p27 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=27) p28 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=28) p29 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=29) p30 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=30) p31 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=31) s01 = _libdevice.add_rn(p0, p1) s23 = _libdevice.add_rn(p2, p3) s45 = _libdevice.add_rn(p4, p5) s67 = _libdevice.add_rn(p6, p7) s89 = _libdevice.add_rn(p8, p9) s1011 = _libdevice.add_rn(p10, p11) s1213 = _libdevice.add_rn(p12, p13) s1415 = _libdevice.add_rn(p14, p15) s1617 = _libdevice.add_rn(p16, p17) s1819 = _libdevice.add_rn(p18, p19) s2021 = _libdevice.add_rn(p20, p21) s2223 = _libdevice.add_rn(p22, p23) s2425 = _libdevice.add_rn(p24, p25) s2627 = _libdevice.add_rn(p26, p27) s2829 = _libdevice.add_rn(p28, p29) s3031 = _libdevice.add_rn(p30, p31) s03 = _libdevice.add_rn(s01, s23) s47 = _libdevice.add_rn(s45, s67) s811 = _libdevice.add_rn(s89, s1011) s1215 = _libdevice.add_rn(s1213, s1415) s1619 = _libdevice.add_rn(s1617, s1819) s2023 = _libdevice.add_rn(s2021, s2223) s2427 = _libdevice.add_rn(s2425, s2627) s2831 = _libdevice.add_rn(s2829, s3031) s07 = _libdevice.add_rn(s03, s47) s815 = _libdevice.add_rn(s811, s1215) s1623 = _libdevice.add_rn(s1619, s2023) s2431 = _libdevice.add_rn(s2427, s2831) s015 = _libdevice.add_rn(s07, s815) s1631 = _libdevice.add_rn(s1623, s2431) sum_squares = _libdevice.add_rn(s015, s1631) else: # 32-pair full-attention layers use 16 lanes. Each pN is the # sequential 12-value fma.rn chain of logical lane N. p0 = _lane_square_sum_192(qkv_ptr, head_base, LANE=0) p1 = _lane_square_sum_192(qkv_ptr, head_base, LANE=1) p2 = _lane_square_sum_192(qkv_ptr, head_base, LANE=2) p3 = _lane_square_sum_192(qkv_ptr, head_base, LANE=3) p4 = _lane_square_sum_192(qkv_ptr, head_base, LANE=4) p5 = _lane_square_sum_192(qkv_ptr, head_base, LANE=5) p6 = _lane_square_sum_192(qkv_ptr, head_base, LANE=6) p7 = _lane_square_sum_192(qkv_ptr, head_base, LANE=7) p8 = _lane_square_sum_192(qkv_ptr, head_base, LANE=8) p9 = _lane_square_sum_192(qkv_ptr, head_base, LANE=9) p10 = _lane_square_sum_192(qkv_ptr, head_base, LANE=10) p11 = _lane_square_sum_192(qkv_ptr, head_base, LANE=11) p12 = _lane_square_sum_192(qkv_ptr, head_base, LANE=12) p13 = _lane_square_sum_192(qkv_ptr, head_base, LANE=13) p14 = _lane_square_sum_192(qkv_ptr, head_base, LANE=14) p15 = _lane_square_sum_192(qkv_ptr, head_base, LANE=15) # Lane 0 after ascending shfl_xor offsets 1, 2, 4, 8. s01 = _libdevice.add_rn(p0, p1) s23 = _libdevice.add_rn(p2, p3) s45 = _libdevice.add_rn(p4, p5) s67 = _libdevice.add_rn(p6, p7) s89 = _libdevice.add_rn(p8, p9) s1011 = _libdevice.add_rn(p10, p11) s1213 = _libdevice.add_rn(p12, p13) s1415 = _libdevice.add_rn(p14, p15) s03 = _libdevice.add_rn(s01, s23) s47 = _libdevice.add_rn(s45, s67) s811 = _libdevice.add_rn(s89, s1011) s1215 = _libdevice.add_rn(s1213, s1415) s07 = _libdevice.add_rn(s03, s47) s815 = _libdevice.add_rn(s811, s1215) sum_squares = _libdevice.add_rn(s07, s815) mean_square = _libdevice.div_rn(sum_squares, 192.0) inverse_rms = _libdevice.rsqrt(_libdevice.add_rn(mean_square, EPS)) dims = tl.arange(0, BLOCK_D) dim_mask = dims < 192 values = tl.load( qkv_ptr + head_base + dims, mask=dim_mask, other=0.0, ).to(tl.float32) q_weight = tl.load( q_weight_ptr + dims, mask=dim_mask, other=0.0, ).to(tl.float32) k_weight = tl.load( k_weight_ptr + dims, mask=dim_mask, other=0.0, ).to(tl.float32) weight = tl.where(is_q, q_weight, k_weight) weight = _libdevice.add_rn(weight, NORM_WEIGHT_BIAS) x_hat = _libdevice.mul_rn(values, inverse_rms) normalized = _libdevice.mul_rn(x_hat, weight).to(tl.bfloat16) q_base = token * stride_q_out_token + packed_head * 192 k_base = token * stride_k_out_token + (packed_head - NUM_Q_HEADS) * 192 tl.store( q_out_ptr + q_base + dims, normalized, mask=dim_mask & is_q, ) tl.store( k_out_ptr + k_base + dims, normalized, mask=dim_mask & ~is_q, ) # Store and reload to round normalized values to the activation dtype # before RoPE. The barrier makes the stores visible to the reloads. tl.debug_barrier() pairs = tl.arange(0, BLOCK_R) pair_mask = pairs < ROTARY_PAIRS dim0 = pairs dim1 = pairs + ROTARY_PAIRS q_value0 = tl.load( q_out_ptr + q_base + dim0, mask=pair_mask & is_q, other=0.0, ).to(tl.float32) q_value1 = tl.load( q_out_ptr + q_base + dim1, mask=pair_mask & is_q, other=0.0, ).to(tl.float32) k_value0 = tl.load( k_out_ptr + k_base + dim0, mask=pair_mask & ~is_q, other=0.0, ).to(tl.float32) k_value1 = tl.load( k_out_ptr + k_base + dim1, mask=pair_mask & ~is_q, other=0.0, ).to(tl.float32) value0 = tl.where(is_q, q_value0, k_value0) value1 = tl.where(is_q, q_value1, k_value1) position = tl.load(positions_ptr + token) cos_value = tl.load( cos_ptr + position * stride_cos_token + pairs, mask=pair_mask, other=0.0, ).to(tl.float32) sin_value = tl.load( sin_ptr + position * stride_sin_token + pairs, mask=pair_mask, other=0.0, ).to(tl.float32) # Step-5 RoPE rounding contract: # rot0 = mul(value0, cos) - mul(value1, sin) # rot1 = fma(value0, sin, mul(value1, cos)) rotated0 = _libdevice.sub_rn( _libdevice.mul_rn(value0, cos_value), _libdevice.mul_rn(value1, sin_value), ).to(tl.bfloat16) rotated1 = _libdevice.fma_rn( value0, sin_value, _libdevice.mul_rn(value1, cos_value), ).to(tl.bfloat16) tl.store( q_out_ptr + q_base + dim0, rotated0, mask=pair_mask & is_q, ) tl.store( q_out_ptr + q_base + dim1, rotated1, mask=pair_mask & is_q, ) tl.store( k_out_ptr + k_base + dim0, rotated0, mask=pair_mask & ~is_q, ) tl.store( k_out_ptr + k_base + dim1, rotated1, mask=pair_mask & ~is_q, ) @triton.jit def _copy_v_192_kernel( qkv_ptr, v_out_ptr, stride_qkv_token, stride_v_out_token, NUM_Q_HEADS: tl.constexpr, NUM_KV_HEADS: tl.constexpr, BLOCK_D: tl.constexpr, ): token = tl.program_id(0) head = tl.program_id(1) dims = tl.arange(0, BLOCK_D) mask = dims < 192 input_base = ( token * stride_qkv_token + (NUM_Q_HEADS + NUM_KV_HEADS + head) * 192 ) output_base = token * stride_v_out_token + head * 192 # BF16 load/store only: V must retain its original bits. value = tl.load(qkv_ptr + input_base + dims, mask=mask) tl.store(v_out_ptr + output_base + dims, value, mask=mask) def _validate_qknorm_rope_inputs( qkv: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, positions: torch.Tensor, *, head_dim: int, num_q_heads: int, num_kv_heads: int, rotary_pairs: int, eps: float, norm_weight_bias: float, ) -> None: tensors = { "qkv": qkv, "q_weight": q_weight, "k_weight": k_weight, "cos": cos, "sin": sin, "positions": positions, } for name, tensor in tensors.items(): if not isinstance(tensor, torch.Tensor): raise TypeError(f"{name} must be a torch.Tensor") if head_dim != HEAD_DIM: raise ValueError(f"head_dim must be {HEAD_DIM}, got {head_dim}") if num_q_heads <= 0 or num_kv_heads <= 0: raise ValueError("num_q_heads and num_kv_heads must be positive") if rotary_pairs not in _SUPPORTED_ROTARY_PAIRS: raise ValueError( "rotary_pairs must be 32 (full-attention layers) or 96 " f"(sliding-attention layers), got {rotary_pairs}" ) if not math.isfinite(eps) or eps <= 0.0: raise ValueError(f"eps must be finite and positive, got {eps}") if not math.isfinite(norm_weight_bias): raise ValueError(f"norm_weight_bias must be finite, got {norm_weight_bias}") packed_width = (num_q_heads + 2 * num_kv_heads) * head_dim if qkv.ndim != 2 or tuple(qkv.shape) != (qkv.shape[0], packed_width): raise ValueError( f"qkv must have shape [tokens, {packed_width}], got {tuple(qkv.shape)}" ) if qkv.stride(1) != 1: raise ValueError("qkv must be contiguous in its last dimension") if qkv.dtype != torch.bfloat16: raise TypeError(f"qkv must have dtype torch.bfloat16, got {qkv.dtype}") for name, weight in (("q_weight", q_weight), ("k_weight", k_weight)): if weight.ndim != 1 or weight.numel() != head_dim: raise ValueError( f"{name} must have shape [{head_dim}], got {tuple(weight.shape)}" ) if not weight.is_contiguous(): raise ValueError(f"{name} must be contiguous") if weight.dtype != torch.float32: raise TypeError(f"{name} must have dtype torch.float32, got {weight.dtype}") if cos.ndim != 2 or sin.ndim != 2: raise ValueError( f"cos and sin must be 2D, got {tuple(cos.shape)} and {tuple(sin.shape)}" ) if cos.shape != sin.shape: raise ValueError( f"cos and sin must have equal shapes, got {tuple(cos.shape)} " f"and {tuple(sin.shape)}" ) if cos.shape[1] < rotary_pairs: raise ValueError( f"cos/sin width must be at least {rotary_pairs}, got {cos.shape[1]}" ) if cos.stride(1) != 1 or sin.stride(1) != 1: raise ValueError("cos and sin must be contiguous in their last dimension") if cos.dtype != torch.bfloat16 or sin.dtype != torch.bfloat16: raise TypeError( "cos and sin must have dtype torch.bfloat16, got " f"{cos.dtype} and {sin.dtype}" ) if positions.ndim != 1 or positions.numel() != qkv.shape[0]: raise ValueError( "positions must have one entry per qkv token, got " f"shape {tuple(positions.shape)} for {qkv.shape[0]} tokens" ) if not positions.is_contiguous(): raise ValueError("positions must be contiguous") if positions.dtype not in (torch.int32, torch.int64): raise TypeError( "positions must have dtype torch.int32 or torch.int64, got " f"{positions.dtype}" ) # Device checks intentionally come last. Shape and dtype are properties of # the public contract too, and validating them first lets CPU-only tooling # diagnose malformed model inputs without needing a CUDA allocation. for name, tensor in tensors.items(): if not tensor.is_cuda: raise ValueError(f"{name} must be a CUDA tensor") devices = {tensor.device for tensor in tensors.values()} if len(devices) != 1: raise ValueError("all inputs must be on the same CUDA device") def fused_qknorm_rope( qkv: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, positions: torch.Tensor, *, head_dim: int, num_q_heads: int, num_kv_heads: int, rotary_pairs: int, eps: float = 1e-5, norm_weight_bias: float = 1.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Normalize and rotate packed TP-local Q/K heads, and copy V exactly. Args: qkv: BF16 ``[tokens, (num_q_heads + 2 * num_kv_heads) * 192]``. A wider physical row stride is allowed, but the last dimension must be contiguous. q_weight, k_weight: FP32 zero-centered RMSNorm weights of shape ``[192]``. The effective scale is ``weight + norm_weight_bias``; Step-5 uses the default bias of ``1.0``. cos, sin: BF16 NeoX tables ``[max_position, >= rotary_pairs]``. positions: contiguous int32 or int64 position for every token. rotary_pairs: ``96`` on sliding layers and ``32`` on full-attention layers. The leading ``2 * rotary_pairs`` dimensions are rotated. Returns: Flat contiguous BF16 tensors ``q``, ``k``, ``v`` with shapes ``[tokens, num_q_heads * 192]`` and ``[tokens, num_kv_heads * 192]`` (for both K and V). The 96-pair sliding path uses 32 lanes and a padded 256-value normalization tile, rotating all 96 pairs. The 32-pair full-attention path uses a 16-lane blocked reduction layout. """ _validate_qknorm_rope_inputs( qkv, q_weight, k_weight, cos, sin, positions, head_dim=head_dim, num_q_heads=num_q_heads, num_kv_heads=num_kv_heads, rotary_pairs=rotary_pairs, eps=eps, norm_weight_bias=norm_weight_bias, ) tokens = qkv.shape[0] q_out = torch.empty( (tokens, num_q_heads * HEAD_DIM), device=qkv.device, dtype=qkv.dtype, ) k_out = torch.empty( (tokens, num_kv_heads * HEAD_DIM), device=qkv.device, dtype=qkv.dtype, ) v_out = torch.empty_like(k_out) if tokens == 0: return q_out, k_out, v_out _qknorm_rope_192_kernel[(tokens, num_q_heads + num_kv_heads)]( qkv, q_out, k_out, q_weight, k_weight, cos, sin, positions, qkv.stride(0), q_out.stride(0), k_out.stride(0), cos.stride(0), sin.stride(0), NUM_Q_HEADS=num_q_heads, NUM_KV_HEADS=num_kv_heads, ROTARY_PAIRS=rotary_pairs, EPS=float(eps), NORM_WEIGHT_BIAS=float(norm_weight_bias), BLOCK_D=256, BLOCK_R=128, num_warps=4, ) _copy_v_192_kernel[(tokens, num_kv_heads)]( qkv, v_out, qkv.stride(0), v_out.stride(0), NUM_Q_HEADS=num_q_heads, NUM_KV_HEADS=num_kv_heads, BLOCK_D=256, num_warps=4, ) return q_out, k_out, v_out __all__ = [ "REGION_BLOCK_SIZE", "PROXY_DIM", "shared_table_stride", "indexer_norm_rope", "csa_region_layout", "csa_compress_regions", "round_activations_e4m3", "indexer_logits", "region_topk_ids", "region_topk_pack", "prefill_sparse_meta", "decode_sparse_meta", "sparse_attention_prefill", "merge_split_states", "sparse_attention_decode", "build_rope_cache", "DSAGeometry", "DSALayerCache", "Step5SparseIndexer", "update_summaries_prefill", "update_summaries_decode", "score_regions", "prefill_metadata", "decode_metadata", "act_quant", "fp8_gemm", "linear_fp8_or_bf16", "clamped_swiglu", "weighted_topk_gather", "HEAD_DIM", "fused_qknorm_rope", ]