| import sys |
| from types import MethodType |
| import torch |
| import triton |
| import triton.language as tl |
| from model import inference_kernels as kernels |
| from transformers.models.qwen3_5 import modeling_qwen3_5 as hf_model |
| from . import b1_bf16_projection as projection_patch |
|
|
| original_rope = hf_model.apply_rotary_pos_emb |
| original_query = {} |
|
|
|
|
| @triton.jit |
| def fast_norm(X, U, G, Y, R, WIDTH: tl.constexpr, EPS: tl.constexpr, |
| BLOCK: tl.constexpr, ADD: tl.constexpr, ROUND_ADD: tl.constexpr): |
| row = tl.program_id(0) |
| column = tl.arange(0, BLOCK) |
| x = tl.load(X + row * WIDTH + column, column < WIDTH, 0).to(tl.float32) |
| if ADD: |
| x += tl.load(U + row * WIDTH + column, column < WIDTH, 0).to(tl.float32) |
| if ROUND_ADD: |
| x = x.to(X.dtype.element_ty).to(tl.float32) |
| tl.store(R + row * WIDTH + column, x, column < WIDTH) |
| inv = tl.rsqrt(tl.sum(x * x, 0) / WIDTH + EPS) |
| gamma = tl.load(G + column, column < WIDTH, 0) |
| tl.store(Y + row * WIDTH + column, x * inv * gamma, column < WIDTH) |
|
|
|
|
| def norm(values, gamma, eps, update=None, round_update=True): |
| values = values.contiguous() |
| width = values.shape[-1] |
| rows = values.numel() // width |
| if rows == 0: |
| return values, values |
| output = torch.empty_like(values) |
| summed = torch.empty_like(values) if update is not None else values |
| fast_norm[rows,](values, values if update is None else update, gamma, output, summed, |
| width, float(eps), triton.next_power_of_2(width), update is not None, round_update, |
| num_warps=4, enable_fp_fusion=False) |
| return summed, output |
|
|
|
|
|
|
| @triton.jit |
| def rope_kernel(X, C, S, Y, TOTAL: tl.constexpr, D: tl.constexpr, |
| ROT: tl.constexpr, STRIDE: tl.constexpr, BLOCK: tl.constexpr): |
| index = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) |
| row, column = index // D, index % D |
| x = tl.load(X + row * STRIDE + column, index < TOTAL, 0).to(tl.float32) |
| paired = tl.where(column < ROT // 2, column + ROT // 2, column - ROT // 2) |
| other = tl.load(X + row * STRIDE + paired, (index < TOTAL) & (column < ROT), 0).to(tl.float32) |
| other = tl.where(column < ROT // 2, -other, other) |
| cos = tl.load(C + column, column < ROT, 0).to(tl.float32) |
| sin = tl.load(S + column, column < ROT, 0).to(tl.float32) |
| first = (x * cos).to(X.dtype.element_ty).to(tl.float32) |
| second = (other * sin).to(X.dtype.element_ty).to(tl.float32) |
| rotated = tl.where(column < ROT, first + second, x) |
| tl.store(Y + index, rotated, index < TOTAL) |
|
|
|
|
| def rotate(values, cos, sin): |
| d = values.shape[-1] |
| flat = values.reshape(-1, d) |
| out = torch.empty(values.shape, device=values.device, dtype=values.dtype) |
| rope_kernel[triton.cdiv(values.numel(), 256),](flat, cos, sin, out, |
| values.numel(), d, cos.shape[-1], flat.stride(0), 256, |
| num_warps=4, enable_fp_fusion=False) |
| return out |
|
|
|
|
| def rope(q, k, cos, sin, unsqueeze_dim=1): |
| if cos.numel() == cos.shape[-1] and cos.dtype == q.dtype and q.stride(-1) == 1 and k.stride(-1) == 1: |
| return rotate(q, cos, sin), rotate(k, cos, sin) |
| return original_rope(q, k, cos, sin, unsqueeze_dim) |
|
|
|
|
|
|
| @triton.jit |
| def norm_gate_query(X, Gamma, W, Bias, Gain, C, S, Q, Gate, |
| D: tl.constexpr, ROT: tl.constexpr, EPS: tl.constexpr): |
| head = tl.program_id(0) |
| column = tl.arange(0, D) |
| raw = tl.load(X + head * D + column).to(tl.float32) |
| gamma = tl.load(Gamma + column) |
| normalized = (raw * tl.rsqrt(tl.sum(raw * raw, 0) / D + EPS) * gamma).to(X.dtype.element_ty) |
| product = (normalized.to(tl.float32) * tl.load(W + head * D + column).to(tl.float32)).to(X.dtype.element_ty) |
| score = tl.sum(product.to(tl.float32), 0).to(X.dtype.element_ty).to(tl.float32) |
| score = (score / tl.sqrt(float(D))).to(X.dtype.element_ty).to(tl.float32) |
| score = (score + tl.load(Bias + head).to(tl.float32)).to(X.dtype.element_ty).to(tl.float32) |
| sigmoid = tl.sigmoid(score).to(X.dtype.element_ty).to(tl.float32) |
| gate = sigmoid * tl.load(Gain + head).to(tl.float32) |
| tl.store(Gate + head, gate) |
| paired = tl.where(column < ROT // 2, column + ROT // 2, column - ROT // 2) |
| paired = tl.maximum(paired, 0) % D |
| other = tl.gather(normalized, paired, 0).to(tl.float32) |
| other = tl.where(column < ROT // 2, -other, other) |
| cos = tl.load(C + column, column < ROT, 0).to(tl.float32) |
| sin = tl.load(S + column, column < ROT, 0).to(tl.float32) |
| first = (normalized.to(tl.float32) * cos).to(X.dtype.element_ty).to(tl.float32) |
| second = (other * sin).to(X.dtype.element_ty).to(tl.float32) |
| query = tl.where(column < ROT, first + second, normalized.to(tl.float32)) |
| tl.store(Q + head * D + column, query) |
|
|
|
|
| @triton.jit |
| def admission_kernel(X, W, Y, RANK: tl.constexpr, GROUPS: tl.constexpr, BLOCK: tl.constexpr): |
| group = tl.program_id(0) |
| column = tl.arange(0, BLOCK) |
| x = tl.load(X + column, column < RANK, 0).to(tl.float32) |
| activation = (x * tl.sigmoid(x)).to(X.dtype.element_ty).to(tl.float32) |
| weight = tl.load(W + group * RANK + column, column < RANK, 0).to(tl.float32) |
| tl.store(Y + group, tl.sum(activation * weight, 0)) |
|
|
|
|
| def fused_query(m, u, position_embeddings): |
| if u.shape[:2] != (1, 1): |
| return original_query[m](u, position_embeddings) |
| projected = m.q_proj(u).contiguous() |
| query = torch.empty_like(projected).unflatten(-1, (m.heads, m.dim)) |
| gate = torch.empty((1, 1, m.heads), device=u.device, dtype=u.dtype) |
| cos, sin = position_embeddings |
| gain = m.read_scale.tanh() * 2 |
| norm_gate_query[m.heads,](projected, m.q_norm._qwen35_native_rms_gamma, |
| m.read_gate_weight, m.read_gate_bias, gain, cos, sin, query, gate, |
| m.dim, cos.shape[-1], float(m.q_norm.eps), num_warps=4, enable_fp_fusion=False) |
| features = m.admission_down(u.detach()).contiguous() |
| scores = torch.empty((1, 1, m.groups), device=u.device, dtype=torch.float32) |
| admission_kernel[m.groups,](features, m.admission_up.weight, scores, |
| features.shape[-1], m.groups, triton.next_power_of_2(features.shape[-1]), |
| num_warps=4, enable_fp_fusion=False) |
| return query, scores, gate |
|
|
|
|
|
|
| def install(model): |
| projection_patch.install() |
| kernels.triton_qwen35_rmsnorm_torch_tree = norm |
| hf_model.apply_rotary_pos_emb = rope |
| modules = [m for m in model.modules() if hasattr(m, '_query_scores') and hasattr(m, '_write_kv')] |
| source_modules = {sys.modules[type(m).__module__] for m in modules} |
| for module in source_modules: |
| previous_memory_rope = module._memory_rope |
| def memory_rope(values, cos, sin, original=previous_memory_rope): |
| if cos.numel() == cos.shape[-1] and cos.dtype == values.dtype and values.stride(-1) == 1: |
| return rotate(values, cos, sin) |
| return original(values, cos, sin) |
| module._memory_rope = memory_rope |
| if hasattr(module, 'apply_rotary_pos_emb'): |
| module.apply_rotary_pos_emb = rope |
| for m in modules: |
| if m not in original_query: |
| original_query[m] = m._query_scores |
| m._query_scores = MethodType(fused_query, m) |
|
|
