[DSA] Q8KV8 FP8 Sparse Prefill on GLM-5.2 & DeepSeek-V3.2: Q8-Path & Shared-Path Optimizations (#31888)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Ho-Ren (Jack) Chuang
co-authored by
Claude Fable 5
Claude Fable 5
parent
4f51dad1da
commit
e4a40a71f8
@@ -0,0 +1,274 @@
|
||||
"""Standalone GPU test for SGLANG_OPT_MOE_QUANT_ONCE (quantize the MoE input
|
||||
once, feed both the fused shared-expert GEMM and the routed triton runner).
|
||||
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0 python test/manual/test_moe_quant_once.py
|
||||
|
||||
Verifies, against the double-quant baseline:
|
||||
(1) quant equivalence: the row-padded quantize-once kernel produces the
|
||||
same q bits / scale values as the routed path's default row-major quant
|
||||
(JIT v2 kernel) on the valid rows;
|
||||
(2) shared consumer: cutlass_w8a8_block_fp8_linear_with_fallback with a
|
||||
pre-quantized (q, s) tuple vs its own internal quant -- expected BITWISE
|
||||
(baseline uses the identical row-padded quant + identical GEMM);
|
||||
(2b) shared consumer under SGLANG_ENABLE_JIT_DEEPGEMM=1 (the recommended JIT-DeepGEMM config):
|
||||
deepgemm_w8a8_block_fp8_linear_with_fallback with the same (q, s) tuple
|
||||
-- expected BITWISE (DG's own quant layout, column-major TMA-aligned
|
||||
fp32 scales, is byte-identical to the row-padded quantize-once layout);
|
||||
skipped cleanly when deep_gemm is unavailable or UE8M0 (Blackwell);
|
||||
(3) routed consumer: fused_experts(a1_q=..., a1_scale=...) vs the in-kernel
|
||||
quant baseline -- expected BITWISE if (1) is bitwise (the fused kernel
|
||||
reads A_scale through explicit strides, so the column-major scale view
|
||||
feeds identical values).
|
||||
|
||||
If (1) is not bitwise (AOT v2 vs JIT v2 quant kernels round differently),
|
||||
(3) falls back to an allclose check at atol=1e-2 and the discrepancy is
|
||||
reported --.
|
||||
"""
|
||||
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import (
|
||||
sglang_per_token_group_quant_fp8,
|
||||
sglang_per_token_group_quant_fp8_row_padded,
|
||||
)
|
||||
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import fused_experts
|
||||
from sglang.srt.layers.moe.topk import StandardTopKOutput
|
||||
from sglang.srt.layers.quantization.fp8_utils import (
|
||||
cutlass_w8a8_block_fp8_linear_with_fallback,
|
||||
)
|
||||
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
|
||||
|
||||
GROUP = 128
|
||||
FAILURES = []
|
||||
|
||||
|
||||
def _report(name, ok, detail=""):
|
||||
status = "PASS" if ok else "FAIL"
|
||||
print(f"[{status}] {name} {detail}")
|
||||
if not ok:
|
||||
FAILURES.append(name)
|
||||
|
||||
|
||||
def _quant_weight_blockwise(w_bf16, block=128):
|
||||
"""Per-[128,128]-block fp8 weight quant (reference, fp32 math)."""
|
||||
n, k = w_bf16.shape
|
||||
w = w_bf16.float().view(n // block, block, k // block, block)
|
||||
amax = w.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-4)
|
||||
scale = amax / torch.finfo(torch.float8_e4m3fn).max
|
||||
q = (w / scale).clamp(-448, 448).to(torch.float8_e4m3fn)
|
||||
return (
|
||||
q.view(n, k),
|
||||
scale.squeeze(1).squeeze(-1).to(torch.float32), # [n/128, k/128]
|
||||
)
|
||||
|
||||
|
||||
def test_quant_equivalence(T, K, device):
|
||||
x = torch.randn(T, K, device=device, dtype=torch.bfloat16) * 3
|
||||
q_ref, s_ref = sglang_per_token_group_quant_fp8(x, GROUP) # routed baseline
|
||||
q_pad, s_pad = sglang_per_token_group_quant_fp8_row_padded(x, GROUP)
|
||||
|
||||
bitwise_q = torch.equal(q_pad[:T].view(torch.uint8), q_ref.view(torch.uint8))
|
||||
bitwise_s = torch.equal(s_pad[:T].contiguous(), s_ref)
|
||||
_report(
|
||||
f"quant-equivalence T={T} K={K}",
|
||||
bitwise_q and bitwise_s,
|
||||
f"(q bitwise={bitwise_q}, s bitwise={bitwise_s})",
|
||||
)
|
||||
return bitwise_q and bitwise_s
|
||||
|
||||
|
||||
def test_shared_consumer(T, K, N, device):
|
||||
torch.manual_seed(T + K)
|
||||
x = torch.randn(T, K, device=device, dtype=torch.bfloat16)
|
||||
w_bf16 = torch.randn(N, K, device=device, dtype=torch.bfloat16) / K**0.5
|
||||
w, ws = _quant_weight_blockwise(w_bf16)
|
||||
|
||||
ref = cutlass_w8a8_block_fp8_linear_with_fallback(x, w, [128, 128], ws)
|
||||
|
||||
q_pad, s_pad = sglang_per_token_group_quant_fp8_row_padded(x, GROUP)
|
||||
out = cutlass_w8a8_block_fp8_linear_with_fallback(
|
||||
q_pad, w, [128, 128], ws, input_scale=s_pad
|
||||
)[:T]
|
||||
|
||||
bitwise = torch.equal(out, ref)
|
||||
close = torch.allclose(out.float(), ref.float(), atol=1e-2, rtol=1e-2)
|
||||
_report(
|
||||
f"shared-consumer T={T} K={K} N={N}",
|
||||
close,
|
||||
f"(bitwise={bitwise}, max|d|={(out.float() - ref.float()).abs().max().item():.3e})",
|
||||
)
|
||||
return bitwise
|
||||
|
||||
|
||||
def test_shared_consumer_deepgemm(T, K, N, device):
|
||||
"""DG branch (SGLANG_ENABLE_JIT_DEEPGEMM=1 recommended JIT-DeepGEMM config): the shared-expert
|
||||
linear resolves to deepgemm_w8a8_block_fp8_linear_with_fallback. Its own
|
||||
quant (column-major + TMA-aligned fp32 scales) has the same buffer layout
|
||||
as the row-padded quantize-once kernel, so this is expected BITWISE."""
|
||||
from sglang.srt.layers.quantization.fp8_utils import (
|
||||
deepgemm_w8a8_block_fp8_linear_with_fallback,
|
||||
)
|
||||
|
||||
torch.manual_seed(T + K + 1)
|
||||
x = torch.randn(T, K, device=device, dtype=torch.bfloat16)
|
||||
w_bf16 = torch.randn(N, K, device=device, dtype=torch.bfloat16) / K**0.5
|
||||
w, ws = _quant_weight_blockwise_n64(w_bf16)
|
||||
|
||||
ref = deepgemm_w8a8_block_fp8_linear_with_fallback(x, w, [128, 128], ws)
|
||||
|
||||
q_pad, s_pad = sglang_per_token_group_quant_fp8_row_padded(x, GROUP)
|
||||
out = deepgemm_w8a8_block_fp8_linear_with_fallback(
|
||||
q_pad, w, [128, 128], ws, input_scale=s_pad
|
||||
)[:T]
|
||||
|
||||
bitwise = torch.equal(out, ref)
|
||||
close = torch.allclose(out.float(), ref.float(), atol=1e-2, rtol=1e-2)
|
||||
_report(
|
||||
f"shared-consumer-deepgemm T={T} K={K} N={N}",
|
||||
close,
|
||||
f"(bitwise={bitwise}, max|d|={(out.float() - ref.float()).abs().max().item():.3e})",
|
||||
)
|
||||
return bitwise
|
||||
|
||||
|
||||
def _quant_weight_blockwise_n64(w_bf16, block=128):
|
||||
"""Like _quant_weight_blockwise but supports N % 64 == 0 (DeepGEMM's
|
||||
minimum): the last (partial) N-block reuses ceil-division block indexing."""
|
||||
n, k = w_bf16.shape
|
||||
if n % block == 0:
|
||||
return _quant_weight_blockwise(w_bf16, block)
|
||||
import math
|
||||
|
||||
n_blocks = math.ceil(n / block)
|
||||
w = w_bf16.float()
|
||||
q = torch.empty(n, k, device=w.device, dtype=torch.float8_e4m3fn)
|
||||
scale = torch.empty(n_blocks, k // block, device=w.device, dtype=torch.float32)
|
||||
for bn in range(n_blocks):
|
||||
rows = slice(bn * block, min((bn + 1) * block, n))
|
||||
wb = w[rows].view(rows.stop - rows.start, k // block, block)
|
||||
amax = wb.abs().amax(dim=(0, 2)).clamp(min=1e-4)
|
||||
s = amax / torch.finfo(torch.float8_e4m3fn).max
|
||||
q[rows] = (
|
||||
(wb / s[None, :, None]).clamp(-448, 448).to(torch.float8_e4m3fn).view(-1, k)
|
||||
)
|
||||
scale[bn] = s
|
||||
return q, scale
|
||||
|
||||
|
||||
def test_routed_consumer(T, K, E, I, topk, device):
|
||||
torch.manual_seed(T * 7 + K)
|
||||
x = torch.randn(T, K, device=device, dtype=torch.bfloat16)
|
||||
w1 = torch.empty(E, 2 * I, K, device=device, dtype=torch.float8_e4m3fn)
|
||||
w1s = torch.empty(E, 2 * I // 128, K // 128, device=device)
|
||||
w2 = torch.empty(E, K, I, device=device, dtype=torch.float8_e4m3fn)
|
||||
w2s = torch.empty(E, K // 128, I // 128, device=device)
|
||||
for e in range(E):
|
||||
w1[e], w1s[e] = _quant_weight_blockwise(
|
||||
torch.randn(2 * I, K, device=device, dtype=torch.bfloat16) / K**0.5
|
||||
)
|
||||
w2[e], w2s[e] = _quant_weight_blockwise(
|
||||
torch.randn(K, I, device=device, dtype=torch.bfloat16) / I**0.5
|
||||
)
|
||||
|
||||
topk_weights = torch.rand(T, topk, device=device)
|
||||
topk_weights = (topk_weights / topk_weights.sum(-1, keepdim=True)).to(torch.float32)
|
||||
topk_ids = torch.stack(
|
||||
[torch.randperm(E, device=device)[:topk] for _ in range(T)]
|
||||
).to(torch.int32)
|
||||
topk_output = StandardTopKOutput(
|
||||
topk_weights=topk_weights, topk_ids=topk_ids, router_logits=None
|
||||
)
|
||||
# num_experts == num_local_experts => filter_expert=False (pure TP layout)
|
||||
cfg = MoeRunnerConfig(
|
||||
num_experts=E,
|
||||
num_local_experts=E,
|
||||
top_k=topk,
|
||||
inplace=False,
|
||||
activation="silu",
|
||||
is_gated=True,
|
||||
)
|
||||
|
||||
kwargs = dict(
|
||||
w1=w1,
|
||||
w2=w2,
|
||||
topk_output=topk_output,
|
||||
moe_runner_config=cfg,
|
||||
use_fp8_w8a8=True,
|
||||
w1_scale=w1s,
|
||||
w2_scale=w2s,
|
||||
block_shape=[128, 128],
|
||||
)
|
||||
ref = fused_experts(hidden_states=x, **kwargs)
|
||||
|
||||
q_pad, s_pad = sglang_per_token_group_quant_fp8_row_padded(x, GROUP)
|
||||
out = fused_experts(hidden_states=x, a1_q=q_pad, a1_scale=s_pad, **kwargs)
|
||||
|
||||
bitwise = torch.equal(out, ref)
|
||||
close = torch.allclose(out.float(), ref.float(), atol=1e-2, rtol=1e-2)
|
||||
_report(
|
||||
f"routed-consumer T={T} K={K} E={E} topk={topk}",
|
||||
close,
|
||||
f"(bitwise={bitwise}, max|d|={(out.float() - ref.float()).abs().max().item():.3e})",
|
||||
)
|
||||
return bitwise
|
||||
|
||||
|
||||
def main():
|
||||
assert torch.cuda.is_available(), "CUDA required"
|
||||
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
|
||||
device = "cuda"
|
||||
torch.manual_seed(0)
|
||||
|
||||
print("== (1) quantize-once vs routed-baseline quant equivalence ==")
|
||||
all_bitwise_q = True
|
||||
for T in (1, 3, 4093, 4096):
|
||||
for K in (6144, 7168):
|
||||
all_bitwise_q &= test_quant_equivalence(T, K, device)
|
||||
|
||||
print("== (2) shared consumer (cutlass w8a8 linear) ==")
|
||||
# N=512 mirrors a tp8 shared expert gate_up (2*2048/8); must be %128==0.
|
||||
for T in (4093, 4096):
|
||||
for K in (6144, 7168):
|
||||
test_shared_consumer(T, K, 512, device)
|
||||
|
||||
print(
|
||||
"== (2b) shared consumer (deepgemm w8a8 linear, JIT DG recommended JIT-DeepGEMM config) =="
|
||||
)
|
||||
from sglang.srt.layers import deep_gemm_wrapper
|
||||
|
||||
if not deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
|
||||
print("SKIP: deep_gemm unavailable or SGLANG_ENABLE_JIT_DEEPGEMM=0")
|
||||
elif deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
|
||||
print("SKIP: Blackwell UE8M0 scale layout (gated ineligible by design)")
|
||||
else:
|
||||
for T in (4093, 4096):
|
||||
for K in (6144, 7168):
|
||||
test_shared_consumer_deepgemm(T, K, 512, device)
|
||||
# DG accepts N % 64 (cutlass needs % 128) -- exercise the DG-only shape.
|
||||
test_shared_consumer_deepgemm(4096, 7168, 320, device)
|
||||
|
||||
print("== (3) routed consumer (triton fused_experts) ==")
|
||||
# Identical in both cutlass and JIT-DG configs: the MoE runner stays
|
||||
# triton with a2a=none (is_deepgemm_moe_runner_backend_enabled() is False
|
||||
# for auto + a2a=none even when SGLANG_ENABLE_JIT_DEEPGEMM=1).
|
||||
for T in (61, 4093, 4096):
|
||||
test_routed_consumer(T, 7168, E=32, I=256, topk=8, device=device)
|
||||
|
||||
if not all_bitwise_q:
|
||||
print(
|
||||
"NOTE: quantize-once q/s not bitwise vs the routed baseline quant "
|
||||
"(AOT v2 vs JIT v2 kernel rounding) -- routed consumer is then "
|
||||
"allclose-only; document this in the PR."
|
||||
)
|
||||
if FAILURES:
|
||||
print(f"FAILED: {FAILURES}")
|
||||
sys.exit(1)
|
||||
print("ALL PASS")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user