kvcache-ai/ktransformers

Fork 0

mirror of https://github.com/kvcache-ai/ktransformers.git synced 2026-11-05 01:07:28 +00:00

[GH-ISSUE #2172] [MXFP4 layerwise prefill] SM 8.6 (Ampere) excluded by capability whitelist — DSV4-Flash on dual A10 falls back to serialized full-GPU path #1295

New issue

Open

opened 2026-10-29 14:50:55 +00:00 by bsakaguchi · 2 comments

bsakaguchi commented

2026-10-29 14:50:55 +00:00

Owner

Originally created by @1stlook-dev on GitHub (Aug 26, 2026).
Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/2172

Summary

While evaluating KTransformers (DSV4-Flash Day-0 support) for a dual-NVIDIA-A10 (Ampere, SM 8.6) deployment, I noticed that the MXFP4 layerwise prefill pipeline is enabled only on Ada/Blackwell.

python/sglang/srt/layers/moe/kt_ep_wrapper.py (in the third_party/sglang submodule, branch main):

def _mxfp4_pipeline_backend_supported(method, layer: torch.nn.Module) -> bool:
    ...
    device = next(layer.parameters()).device
    return torch.cuda.get_device_capability(device) in ((8, 9), (12, 0))

SM 8.6 (A10/A30/A40/A6000/L40) is excluded, and there is no env override to opt in (the SGLANG_V4_USE_TRITON_KERNELS diagnostic only disables the pipeline).

Impact

On Ampere, when --kt-gpu-prefill-token-threshold > 0 (the DSV4-Flash doc uses 4096), a prefill request at/above the threshold enters the _full_gpu_gate branch, but since _mxfp4_pipeline_runtime_supported() is False it falls through to:

# Non-MXFP4 and unsupported MXFP4 backends retain the existing
# serialized full-GPU fallback.
ctx = self._build_full_context(layer)
...
result = ctx.gpu_method.apply(ctx.gpu_layer, dispatch_output)
if torch.cuda.is_available():
    torch.cuda.synchronize()

So each layer does: H2D-load all 256 experts -> compute -> hard sync -> next layer. That is a serialized per-layer transfer chain — the exact bottleneck we are trying to escape (see below). As the gate stands, Ampere deployments get no benefit from the layerwise pipeline.

Background / motivation

Our current DSV4-Flash deployment (vLLM + CPU-MoE backend, TP2 on 2x A10) plateaus at ~650 tok/s cold prefill, independent of prompt length (8k-64k). Thread-level diagnosis attributes it to a latency-bound serialized per-layer weight-transfer chain (1-3 runnable threads on the transfer side, zero threads blocked on bandwidth; GPU memory util 14-44% during prefill — i.e. the GPU idles waiting on H2D, not on compute).

We want to evaluate whether KTransformers' double-buffered layerwise prefill can break that plateau, but with the current capability gate we would land on the same serialized-chain shape on SM 8.6.

Questions

Is SM 8.6 excluded because the MXFP4 Marlin path is unverified/incorrect on Ampere, or simply untested?
Are there plans to enable the layerwise pipeline on Ampere, or to add an opt-in env override (e.g. SGLANG_V4_MXFP4_LAYERWISE_FORCE=1)?
If it is a verification gap: we can run the DSV4-Flash benchmark on our dual-A10 machine and report prefill throughput + correctness results. Please tell us which checks are needed to certify SM 8.6.

Environment

GPU: 2x NVIDIA A10, 24 GB each, SM 8.6, PCIe 4.0 x16, no NVLink
CPU: 2x AMD EPYC 7H12 (AVX2 only, no AVX-512/AMX)
RAM: 512 GB
Model: DeepSeek-V4-Flash (284B MoE, 13B active, 256 experts / top-6, native MXFP4, 43 layers, 160K context)
KTransformers: approachingai/ktransformers:DSV4-specific / main branch

Originally created by @1stlook-dev on GitHub (Aug 26, 2026). Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/2172 ## Summary While evaluating KTransformers (DSV4-Flash Day-0 support) for a **dual-NVIDIA-A10 (Ampere, SM 8.6)** deployment, I noticed that the **MXFP4 layerwise prefill pipeline** is enabled only on Ada/Blackwell. `python/sglang/srt/layers/moe/kt_ep_wrapper.py` (in the `third_party/sglang` submodule, branch `main`): ```python def _mxfp4_pipeline_backend_supported(method, layer: torch.nn.Module) -> bool: ... device = next(layer.parameters()).device return torch.cuda.get_device_capability(device) in ((8, 9), (12, 0)) ``` SM 8.6 (A10/A30/A40/A6000/L40) is excluded, and there is no env override to opt in (the `SGLANG_V4_USE_TRITON_KERNELS` diagnostic only *disables* the pipeline). ## Impact On Ampere, when `--kt-gpu-prefill-token-threshold > 0` (the DSV4-Flash doc uses `4096`), a prefill request at/above the threshold enters the `_full_gpu_gate` branch, but since `_mxfp4_pipeline_runtime_supported()` is False it falls through to: ```python # Non-MXFP4 and unsupported MXFP4 backends retain the existing # serialized full-GPU fallback. ctx = self._build_full_context(layer) ... result = ctx.gpu_method.apply(ctx.gpu_layer, dispatch_output) if torch.cuda.is_available(): torch.cuda.synchronize() ``` So each layer does: H2D-load all 256 experts -> compute -> hard sync -> next layer. That is a **serialized per-layer transfer chain** — the exact bottleneck we are trying to escape (see below). As the gate stands, Ampere deployments get no benefit from the layerwise pipeline. ## Background / motivation Our current DSV4-Flash deployment (vLLM + CPU-MoE backend, TP2 on 2x A10) **plateaus at ~650 tok/s cold prefill, independent of prompt length (8k-64k)**. Thread-level diagnosis attributes it to a latency-bound serialized per-layer weight-transfer chain (1-3 runnable threads on the transfer side, zero threads blocked on bandwidth; GPU memory util 14-44% during prefill — i.e. the GPU idles waiting on H2D, not on compute). We want to evaluate whether KTransformers' double-buffered layerwise prefill can break that plateau, but with the current capability gate we would land on the same serialized-chain shape on SM 8.6. ## Questions 1. Is SM 8.6 excluded because the MXFP4 Marlin path is **unverified/incorrect on Ampere**, or simply untested? 2. Are there plans to enable the layerwise pipeline on Ampere, or to add an opt-in env override (e.g. `SGLANG_V4_MXFP4_LAYERWISE_FORCE=1`)? 3. If it is a verification gap: we can run the DSV4-Flash benchmark on our dual-A10 machine and report prefill throughput + correctness results. Please tell us which checks are needed to certify SM 8.6. ## Environment - GPU: 2x NVIDIA A10, 24 GB each, SM 8.6, PCIe 4.0 x16, no NVLink - CPU: 2x AMD EPYC 7H12 (AVX2 only, no AVX-512/AMX) - RAM: 512 GB - Model: DeepSeek-V4-Flash (284B MoE, 13B active, 256 experts / top-6, native MXFP4, 43 layers, 160K context) - KTransformers: `approachingai/ktransformers:DSV4-specific` / main branch

bsakaguchi commented

2026-10-29 14:50:56 +00:00

Author

Owner

@wenmin-wu commented on GitHub (Aug 26, 2026):

Regarding Q1 ("unverified/incorrect on Ampere, or simply untested?") — I went through the kernel sources rather than the Python gate, and everything I can find points to untested rather than unsupported. Sharing the trace in case it saves the maintainers a pass; all of it is verifiable from main.

The Python whitelist is narrower than the kernel's own requirement.

python/sglang/srt/layers/moe/kt_ep_wrapper.py:3228

return torch.cuda.get_device_capability(device) in ((8, 9), (12, 0))

while the CUDA side gates only on:

python/sglang/jit_kernel/csrc/gemm/marlin_moe/moe_wna16_marlin.cuh:37 → #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 (stub below 8.0, real implementation otherwise)
python/sglang/jit_kernel/csrc/gemm/marlin_moe/marlin_template.h:37 → same guard
python/sglang/jit_kernel/csrc/gemm/marlin/gptq_marlin_repack.cuh:32 → same guard

So the compiled artifacts declare sm_80+, and SM 8.6 sits inside that.

The tensor-core ISA in the MoE template is Ampere-era. Grepping marlin_template.h for every PTX/intrinsic it emits returns exactly:

mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32     (sm_80+)
mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32   (sm_80+)
ldmatrix.sync.aligned.m8n8.x{1,2,4}.shared.b16        (sm_75+)

No wgmma, no TMA, no cluster barriers, no cvt.* to/from FP8/FP4, and no __CUDA_ARCH__ branch on 900/1000/1200. MXFP4 is dequantized into 16-bit and multiplied on tensor cores that Ampere already has — which is consistent with Marlin's original design target.

The repack path looks arch-neutral too. mxfp4_marlin_repack carries the same < 800 guard and emits no FP4/FP8 conversion instructions; it treats MXFP4 as packed bytes, so there is no Ada/Blackwell-only numerics step I can see.

Shared memory is already handled dynamically, which was my main a-priori worry for GA10x (100 KB opt-in per block vs A100's 164 KB):

moe_wna16_marlin.cuh:692 → cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev)
marlin_template.h:804 → sh_a_max_row = ((max_shared_mem - 1024) / 16 - shm_size_used) / (thread_k_blocks * 2)

i.e. the tiling adapts to whatever budget the device reports rather than assuming a fixed size.

And the GEMM is JIT-compiled (load_jit("moe_wna16_marlin", ..., cuda_files=["gemm/marlin_moe/moe_wna16_marlin.cuh"])), so there is no prebuilt sm_89 cubin that SM 8.6 would be missing.

To be explicit about what this is and is not: the above is a source read, not a benchmark. I have not run the layerwise pipeline on Ampere — the gate stops it before it can be observed, which is exactly why the question is worth asking.

Offer for Q3. If it is a verification gap, we can contribute a second, independent SM 8.6 data point that is deliberately different from the OP's box on the CPU side:

GPU: 2x RTX A6000 48 GB, SM 8.6, PCIe 4.0 x16, no NVLink
CPU: 2x Intel Xeon Platinum 8336C, 64 cores / 128 threads, 2 NUMA nodes, AVX-512F + AVX512_VNNI (no AMX)
RAM: 503 GB
The OP's machine is dual A10 + EPYC 7H12 (AVX2 only), so between us you would get SM 8.6 covered across both the AVX2 and AVX-512-without-AMX CPU backends.

Happy to run whatever certification matrix you specify — correctness (logit/output parity against the serialized full-GPU fallback), prefill throughput vs prompt length, and MXFP4 numerical checks. An opt-in escape hatch such as SGLANG_V4_MXFP4_LAYERWISE_FORCE=1 (as the OP suggests) would be enough for us to produce those numbers without patching the submodule; today the only knob, SGLANG_V4_USE_TRITON_KERNELS, can just disable the pipeline.

One narrower question, in case the whitelist is deliberate rather than conservative: is there a known correctness issue specific to mma.m16n8k16 without the Ada/Blackwell FP8 path — e.g. an accumulate-order or scale-swizzle assumption in _swizzle_e8m0_scales that only holds on 8.9+? That is the one thing a source read cannot rule out, and it would explain the gate.

@wenmin-wu commented on GitHub (Aug 26, 2026): Regarding Q1 ("unverified/incorrect on Ampere, or simply untested?") — I went through the kernel sources rather than the Python gate, and everything I can find points to **untested rather than unsupported**. Sharing the trace in case it saves the maintainers a pass; all of it is verifiable from `main`. **The Python whitelist is narrower than the kernel's own requirement.** `python/sglang/srt/layers/moe/kt_ep_wrapper.py:3228` ```python return torch.cuda.get_device_capability(device) in ((8, 9), (12, 0)) ``` while the CUDA side gates only on: - `python/sglang/jit_kernel/csrc/gemm/marlin_moe/moe_wna16_marlin.cuh:37` → `#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800` (stub below 8.0, real implementation otherwise) - `python/sglang/jit_kernel/csrc/gemm/marlin_moe/marlin_template.h:37` → same guard - `python/sglang/jit_kernel/csrc/gemm/marlin/gptq_marlin_repack.cuh:32` → same guard So the compiled artifacts declare **sm_80+**, and SM 8.6 sits inside that. **The tensor-core ISA in the MoE template is Ampere-era.** Grepping `marlin_template.h` for every PTX/intrinsic it emits returns exactly: ``` mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 (sm_80+) mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 (sm_80+) ldmatrix.sync.aligned.m8n8.x{1,2,4}.shared.b16 (sm_75+) ``` No `wgmma`, no TMA, no cluster barriers, no `cvt.*` to/from FP8/FP4, and no `__CUDA_ARCH__` branch on 900/1000/1200. MXFP4 is dequantized into 16-bit and multiplied on tensor cores that Ampere already has — which is consistent with Marlin's original design target. **The repack path looks arch-neutral too.** `mxfp4_marlin_repack` carries the same `< 800` guard and emits no FP4/FP8 conversion instructions; it treats MXFP4 as packed bytes, so there is no Ada/Blackwell-only numerics step I can see. **Shared memory is already handled dynamically**, which was my main a-priori worry for GA10x (100 KB opt-in per block vs A100's 164 KB): - `moe_wna16_marlin.cuh:692` → `cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev)` - `marlin_template.h:804` → `sh_a_max_row = ((max_shared_mem - 1024) / 16 - shm_size_used) / (thread_k_blocks * 2)` i.e. the tiling adapts to whatever budget the device reports rather than assuming a fixed size. **And the GEMM is JIT-compiled** (`load_jit("moe_wna16_marlin", ..., cuda_files=["gemm/marlin_moe/moe_wna16_marlin.cuh"])`), so there is no prebuilt sm_89 cubin that SM 8.6 would be missing. **To be explicit about what this is and is not:** the above is a source read, not a benchmark. I have not run the layerwise pipeline on Ampere — the gate stops it before it can be observed, which is exactly why the question is worth asking. **Offer for Q3.** If it is a verification gap, we can contribute a second, independent SM 8.6 data point that is deliberately different from the OP's box on the CPU side: - GPU: 2x RTX A6000 48 GB, SM 8.6, PCIe 4.0 x16, no NVLink - CPU: 2x Intel Xeon Platinum 8336C, 64 cores / 128 threads, 2 NUMA nodes, **AVX-512F + AVX512_VNNI (no AMX)** - RAM: 503 GB - The OP's machine is dual A10 + EPYC 7H12 (AVX2 only), so between us you would get SM 8.6 covered across both the AVX2 and AVX-512-without-AMX CPU backends. Happy to run whatever certification matrix you specify — correctness (logit/output parity against the serialized full-GPU fallback), prefill throughput vs prompt length, and MXFP4 numerical checks. An opt-in escape hatch such as `SGLANG_V4_MXFP4_LAYERWISE_FORCE=1` (as the OP suggests) would be enough for us to produce those numbers without patching the submodule; today the only knob, `SGLANG_V4_USE_TRITON_KERNELS`, can just disable the pipeline. One narrower question, in case the whitelist is deliberate rather than conservative: is there a known correctness issue specific to `mma.m16n8k16` **without** the Ada/Blackwell FP8 path — e.g. an accumulate-order or scale-swizzle assumption in `_swizzle_e8m0_scales` that only holds on 8.9+? That is the one thing a source read cannot rule out, and it would explain the gate.

bsakaguchi commented

2026-10-29 14:50:56 +00:00

Author

Owner

@1stlook-dev commented on GitHub (Aug 27, 2026):

Good catch — and this matches what I can verify in the sglang submodule pinned by KTransformers: the same ((8, 9), (12, 0)) check (kt_ep_wrapper.py:2450 in our checkout, :3228 on main), and the same < 800 stub/real split in moe_wna16_marlin.cuh (empty stub at :37, real kernel under #else :51–:822) and marlin_template.h — only Ampere tensor-core intrinsics (mma.sync.aligned.m16n8k16, ldmatrix), no wgmma/TMA/FP8 cvt. So: the gate is narrower than the kernel requires, and the only reason it has never been observed on Ampere is that the gate stops it first.

One clarification for triage on where the OP's box actually dies — the GPU gate was not the blocker on our machine. 2x A10 (sm86), 2x EPYC 7H12 AVX2-only (no AVX-512/AMX), NPS4: 32 NUMA nodes x ~32GB:

OOM on the default topology. Default is numa_ids=[0,1] + process-wide numa_bind() (MPOL_BIND), so each rank's CPU-expert working set (raw + swizzled MXFP4 both resident while --kt-gpu-prefill-token-threshold > 0, ~140GB/rank) is hard-pinned into node0's 32GB → OOM-kill at layer 17 (sglang::scheduler, anon-rss ~60GB and still growing). We relaxed the bind to numa_set_localalloc() and spread 16 pools across nodes 16–31: the OOM went away and TP1 completed the full 43-layer load (28.75s).
Then the AVX2 worker pool crashes, before the GPU MoE path is ever reached. With 16 pools: segfault in layer-0 load_weights (kt-kernel amx.py:936, work-stealing) — a worker thread jumps to 0xffffffffffffffff. With a 2-pool control: it dies earlier, in pool construction — Failed to set thread name: Numerical result out of range, Core 27 inside NUMA node 16 not found, malloc(): unsorted double linked list corrupted → abort. The control proves it is not our memory-policy change: the WorkerPool machinery itself (thread naming / core enumeration / heap) assumes few NUMA nodes and breaks on 32.

So the current data point for sm86 + AVX2-only + NPS4/32-node is: weights load fine (43/43 layers on one rank), but the box cannot get past layer-0 CPU MoE.

Your offer complements ours nicely: between the two boxes, sm86 would be covered for both AVX-512 (your 2-node Xeon) and AVX2-only (our 32-node EPYC) CPU backends, across both 2-node and 32-node topologies. Happy to share the raw crash logs or the patched worker_pool.h if it helps.

On your narrower question — whether mma.m16n8k16 without the Ada/Blackwell FP8 path has a correctness issue specific to 8.9+ — we can't speak to it: our box dies before any GPU MoE kernel runs, so there is no output to compare against.

@1stlook-dev commented on GitHub (Aug 27, 2026): Good catch — and this matches what I can verify in the sglang submodule pinned by KTransformers: the same `((8, 9), (12, 0))` check (`kt_ep_wrapper.py:2450` in our checkout, `:3228` on `main`), and the same `< 800` stub/real split in `moe_wna16_marlin.cuh` (empty stub at :37, real kernel under `#else` :51–:822) and `marlin_template.h` — only Ampere tensor-core intrinsics (`mma.sync.aligned.m16n8k16`, `ldmatrix`), no wgmma/TMA/FP8 cvt. So: the gate is narrower than the kernel requires, and the only reason it has never been observed on Ampere is that the gate stops it first. One clarification for triage on where the OP's box actually dies — **the GPU gate was not the blocker on our machine.** 2x A10 (sm86), 2x EPYC 7H12 **AVX2-only** (no AVX-512/AMX), **NPS4: 32 NUMA nodes x ~32GB**: 1. **OOM on the default topology.** Default is `numa_ids=[0,1]` + process-wide `numa_bind()` (MPOL_BIND), so each rank's CPU-expert working set (raw + swizzled MXFP4 both resident while `--kt-gpu-prefill-token-threshold > 0`, ~140GB/rank) is hard-pinned into node0's 32GB → OOM-kill at layer 17 (`sglang::scheduler`, anon-rss ~60GB and still growing). We relaxed the bind to `numa_set_localalloc()` and spread 16 pools across nodes 16–31: the OOM went away and **TP1 completed the full 43-layer load** (28.75s). 2. **Then the AVX2 worker pool crashes, before the GPU MoE path is ever reached.** With 16 pools: segfault in layer-0 `load_weights` (kt-kernel `amx.py:936`, work-stealing) — a worker thread jumps to `0xffffffffffffffff`. With a 2-pool control: it dies earlier, in pool construction — `Failed to set thread name: Numerical result out of range`, `Core 27 inside NUMA node 16 not found`, `malloc(): unsorted double linked list corrupted` → abort. The control proves it is not our memory-policy change: the `WorkerPool` machinery itself (thread naming / core enumeration / heap) assumes few NUMA nodes and breaks on 32. So the current data point for **sm86 + AVX2-only + NPS4/32-node** is: weights load fine (43/43 layers on one rank), but the box cannot get past layer-0 CPU MoE. Your offer complements ours nicely: between the two boxes, sm86 would be covered for both AVX-512 (your 2-node Xeon) and AVX2-only (our 32-node EPYC) CPU backends, across both 2-node and 32-node topologies. Happy to share the raw crash logs or the patched `worker_pool.h` if it helps. On your narrower question — whether `mma.m16n8k16` without the Ada/Blackwell FP8 path has a correctness issue specific to 8.9+ — we can't speak to it: our box dies before any GPU MoE kernel runs, so there is no output to compare against.