TBNilles/qwen3.5-122B-A10B-on-spark
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

int8 lm-head: output bf16 (match stock lm-head dtype) instead of fp32 — halves the logits buffer that was inflating vLLM's profiled activation reserve and shrinking the dense KV pool

Browse Source
This commit is contained in:
ent
2026-06-24 15:53:58 +10:00
parent a0b661b016
commit 9b7be16bee
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
runtime/patch_int8_lmhead_v3.py
+4 -2
View File
@@ -59,7 +59,9 @@ def _spark_k_int8(x_ptr, w_ptr, s_ptr, o_ptr, B, N, K,
s = _spark_tl.load(s_ptr + offs_n, mask=offs_n < N, other=0.0).to(_spark_tl.float32) s = _spark_tl.load(s_ptr + offs_n, mask=offs_n < N, other=0.0).to(_spark_tl.float32)
acc = acc * s[None, :] acc = acc * s[None, :]
o_ptrs = o_ptr + offs_b[:, None] * sob + offs_n[None, :] * son o_ptrs = o_ptr + offs_b[:, None] * sob + offs_n[None, :] * son
_spark_tl.store(o_ptrs, acc, mask=(offs_b[:, None] < B) & (offs_n[None, :] < N)) # bf16 output matches the stock bf16 lm-head's logits dtype (F.linear keeps bf16);
# fp32 here would ~2x the logits buffer and inflate vLLM's profiled activation reserve.
_spark_tl.store(o_ptrs, acc.to(_spark_tl.bfloat16), mask=(offs_b[:, None] < B) & (offs_n[None, :] < N))
def _spark_int8_gemm(hidden, w_int8, w_scale): def _spark_int8_gemm(hidden, w_int8, w_scale):
@@ -68,7 +70,7 @@ def _spark_int8_gemm(hidden, w_int8, w_scale):
x = hidden.reshape(-1, K) x = hidden.reshape(-1, K)
B = x.shape[0] B = x.shape[0]
BLOCK_B = max(16, _spark_triton.next_power_of_2(B)) BLOCK_B = max(16, _spark_triton.next_power_of_2(B))
out = torch.empty(B, N, dtype=torch.float32, device=x.device) out = torch.empty(B, N, dtype=torch.bfloat16, device=x.device) # match stock bf16 logits
xf = x.to(torch.float16) xf = x.to(torch.float16)
grid = ((N + 127) // 128,) grid = ((N + 127) // 128,)
_spark_k_int8[grid](xf, w_int8, w_scale, out, B, N, K, _spark_k_int8[grid](xf, w_int8, w_scale, out, B, N, K,