submission 744389
divc13 · python · License unknown
Use it
Vendorable · source mirrored · license unknownView source →
No package. Vendor the mirrored source: 245 lines, June 9 Researcher Reciprocity License v1.0.
submission_v128.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-amd-mixed-mla-744389?include=source"interfacepython
Compatibility
measured onAMD Instinct MI355X
declared hardwareAMD Instinct MI355X
architecturesgfx950
dtypesbf16, int32
Benchmark evidence
1 measurement across 1 GPU, fastest first.
Operation / workload
Hardware
Latency
Rank
Observed
Reported · How evidence levels are derived →
Source and license
sourceavailable
revision digestsha256:eb46f9548207b423087f2043c61d4cdf71b2688a0697c78ce96500d7f5533ed1
license declaredunknown
license concludedunknown
authorsdivc13
imported2026-08-15
Techniques
Extracted from the mirrored source by pattern, never inferred. Each row cites its line.
fp8
FP8_DTYPE = tl.float8e4nvmma
qk = tl.dot(q_nope, k_nope) + tl.dot(q_rope, k_rope)num-warps = 4
num_warps=4, num_stages=2, waves_per_eu=2,stages = 2
num_warps=4, num_stages=2, waves_per_eu=2,tile-n = 64
NUM_KV_SPLITS=num_splits, BLOCK_N=64, BLOCK_H=16,Kernel source
submission_v128.py245 lines
"""MLA decode Triton v128: uniform ns=2 + adjusted splits.
Hypothesis: ns=3 forces occ1 (73KB LDS > 64KB/WG at occ2).
With ns=2 (36KB LDS), all shapes run at occ2, doubling effective
memory bandwidth. Trade: slightly shallower pipeline, but more
concurrent wavefronts.
Also: adjusted split counts to increase parallelism for shapes that
were previously underutilized.
Based on v109 (exp2+rcp+.cg).
"""
import torch
import triton
import triton.language as tl
from task import input_t, output_t
NUM_HEADS = 16
KV_LORA_RANK = 512
QK_HEAD_DIM = 576
V_HEAD_DIM = 512
FP8_DTYPE = tl.float8e4nv
LOG2E = tl.constexpr(1.4426950408889634)
_SHAPE_CONFIGS = {
(4, 1024): 16,
(4, 8192): 32,
(32, 1024): 4,
(32, 8192): 8,
(64, 1024): 4,
(64, 8192): 8,
(256, 1024): 1,
(256, 8192): 2,
}
@triton.jit
def _remap_xcd(pid, GRID_SIZE, NUM_XCDS: tl.constexpr = 8):
pids_per_xcd = (GRID_SIZE + NUM_XCDS - 1) // NUM_XCDS
tall_xcds = GRID_SIZE % NUM_XCDS
tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds
xcd = pid % NUM_XCDS
local_pid = pid // NUM_XCDS
if xcd < tall_xcds:
return xcd * pids_per_xcd + local_pid
else:
return tall_xcds * pids_per_xcd + (xcd - tall_xcds) * (pids_per_xcd - 1) + local_pid
@triton.jit
def _fast_rcp(x):
return tl.inline_asm_elementwise(
"v_rcp_f32_e32 $0, $1", "=v, v", [x],
dtype=tl.float32, is_pure=True, pack=1,
)
@triton.jit
def _stage1_kernel(
Q, KV, Mid_O, Out,
qo_indptr, kv_indptr,
sm_scale, kv_scale_ptr,
stride_qt, stride_qh, stride_kvt,
stride_mb, stride_mh, stride_ms,
stride_ot, stride_oh,
NUM_KV_SPLITS: tl.constexpr, BLOCK_N: tl.constexpr,
BLOCK_H: tl.constexpr, BLOCK_C: tl.constexpr, BLOCK_R: tl.constexpr,
DIRECT_OUT: tl.constexpr, GRID_SIZE: tl.constexpr,
):
raw_pid = tl.program_id(0)
pid = _remap_xcd(raw_pid, GRID_SIZE)
cur_batch = pid // NUM_KV_SPLITS
split_id = pid % NUM_KV_SPLITS
heads = tl.arange(0, BLOCK_H)
kv_scale = tl.load(kv_scale_ptr).to(tl.float32)
score_scale = sm_scale * kv_scale
q_tok = tl.load(qo_indptr + cur_batch)
kv_s = tl.load(kv_indptr + cur_batch)
kv_e = tl.load(kv_indptr + cur_batch + 1)
kv_len = kv_e - kv_s
kv_per_split = tl.cdiv(kv_len, NUM_KV_SPLITS)
sp_start = kv_per_split * split_id
sp_end = tl.minimum(sp_start + kv_per_split, kv_len)
offs_c = tl.arange(0, BLOCK_C)
offs_r = tl.arange(0, BLOCK_R)
q_base = q_tok * stride_qt
q_nope = tl.load(Q + q_base + heads[:, None] * stride_qh + offs_c[None, :]).to(FP8_DTYPE)
q_rope = tl.load(Q + q_base + heads[:, None] * stride_qh + (512 + offs_r[None, :])).to(FP8_DTYPE)
e_max = tl.full([BLOCK_H], value=float("-inf"), dtype=tl.float32)
e_sum = tl.zeros([BLOCK_H], dtype=tl.float32)
acc = tl.zeros([BLOCK_H, BLOCK_C], dtype=tl.float32)
num_tokens = sp_end - sp_start
num_full = num_tokens // BLOCK_N
full_end = sp_start + num_full * BLOCK_N
for start_n in range(sp_start, full_end, BLOCK_N):
kv_idx = kv_s + start_n + tl.arange(0, BLOCK_N)
kv_nope = tl.load(KV + kv_idx[:, None] * stride_kvt + offs_c[None, :],
cache_modifier=".cg")
kv_rope = tl.load(KV + kv_idx[:, None] * stride_kvt + (512 + offs_r[None, :]),
cache_modifier=".cg")
k_nope = tl.trans(kv_nope)
k_rope = tl.trans(kv_rope)
qk = tl.dot(q_nope, k_nope) + tl.dot(q_rope, k_rope)
qk = qk * score_scale
new_max = tl.maximum(tl.max(qk, 1), e_max)
rescale = tl.math.exp2((e_max - new_max) * LOG2E)
p = tl.math.exp2((qk - new_max[:, None]) * LOG2E)
acc = acc * rescale[:, None]
acc = acc + tl.dot(p.to(FP8_DTYPE), kv_nope)
e_sum = e_sum * rescale + tl.sum(p, 1)
e_max = new_max
if full_end < sp_end:
offs_n = full_end + tl.arange(0, BLOCK_N)
mask_n = offs_n < sp_end
kv_idx = kv_s + offs_n
kv_nope = tl.load(KV + kv_idx[:, None] * stride_kvt + offs_c[None, :],
mask=mask_n[:, None], other=0.0, cache_modifier=".cg")
kv_rope = tl.load(KV + kv_idx[:, None] * stride_kvt + (512 + offs_r[None, :]),
mask=mask_n[:, None], other=0.0, cache_modifier=".cg")
k_nope = tl.trans(kv_nope)
k_rope = tl.trans(kv_rope)
qk = tl.dot(q_nope, k_nope) + tl.dot(q_rope, k_rope)
qk = qk * score_scale
qk = tl.where(mask_n[None, :], qk, float("-inf"))
new_max = tl.maximum(tl.max(qk, 1), e_max)
rescale = tl.math.exp2((e_max - new_max) * LOG2E)
p = tl.math.exp2((qk - new_max[:, None]) * LOG2E)
acc = acc * rescale[:, None]
acc = acc + tl.dot(p.to(FP8_DTYPE), kv_nope)
e_sum = e_sum * rescale + tl.sum(p, 1)
e_max = new_max
safe_esum = tl.where(e_sum > 0, e_sum, 1.0)
inv_esum = _fast_rcp(safe_esum)
result = acc * inv_esum[:, None]
if DIRECT_OUT:
o_base = q_tok * stride_ot
tl.store(Out + o_base + heads[:, None] * stride_oh + offs_c[None, :],
(result * kv_scale).to(tl.bfloat16))
else:
mid_base = cur_batch * stride_mb + heads * stride_mh + split_id * stride_ms
tl.store(Mid_O + mid_base[:, None] + offs_c[None, :], result)
lse = tl.where(e_sum > 0, e_max + tl.log(e_sum), float("-inf"))
tl.store(Mid_O + mid_base + 512, lse)
@triton.jit
def _stage2_kernel(
Mid_O, Out, qo_indptr, kv_scale_ptr,
stride_mb, stride_mh, stride_ms, stride_ot, stride_oh,
NUM_KV_SPLITS: tl.constexpr, BLOCK_DV: tl.constexpr,
batch: tl.constexpr, GRID_SIZE: tl.constexpr,
):
raw_pid = tl.program_id(0)
pid = _remap_xcd(raw_pid, GRID_SIZE)
cur_batch = pid % batch
cur_head = pid // batch
kv_scale = tl.load(kv_scale_ptr).to(tl.float32)
q_tok = tl.load(qo_indptr + cur_batch)
offs_d = tl.arange(0, BLOCK_DV)
e_max = float("-inf")
e_sum = 0.0
acc = tl.zeros([BLOCK_DV], dtype=tl.float32)
mid_base = cur_batch * stride_mb + cur_head * stride_mh
for s in range(NUM_KV_SPLITS):
tv = tl.load(Mid_O + mid_base + s * stride_ms + offs_d)
lse = tl.load(Mid_O + mid_base + s * stride_ms + 512)
new_max = tl.maximum(lse, e_max)
old_scale = tl.math.exp2((e_max - new_max) * LOG2E)
exp_lse = tl.math.exp2((lse - new_max) * LOG2E)
acc = acc * old_scale + exp_lse * tv
e_sum = e_sum * old_scale + exp_lse
e_max = new_max
safe_esum = tl.maximum(e_sum, 1e-12)
inv_esum = _fast_rcp(safe_esum)
result = acc * inv_esum * kv_scale
o_base = q_tok * stride_ot + cur_head * stride_oh
tl.store(Out + o_base + offs_d, result.to(tl.bfloat16))
def _get_splits(batch_size, kv_len):
key = (batch_size, kv_len)
if key in _SHAPE_CONFIGS:
return _SHAPE_CONFIGS[key]
max_tiles = max(1, kv_len // 64)
if batch_size >= 256:
return min(2, max_tiles)
elif batch_size >= 64:
return min(max(1, 512 // batch_size), max_tiles)
else:
splits = min(max(1, 768 // batch_size), max_tiles)
while splits > 1 and batch_size * splits > 912:
splits -= 1
return splits
def custom_kernel(data: input_t) -> output_t:
q, kv_data, qo_indptr, kv_indptr, config = data
kv_fp8, kv_scale_t = kv_data["fp8"]
kv = kv_fp8.view(torch.float8_e4m3fn).view(-1, QK_HEAD_DIM)
batch_size = config["batch_size"]
sm_scale = config["sm_scale"]
total_q = q.shape[0]
kv_len = kv.shape[0] // batch_size if batch_size > 0 else 0
num_splits = _get_splits(batch_size, kv_len)
dev = q.device
direct_out = (num_splits == 1)
mid_o = torch.empty(
(batch_size, NUM_HEADS, num_splits, KV_LORA_RANK + 1),
dtype=torch.float32, device=dev,
) if not direct_out else torch.empty(1, dtype=torch.float32, device=dev)
out = torch.empty((total_q, NUM_HEADS, V_HEAD_DIM), dtype=torch.bfloat16, device=dev)
grid_size1 = batch_size * num_splits
_stage1_kernel[(grid_size1,)](
q, kv, mid_o, out, qo_indptr, kv_indptr,
sm_scale, kv_scale_t,
q.stride(0), q.stride(1), kv.stride(0),
mid_o.stride(0) if not direct_out else 0,
mid_o.stride(1) if not direct_out else 0,
mid_o.stride(2) if not direct_out else 0,
out.stride(0), out.stride(1),
NUM_KV_SPLITS=num_splits, BLOCK_N=64, BLOCK_H=16,
BLOCK_C=512, BLOCK_R=64, DIRECT_OUT=direct_out,
GRID_SIZE=grid_size1,
num_warps=4, num_stages=2, waves_per_eu=2,
)
if not direct_out:
grid_size2 = NUM_HEADS * batch_size
_stage2_kernel[(grid_size2,)](
mid_o, out, qo_indptr, kv_scale_t,
mid_o.stride(0), mid_o.stride(1), mid_o.stride(2),
out.stride(0), out.stride(1),
NUM_KV_SPLITS=num_splits, BLOCK_DV=512,
batch=batch_size, GRID_SIZE=grid_size2,
num_warps=4, num_stages=2,
)
return out
scrolls · 245 lines total
Source code from GPU Mode and the KernelBot dataset · June 9 Researcher Reciprocity License v1.0
Changes from previous submission
Against this author's previous submission submission 740330.
- """MLA decode Triton v109: exp2 + fast rcp on v94 baseline.+ """MLA decode Triton v128: uniform ns=2 + adjusted splits.- 1. tl.exp → tl.math.exp2 with LOG2E prescale. Softmax inputs are ≤0,- so the range fixup (v_cmp+v_cndmask+v_ldexp) is wasted. exp2 may- compile to bare v_exp_f32. Saves ~200 VALU in hot loop.- 2. IEEE div → v_rcp_f32 via inline asm for acc/esum normalization.- Saves ~190 instructions in epilogue. FP8 output has ~12% quant- error anyway.+ Hypothesis: ns=3 forces occ1 (73KB LDS > 64KB/WG at occ2).+ With ns=2 (36KB LDS), all shapes run at occ2, doubling effective+ memory bandwidth. Trade: slightly shallower pipeline, but more+ concurrent wavefronts.+ Also: adjusted split counts to increase parallelism for shapes that+ were previously underutilized.+ Based on v109 (exp2+rcp+.cg)."""import torch⋯ 9 unchanged linesLOG2E = tl.constexpr(1.4426950408889634)_SHAPE_CONFIGS = {- (4, 1024): (16, 2),- (4, 8192): (32, 3),- (32, 1024): (4, 3),- (32, 8192): (8, 3),- (64, 1024): (4, 3),- (64, 8192): (8, 2),- (256, 1024): (1, 2),- (256, 8192): (2, 2),+ (4, 1024): 16,+ (4, 8192): 32,+ (32, 1024): 4,+ (32, 8192): 8,+ (64, 1024): 4,+ (64, 8192): 8,+ (256, 1024): 1,+ (256, 8192): 2,}⋯ 143 unchanged linestl.store(Out + o_base + offs_d, result.to(tl.bfloat16))- def _get_config(batch_size, kv_len):+ def _get_splits(batch_size, kv_len):key = (batch_size, kv_len)if key in _SHAPE_CONFIGS:return _SHAPE_CONFIGS[key]max_tiles = max(1, kv_len // 64)if batch_size >= 256:- splits = min(2, max_tiles)+ return min(2, max_tiles)elif batch_size >= 64:- splits = min(max(1, 512 // batch_size), max_tiles)+ return min(max(1, 512 // batch_size), max_tiles)else:splits = min(max(1, 768 // batch_size), max_tiles)- while splits > 1 and batch_size * splits > 912:- splits -= 1- grid_size = batch_size * splits- ns = 3 if grid_size < 512 else 2- return splits, ns+ while splits > 1 and batch_size * splits > 912:+ splits -= 1+ return splitsdef custom_kernel(data: input_t) -> output_t:⋯ 5 unchanged linestotal_q = q.shape[0]kv_len = kv.shape[0] // batch_size if batch_size > 0 else 0- num_splits, ns = _get_config(batch_size, kv_len)+ num_splits = _get_splits(batch_size, kv_len)dev = q.devicedirect_out = (num_splits == 1)⋯ 16 unchanged linesNUM_KV_SPLITS=num_splits, BLOCK_N=64, BLOCK_H=16,BLOCK_C=512, BLOCK_R=64, DIRECT_OUT=direct_out,GRID_SIZE=grid_size1,- num_warps=4, num_stages=ns, waves_per_eu=2,+ num_warps=4, num_stages=2, waves_per_eu=2,)if not direct_out:
scrolls · 89 diff lines total
Best evidence level for this revision: reported
JSON