[Kernel] Phase 4 batch-3: migrate tangled JIT subsystems + new groups into kernels.ops (RFC #29630) (#32045)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
71fe649d68
commit
74338e94f1
@@ -5,7 +5,7 @@ from typing import Callable
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES, RealKvSource
|
||||
from sglang.kernels.ops.kv_canary.verify import CANARY_SLOT_BYTES, RealKvSource
|
||||
|
||||
BS_AXIS: list[int] = [1, 4, 32, 128, 256, 1024]
|
||||
PREFIX_AXIS: list[int] = [0, 128, 1024, 4096, 10240, 16384]
|
||||
|
||||
@@ -1,30 +0,0 @@
|
||||
from .paged_mqa_logits import (
|
||||
aiter_paged_mqa_logits,
|
||||
cutedsl_paged_mqa_logits,
|
||||
deepgemm_paged_mqa_logits_native,
|
||||
deepgemm_paged_mqa_logits_split,
|
||||
)
|
||||
|
||||
|
||||
def pick_dsl_expand(*args, **kwargs):
|
||||
from .cutedsl_paged_mqa_logits import pick_dsl_expand as _pick_dsl_expand
|
||||
|
||||
return _pick_dsl_expand(*args, **kwargs)
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
if name == "CuteDSLPagedMQALogitsRunner":
|
||||
from .cutedsl_paged_mqa_logits import CuteDSLPagedMQALogitsRunner
|
||||
|
||||
return CuteDSLPagedMQALogitsRunner
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CuteDSLPagedMQALogitsRunner",
|
||||
"pick_dsl_expand",
|
||||
"aiter_paged_mqa_logits",
|
||||
"cutedsl_paged_mqa_logits",
|
||||
"deepgemm_paged_mqa_logits_native",
|
||||
"deepgemm_paged_mqa_logits_split",
|
||||
]
|
||||
@@ -1,63 +0,0 @@
|
||||
from .attn import (
|
||||
fused_store_cache,
|
||||
get_paged_mqa_logits_metadata,
|
||||
triton_create_paged_compress_data,
|
||||
)
|
||||
from .c128_cleanup import clear_unaccepted_c128_draft_states
|
||||
from .compress import (
|
||||
CompressorDecodePlan,
|
||||
CompressorPrefillPlan,
|
||||
compress_forward,
|
||||
compress_norm_rope_store,
|
||||
)
|
||||
from .compress_old import fused_norm_rope_inplace
|
||||
from .elementwise import (
|
||||
fused_k_norm_rope_flashmla,
|
||||
fused_q_indexer_rope_first_quant,
|
||||
fused_q_indexer_rope_hadamard_fp4_quant,
|
||||
fused_q_indexer_rope_hadamard_quant,
|
||||
fused_q_norm_rope,
|
||||
fused_rope_inplace,
|
||||
)
|
||||
from .fp8_wo_a import sglang_per_token_group_quant_fp8_dsv4_wo_a
|
||||
from .gemm import linear_bf16_fp32
|
||||
from .moe import (
|
||||
hash_topk,
|
||||
mask_topk_ids,
|
||||
mega_moe_pre_dispatch,
|
||||
silu_and_mul_clamp,
|
||||
silu_and_mul_contig_post_quant,
|
||||
silu_and_mul_masked_post_quant,
|
||||
)
|
||||
from .topk import plan_topk_v2, topk_transform_512, topk_transform_512_v2
|
||||
from .utils import make_name
|
||||
|
||||
__all__ = [
|
||||
"CompressorDecodePlan",
|
||||
"CompressorPrefillPlan",
|
||||
"compress_forward",
|
||||
"compress_norm_rope_store",
|
||||
"clear_unaccepted_c128_draft_states",
|
||||
"fused_norm_rope_inplace",
|
||||
"fused_store_cache",
|
||||
"fused_rope_inplace",
|
||||
"fused_q_norm_rope",
|
||||
"fused_q_indexer_rope_first_quant",
|
||||
"fused_q_indexer_rope_hadamard_fp4_quant",
|
||||
"fused_q_indexer_rope_hadamard_quant",
|
||||
"fused_k_norm_rope_flashmla",
|
||||
"sglang_per_token_group_quant_fp8_dsv4_wo_a",
|
||||
"make_name",
|
||||
"linear_bf16_fp32",
|
||||
"get_paged_mqa_logits_metadata",
|
||||
"triton_create_paged_compress_data",
|
||||
"topk_transform_512",
|
||||
"topk_transform_512_v2",
|
||||
"plan_topk_v2",
|
||||
"hash_topk",
|
||||
"mega_moe_pre_dispatch",
|
||||
"mask_topk_ids",
|
||||
"silu_and_mul_clamp",
|
||||
"silu_and_mul_masked_post_quant",
|
||||
"silu_and_mul_contig_post_quant",
|
||||
]
|
||||
@@ -1 +0,0 @@
|
||||
from sglang.jit_kernel.kv_canary.plan.api import launch_canary_plan_kernels
|
||||
@@ -1,25 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Fused Triton kernels for MiniMax-M3 on AMD ROCm (gfx94x / gfx95x).
|
||||
|
||||
Model-scoped JIT kernels (mirrors ``jit_kernel/dsv4``), split by op type:
|
||||
* ``rmsnorm`` -- fused fp32 Gemma RMSNorm (plain + fused-add-residual)
|
||||
* ``swiglu`` -- fused fp32 SwiGLU-OAI (split layout)
|
||||
"""
|
||||
|
||||
from sglang.jit_kernel.minimax_m3.rmsnorm import (
|
||||
_num_warps,
|
||||
gemma_fused_add_rmsnorm,
|
||||
gemma_rmsnorm,
|
||||
)
|
||||
from sglang.jit_kernel.minimax_m3.swiglu import (
|
||||
swiglu_oai_mxfp8_quant,
|
||||
swiglu_oai_split,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"gemma_rmsnorm",
|
||||
"gemma_fused_add_rmsnorm",
|
||||
"swiglu_oai_split",
|
||||
"swiglu_oai_mxfp8_quant",
|
||||
"_num_warps",
|
||||
]
|
||||
@@ -5,7 +5,10 @@ from typing import List, Literal, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.dsv4 import CompressorDecodePlan, CompressorPrefillPlan
|
||||
from sglang.kernels.ops.attention.dsv4 import (
|
||||
CompressorDecodePlan,
|
||||
CompressorPrefillPlan,
|
||||
)
|
||||
from sglang.srt.utils import get_device
|
||||
|
||||
|
||||
|
||||
@@ -5,10 +5,6 @@ from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary import consts
|
||||
from sglang.jit_kernel.kv_canary.consts import splitmix64, splitmix64_mix3
|
||||
from sglang.jit_kernel.kv_canary.verify import VerifyPlan
|
||||
from sglang.jit_kernel.kv_canary.write import WritePlan
|
||||
from sglang.jit_kernel.tests.kv_canary._constants import (
|
||||
_I64_SIGN_BIT,
|
||||
_U64_MASK,
|
||||
@@ -20,6 +16,10 @@ from sglang.jit_kernel.tests.kv_canary._fixtures import (
|
||||
make_real_kv_source,
|
||||
make_real_kv_sources,
|
||||
)
|
||||
from sglang.kernels.ops.kv_canary import consts
|
||||
from sglang.kernels.ops.kv_canary.consts import splitmix64, splitmix64_mix3
|
||||
from sglang.kernels.ops.kv_canary.verify import VerifyPlan
|
||||
from sglang.kernels.ops.kv_canary.write import WritePlan
|
||||
|
||||
__all__ = [
|
||||
"FakeViolationLog",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES
|
||||
from sglang.kernels.ops.kv_canary.verify import CANARY_SLOT_BYTES
|
||||
|
||||
# Default fixture sizes — small enough for fast tests, large enough that ring overflow / multi-req cases
|
||||
# stay realistic without bloating the assertion surface.
|
||||
|
||||
@@ -5,31 +5,31 @@ from typing import Any, Callable, Iterator, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary import consts
|
||||
from sglang.jit_kernel.kv_canary.plan import launch_canary_plan_kernels
|
||||
from sglang.jit_kernel.kv_canary.plan_ref import (
|
||||
from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
|
||||
FakeViolationLog,
|
||||
assert_canary_buf_equal,
|
||||
assert_canary_state_equal,
|
||||
make_log_pair,
|
||||
)
|
||||
from sglang.kernels.ops.kv_canary import consts
|
||||
from sglang.kernels.ops.kv_canary.plan import launch_canary_plan_kernels
|
||||
from sglang.kernels.ops.kv_canary.plan_ref import (
|
||||
launch_canary_plan_kernels_torch_reference,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.verify import (
|
||||
from sglang.kernels.ops.kv_canary.verify import (
|
||||
CanaryLaunchTag,
|
||||
RealKvSource,
|
||||
VerifyOrWriteContext,
|
||||
VerifyPlan,
|
||||
launch_canary_verify_kernel,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.verify_ref import (
|
||||
from sglang.kernels.ops.kv_canary.verify_ref import (
|
||||
launch_canary_verify_kernel_torch_reference,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.write import WritePlan, launch_canary_write_kernel
|
||||
from sglang.jit_kernel.kv_canary.write_ref import (
|
||||
from sglang.kernels.ops.kv_canary.write import WritePlan, launch_canary_write_kernel
|
||||
from sglang.kernels.ops.kv_canary.write_ref import (
|
||||
launch_canary_write_kernel_torch_reference,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
|
||||
FakeViolationLog,
|
||||
assert_canary_buf_equal,
|
||||
assert_canary_state_equal,
|
||||
make_log_pair,
|
||||
)
|
||||
|
||||
_DEVICE = torch.device("cuda")
|
||||
|
||||
|
||||
@@ -5,12 +5,12 @@ from typing import Literal, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary.verify import (
|
||||
from sglang.jit_kernel.tests.kv_canary._constants import DEFAULT_NUM_SLOTS
|
||||
from sglang.kernels.ops.kv_canary.verify import (
|
||||
RealKvSource,
|
||||
VerifyPlan,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.write import WritePlan
|
||||
from sglang.jit_kernel.tests.kv_canary._constants import DEFAULT_NUM_SLOTS
|
||||
from sglang.kernels.ops.kv_canary.write import WritePlan
|
||||
|
||||
_DEVICE = torch.device("cuda")
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ diff comparison."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from sglang.jit_kernel.kv_canary.consts import splitmix64
|
||||
from sglang.kernels.ops.kv_canary.consts import splitmix64
|
||||
|
||||
|
||||
def _fold_words(padded: bytes) -> int:
|
||||
|
||||
@@ -11,10 +11,10 @@ from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary import consts
|
||||
from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag, VerifyPlan
|
||||
from sglang.jit_kernel.kv_canary.write import WritePlan
|
||||
from sglang.jit_kernel.tests.kv_canary._canary_helpers import FakeViolationLog
|
||||
from sglang.kernels.ops.kv_canary import consts
|
||||
from sglang.kernels.ops.kv_canary.verify import CanaryLaunchTag, VerifyPlan
|
||||
from sglang.kernels.ops.kv_canary.write import WritePlan
|
||||
|
||||
|
||||
class PlanInvariants:
|
||||
|
||||
@@ -86,7 +86,7 @@ def test_minimax_swiglu_mxfp8_quant_matches_unfused_fp32(m, inter):
|
||||
# the reference is the unfused fp32 swiglu followed by MXFP8 quant. Not
|
||||
# bit-identical because the reference quant runs in torch vs the fused triton
|
||||
# path, but numerically equivalent (tight relerr, scales agree within 1 ulp).
|
||||
from sglang.jit_kernel.minimax_m3 import (
|
||||
from sglang.kernels.ops.moe.minimax_m3_swiglu import (
|
||||
swiglu_oai_mxfp8_quant,
|
||||
swiglu_oai_split,
|
||||
)
|
||||
|
||||
@@ -14,7 +14,7 @@ if not is_hip():
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("Requires a GPU.", allow_module_level=True)
|
||||
|
||||
from sglang.jit_kernel.minimax_m3.rmsnorm import ( # noqa: E402
|
||||
from sglang.kernels.ops.layernorm.minimax_m3_rmsnorm import ( # noqa: E402
|
||||
gemma_fused_add_rmsnorm,
|
||||
gemma_rmsnorm,
|
||||
)
|
||||
|
||||
@@ -13,7 +13,7 @@ from typing import Optional
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.moe_topk_sigmoid import topk_sigmoid
|
||||
from sglang.kernels.ops.moe.moe_topk_sigmoid import topk_sigmoid
|
||||
|
||||
try:
|
||||
from sgl_kernel import topk_sigmoid as topk_sigmoid_aot
|
||||
|
||||
@@ -1,19 +0,0 @@
|
||||
from sglang.jit_kernel.trtllm_lora_temp.core import (
|
||||
trtllm_bf16_routed_moe_lora,
|
||||
trtllm_fp4_block_scale_moe_lora_finalize,
|
||||
trtllm_fp4_block_scale_routed_moe_lora,
|
||||
trtllm_fp8_block_scale_moe,
|
||||
trtllm_fp8_block_scale_moe_lora_finalize,
|
||||
trtllm_fp8_block_scale_routed_moe,
|
||||
trtllm_fp8_block_scale_routed_moe_lora,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"trtllm_bf16_routed_moe_lora",
|
||||
"trtllm_fp4_block_scale_moe_lora_finalize",
|
||||
"trtllm_fp4_block_scale_routed_moe_lora",
|
||||
"trtllm_fp8_block_scale_moe_lora_finalize",
|
||||
"trtllm_fp8_block_scale_moe",
|
||||
"trtllm_fp8_block_scale_routed_moe",
|
||||
"trtllm_fp8_block_scale_routed_moe_lora",
|
||||
]
|
||||
@@ -1,246 +0,0 @@
|
||||
// Fused SwiGLU+LoRA activation -> NVFP4 per-token quant for the FP4 MoE LoRA path.
|
||||
//
|
||||
// Modeled on tensorrt_llm::kernels::nvfp4QuantAndPerTokenScaleKernel (flashinfer
|
||||
// quantization.cuh): one block per expanded row, per-token amax via cub::BlockReduce,
|
||||
// cvt_warp_fp16_to_fp4 for the e4m3 block scale + e2m1 packing + swizzled SF layout.
|
||||
// The ONLY change vs that kernel is the pass-1 input: instead of reading activated_bf16
|
||||
// from gmem, it reads the interleaved gate/up GEMM1 output + the LoRA delta and computes
|
||||
// silu(up)*gate on the fly, rounds to bf16 (matching the standalone activation kernel's
|
||||
// bf16 output exactly), caches in smem, and also writes activation_lora_input. The down
|
||||
// GEMM input (activated_bf16) is therefore never materialized to HBM.
|
||||
//
|
||||
// Because the activated value is rounded to bf16 before quantization (same as the separate
|
||||
// activation kernel) and silu matches, the fp4 / SF / per_token_sf / activation_lora_input
|
||||
// outputs are BITWISE-identical to the unfused activation -> quant#2 chain.
|
||||
#pragma once
|
||||
|
||||
#include <cub/block/block_reduce.cuh>
|
||||
#include <cuda/functional>
|
||||
|
||||
#include "nv_internal/tensorrt_llm/kernels/quantization_utils.cuh"
|
||||
#include <cstdint>
|
||||
#include <cuda_bf16.h>
|
||||
#include <optional>
|
||||
|
||||
namespace flashinfer {
|
||||
namespace sgl_fused_act_quant {
|
||||
|
||||
namespace tk = tensorrt_llm::kernels;
|
||||
|
||||
// Same silu as moe::dev::activation (trtllm_fused_moe_dev_kernel.cu:55): x / (1 + exp(-x)).
|
||||
inline __device__ float fused_silu(float x) {
|
||||
return x / (1.0f + expf(-x));
|
||||
}
|
||||
|
||||
// One block per expanded row. gateUp is the column-interleaved GEMM1 output (g0,u0,g1,u1,...)
|
||||
// indexed by permutedIdx; loraDelta is the contiguous [gate|up] delta indexed by expandedIdx.
|
||||
template <uint32_t BLOCK_SIZE, tensorrt_llm::QuantizationSFLayout SF_LAYOUT, bool DISABLE_FP4_FAST_MATH>
|
||||
__global__ void fusedActivationQuantKernel(
|
||||
int m, // numTokens * topK (number of expanded rows)
|
||||
int innerHalf, // inter == n (output width per row); must be a multiple of 16
|
||||
int innerDim, // gate_up_n == 2 * innerHalf
|
||||
__nv_bfloat16 const* __restrict__ gateUp, // interleaved gate/up, [.., innerDim] by permutedIdx
|
||||
__nv_bfloat16 const* __restrict__ loraDelta, // [.., innerDim] by expandedIdx, may be null
|
||||
__nv_bfloat16* __restrict__ loraInputOut, // [.., innerHalf] by expandedIdx, may be null
|
||||
int32_t const* __restrict__ expandedIdxToPermutedIdx,
|
||||
float globalScaleInv,
|
||||
uint8_t* __restrict__ weightOutput, // fp4 [.., innerHalf/2] by permutedIdx
|
||||
uint8_t* __restrict__ scaleOutput, // swizzled e4m3 SF
|
||||
float* __restrict__ perTokenScaleOutput) {
|
||||
constexpr int SF_VEC_SIZE = 16;
|
||||
using InType = tk::PackedVec<__nv_bfloat16, SF_VEC_SIZE>; // 16 bf16 == 8 __nv_bfloat162
|
||||
using PackedFp4Type = uint64_t; // SF_VEC_SIZE == 16
|
||||
|
||||
int const expandedIdx = blockIdx.x;
|
||||
if (expandedIdx >= m) return;
|
||||
int const permutedIdx = expandedIdxToPermutedIdx[expandedIdx];
|
||||
int const num_vecs_per_row = innerHalf / SF_VEC_SIZE;
|
||||
int64_t const liBaseRow = (int64_t)expandedIdx * innerHalf;
|
||||
|
||||
// Padding row: the separate activation kernel writes 0 to activation_lora_input and skips
|
||||
// the quant outputs. Mirror that, then return.
|
||||
if (permutedIdx < 0) {
|
||||
if (loraInputOut != nullptr) {
|
||||
InType z;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < SF_VEC_SIZE / 2; ++i)
|
||||
z.elts[i] = __float2bfloat162_rn(0.0f);
|
||||
for (int vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
|
||||
*reinterpret_cast<InType*>(&loraInputOut[liBaseRow + (int64_t)vecIdx * SF_VEC_SIZE]) = z;
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
int64_t const permBase = (int64_t)permutedIdx * innerDim; // gate_up row (interleaved)
|
||||
int64_t const expBase = (int64_t)expandedIdx * innerDim; // delta row (contiguous gate|up)
|
||||
(void)DISABLE_FP4_FAST_MATH;
|
||||
|
||||
// 1 SF block (16 outputs) per thread, held in registers across the amax barrier (no smem cache):
|
||||
// requires num_vecs_per_row <= BLOCK_SIZE (inter=2048 -> 128 == BLOCK_SIZE). With
|
||||
// CVT_ELTS_PER_THREAD == SF_VEC_SIZE the cvt needs no cross-thread shuffle, so masking is safe.
|
||||
int const vecIdx = threadIdx.x;
|
||||
bool const active = vecIdx < num_vecs_per_row;
|
||||
|
||||
InType vec;
|
||||
float localAmax = 0.f;
|
||||
if (active) {
|
||||
int const h0 = vecIdx * SF_VEC_SIZE;
|
||||
__nv_bfloat16 const* g = gateUp + permBase + (int64_t)2 * h0; // 32 interleaved bf16
|
||||
__nv_bfloat16 const* dlo = loraDelta + expBase + h0; // silu-arg delta (lower half)
|
||||
__nv_bfloat16 const* dhi = loraDelta + expBase + innerHalf + h0; // multiplier delta (upper half)
|
||||
__nv_bfloat162 amax2 = __float2bfloat162_rn(0.0f);
|
||||
union {
|
||||
int4 v[4];
|
||||
__nv_bfloat16 b[32];
|
||||
} gu;
|
||||
union {
|
||||
int4 v[2];
|
||||
__nv_bfloat16 b[16];
|
||||
} dl, dh;
|
||||
int4 const* gp = reinterpret_cast<int4 const*>(g);
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
gu.v[k] = gp[k];
|
||||
if (loraDelta != nullptr) {
|
||||
int4 const* dlp = reinterpret_cast<int4 const*>(dlo);
|
||||
int4 const* dhp = reinterpret_cast<int4 const*>(dhi);
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2; ++k) {
|
||||
dl.v[k] = dlp[k];
|
||||
dh.v[k] = dhp[k];
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < SF_VEC_SIZE / 2; ++i) { // 8 bf162 = 16 output elements
|
||||
int const j0 = 2 * i, j1 = 2 * i + 1;
|
||||
float even0 = (float)gu.b[2 * j0], odd0 = (float)gu.b[2 * j0 + 1];
|
||||
float even1 = (float)gu.b[2 * j1], odd1 = (float)gu.b[2 * j1 + 1];
|
||||
float a0 = odd0, b0 = even0, a1 = odd1, b1 = even1;
|
||||
if (loraDelta != nullptr) {
|
||||
a0 += (float)dl.b[j0];
|
||||
b0 += (float)dh.b[j0];
|
||||
a1 += (float)dl.b[j1];
|
||||
b1 += (float)dh.b[j1];
|
||||
}
|
||||
float act0 = fused_silu(a0) * b0;
|
||||
float act1 = fused_silu(a1) * b1;
|
||||
__nv_bfloat162 e = __float22bfloat162_rn(make_float2(act0, act1));
|
||||
vec.elts[i] = e;
|
||||
amax2 = __hmax2(amax2, __habs2(e));
|
||||
}
|
||||
localAmax = (float)__hmax(amax2.x, amax2.y);
|
||||
if (loraInputOut != nullptr) {
|
||||
*reinterpret_cast<InType*>(&loraInputOut[liBaseRow + h0]) = vec;
|
||||
}
|
||||
}
|
||||
|
||||
// ---- per-token scale: blockReduce amax, broadcast via smem (no gmem round-trip) ----
|
||||
using BlockReduce = cub::BlockReduce<float, BLOCK_SIZE>;
|
||||
__shared__ typename BlockReduce::TempStorage tempStorage;
|
||||
__shared__ float sScale;
|
||||
float const globalAmax = BlockReduce(tempStorage).Reduce(localAmax, cuda::maximum<>{});
|
||||
if (threadIdx.x == 0) {
|
||||
float const pts = globalAmax * globalScaleInv;
|
||||
perTokenScaleOutput[permutedIdx] = pts;
|
||||
sScale = pts;
|
||||
}
|
||||
__syncthreads();
|
||||
float const globalEncodeScale = tk::reciprocal_approximate_ftz(sScale);
|
||||
|
||||
// ---- quantize from registers (cvt computes the per-16 e4m3 block scale internally) ----
|
||||
if (active) {
|
||||
uint8_t fp8Scale;
|
||||
// 5 template args on this flashinfer build: Type, SF_VEC_SIZE, CVT_ELTS_PER_THREAD,
|
||||
// UE8M0_SF=false, TE_EXACT_NVFP4=false (the default nvfp4 quant path).
|
||||
auto fp4Vals = tk::cvt_warp_fp16_to_fp4<__nv_bfloat16, SF_VEC_SIZE, SF_VEC_SIZE, false, false>(
|
||||
vec, globalEncodeScale, &fp8Scale);
|
||||
int64_t const vecOffset = (int64_t)permutedIdx * num_vecs_per_row + vecIdx;
|
||||
reinterpret_cast<PackedFp4Type*>(weightOutput)[vecOffset] = fp4Vals;
|
||||
|
||||
// Match nvfp4QuantAndPerTokenScaleKernel exactly (it passes the kernel's `m` as numRows).
|
||||
int64_t sfOffset;
|
||||
if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::LINEAR) {
|
||||
sfOffset = vecOffset;
|
||||
} else if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
|
||||
sfOffset = tk::get_sf_out_offset_128x4(/*batchIdx=*/0, permutedIdx, vecIdx, m, num_vecs_per_row);
|
||||
} else {
|
||||
sfOffset = tk::get_sf_out_offset_8x4(/*batchIdx=*/0, permutedIdx, vecIdx, m, num_vecs_per_row);
|
||||
}
|
||||
scaleOutput[sfOffset] = fp8Scale;
|
||||
}
|
||||
}
|
||||
|
||||
// Host launch: globalScaleInv = 1/448/6. BLOCK_SIZE must be >= innerHalf/16 (one SF block/thread).
|
||||
inline void launchFusedActivationQuant(
|
||||
int m,
|
||||
int innerHalf,
|
||||
int innerDim,
|
||||
__nv_bfloat16 const* gateUp,
|
||||
__nv_bfloat16 const* loraDelta,
|
||||
__nv_bfloat16* loraInputOut,
|
||||
int32_t const* expandedIdxToPermutedIdx,
|
||||
float globalScaleInv,
|
||||
uint8_t* weightOutput,
|
||||
uint8_t* scaleOutput,
|
||||
float* perTokenScaleOutput,
|
||||
tensorrt_llm::QuantizationSFLayout sfLayout,
|
||||
bool disableFp4FastMath,
|
||||
cudaStream_t stream) {
|
||||
// One SF block per thread, no stride loop: BLOCK_SIZE must cover innerHalf/16 (a fixed
|
||||
// 128 left cols [2048,inter) unwritten at Inkling EP8's inter=3072 -> NaN from the down GEMM).
|
||||
uint32_t const numVecs = static_cast<uint32_t>(innerHalf) / 16;
|
||||
auto dispatchBlock = [&](auto blockTag) {
|
||||
constexpr uint32_t BLOCK_SIZE = decltype(blockTag)::value;
|
||||
dim3 const grid(m), block(BLOCK_SIZE);
|
||||
auto launch = [&](auto layoutTag, auto fastMathTag) {
|
||||
fusedActivationQuantKernel<BLOCK_SIZE, decltype(layoutTag)::value, decltype(fastMathTag)::value>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
m,
|
||||
innerHalf,
|
||||
innerDim,
|
||||
gateUp,
|
||||
loraDelta,
|
||||
loraInputOut,
|
||||
expandedIdxToPermutedIdx,
|
||||
globalScaleInv,
|
||||
weightOutput,
|
||||
scaleOutput,
|
||||
perTokenScaleOutput);
|
||||
};
|
||||
auto withFastMath = [&](auto layoutTag) {
|
||||
if (disableFp4FastMath) {
|
||||
launch(layoutTag, std::integral_constant<bool, true>{});
|
||||
} else {
|
||||
launch(layoutTag, std::integral_constant<bool, false>{});
|
||||
}
|
||||
};
|
||||
if (sfLayout == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
|
||||
withFastMath(
|
||||
std::integral_constant<
|
||||
tensorrt_llm::QuantizationSFLayout,
|
||||
tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4>{});
|
||||
} else if (sfLayout == tensorrt_llm::QuantizationSFLayout::LINEAR) {
|
||||
withFastMath(
|
||||
std::integral_constant<tensorrt_llm::QuantizationSFLayout, tensorrt_llm::QuantizationSFLayout::LINEAR>{});
|
||||
} else {
|
||||
withFastMath(
|
||||
std::integral_constant<
|
||||
tensorrt_llm::QuantizationSFLayout,
|
||||
tensorrt_llm::QuantizationSFLayout::SWIZZLED_8x4>{});
|
||||
}
|
||||
};
|
||||
if (numVecs <= 128) {
|
||||
dispatchBlock(std::integral_constant<uint32_t, 128>{});
|
||||
} else if (numVecs <= 256) {
|
||||
dispatchBlock(std::integral_constant<uint32_t, 256>{});
|
||||
} else if (numVecs <= 512) {
|
||||
dispatchBlock(std::integral_constant<uint32_t, 512>{});
|
||||
} else {
|
||||
// Callers guard on numVecs <= 512 and fall back to the unfused chain.
|
||||
dispatchBlock(std::integral_constant<uint32_t, 1024>{});
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace sgl_fused_act_quant
|
||||
} // namespace flashinfer
|
||||
@@ -1,307 +0,0 @@
|
||||
// Fused permute + NvFP4-per-token-quant for the FP4 MoE-LoRA gate_up path.
|
||||
//
|
||||
// Background (decode bs64, EP8): the plain path runs `permuteKernel` (gather bf16 hidden into the
|
||||
// padded [max_padded, hidden] permuted buffer) then `nvfp4QuantAndPerTokenScaleKernel` over ALL
|
||||
// max_padded rows. At decode only num_tokens*top_k of those rows are real (the rest are padding),
|
||||
// so both kernels waste ~6x of their work on padding, and the bf16 permuted buffer is a full
|
||||
// HBM round-trip (written by permute, read back by quant).
|
||||
//
|
||||
// This kernel fuses the two: it reads the UN-permuted hidden, NvFP4-quantizes each (token,expert)
|
||||
// pair's row, and scatter-writes fp4 + swizzled block-sf + per-token-sf directly to that pair's
|
||||
// permuted position. It iterates only the num_tokens*top_k real pairs (skipping pad), and never
|
||||
// materializes the bf16 permuted buffer.
|
||||
//
|
||||
// It mirrors `nvfp4QuantAndPerTokenScaleKernel` (quantization.cuh) — same amax, same per-token-scale
|
||||
// recipe, same `cvt_warp_fp16_to_fp4`, same swizzled-sf offset (`get_sf_out_offset_8x4`) — with the
|
||||
// single `rowIdx` split into a READ row (the unpermuted source token) and a WRITE row (the permuted
|
||||
// destination). For the valid rows the result is BITWISE-identical to the plain permute->quant chain
|
||||
// (the chain's quant reads permuted_hidden[writeRow], filled by permute from hidden[readRow]; we
|
||||
// read hidden[readRow] directly), verified by the bench's fused-vs-old guard.
|
||||
//
|
||||
// PER-TOKEN-SCALE BRANCH: uses TE_EXACT (globalEncodeScale = __fdiv_rn(globalScale, globalAmax),
|
||||
// stored scale = 1/globalEncodeScale, cvt TE_EXACT_NVFP4=true). This matches the installed
|
||||
// flashinfer 0.6.11.post1, whose DISPATCH macro hard-codes the bf16 kernel to TE_EXACT_NVFP4=true
|
||||
// (quantization.cu DISPATCH_NVP4_QUANT_AND_PER_TOKEN_SCALE_KERNEL). NOTE: if a future flashinfer
|
||||
// reverts the bf16 path to the fast-math branch (reciprocal_approximate / TE_EXACT=false), or if
|
||||
// FLASHINFER_NVFP4_4OVER6 is enabled, this fused kernel would diverge — re-validate the bench's
|
||||
// fused-vs-old bitwise guard against the deployed flashinfer before trusting it there.
|
||||
//
|
||||
// Two variants (both kept, selectable, for cross-scenario perf comparison):
|
||||
// - no-dedup: grid over the num_tokens*top_k pairs; each block re-reads+re-quantizes its source
|
||||
// token and writes 1 destination (more blocks -> better occupancy at tiny decode sizes).
|
||||
// - dedup: grid over num_tokens; each block reads+quantizes its token once and scatter-writes
|
||||
// to all of that token's (valid) permuted destinations (no redundant quant, fewer blocks).
|
||||
//
|
||||
// Helpers are pulled from quantization_utils.cuh (cvt_warp_fp16_to_fp4 / get_sf_out_offset_* /
|
||||
// PackedVec / reciprocal_approximate_ftz) rather than quantization.cuh, because the latter pulls in
|
||||
// nv_internal/.../common/cudaUtils.h, which ODR-conflicts with the flashinfer/trtllm/common twin
|
||||
// already in trtllm_fused_moe_kernel_launcher.cu's TU. loadPackedVec lives in quantization.cuh, so
|
||||
// we do a direct aligned PackedVec load instead.
|
||||
#pragma once
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
#include "nv_internal/tensorrt_llm/kernels/quantization_utils.cuh"
|
||||
#include <cstdint>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <optional>
|
||||
#include <type_traits>
|
||||
|
||||
namespace sgl_fused_permute_quant {
|
||||
|
||||
namespace tk = tensorrt_llm::kernels;
|
||||
|
||||
// Quantize source row `readRow` of `input` (unpermuted) and write fp4 + block-sf + per-token-sf to
|
||||
// destination row `writeRow` of the permuted outputs. `numRowsSf` is the SF buffer's row count
|
||||
// (= max_padded), matching the plain quant's `m` arg to get_sf_out_offset_*.
|
||||
template <typename T, uint32_t BLOCK_SIZE, tensorrt_llm::QuantizationSFLayout SF_LAYOUT>
|
||||
__device__ __forceinline__ void fused_quant_one_row(
|
||||
uint32_t n,
|
||||
T const* input,
|
||||
int readRow,
|
||||
int writeRow,
|
||||
int numRowsSf,
|
||||
float globalScaleInv,
|
||||
uint8_t* weightOutput,
|
||||
uint8_t* scaleOutput,
|
||||
float* perTokenScaleOutput) {
|
||||
constexpr int SF_VEC_SIZE = 16;
|
||||
constexpr int ELTS_PER_THREAD = 16;
|
||||
using InType = tk::PackedVec<T, ELTS_PER_THREAD>;
|
||||
using PackedFp4Type = std::conditional_t<ELTS_PER_THREAD == 16, uint64_t, uint32_t>;
|
||||
uint32_t const num_vecs_per_row = (n + ELTS_PER_THREAD - 1) / ELTS_PER_THREAD;
|
||||
uint32_t const num_sf_vecs_per_row = (n + SF_VEC_SIZE - 1) / SF_VEC_SIZE;
|
||||
InType const* inBase = reinterpret_cast<InType const*>(input);
|
||||
|
||||
// ---- pass 1: per-row amax over the (unpermuted) source row ----
|
||||
float localAmax = 0.f;
|
||||
for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
|
||||
InType vec_in = inBase[static_cast<int64_t>(readRow) * num_vecs_per_row + vecIdx];
|
||||
std::remove_reference_t<decltype(vec_in.elts[0])> a(0.f, 0.f);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < ELTS_PER_THREAD / 2; ++i) {
|
||||
a = __hmax2(a, __habs2(vec_in.elts[i]));
|
||||
}
|
||||
localAmax = fmaxf(localAmax, static_cast<float>(__hmax(a.x, a.y)));
|
||||
}
|
||||
|
||||
using BlockReduce = cub::BlockReduce<float, BLOCK_SIZE>;
|
||||
__shared__ typename BlockReduce::TempStorage tempStorage;
|
||||
float const globalAmax = BlockReduce(tempStorage).Reduce(localAmax, cuda::maximum<>{});
|
||||
|
||||
// ---- per-token scale (TE_EXACT branch — production instantiates TE_EXACT_NVFP4=true for bf16,
|
||||
// quantization.cu:247): globalEncodeScale = globalScale/globalAmax (exact __fdiv_rn), stored
|
||||
// per-token scale = 1/globalEncodeScale. __shared__ scalar replaces the gmem round-trip
|
||||
// (bit-identical: an fp32 store->load doesn't change the value). ----
|
||||
__shared__ float sEncodeScale;
|
||||
if (threadIdx.x == 0) {
|
||||
float const globalScale = __fdiv_rn(1.0f, globalScaleInv);
|
||||
float const rowEncodeScale = globalAmax != 0.0f ? fminf(__fdiv_rn(globalScale, globalAmax), FLT_MAX) : FLT_MAX;
|
||||
sEncodeScale = rowEncodeScale != 0.0f ? rowEncodeScale : 1.0f;
|
||||
}
|
||||
__syncthreads();
|
||||
float const globalEncodeScale = sEncodeScale;
|
||||
float const perTokenScale = __fdiv_rn(1.0f, globalEncodeScale);
|
||||
if (threadIdx.x == 0) perTokenScaleOutput[writeRow] = perTokenScale;
|
||||
|
||||
// ---- pass 2: quantize + scatter-write to the permuted destination ----
|
||||
for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
|
||||
InType vec_in = inBase[static_cast<int64_t>(readRow) * num_vecs_per_row + vecIdx];
|
||||
uint8_t fp8Scale;
|
||||
auto fp4Vals = tk::cvt_warp_fp16_to_fp4<
|
||||
T,
|
||||
SF_VEC_SIZE,
|
||||
ELTS_PER_THREAD,
|
||||
/*UE8M0_SF=*/false,
|
||||
/*TE_EXACT_NVFP4=*/true>(vec_in, globalEncodeScale, &fp8Scale);
|
||||
reinterpret_cast<PackedFp4Type*>(weightOutput)[static_cast<int64_t>(writeRow) * num_vecs_per_row + vecIdx] =
|
||||
fp4Vals;
|
||||
|
||||
int64_t sfOffset;
|
||||
if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::LINEAR) {
|
||||
sfOffset = static_cast<int64_t>(writeRow) * num_sf_vecs_per_row + vecIdx;
|
||||
} else if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
|
||||
sfOffset = tk::get_sf_out_offset_128x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
|
||||
} else {
|
||||
sfOffset = tk::get_sf_out_offset_8x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
|
||||
}
|
||||
scaleOutput[sfOffset] = fp8Scale;
|
||||
}
|
||||
}
|
||||
|
||||
// no-dedup: grid.x = num_tokens*top_k (one block per (token,expert) pair).
|
||||
template <typename T, uint32_t BLOCK_SIZE, tensorrt_llm::QuantizationSFLayout SF_LAYOUT>
|
||||
__global__ void fusedPermuteNvfp4QuantKernel(
|
||||
uint32_t numPairs,
|
||||
uint32_t n,
|
||||
uint32_t topK,
|
||||
int numRowsSf,
|
||||
T const* input,
|
||||
float globalScaleInv,
|
||||
int32_t const* expandedIdxToPermutedIdx,
|
||||
uint8_t* weightOutput,
|
||||
uint8_t* scaleOutput,
|
||||
float* perTokenScaleOutput) {
|
||||
uint32_t const expandedIdx = blockIdx.x;
|
||||
if (expandedIdx >= numPairs) return;
|
||||
int const writeRow = expandedIdxToPermutedIdx[expandedIdx];
|
||||
if (writeRow < 0) return;
|
||||
int const readRow = static_cast<int>(expandedIdx / topK);
|
||||
fused_quant_one_row<T, BLOCK_SIZE, SF_LAYOUT>(
|
||||
n, input, readRow, writeRow, numRowsSf, globalScaleInv, weightOutput, scaleOutput, perTokenScaleOutput);
|
||||
}
|
||||
|
||||
// dedup: grid.x = num_tokens (one block per source token, scatter to its top_k destinations).
|
||||
template <typename T, uint32_t BLOCK_SIZE, tensorrt_llm::QuantizationSFLayout SF_LAYOUT>
|
||||
__global__ void fusedPermuteNvfp4QuantDedupKernel(
|
||||
uint32_t numTokens,
|
||||
uint32_t n,
|
||||
uint32_t topK,
|
||||
int numRowsSf,
|
||||
T const* input,
|
||||
float globalScaleInv,
|
||||
int32_t const* expandedIdxToPermutedIdx,
|
||||
uint8_t* weightOutput,
|
||||
uint8_t* scaleOutput,
|
||||
float* perTokenScaleOutput) {
|
||||
constexpr int SF_VEC_SIZE = 16;
|
||||
constexpr int ELTS_PER_THREAD = 16;
|
||||
using InType = tk::PackedVec<T, ELTS_PER_THREAD>;
|
||||
using PackedFp4Type = std::conditional_t<ELTS_PER_THREAD == 16, uint64_t, uint32_t>;
|
||||
uint32_t const token = blockIdx.x;
|
||||
if (token >= numTokens) return;
|
||||
uint32_t const num_vecs_per_row = (n + ELTS_PER_THREAD - 1) / ELTS_PER_THREAD;
|
||||
uint32_t const num_sf_vecs_per_row = (n + SF_VEC_SIZE - 1) / SF_VEC_SIZE;
|
||||
InType const* inBase = reinterpret_cast<InType const*>(input);
|
||||
|
||||
// pass 1: amax over the source token row (read once).
|
||||
float localAmax = 0.f;
|
||||
for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
|
||||
InType vec_in = inBase[static_cast<int64_t>(token) * num_vecs_per_row + vecIdx];
|
||||
std::remove_reference_t<decltype(vec_in.elts[0])> a(0.f, 0.f);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < ELTS_PER_THREAD / 2; ++i) {
|
||||
a = __hmax2(a, __habs2(vec_in.elts[i]));
|
||||
}
|
||||
localAmax = fmaxf(localAmax, static_cast<float>(__hmax(a.x, a.y)));
|
||||
}
|
||||
using BlockReduce = cub::BlockReduce<float, BLOCK_SIZE>;
|
||||
__shared__ typename BlockReduce::TempStorage tempStorage;
|
||||
float const globalAmax = BlockReduce(tempStorage).Reduce(localAmax, cuda::maximum<>{});
|
||||
|
||||
// TE_EXACT per-token scale (matches production; see fused_quant_one_row).
|
||||
__shared__ float sEncodeScale;
|
||||
if (threadIdx.x == 0) {
|
||||
float const globalScale = __fdiv_rn(1.0f, globalScaleInv);
|
||||
float const rowEncodeScale = globalAmax != 0.0f ? fminf(__fdiv_rn(globalScale, globalAmax), FLT_MAX) : FLT_MAX;
|
||||
sEncodeScale = rowEncodeScale != 0.0f ? rowEncodeScale : 1.0f;
|
||||
}
|
||||
__syncthreads();
|
||||
float const globalEncodeScale = sEncodeScale;
|
||||
float const perTokenScale = __fdiv_rn(1.0f, globalEncodeScale);
|
||||
|
||||
// per-token scale -> each (valid) destination (top_k small; first top_k threads write).
|
||||
if (threadIdx.x < topK) {
|
||||
int const writeRow = expandedIdxToPermutedIdx[token * topK + threadIdx.x];
|
||||
if (writeRow >= 0) perTokenScaleOutput[writeRow] = perTokenScale;
|
||||
}
|
||||
|
||||
// pass 2: quantize each vec once, scatter to all valid destinations.
|
||||
for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
|
||||
InType vec_in = inBase[static_cast<int64_t>(token) * num_vecs_per_row + vecIdx];
|
||||
uint8_t fp8Scale;
|
||||
auto fp4Vals = tk::cvt_warp_fp16_to_fp4<
|
||||
T,
|
||||
SF_VEC_SIZE,
|
||||
ELTS_PER_THREAD,
|
||||
/*UE8M0_SF=*/false,
|
||||
/*TE_EXACT_NVFP4=*/true>(vec_in, globalEncodeScale, &fp8Scale);
|
||||
#pragma unroll 1
|
||||
for (uint32_t k = 0; k < topK; ++k) {
|
||||
int const writeRow = expandedIdxToPermutedIdx[token * topK + k];
|
||||
if (writeRow < 0) continue;
|
||||
reinterpret_cast<PackedFp4Type*>(weightOutput)[static_cast<int64_t>(writeRow) * num_vecs_per_row + vecIdx] =
|
||||
fp4Vals;
|
||||
int64_t sfOffset;
|
||||
if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::LINEAR) {
|
||||
sfOffset = static_cast<int64_t>(writeRow) * num_sf_vecs_per_row + vecIdx;
|
||||
} else if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
|
||||
sfOffset = tk::get_sf_out_offset_128x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
|
||||
} else {
|
||||
sfOffset = tk::get_sf_out_offset_8x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
|
||||
}
|
||||
scaleOutput[sfOffset] = fp8Scale;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Launcher. `dedup` picks the variant. `n` (= hidden) must be a multiple of 16. `numRowsSf` is the
|
||||
// SF buffer's row count (max_padded).
|
||||
template <typename T>
|
||||
void invokeFusedPermuteNvfp4Quant(
|
||||
uint32_t numTokens,
|
||||
uint32_t topK,
|
||||
uint32_t n,
|
||||
int numRowsSf,
|
||||
T const* input,
|
||||
float globalScaleInv,
|
||||
int32_t const* expandedIdxToPermutedIdx,
|
||||
uint8_t* weightOutput,
|
||||
uint8_t* scaleOutput,
|
||||
float* perTokenScaleOutput,
|
||||
tensorrt_llm::QuantizationSFLayout sfLayout,
|
||||
bool dedup,
|
||||
cudaStream_t stream) {
|
||||
// [opt] Occupancy tuning (ncu: kernel is occupancy-bound, not DRAM-bound — DRAM <1%, achieved
|
||||
// occupancy was 19.8% no-dedup / 5.5% dedup at BLOCK_SIZE=128). The dedup variant launches only
|
||||
// num_tokens CTAs (=64 at decode bs64), so it is the most CTA-starved; widening the block raises
|
||||
// threads/CTA and hides the per-row amax-reduction + scatter latency. Decode bs64 dedup sweep:
|
||||
// 128 -> 5.52us, 256 -> 4.09us, 512 -> 3.71us. 512 is the chosen default (the prod path uses
|
||||
// dedup). (7168/16 = 448 vecs/row, so >448 threads idle on the tail, but the win dominates.)
|
||||
constexpr uint32_t BLOCK_SIZE = 512;
|
||||
dim3 const block(BLOCK_SIZE);
|
||||
|
||||
auto dispatch = [&](auto layoutTag) {
|
||||
constexpr tensorrt_llm::QuantizationSFLayout LAYOUT = decltype(layoutTag)::value;
|
||||
if (dedup) {
|
||||
dim3 const grid(numTokens);
|
||||
fusedPermuteNvfp4QuantDedupKernel<T, BLOCK_SIZE, LAYOUT><<<grid, block, 0, stream>>>(
|
||||
numTokens,
|
||||
n,
|
||||
topK,
|
||||
numRowsSf,
|
||||
input,
|
||||
globalScaleInv,
|
||||
expandedIdxToPermutedIdx,
|
||||
weightOutput,
|
||||
scaleOutput,
|
||||
perTokenScaleOutput);
|
||||
} else {
|
||||
dim3 const grid(numTokens * topK);
|
||||
fusedPermuteNvfp4QuantKernel<T, BLOCK_SIZE, LAYOUT><<<grid, block, 0, stream>>>(
|
||||
numTokens * topK,
|
||||
n,
|
||||
topK,
|
||||
numRowsSf,
|
||||
input,
|
||||
globalScaleInv,
|
||||
expandedIdxToPermutedIdx,
|
||||
weightOutput,
|
||||
scaleOutput,
|
||||
perTokenScaleOutput);
|
||||
}
|
||||
};
|
||||
|
||||
if (sfLayout == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
|
||||
dispatch(
|
||||
std::integral_constant<
|
||||
tensorrt_llm::QuantizationSFLayout,
|
||||
tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4>{});
|
||||
} else {
|
||||
dispatch(
|
||||
std::integral_constant<tensorrt_llm::QuantizationSFLayout, tensorrt_llm::QuantizationSFLayout::SWIZZLED_8x4>{});
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace sgl_fused_permute_quant
|
||||
-4380
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
-472
@@ -1,472 +0,0 @@
|
||||
/*
|
||||
* Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h"
|
||||
#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/SfLayoutDecl.h"
|
||||
#include <cuda.h>
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
// #include <cuda_runtime_api.h>
|
||||
#include <cutlass/cutlass.h>
|
||||
#include <cutlass/numeric_size.h>
|
||||
#include <cutlass/numeric_types.h>
|
||||
|
||||
#include "flashinfer/exception.h"
|
||||
// #include <tensorrt_llm/common/assert.h>
|
||||
#include "flashinfer/trtllm/common/cudaUtils.h"
|
||||
#include "tensorrt_llm/common/logger.h"
|
||||
|
||||
namespace moe::dev {
|
||||
|
||||
#define CHECK_CUDA_ERROR(cmd) \
|
||||
do { \
|
||||
cudaError_t e = cmd; \
|
||||
if (e != cudaSuccess) { \
|
||||
std::cout << "CUDA error in " << __FILE__ << ":" << __LINE__ << " executing '" << #cmd \
|
||||
<< "': " << cudaGetErrorString(e); \
|
||||
} \
|
||||
FLASHINFER_CHECK(e == cudaSuccess, "Got CUDA error. See above for details."); \
|
||||
} while (0)
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#define LAUNCH_ESC(...) __VA_ARGS__
|
||||
|
||||
#define LAUNCH_PDL(data, coopLaunch, types, kernel, numBlocks, numThreads, smemSize, stream) \
|
||||
cudaLaunchConfig_t config{}; \
|
||||
config.gridDim = numBlocks; \
|
||||
config.blockDim = numThreads; \
|
||||
config.dynamicSmemBytes = smemSize; \
|
||||
config.stream = (cudaStream_t)stream; \
|
||||
\
|
||||
cudaLaunchAttribute attributes[2] = {}; \
|
||||
attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \
|
||||
attributes[0].val.programmaticStreamSerializationAllowed = int(data.mUsePdl); \
|
||||
attributes[1].id = cudaLaunchAttributeCooperative; \
|
||||
attributes[1].val.cooperative = int(coopLaunch); \
|
||||
config.attrs = attributes; \
|
||||
config.numAttrs = 2; \
|
||||
if (data.mUsePdl) { \
|
||||
auto params = KernelParams<types, true>::setKernelParams(data); \
|
||||
auto kernelTyped = kernel<KernelParams<types, true>>; \
|
||||
if (smemSize > 48 * 1024) \
|
||||
CHECK_CUDA_ERROR(cudaFuncSetAttribute(kernelTyped, cudaFuncAttributeMaxDynamicSharedMemorySize, smemSize)); \
|
||||
CHECK_CUDA_ERROR(cudaLaunchKernelEx(&config, kernelTyped, params)); \
|
||||
} else { \
|
||||
auto params = KernelParams<types, false>::setKernelParams(data); \
|
||||
auto kernelTyped = kernel<KernelParams<types, false>>; \
|
||||
if (smemSize > 48 * 1024) \
|
||||
CHECK_CUDA_ERROR(cudaFuncSetAttribute(kernelTyped, cudaFuncAttributeMaxDynamicSharedMemorySize, smemSize)); \
|
||||
CHECK_CUDA_ERROR(cudaLaunchKernelEx(&config, kernelTyped, params)); \
|
||||
}
|
||||
|
||||
#define LAUNCH(data, kernel, numBlocks, numThreads, smemSize, stream) \
|
||||
if (data.mDtypeElt == tg::Dtype::Fp16) { \
|
||||
LAUNCH_PDL(data, false, cutlass::half_t, kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (data.mDtypeElt == tg::Dtype::E4m3) { \
|
||||
LAUNCH_PDL(data, false, cutlass::float_e4m3_t, kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (data.mDtypeElt == tg::Dtype::Bfloat16) { \
|
||||
LAUNCH_PDL(data, false, cutlass::bfloat16_t, kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else { \
|
||||
FLASHINFER_WARN("Unsupported dtypeElt"); \
|
||||
}
|
||||
|
||||
#define LAUNCH_NUM_TOKENS_PER_CTA(data, type, numTokensPerCta, kernel, numBlocks, numThreads, smemSize, stream) \
|
||||
if (numTokensPerCta == 4) { \
|
||||
LAUNCH_PDL(data, false, LAUNCH_ESC(type, 4), kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (numTokensPerCta == 2) { \
|
||||
LAUNCH_PDL(data, false, LAUNCH_ESC(type, 2), kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (numTokensPerCta == 1) { \
|
||||
LAUNCH_PDL(data, false, LAUNCH_ESC(type, 1), kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else { \
|
||||
FLASHINFER_WARN("Unsupported numTokensPerCta"); \
|
||||
}
|
||||
|
||||
#define LAUNCH_ACTIVATION(data, kernel, numTokensPerCta, numBlocks, numThreads, smemSize, stream) \
|
||||
if (data.mDtypeElt == tg::Dtype::Fp16) { \
|
||||
LAUNCH_NUM_TOKENS_PER_CTA( \
|
||||
data, cutlass::half_t, numTokensPerCta, kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (data.mDtypeElt == tg::Dtype::E4m3) { \
|
||||
LAUNCH_NUM_TOKENS_PER_CTA( \
|
||||
data, cutlass::float_e4m3_t, numTokensPerCta, kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (data.mDtypeElt == tg::Dtype::Bfloat16) { \
|
||||
LAUNCH_NUM_TOKENS_PER_CTA( \
|
||||
data, cutlass::bfloat16_t, numTokensPerCta, kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else { \
|
||||
FLASHINFER_WARN("Unsupported dtypeElt"); \
|
||||
}
|
||||
|
||||
#define LAUNCH_EXPW(data, kernel, topK, numBlocks, numThreads, smemSize, stream) \
|
||||
if (data.mDtypeElt == tg::Dtype::Fp16 && data.mDtypeExpW == tg::Dtype::Fp32) { \
|
||||
LAUNCH_PDL( \
|
||||
data, false, LAUNCH_ESC(cutlass::half_t, float, topK), kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (data.mDtypeElt == tg::Dtype::E4m3 && data.mDtypeExpW == tg::Dtype::Fp32) { \
|
||||
LAUNCH_PDL( \
|
||||
data, false, LAUNCH_ESC(cutlass::float_e4m3_t, float, topK), kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (data.mDtypeElt == tg::Dtype::Bfloat16 && data.mDtypeExpW == tg::Dtype::Fp32) { \
|
||||
LAUNCH_PDL( \
|
||||
data, false, LAUNCH_ESC(cutlass::bfloat16_t, float, topK), kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (data.mDtypeElt == tg::Dtype::Fp16 && data.mDtypeExpW == tg::Dtype::Bfloat16) { \
|
||||
LAUNCH_PDL( \
|
||||
data, \
|
||||
false, \
|
||||
LAUNCH_ESC(cutlass::half_t, cutlass::bfloat16_t, topK), \
|
||||
kernel, \
|
||||
numBlocks, \
|
||||
numThreads, \
|
||||
smemSize, \
|
||||
stream); \
|
||||
} else if (data.mDtypeElt == tg::Dtype::E4m3 && data.mDtypeExpW == tg::Dtype::Bfloat16) { \
|
||||
LAUNCH_PDL( \
|
||||
data, \
|
||||
false, \
|
||||
LAUNCH_ESC(cutlass::float_e4m3_t, cutlass::bfloat16_t, topK), \
|
||||
kernel, \
|
||||
numBlocks, \
|
||||
numThreads, \
|
||||
smemSize, \
|
||||
stream); \
|
||||
} else if (data.mDtypeElt == tg::Dtype::Bfloat16 && data.mDtypeExpW == tg::Dtype::Bfloat16) { \
|
||||
LAUNCH_PDL( \
|
||||
data, \
|
||||
false, \
|
||||
LAUNCH_ESC(cutlass::bfloat16_t, cutlass::bfloat16_t, topK), \
|
||||
kernel, \
|
||||
numBlocks, \
|
||||
numThreads, \
|
||||
smemSize, \
|
||||
stream); \
|
||||
} else { \
|
||||
FLASHINFER_WARN("Unsupported pair"); \
|
||||
}
|
||||
|
||||
#define LAUNCH_TOPK_EXPW(data, kernel, numBlocks, numThreads, smemSize, stream) \
|
||||
if (data.topK % 4 == 0) { \
|
||||
LAUNCH_EXPW(data, kernel, 4, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (data.topK % 2 == 0) { \
|
||||
LAUNCH_EXPW(data, kernel, 2, numBlocks, numThreads, smemSize, stream); \
|
||||
} else { \
|
||||
LAUNCH_EXPW(data, kernel, 1, numBlocks, numThreads, smemSize, stream); \
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// NOTE: Old routing-specific macros (LAUNCH_TILEN, LAUNCH_ROUTING_LLAMA4,
|
||||
// LAUNCH_ROUTING_DEEPSEEK_*, LAUNCH_ROUTING_WITH_NUM_EXPERTS) have been moved to
|
||||
// RoutingDevKernel.h which uses the new template signature with runtime isPow2/UsePdl.
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace activation {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace tg = batchedGemm::trtllm::gen;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct Data {
|
||||
tg::Dtype mDtypeElt{tg::Dtype::Fp16};
|
||||
bool mUsePdl{false};
|
||||
bool mUseDeepSeekFp8{false};
|
||||
|
||||
void* inPtr;
|
||||
void* outPtr;
|
||||
float* inDqSfsPtr = nullptr;
|
||||
float* outDqSfsPtr = nullptr;
|
||||
cutlass::bfloat16_t const* gateUpLoraDeltaPtr = nullptr;
|
||||
cutlass::bfloat16_t* activationLoraInputOutPtr = nullptr;
|
||||
|
||||
// When true, inPtr holds the column-interleaved gate/up GEMM1 output
|
||||
// (g0,u0,g1,u1,...) and the kernel de-interleaves on read (x1=col 2k, x2=col 2k+1);
|
||||
// when false, inPtr is the contiguous [gate | up] layout. Default false keeps the
|
||||
// FP8 / non-LoRA semantics; the FP4 LoRA path sets it true to fuse the standalone
|
||||
// de-interleave kernel into this activation read.
|
||||
bool interleavedGateUpInput = false;
|
||||
|
||||
int32_t innerDim;
|
||||
int32_t numTokens;
|
||||
int32_t topK;
|
||||
int32_t* expandedIdxToPermutedIdx;
|
||||
|
||||
int32_t const* totalNumPaddedTokens;
|
||||
|
||||
// Bench/tuning knob: when > 0, overrides the activation launch grid.x (default
|
||||
// innerDim/128). The kernel's grid-stride hidden-dim loop makes any grid.x
|
||||
// produce bitwise-identical output, so a smaller grid.x removes empty blocks and
|
||||
// gives each thread a longer strip (more in-flight loads) without changing math.
|
||||
int32_t actGridXOverride = 0;
|
||||
|
||||
// Activation kernel variant: 0 = scalar activationKernel, 1 = vectorized
|
||||
// activationKernelOpt (128-bit gate/up + 64-bit delta/store, 4 pairs/thread).
|
||||
// Variant 1 applies only to the bf16 interleaved path with (innerDim/2)%4==0;
|
||||
// run() falls back to the scalar kernel otherwise. Output is bitwise-identical.
|
||||
int32_t actOptMode = 0;
|
||||
};
|
||||
|
||||
template <typename Type_, int32_t NumTokensPerCta_, bool UsePdl_>
|
||||
struct KernelParams {
|
||||
using Type = Type_;
|
||||
static constexpr int32_t NumTokensPerCta = NumTokensPerCta_;
|
||||
static constexpr bool UsePdl = UsePdl_;
|
||||
|
||||
Type const* inPtr;
|
||||
Type* outPtr;
|
||||
|
||||
float* inDqSfsPtr = nullptr;
|
||||
float* outDqSfsPtr = nullptr;
|
||||
cutlass::bfloat16_t const* gateUpLoraDeltaPtr = nullptr;
|
||||
cutlass::bfloat16_t* activationLoraInputOutPtr = nullptr;
|
||||
|
||||
bool interleavedGateUpInput = false;
|
||||
|
||||
int32_t innerDim;
|
||||
int32_t numTokens;
|
||||
int32_t topK;
|
||||
int32_t* expandedIdxToPermutedIdx;
|
||||
|
||||
int32_t const* totalNumPaddedTokens;
|
||||
|
||||
static KernelParams setKernelParams(Data const& data) {
|
||||
KernelParams params;
|
||||
|
||||
params.inPtr = (Type*)data.inPtr;
|
||||
params.outPtr = (Type*)data.outPtr;
|
||||
params.inDqSfsPtr = data.inDqSfsPtr;
|
||||
params.outDqSfsPtr = data.outDqSfsPtr;
|
||||
params.gateUpLoraDeltaPtr = data.gateUpLoraDeltaPtr;
|
||||
params.activationLoraInputOutPtr = data.activationLoraInputOutPtr;
|
||||
params.interleavedGateUpInput = data.interleavedGateUpInput;
|
||||
|
||||
params.expandedIdxToPermutedIdx = data.expandedIdxToPermutedIdx;
|
||||
|
||||
params.innerDim = data.innerDim;
|
||||
params.numTokens = data.numTokens;
|
||||
params.topK = data.topK;
|
||||
params.totalNumPaddedTokens = data.totalNumPaddedTokens;
|
||||
|
||||
return params;
|
||||
}
|
||||
};
|
||||
|
||||
void run(Data const& data, void* stream);
|
||||
|
||||
} // namespace activation
|
||||
|
||||
namespace convertsf {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace tg = batchedGemm::trtllm::gen;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct Data {
|
||||
bool mUsePdl{false};
|
||||
|
||||
void* inSfPtr = nullptr;
|
||||
void* outSfPtr = nullptr;
|
||||
int32_t hiddenDimSf;
|
||||
int32_t numTokens;
|
||||
tg::SfLayout sfLayoutSrc;
|
||||
tg::SfLayout sfLayoutDst;
|
||||
};
|
||||
|
||||
template <typename Type_, bool UsePdl_>
|
||||
struct KernelParams {
|
||||
using Type = Type_;
|
||||
static constexpr bool UsePdl = UsePdl_;
|
||||
|
||||
void const* inSfPtr = nullptr;
|
||||
void* outSfPtr = nullptr;
|
||||
int32_t hiddenDimSf;
|
||||
int32_t numTokens;
|
||||
tg::SfLayout sfLayoutSrc;
|
||||
tg::SfLayout sfLayoutDst;
|
||||
|
||||
static KernelParams setKernelParams(Data const& data) {
|
||||
KernelParams params;
|
||||
|
||||
params.inSfPtr = data.inSfPtr;
|
||||
params.outSfPtr = data.outSfPtr;
|
||||
params.hiddenDimSf = data.hiddenDimSf;
|
||||
params.numTokens = data.numTokens;
|
||||
params.sfLayoutSrc = data.sfLayoutSrc;
|
||||
params.sfLayoutDst = data.sfLayoutDst;
|
||||
|
||||
return params;
|
||||
}
|
||||
};
|
||||
|
||||
void run(Data const& data, void* stream);
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace convertsf
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace permute {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace tg = batchedGemm::trtllm::gen;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct Data {
|
||||
tg::Dtype mDtypeElt{tg::Dtype::Fp16};
|
||||
bool mUsePdl{false};
|
||||
bool mUseDeepSeekFp8{false};
|
||||
|
||||
void* inPtr;
|
||||
void* outPtr;
|
||||
float* inDqSfsPtr = nullptr;
|
||||
float* outDqSfsPtr = nullptr;
|
||||
int32_t* expandedIdxToPermutedIdx;
|
||||
int32_t hiddenDim;
|
||||
int32_t numTokens;
|
||||
int32_t topK;
|
||||
int32_t const* totalNumPaddedTokens;
|
||||
};
|
||||
|
||||
template <typename Type_, bool UsePdl_>
|
||||
struct KernelParams {
|
||||
using Type = Type_;
|
||||
static constexpr bool UsePdl = UsePdl_;
|
||||
|
||||
Type const* inPtr;
|
||||
Type* outPtr;
|
||||
float const* inDqSfsPtr;
|
||||
float* outDqSfsPtr;
|
||||
int32_t* expandedIdxToPermutedIdx;
|
||||
int32_t hiddenDim;
|
||||
int32_t numTokens;
|
||||
int32_t topK;
|
||||
int32_t const* totalNumPaddedTokens;
|
||||
bool useDeepSeekFp8;
|
||||
|
||||
static KernelParams setKernelParams(Data const& data) {
|
||||
KernelParams params;
|
||||
|
||||
params.inPtr = (Type*)data.inPtr;
|
||||
params.outPtr = (Type*)data.outPtr;
|
||||
params.inDqSfsPtr = data.inDqSfsPtr;
|
||||
params.outDqSfsPtr = data.outDqSfsPtr;
|
||||
params.expandedIdxToPermutedIdx = data.expandedIdxToPermutedIdx;
|
||||
params.hiddenDim = data.hiddenDim;
|
||||
params.numTokens = data.numTokens;
|
||||
params.topK = data.topK;
|
||||
params.totalNumPaddedTokens = data.totalNumPaddedTokens;
|
||||
params.useDeepSeekFp8 = data.mUseDeepSeekFp8;
|
||||
|
||||
return params;
|
||||
}
|
||||
};
|
||||
|
||||
void run(Data const& data, void* stream);
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace permute
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace finalize {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace tg = batchedGemm::trtllm::gen;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct Data {
|
||||
tg::Dtype mDtypeElt{tg::Dtype::Fp16};
|
||||
tg::Dtype mDtypeExpW{tg::Dtype::Bfloat16};
|
||||
bool mUsePdl{false};
|
||||
bool mUseDeepSeekFp8{false};
|
||||
|
||||
void* inPtr;
|
||||
void* outPtr;
|
||||
float* inDqSfsPtr = nullptr;
|
||||
float* outDqSfsPtr = nullptr;
|
||||
|
||||
void* expertWeightsPtr;
|
||||
int32_t* expandedIdxToPermutedIdx;
|
||||
|
||||
int32_t numTokens;
|
||||
int32_t numExperts;
|
||||
int32_t topK;
|
||||
// Hidden dimension output of MoE block. It is not padded.
|
||||
int32_t hiddenDim;
|
||||
// Hidden dimension output of FC2. It might be padded.
|
||||
int32_t hiddenDimPadded;
|
||||
int32_t const* totalNumPaddedTokens;
|
||||
};
|
||||
|
||||
template <typename Type_, typename TypeExpW_, int TopKUnrollFactor_, bool UsePdl_>
|
||||
struct KernelParams {
|
||||
using Type = Type_;
|
||||
using TypeExpW = TypeExpW_;
|
||||
static constexpr int TopKUnrollFactor = TopKUnrollFactor_;
|
||||
static constexpr bool UsePdl = UsePdl_;
|
||||
|
||||
Type const* inPtr;
|
||||
TypeExpW const* expertWeightsPtr;
|
||||
Type* outPtr;
|
||||
|
||||
float* inDqSfsPtr = nullptr;
|
||||
float* outDqSfsPtr = nullptr;
|
||||
|
||||
int32_t* expandedIdxToPermutedIdx;
|
||||
|
||||
int32_t hiddenDim;
|
||||
int32_t hiddenDimPadded;
|
||||
int32_t numTokens;
|
||||
int32_t numExperts;
|
||||
int32_t topK;
|
||||
int32_t const* totalNumPaddedTokens;
|
||||
|
||||
static KernelParams setKernelParams(Data const& data) {
|
||||
KernelParams params;
|
||||
|
||||
params.inPtr = (Type*)data.inPtr;
|
||||
params.expertWeightsPtr = (TypeExpW*)data.expertWeightsPtr;
|
||||
params.outPtr = (Type*)data.outPtr;
|
||||
params.inDqSfsPtr = data.inDqSfsPtr;
|
||||
params.outDqSfsPtr = data.outDqSfsPtr;
|
||||
|
||||
params.expandedIdxToPermutedIdx = data.expandedIdxToPermutedIdx;
|
||||
|
||||
params.hiddenDim = data.hiddenDim;
|
||||
params.hiddenDimPadded = data.hiddenDimPadded;
|
||||
params.numTokens = data.numTokens;
|
||||
params.numExperts = data.numExperts;
|
||||
params.topK = data.topK;
|
||||
params.totalNumPaddedTokens = data.totalNumPaddedTokens;
|
||||
|
||||
return params;
|
||||
}
|
||||
};
|
||||
|
||||
void run(Data const& data, void* stream);
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace finalize
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace moe::dev
|
||||
-586
@@ -1,586 +0,0 @@
|
||||
/*
|
||||
* Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "DevKernel.h"
|
||||
#include "flashinfer/trtllm/fused_moe/RoutingKernel.h"
|
||||
#include <string>
|
||||
// #include "flashinfer/trtllm/common/cudaDriverWrapper.h"
|
||||
#include "flashinfer/trtllm/batched_gemm/KernelRunner.h"
|
||||
#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h"
|
||||
#include "flashinfer/trtllm/common/cudaUtils.h"
|
||||
|
||||
namespace tensorrt_llm {
|
||||
namespace kernels {
|
||||
namespace trtllmgen_moe {
|
||||
|
||||
namespace MoE {
|
||||
class Runner;
|
||||
} // namespace MoE
|
||||
|
||||
namespace Routing {
|
||||
|
||||
// The type of method in top-K routing, for use in torch custom op
|
||||
// Please keep this in sync with the counterpart defined in
|
||||
// flashinfer/fused_moe/core.py
|
||||
enum class RoutingMethodType : int64_t {
|
||||
// Default: Softmax -> TopK
|
||||
Default = 0,
|
||||
// Renormalize: TopK -> Softmax
|
||||
Renormalize = 1,
|
||||
// DeepSeekV3: Sigmoid -> RoutingBiasAdd -> Top2 in group -> Top4 groups -> Top8 experts from the
|
||||
// Top4 groups
|
||||
DeepSeekV3 = 2,
|
||||
// Llama4: Top1 -> Sigmoid
|
||||
Llama4 = 3,
|
||||
// RenormalizeNaive: Softmax -> TopK -> Renormalize
|
||||
RenormalizeNaive = 4,
|
||||
// TopK only (no softmax)
|
||||
TopK = 5,
|
||||
// SigmoidRenorm: Sigmoid -> TopK -> Renormalize (divide by sum of top-K weights)
|
||||
SigmoidRenorm = 6,
|
||||
// MiniMax2: Sigmoid + Bias -> TopK -> ScaledSumNormalize (routeScale=1.0, epsilon=1e-20)
|
||||
MiniMax2 = 7,
|
||||
// Sigmoid: Sigmoid -> TopK (no renormalization)
|
||||
Sigmoid = 8,
|
||||
// Unspecified
|
||||
Unspecified = 9,
|
||||
};
|
||||
|
||||
inline int32_t maybeGetMinTokenCount(int32_t numPaddedTokens, int32_t hiddenSize, int32_t dtypeSizeBits) {
|
||||
// Pad so total size exceeds 128KiB for performance reasons
|
||||
int32_t minNumTokensRequired = common::divUp(128 * 1024 * 8, hiddenSize * dtypeSizeBits);
|
||||
return std::max(numPaddedTokens, minNumTokensRequired);
|
||||
}
|
||||
|
||||
inline std::string serializeMoeRoutingMethodType(RoutingMethodType routingMethodType) {
|
||||
switch (routingMethodType) {
|
||||
case RoutingMethodType::Default:
|
||||
return "Default";
|
||||
case RoutingMethodType::Renormalize:
|
||||
return "Renormalize";
|
||||
case RoutingMethodType::DeepSeekV3:
|
||||
return "DeepSeekV3";
|
||||
case RoutingMethodType::Llama4:
|
||||
return "Llama4";
|
||||
case RoutingMethodType::RenormalizeNaive:
|
||||
return "RenormalizeNaive";
|
||||
case RoutingMethodType::TopK:
|
||||
return "TopK";
|
||||
case RoutingMethodType::SigmoidRenorm:
|
||||
return "SigmoidRenorm";
|
||||
case RoutingMethodType::MiniMax2:
|
||||
return "MiniMax2";
|
||||
case RoutingMethodType::Sigmoid:
|
||||
return "Sigmoid";
|
||||
default:
|
||||
return "InvalidRountingMethod"; // TODO throw error
|
||||
};
|
||||
}
|
||||
|
||||
inline int32_t getMaxNumCtasInBatchDim(int32_t numTokens, int32_t topK, int32_t numExperts, int32_t tileTokensDim) {
|
||||
// For MoE, mNumTokens != 0 and the number of CTAs is known only at runtime.
|
||||
// We launch maximally possible number of CTAs and use ptrNumNonExitingCtas to determine
|
||||
// the actual number of CTAs to run.
|
||||
|
||||
// Initialize number of tokens with the number of expanded tokens after routing.
|
||||
int32_t numRemainingTokens = numTokens * topK;
|
||||
int32_t maxNumCtasInBatchDim = 0;
|
||||
// First, distribute one token each expert until token depletion to maximize CTA tile count.
|
||||
int32_t numExpertsFilled = std::min(numExperts, numRemainingTokens);
|
||||
maxNumCtasInBatchDim += numExpertsFilled;
|
||||
numRemainingTokens -= numExpertsFilled;
|
||||
// Next, greedily pour all remaining tokens to one expert to maximize CTA tile count.
|
||||
// E.g., at this point tokens over 4 experts are [1, 1, 1, 1], and we have 4 tokens left.
|
||||
// If each CTA handles 4 tokens/expert, the greedy strategy is to pour all remaining tokens
|
||||
// to any one expert to get to the 5th CTA tile. Otherwise, we can only get 4 tiles in total.
|
||||
//
|
||||
// Another way to reason about this is to pour the remaining tokens into buckets of some fixed
|
||||
// capacity. These buckets, if full, can then be attributed to any expert; it does not have to
|
||||
// belong to the same expert every time.
|
||||
if (numRemainingTokens > 0) {
|
||||
// For every tileTokenDim tokens, we add an extra CTA tile in the token dimension.
|
||||
// The number of CTA tiles is given by divDown(numRemainingTokens, tokenTileDim).
|
||||
maxNumCtasInBatchDim += (numRemainingTokens / tileTokensDim);
|
||||
}
|
||||
return maxNumCtasInBatchDim;
|
||||
}
|
||||
|
||||
inline int32_t
|
||||
getMaxPermutedPaddedCount(int32_t numTokens, int32_t expertsPerToken, int32_t numExperts, int32_t padding) {
|
||||
int32_t maxCtas = getMaxNumCtasInBatchDim(numTokens, expertsPerToken, numExperts, padding);
|
||||
return maxCtas * padding;
|
||||
}
|
||||
|
||||
class Runner {
|
||||
public:
|
||||
explicit Runner();
|
||||
|
||||
explicit Runner(int32_t tileTokensDim);
|
||||
|
||||
void
|
||||
run(void* routingLogits,
|
||||
void* routingBias,
|
||||
int32_t numTokens,
|
||||
int32_t numExperts,
|
||||
int32_t topK,
|
||||
int32_t nGroups,
|
||||
int32_t topkGroups,
|
||||
int32_t localExpertOffset,
|
||||
int32_t localNumExperts,
|
||||
float routedScalingFactor,
|
||||
int32_t* routingExpertIndexes,
|
||||
int32_t* expertCountHistogram,
|
||||
int32_t* permutedIdxSize,
|
||||
int32_t* expandedIdxToPermutedIdx,
|
||||
int32_t* permutedIdxToExpandedIdx,
|
||||
int32_t* permutedIdxToTokenIdx,
|
||||
int32_t* expertIds,
|
||||
void* expertWeights,
|
||||
int32_t* numTokensPerExpert,
|
||||
int32_t* ctaIdxXyToBatchIdx,
|
||||
int32_t* ctaIdxXyToMnLimit,
|
||||
int32_t* numNonExitingCtas,
|
||||
batchedGemm::trtllm::gen::Dtype dtypeElt,
|
||||
batchedGemm::trtllm::gen::Dtype dtypeBias,
|
||||
bool useRoutingScalesOnInput,
|
||||
bool useDeepSeekFp8,
|
||||
RoutingMethodType routingMethodType,
|
||||
cudaStream_t stream,
|
||||
batchedGemm::trtllm::gen::Dtype dtypeLogits,
|
||||
bool normTopkProb = true,
|
||||
int16_t* routing_replay_out = nullptr);
|
||||
|
||||
private:
|
||||
friend class MoE::Runner;
|
||||
int32_t mTileTokensDim{8};
|
||||
};
|
||||
} // namespace Routing
|
||||
|
||||
namespace MoE {
|
||||
// The type of activation function
|
||||
// Please keep this in sync with the counterpart defined in flashinfer/flashinfer/fused_moe/core.py
|
||||
enum class ActivationType : int64_t {
|
||||
Gelu = 0,
|
||||
Relu = 1,
|
||||
Silu = 2,
|
||||
Swiglu = 3,
|
||||
Geglu = 4,
|
||||
SwigluBias = 5,
|
||||
Relu2 = 6,
|
||||
Identity = 7,
|
||||
InvalidType = 8, // Must be last
|
||||
};
|
||||
|
||||
inline std::string serializeActivationType(ActivationType activationType) {
|
||||
switch (activationType) {
|
||||
case ActivationType::Gelu:
|
||||
return "Gelu";
|
||||
case ActivationType::Relu:
|
||||
return "Relu";
|
||||
case ActivationType::Silu:
|
||||
return "Silu";
|
||||
case ActivationType::Swiglu:
|
||||
return "Swiglu";
|
||||
case ActivationType::Geglu:
|
||||
return "Geglu";
|
||||
case ActivationType::SwigluBias:
|
||||
return "SwigluBias";
|
||||
case ActivationType::Relu2:
|
||||
return "Relu2";
|
||||
case ActivationType::Identity:
|
||||
return "Identity";
|
||||
default:
|
||||
return "InvalidActivationType"; // TODO throw error
|
||||
};
|
||||
}
|
||||
|
||||
inline bool isGatedActivation(ActivationType activationType) {
|
||||
return activationType == ActivationType::Swiglu || activationType == ActivationType::Geglu ||
|
||||
activationType == ActivationType::SwigluBias;
|
||||
}
|
||||
|
||||
} // namespace MoE
|
||||
|
||||
namespace PermuteGemm1 {
|
||||
class Runner {
|
||||
public:
|
||||
explicit Runner(
|
||||
batchedGemm::trtllm::gen::Dtype dtypeAct,
|
||||
batchedGemm::trtllm::gen::Dtype dtypeWeights,
|
||||
batchedGemm::trtllm::gen::Dtype dtypeOutput,
|
||||
bool useDeepSeekFp8,
|
||||
int tileTokensDim,
|
||||
MoE::ActivationType activationType,
|
||||
bool useShuffledMatrix,
|
||||
batchedGemm::gemm::MatrixLayout weight_layout,
|
||||
bool usePerTokenScaling,
|
||||
bool usePerChannelScaling,
|
||||
bool forceUnfusedAct = false);
|
||||
|
||||
size_t getWorkspaceSizeInBytes(
|
||||
int32_t topK,
|
||||
int32_t hiddenSize,
|
||||
int32_t intermediateSize,
|
||||
int32_t numExperts,
|
||||
int32_t numTokens,
|
||||
int32_t configIndex) const;
|
||||
|
||||
[[nodiscard]] int32_t getDefaultValidConfigIndex(
|
||||
int32_t topK, int32_t hiddenSize, int32_t intermediateSize, int32_t numExperts, int32_t numTokens) const;
|
||||
|
||||
[[nodiscard]] bool isValidConfigIndex(
|
||||
int32_t configIndex,
|
||||
int32_t topK,
|
||||
int32_t hiddenSize,
|
||||
int32_t intermediateSize,
|
||||
int32_t numExperts,
|
||||
int32_t numTokens) const;
|
||||
|
||||
[[nodiscard]] std::vector<int64_t> getPassingConfigIndices() const;
|
||||
|
||||
void
|
||||
run(void* hiddenState,
|
||||
void* hiddenStateScale,
|
||||
void* weight,
|
||||
void* weightScale,
|
||||
void* perTokenScales,
|
||||
void* perChannelScales,
|
||||
float* outputScalesScalar,
|
||||
float* outputScalesGateScalar,
|
||||
float* ptrBias,
|
||||
float* ptrGatedActAlpha,
|
||||
float* ptrGatedActBeta,
|
||||
float* ptrClampLimit,
|
||||
void* output,
|
||||
void* outputScale,
|
||||
int32_t topK,
|
||||
int32_t hiddenSize,
|
||||
int32_t intermediateSize,
|
||||
int32_t numExperts,
|
||||
int32_t numTokens,
|
||||
int32_t* permutedIdxToTokenIdx,
|
||||
int32_t* ptrNumNonExitingCtas,
|
||||
int32_t* ptrTotalNumPaddedTokens,
|
||||
int32_t* ptrCtaIdxXyToBatchIdx,
|
||||
int32_t* ptrCtaIdxXyToMnLimit,
|
||||
void* bmm1Workspace,
|
||||
bool useRoutingScalesOnInput,
|
||||
int device,
|
||||
cudaStream_t stream,
|
||||
int32_t configIndex,
|
||||
bool enable_pdl);
|
||||
|
||||
private:
|
||||
friend class MoE::Runner;
|
||||
batchedGemm::trtllm::gen::Dtype mDtypeAct;
|
||||
batchedGemm::trtllm::gen::Dtype mDtypeWeights;
|
||||
batchedGemm::trtllm::gen::Dtype mDtypeOutput;
|
||||
int32_t mTileTokensDim;
|
||||
tensorrt_llm::kernels::TrtllmGenBatchedGemmRunner mRunner;
|
||||
tensorrt_llm::kernels::trtllmgen_moe::MoE::ActivationType mActType;
|
||||
};
|
||||
} // namespace PermuteGemm1
|
||||
|
||||
namespace Gemm2 {
|
||||
class Runner {
|
||||
public:
|
||||
explicit Runner(
|
||||
batchedGemm::trtllm::gen::Dtype dtypeAct,
|
||||
batchedGemm::trtllm::gen::Dtype dtypeWeights,
|
||||
batchedGemm::trtllm::gen::Dtype outputDtype,
|
||||
bool useDeepSeekFp8,
|
||||
int tileTokensDim,
|
||||
bool useShuffledMatrix,
|
||||
batchedGemm::gemm::MatrixLayout weight_layout,
|
||||
bool usePerTokenScaling,
|
||||
bool usePerChannelScaling);
|
||||
|
||||
size_t getWorkspaceSizeInBytes(
|
||||
int32_t topK,
|
||||
int32_t hiddenSize,
|
||||
int32_t intermediateSize,
|
||||
int32_t numExperts,
|
||||
int32_t numTokens,
|
||||
int32_t configIndex) const;
|
||||
|
||||
[[nodiscard]] int32_t getDefaultValidConfigIndex(
|
||||
int32_t topK, int32_t hiddenSize, int32_t intermediateSize, int32_t numExperts, int32_t numTokens) const;
|
||||
|
||||
[[nodiscard]] bool isValidConfigIndex(
|
||||
int32_t configIndex,
|
||||
int32_t topK,
|
||||
int32_t hiddenSize,
|
||||
int32_t intermediateSize,
|
||||
int32_t numExperts,
|
||||
int32_t numTokens) const;
|
||||
|
||||
[[nodiscard]] std::vector<int64_t> getPassingConfigIndices() const;
|
||||
|
||||
void
|
||||
run(void* permutedHiddenState,
|
||||
void* permutedHiddenStateScale,
|
||||
void* weight,
|
||||
void* weightScale,
|
||||
void* perTokenScales,
|
||||
void* perChannelScales,
|
||||
float* outputScalesScalar,
|
||||
float* ptrBias,
|
||||
void* output,
|
||||
void* outputScale,
|
||||
int32_t topK,
|
||||
int32_t hiddenSize,
|
||||
int32_t intermediateSize,
|
||||
int32_t numExperts,
|
||||
int32_t numTokens,
|
||||
int32_t* ptrNumNonExitingCtas,
|
||||
int32_t* ptrTotalNumPaddedTokens,
|
||||
int32_t* ptrCtaIdxXyToBatchIdx,
|
||||
int32_t* ptrCtaIdxXyToMnLimit,
|
||||
void* bmm2Workspace,
|
||||
int device,
|
||||
cudaStream_t stream,
|
||||
int32_t configIndex,
|
||||
bool enable_pdl);
|
||||
|
||||
private:
|
||||
friend class MoE::Runner;
|
||||
batchedGemm::trtllm::gen::Dtype mDtypeAct;
|
||||
batchedGemm::trtllm::gen::Dtype mDtypeWeights;
|
||||
batchedGemm::trtllm::gen::Dtype mDtypeOut;
|
||||
int32_t mTileTokensDim;
|
||||
tensorrt_llm::kernels::TrtllmGenBatchedGemmRunner mRunner;
|
||||
};
|
||||
} // namespace Gemm2
|
||||
|
||||
namespace MoE {
|
||||
namespace btg = batchedGemm::trtllm::gen;
|
||||
|
||||
struct MoERunnerArgs {
|
||||
void* routing_logits = nullptr; // [num_tokens, num_experts] in float, generated after
|
||||
// gemm(hidden_state, routing_weights)
|
||||
void* routing_bias = nullptr; // [num_experts] in bfloat16 for now = mDtypeExpW
|
||||
void* hidden_states = nullptr; // [num_tokens, hidden_size] in fp8 = mDtypeElt
|
||||
// [hidden_size/128, num_tokens] in float for e4m3 DS recipe
|
||||
// and [num_tokens, hidden_size/16] in float for e2m1
|
||||
void* hidden_states_scale = nullptr;
|
||||
|
||||
// Gemm input:
|
||||
void* gemm1_weights = nullptr;
|
||||
void* gemm1_weights_scale = nullptr;
|
||||
void* gemm2_weights = nullptr;
|
||||
void* gemm2_weights_scale = nullptr;
|
||||
|
||||
float* gemm1_bias = nullptr;
|
||||
float* gemm1_alpha = nullptr;
|
||||
float* gemm1_beta = nullptr;
|
||||
float* gemm1_clamp_limit = nullptr;
|
||||
float* gemm2_bias = nullptr;
|
||||
|
||||
ActivationType activation_type = ActivationType::Swiglu;
|
||||
|
||||
int32_t num_tokens{0};
|
||||
int32_t num_experts{0};
|
||||
// Hidden dimension input of MoE block. It might be padded.
|
||||
int32_t hidden_size{0};
|
||||
// Hidden dimension output of MoE block. It is not padded.
|
||||
// If not provided it is the same as hidden_size.
|
||||
std::optional<int32_t> hidden_size_output;
|
||||
// TODO: only compiled routing kernel supports top_k = 8
|
||||
int32_t top_k{0};
|
||||
int32_t n_group{0};
|
||||
// TODO: only compiled routing kernel supports topk_group = 4
|
||||
int32_t topk_group{0};
|
||||
float routed_scaling_factor{0.0f};
|
||||
int32_t intermediate_size{0};
|
||||
int32_t local_expert_offset{0};
|
||||
int32_t local_num_experts{0};
|
||||
// TODO: support other types
|
||||
btg::Dtype mDtypeElt{btg::Dtype::Void};
|
||||
btg::Dtype mDtypeExpW{btg::Dtype::Bfloat16};
|
||||
btg::Dtype mDtypeOut{btg::Dtype::Bfloat16};
|
||||
|
||||
// Apply routing scale factors to input activations
|
||||
bool mUseRoutingScalesOnInput{false};
|
||||
bool mUseDeepSeekFp8{false};
|
||||
float* output1_scales_scalar = nullptr;
|
||||
float* output1_scales_gate_scalar = nullptr;
|
||||
float* output2_scales_scalar = nullptr;
|
||||
|
||||
// Optional LoRA bridge buffers used by the copied SGLang TRTLLM FP8 path.
|
||||
// gate_up_lora_delta: [num_tokens * top_k, 2 * intermediate_size], bf16,
|
||||
// in FlashInfer gate/up order (up first, gate second).
|
||||
// activation_lora_input: [num_tokens * top_k, intermediate_size], bf16,
|
||||
// populated with the post-activation intermediate for down-proj LoRA.
|
||||
void* gate_up_lora_delta = nullptr;
|
||||
void* activation_lora_input = nullptr;
|
||||
|
||||
// Optional CUDA event (cudaEvent_t) recorded on the LoRA side stream. When set, the
|
||||
// runner waits on it right before the activation kernel (which consumes
|
||||
// gate_up_lora_delta), so permute+GEMM1 overlap the side-stream LoRA shrink/expand
|
||||
// instead of joining before the whole MoE op. nullptr = no wait (serial behavior).
|
||||
void* lora_ready_event = nullptr;
|
||||
|
||||
// Down-LoRA/finalize overlap: optional CUDA event (cudaEvent_t) the runner records on the
|
||||
// MoE stream right after GEMM2 (the base down GEMM), before finalize. The LoRA side stream
|
||||
// waits on it to run the down-proj LoRA shrink/expand concurrent with the finalize kernel.
|
||||
// nullptr = no record (serial behavior).
|
||||
void* gemm2_done_event = nullptr;
|
||||
|
||||
// Output:
|
||||
void* output = nullptr;
|
||||
float* output_scale = nullptr;
|
||||
|
||||
// finalize
|
||||
bool do_finalize{true};
|
||||
};
|
||||
|
||||
struct MoEWorkspace {
|
||||
// Routing intermediate outputs:
|
||||
int32_t* routing_expert_indexes = nullptr;
|
||||
int32_t* permuted_idx_size = nullptr;
|
||||
int32_t* total_num_padded_tokens = nullptr; // TODO: duplicate of permuted_idx_size
|
||||
int32_t total_max_padded_tokens{0};
|
||||
|
||||
int32_t* expanded_idx_to_permuted_idx = nullptr;
|
||||
int32_t* permuted_idx_to_expanded_idx = nullptr;
|
||||
int32_t* permuted_idx_to_token_idx = nullptr;
|
||||
|
||||
// consumed by finalize kernel
|
||||
void* expert_weights = nullptr; // [num_tokens, top_k] in bfloat16 = mDtypeExpW
|
||||
// consumed by permuteGemm1 kernel
|
||||
void* token_scales = nullptr;
|
||||
// consumed by Gemm2 kernel
|
||||
void* token_scales_fc2 = nullptr;
|
||||
|
||||
int32_t* cta_idx_xy_to_batch_idx = nullptr;
|
||||
int32_t* cta_idx_xy_to_mn_limit = nullptr;
|
||||
int32_t* num_non_exiting_ctas = nullptr;
|
||||
|
||||
void* hidden_states_scale_linear = nullptr;
|
||||
|
||||
// Permute intermediate outputs:
|
||||
void* permuted_hidden_states = nullptr;
|
||||
float* permuted_hidden_states_scale = nullptr;
|
||||
|
||||
// Gemm1 intermediate outputs:
|
||||
int32_t ProjUpTileN{0};
|
||||
void* gemm1_output = nullptr;
|
||||
float* gemm1_output_scale = nullptr;
|
||||
|
||||
// Activation intermediate outputs:
|
||||
void* activation_output = nullptr;
|
||||
float* activation_output_scale = nullptr;
|
||||
// Unfused FP4 LoRA: bf16 [max_padded_tokens, intermediate_size] activation output written by the
|
||||
// standalone activation kernel (gate_up LoRA added pre-SwiGLU), then NvFP4-quantized for GEMM2.
|
||||
void* activated_lora_bf16 = nullptr;
|
||||
|
||||
// Gemm2 intermediate outputs:
|
||||
void* gemm2_output = nullptr;
|
||||
float* gemm2_output_scale = nullptr;
|
||||
|
||||
// Finalize intermediate outputs (placeholder not used)
|
||||
void* finalize_output = nullptr;
|
||||
float* finalize_output_scale = nullptr;
|
||||
|
||||
// FC1 workspace:
|
||||
void* bmm1_workspace = nullptr;
|
||||
|
||||
// FC2 workspace:
|
||||
void* bmm2_workspace = nullptr;
|
||||
};
|
||||
|
||||
// Config indices to be used with Batched GEMM runners
|
||||
struct MoEConfig {
|
||||
int64_t gemm1Config;
|
||||
int64_t gemm2Config;
|
||||
};
|
||||
|
||||
class Runner {
|
||||
public:
|
||||
// FIXME: tileTokensDim is hardcoded for now
|
||||
Runner(
|
||||
batchedGemm::trtllm::gen::Dtype dtypeAct,
|
||||
batchedGemm::trtllm::gen::Dtype dtypeWeights,
|
||||
bool useDeepSeekFp8,
|
||||
int tileTokensDim = 8,
|
||||
ActivationType activationType = ActivationType::Swiglu,
|
||||
bool useShuffledMatrix = false,
|
||||
batchedGemm::gemm::MatrixLayout weight_layout = batchedGemm::gemm::MatrixLayout::MajorK,
|
||||
bool usePerTokenScalingGemm1 = false,
|
||||
bool usePerTokenScalingGemm2 = false,
|
||||
bool usePerChannelScalingGemm1 = false,
|
||||
bool usePerChannelScalingGemm2 = false,
|
||||
bool unfuseActForLora = false);
|
||||
Runner(
|
||||
batchedGemm::trtllm::gen::Dtype dtypeElt,
|
||||
bool useDeepSeekFp8,
|
||||
int tileTokensDim = 8,
|
||||
bool useShuffledMatrix = false,
|
||||
batchedGemm::gemm::MatrixLayout weight_layout = batchedGemm::gemm::MatrixLayout::MajorK,
|
||||
bool usePerTokenScalingGemm1 = false,
|
||||
bool usePerTokenScalingGemm2 = false,
|
||||
bool usePerChannelScalingGemm1 = false,
|
||||
bool usePerChannelScalingGemm2 = false);
|
||||
|
||||
void
|
||||
run(MoERunnerArgs const& args,
|
||||
MoEWorkspace const& workspace,
|
||||
int device,
|
||||
cudaStream_t stream,
|
||||
int64_t configIndex,
|
||||
bool enable_pdl);
|
||||
|
||||
[[nodiscard]] std::tuple<int32_t, int32_t>
|
||||
getWorkspaceSizeInBytes(MoERunnerArgs const& args, int64_t configIndex) const;
|
||||
|
||||
[[nodiscard]] std::vector<int64_t> getValidConfigIndices(
|
||||
int32_t topK, int32_t hiddenSize, int32_t intermediateSize, int32_t numLocalExperts, int32_t numTokens) const;
|
||||
|
||||
[[nodiscard]] int64_t getDefaultValidConfigIndex(
|
||||
int32_t topK, int32_t hiddenSize, int32_t intermediateSize, int32_t numLocalExperts, int32_t numTokens) const;
|
||||
|
||||
private:
|
||||
void setOpsData(
|
||||
MoERunnerArgs const& args,
|
||||
MoEWorkspace const& workspace,
|
||||
moe::dev::convertsf::Data& convertSfData,
|
||||
moe::dev::activation::Data& activationData,
|
||||
moe::dev::finalize::Data& finalizeData);
|
||||
|
||||
private:
|
||||
bool mUsePerTokenScalingGemm1;
|
||||
bool mUsePerTokenScalingGemm2;
|
||||
bool mUsePerChannelScalingGemm1;
|
||||
bool mUsePerChannelScalingGemm2;
|
||||
// When true (FP4 LoRA path), GEMM1 emits the raw gate_up projection (fusedAct=false) so the
|
||||
// standalone activation kernel can inject the gate_up LoRA delta pre-SwiGLU and capture the
|
||||
// post-activation input for the down LoRA — mirroring the DeepSeek-FP8 unfused activation.
|
||||
bool mUnfuseActForLora;
|
||||
PermuteGemm1::Runner mPermuteGemm1;
|
||||
Gemm2::Runner mGemm2;
|
||||
|
||||
// This will be the cartesian product of the passing configs for gemm1 and gemm2
|
||||
// This allows us to autotune the MoE as one operation instead of tuning gemm1 and gemm2
|
||||
// separately
|
||||
std::vector<MoEConfig> mPassingConfigs;
|
||||
};
|
||||
} // namespace MoE
|
||||
|
||||
} // namespace trtllmgen_moe
|
||||
} // namespace kernels
|
||||
} // namespace tensorrt_llm
|
||||
@@ -31,6 +31,9 @@ _GROUPS = (
|
||||
"sampling",
|
||||
"spatial",
|
||||
"speculative",
|
||||
"lplb",
|
||||
"kv_canary",
|
||||
"model",
|
||||
)
|
||||
|
||||
for _group in _GROUPS:
|
||||
|
||||
@@ -1 +1,33 @@
|
||||
"""DeepSeek DSA kernels (RFC #29630, Phase 2.5)."""
|
||||
|
||||
# --- merged from sglang.kernels.ops.attention.dsa (RFC #29630 Phase 4) ---
|
||||
from .paged_mqa_logits import (
|
||||
aiter_paged_mqa_logits,
|
||||
cutedsl_paged_mqa_logits,
|
||||
deepgemm_paged_mqa_logits_native,
|
||||
deepgemm_paged_mqa_logits_split,
|
||||
)
|
||||
|
||||
|
||||
def pick_dsl_expand(*args, **kwargs):
|
||||
from .cutedsl_paged_mqa_logits import pick_dsl_expand as _pick_dsl_expand
|
||||
|
||||
return _pick_dsl_expand(*args, **kwargs)
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
if name == "CuteDSLPagedMQALogitsRunner":
|
||||
from .cutedsl_paged_mqa_logits import CuteDSLPagedMQALogitsRunner
|
||||
|
||||
return CuteDSLPagedMQALogitsRunner
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CuteDSLPagedMQALogitsRunner",
|
||||
"pick_dsl_expand",
|
||||
"aiter_paged_mqa_logits",
|
||||
"cutedsl_paged_mqa_logits",
|
||||
"deepgemm_paged_mqa_logits_native",
|
||||
"deepgemm_paged_mqa_logits_split",
|
||||
]
|
||||
|
||||
+1
-1
@@ -17,7 +17,7 @@ import cutlass.cute as cute
|
||||
import torch
|
||||
from cutlass.utils import HardwareInfo
|
||||
|
||||
from sglang.jit_kernel.cutedsl_fp8_paged_mqa_logits import FP8MQALogitsKernel
|
||||
from sglang.kernels.ops.attention.cutedsl_fp8_paged_mqa_logits import FP8MQALogitsKernel
|
||||
from sglang.srt.utils import is_sm100_supported
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
+1
-1
@@ -113,7 +113,7 @@ def cutedsl_paged_mqa_logits(
|
||||
sm_count: int,
|
||||
get_paged_mqa_logits_metadata_fn: Callable[..., torch.Tensor],
|
||||
) -> torch.Tensor:
|
||||
from sglang.jit_kernel.dsa.cutedsl_paged_mqa_logits import (
|
||||
from sglang.kernels.ops.attention.dsa.cutedsl_paged_mqa_logits import (
|
||||
CuteDSLPagedMQALogitsRunner,
|
||||
)
|
||||
|
||||
@@ -1 +1,66 @@
|
||||
"""DeepSeek-V4 attention kernels (RFC #29630, Phase 2.5)."""
|
||||
|
||||
# --- merged from sglang.kernels.ops.attention.dsv4 (RFC #29630 Phase 4) ---
|
||||
from .attn import (
|
||||
fused_store_cache,
|
||||
get_paged_mqa_logits_metadata,
|
||||
triton_create_paged_compress_data,
|
||||
)
|
||||
from .c128_cleanup import clear_unaccepted_c128_draft_states
|
||||
from .compress import (
|
||||
CompressorDecodePlan,
|
||||
CompressorPrefillPlan,
|
||||
compress_forward,
|
||||
compress_norm_rope_store,
|
||||
)
|
||||
from .compress_old import fused_norm_rope_inplace
|
||||
from .elementwise import (
|
||||
fused_k_norm_rope_flashmla,
|
||||
fused_q_indexer_rope_first_quant,
|
||||
fused_q_indexer_rope_hadamard_fp4_quant,
|
||||
fused_q_indexer_rope_hadamard_quant,
|
||||
fused_q_norm_rope,
|
||||
fused_rope_inplace,
|
||||
)
|
||||
from .fp8_wo_a import sglang_per_token_group_quant_fp8_dsv4_wo_a
|
||||
from .gemm import linear_bf16_fp32
|
||||
from .moe import (
|
||||
hash_topk,
|
||||
mask_topk_ids,
|
||||
mega_moe_pre_dispatch,
|
||||
silu_and_mul_clamp,
|
||||
silu_and_mul_contig_post_quant,
|
||||
silu_and_mul_masked_post_quant,
|
||||
)
|
||||
from .topk import plan_topk_v2, topk_transform_512, topk_transform_512_v2
|
||||
from .utils import make_name
|
||||
|
||||
__all__ = [
|
||||
"CompressorDecodePlan",
|
||||
"CompressorPrefillPlan",
|
||||
"compress_forward",
|
||||
"compress_norm_rope_store",
|
||||
"clear_unaccepted_c128_draft_states",
|
||||
"fused_norm_rope_inplace",
|
||||
"fused_store_cache",
|
||||
"fused_rope_inplace",
|
||||
"fused_q_norm_rope",
|
||||
"fused_q_indexer_rope_first_quant",
|
||||
"fused_q_indexer_rope_hadamard_fp4_quant",
|
||||
"fused_q_indexer_rope_hadamard_quant",
|
||||
"fused_k_norm_rope_flashmla",
|
||||
"sglang_per_token_group_quant_fp8_dsv4_wo_a",
|
||||
"make_name",
|
||||
"linear_bf16_fp32",
|
||||
"get_paged_mqa_logits_metadata",
|
||||
"triton_create_paged_compress_data",
|
||||
"topk_transform_512",
|
||||
"topk_transform_512_v2",
|
||||
"plan_topk_v2",
|
||||
"hash_topk",
|
||||
"mega_moe_pre_dispatch",
|
||||
"mask_topk_ids",
|
||||
"silu_and_mul_clamp",
|
||||
"silu_and_mul_masked_post_quant",
|
||||
"silu_and_mul_contig_post_quant",
|
||||
]
|
||||
|
||||
@@ -10,11 +10,11 @@ from typing import Union
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.dsv4 import (
|
||||
from sglang.kernels.jit.utils import is_hip_runtime
|
||||
from sglang.kernels.ops.attention.dsv4 import (
|
||||
CompressorDecodePlan,
|
||||
CompressorPrefillPlan,
|
||||
)
|
||||
from sglang.kernels.jit.utils import is_hip_runtime
|
||||
|
||||
_is_hip = is_hip_runtime()
|
||||
|
||||
|
||||
@@ -18,7 +18,7 @@ import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.jit_kernel.dsv4.compress_old import (
|
||||
from sglang.kernels.ops.attention.dsv4.compress_old import (
|
||||
CompressorDecodePlan,
|
||||
CompressorPrefillPlan,
|
||||
)
|
||||
|
||||
+1
-1
@@ -119,7 +119,7 @@ def hash_topk(
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
assert scoring_func == "sqrtsoftplus"
|
||||
if is_hip_runtime():
|
||||
from sglang.jit_kernel.triton.hash_topk import hash_topk_triton
|
||||
from sglang.kernels.ops.moe.triton_hash_topk import hash_topk_triton
|
||||
|
||||
return hash_topk_triton(
|
||||
router_logits,
|
||||
+1
-1
@@ -5,8 +5,8 @@ from typing import TYPE_CHECKING, Any, List, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.dsv4.utils import make_name
|
||||
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
|
||||
from sglang.kernels.ops.attention.dsv4.utils import make_name
|
||||
from sglang.srt.environ import envs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
+1
-1
@@ -13,7 +13,7 @@ try:
|
||||
# (dev's stack). rel_bias is vendored-only, so SHEARED must be 0.
|
||||
from flash_attn.cute import flash_attn_varlen_func as _flash_attn_varlen_func
|
||||
else:
|
||||
from sglang.jit_kernel.flash_attn.cute import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute import (
|
||||
flash_attn_varlen_func as _flash_attn_varlen_func,
|
||||
)
|
||||
except Exception as _e: # pragma: no cover
|
||||
+2
-2
@@ -21,8 +21,8 @@ from typing import Iterable
|
||||
import torch
|
||||
from einops import rearrange
|
||||
|
||||
from sglang.jit_kernel.flash_attn.cute.benchmark import benchmark_forward
|
||||
from sglang.jit_kernel.flash_attn.cute.interface import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.benchmark import benchmark_forward
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.interface import (
|
||||
_flash_attn_fwd as flash_attn_cute_fwd,
|
||||
)
|
||||
|
||||
+1
-1
@@ -7,7 +7,7 @@ from cutlass import Boolean, Int32, const_expr
|
||||
from cutlass._mlir.dialects import llvm
|
||||
from cutlass.cute.nvgpu import tcgen05
|
||||
|
||||
import sglang.jit_kernel.flash_attn.cute.mma_sm100_desc as sm100_desc
|
||||
import sglang.kernels.ops.attention.flash_attn.cute.mma_sm100_desc as sm100_desc
|
||||
|
||||
|
||||
def _tcgen05_mma_kind(op: cute.nvgpu.tcgen05.mma.MmaOp) -> str:
|
||||
+4
-1
@@ -6,7 +6,10 @@ import cutlass
|
||||
import cutlass.cute as cute
|
||||
from cutlass import Int32, const_expr
|
||||
|
||||
from sglang.jit_kernel.flash_attn.cute.seqlen_info import SeqlenInfoQK, SeqlenInfoQKNewK
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.seqlen_info import (
|
||||
SeqlenInfoQK,
|
||||
SeqlenInfoQKNewK,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
+6
-4
@@ -15,10 +15,12 @@ from cutlass import Float32, Int32, const_expr
|
||||
from quack import copy_utils
|
||||
|
||||
# Import data structures from block_sparsity
|
||||
from sglang.jit_kernel.flash_attn.cute.block_sparsity import BlockSparseTensors
|
||||
from sglang.jit_kernel.flash_attn.cute.named_barrier import NamedBarrierBwd
|
||||
from sglang.jit_kernel.flash_attn.cute.seqlen_info import SeqlenInfoQK
|
||||
from sglang.jit_kernel.flash_attn.cute.utils import AuxData
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.block_sparsity import (
|
||||
BlockSparseTensors,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.named_barrier import NamedBarrierBwd
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.seqlen_info import SeqlenInfoQK
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.utils import AuxData
|
||||
|
||||
|
||||
@cute.jit
|
||||
+1
-1
@@ -7,7 +7,7 @@ from typing import Callable, NamedTuple, Tuple
|
||||
import cutlass.cute as cute
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.flash_attn.cute.cute_dsl_utils import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.cute_dsl_utils import (
|
||||
get_broadcast_dims,
|
||||
to_cute_tensor,
|
||||
)
|
||||
+1
-1
@@ -17,7 +17,7 @@ import cutlass.cute as cute
|
||||
import tvm_ffi
|
||||
from cutlass.cutlass_dsl import JitCompiledFunction
|
||||
|
||||
from sglang.jit_kernel.flash_attn.cute.fa_logging import fa_log
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.fa_logging import fa_log
|
||||
|
||||
# Pre-load cute DSL runtime libraries with RTLD_GLOBAL so that their symbols
|
||||
# (e.g. _cudaLibraryLoadData) are visible to .so modules loaded later via dlopen.
|
||||
+7
-7
@@ -6,23 +6,23 @@ import cutlass.cute as cute
|
||||
import torch
|
||||
from cutlass import Boolean, Int8, Int32, const_expr
|
||||
|
||||
from sglang.jit_kernel.flash_attn.cute.block_sparse_utils import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.block_sparse_utils import (
|
||||
get_curr_blocksparse_tensors,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.block_sparsity import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.block_sparsity import (
|
||||
BlockSparseTensors,
|
||||
BlockSparseTensorsTorch,
|
||||
to_cute_block_sparse_tensors,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.cute_dsl_utils import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.cute_dsl_utils import (
|
||||
get_aux_tensor_metadata,
|
||||
to_cute_aux_tensor,
|
||||
to_cute_tensor,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.mask import call_mask_mod
|
||||
from sglang.jit_kernel.flash_attn.cute.seqlen_info import SeqlenInfoQK
|
||||
from sglang.jit_kernel.flash_attn.cute.testing import is_fake_mode
|
||||
from sglang.jit_kernel.flash_attn.cute.utils import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.mask import call_mask_mod
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.seqlen_info import SeqlenInfoQK
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.testing import is_fake_mode
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.utils import (
|
||||
AuxData,
|
||||
get_batch_from_cu_tensor,
|
||||
hash_callable,
|
||||
+20
-13
@@ -20,22 +20,29 @@ from cutlass.cute.nvgpu import cpasync, warp
|
||||
from cutlass.cutlass_dsl import BaseDSL
|
||||
from quack import copy_utils, layout_utils
|
||||
|
||||
from sglang.jit_kernel.flash_attn.cute import ampere_helpers as sm80_utils
|
||||
from sglang.jit_kernel.flash_attn.cute import utils
|
||||
from sglang.jit_kernel.flash_attn.cute.block_info import BlockInfo
|
||||
from sglang.jit_kernel.flash_attn.cute.block_sparsity import BlockSparseTensors
|
||||
from sglang.jit_kernel.flash_attn.cute.cute_dsl_utils import assume_tensor_aligned
|
||||
from sglang.jit_kernel.flash_attn.cute.mask import AttentionMask
|
||||
from sglang.jit_kernel.flash_attn.cute.named_barrier import NamedBarrierFwd
|
||||
from sglang.jit_kernel.flash_attn.cute.pack_gqa import PackGQA
|
||||
from sglang.jit_kernel.flash_attn.cute.seqlen_info import SeqlenInfoQK
|
||||
from sglang.jit_kernel.flash_attn.cute.softmax import Softmax, apply_score_mod_inner
|
||||
from sglang.jit_kernel.flash_attn.cute.tile_scheduler import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute import ampere_helpers as sm80_utils
|
||||
from sglang.kernels.ops.attention.flash_attn.cute import utils
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.block_info import BlockInfo
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.block_sparsity import (
|
||||
BlockSparseTensors,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.cute_dsl_utils import (
|
||||
assume_tensor_aligned,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.mask import AttentionMask
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.named_barrier import NamedBarrierFwd
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.pack_gqa import PackGQA
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.seqlen_info import SeqlenInfoQK
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.softmax import (
|
||||
Softmax,
|
||||
apply_score_mod_inner,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.tile_scheduler import (
|
||||
SingleTileScheduler,
|
||||
SingleTileVarlenScheduler,
|
||||
TileSchedulerArguments,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.utils import AuxData
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.utils import AuxData
|
||||
|
||||
|
||||
class FlashAttentionForwardBase:
|
||||
@@ -1523,7 +1530,7 @@ class FlashAttentionForwardSm80(FlashAttentionForwardBase):
|
||||
# SM90 forward pass moved to flash_fwd_sm90.py; re-export for backward compatibility
|
||||
def __getattr__(name):
|
||||
if name == "FlashAttentionForwardSm90":
|
||||
from sglang.jit_kernel.flash_attn.cute.flash_fwd_sm90 import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.flash_fwd_sm90 import (
|
||||
FlashAttentionForwardSm90,
|
||||
)
|
||||
|
||||
+5
-3
@@ -12,9 +12,11 @@ from cutlass import Boolean, Float32, Int32, const_expr
|
||||
from cutlass.cute import FastDivmodDivisor
|
||||
from cutlass.cute.nvgpu import cpasync
|
||||
|
||||
from sglang.jit_kernel.flash_attn.cute import utils
|
||||
from sglang.jit_kernel.flash_attn.cute.cute_dsl_utils import assume_tensor_aligned
|
||||
from sglang.jit_kernel.flash_attn.cute.seqlen_info import SeqlenInfo
|
||||
from sglang.kernels.ops.attention.flash_attn.cute import utils
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.cute_dsl_utils import (
|
||||
assume_tensor_aligned,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.seqlen_info import SeqlenInfo
|
||||
|
||||
|
||||
class FlashAttentionForwardCombine:
|
||||
+21
-15
@@ -19,22 +19,26 @@ from cutlass.cute.runtime import from_dlpack
|
||||
from cutlass.utils import ClcDynamicPersistentTileScheduler
|
||||
from quack import copy_utils
|
||||
|
||||
import sglang.jit_kernel.flash_attn.cute.blackwell_helpers as fa_sm100_utils
|
||||
from sglang.jit_kernel.flash_attn.cute import utils as fa_utils
|
||||
from sglang.jit_kernel.flash_attn.cute.block_info import BlockInfo
|
||||
from sglang.jit_kernel.flash_attn.cute.cute_dsl_utils import dump_kernel_attributes
|
||||
from sglang.jit_kernel.flash_attn.cute.fa_logging import fa_log, fa_printf
|
||||
from sglang.jit_kernel.flash_attn.cute.mask import AttentionMask
|
||||
from sglang.jit_kernel.flash_attn.cute.named_barrier import NamedBarrierFwdSm100_MLA2CTA
|
||||
from sglang.jit_kernel.flash_attn.cute.pack_gqa import (
|
||||
import sglang.kernels.ops.attention.flash_attn.cute.blackwell_helpers as fa_sm100_utils
|
||||
from sglang.kernels.ops.attention.flash_attn.cute import utils as fa_utils
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.block_info import BlockInfo
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.cute_dsl_utils import (
|
||||
dump_kernel_attributes,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.fa_logging import fa_log, fa_printf
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.mask import AttentionMask
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.named_barrier import (
|
||||
NamedBarrierFwdSm100_MLA2CTA,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.pack_gqa import (
|
||||
make_packgqa_tiled_tma_atom,
|
||||
pack_gqa_layout,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.paged_kv import PagedKVManager
|
||||
from sglang.jit_kernel.flash_attn.cute.seqlen_info import SeqlenInfoQK
|
||||
from sglang.jit_kernel.flash_attn.cute.softmax import SoftmaxSm100
|
||||
from sglang.jit_kernel.flash_attn.cute.testing import attention_ref
|
||||
from sglang.jit_kernel.flash_attn.cute.tile_scheduler import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.paged_kv import PagedKVManager
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.seqlen_info import SeqlenInfoQK
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.softmax import SoftmaxSm100
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.testing import attention_ref
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.tile_scheduler import (
|
||||
ClcState,
|
||||
ParamsBase,
|
||||
SchedulingMode,
|
||||
@@ -44,8 +48,10 @@ from sglang.jit_kernel.flash_attn.cute.tile_scheduler import (
|
||||
TileSchedulerArguments,
|
||||
TileSchedulerProtocol,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.topk_gather_kv import CpasyncGatherKVManager
|
||||
from sglang.jit_kernel.flash_attn.cute.utils import smid
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.topk_gather_kv import (
|
||||
CpasyncGatherKVManager,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.utils import smid
|
||||
|
||||
|
||||
class FlashAttentionMLAForwardSm100:
|
||||
+30
-17
@@ -43,30 +43,43 @@ from cutlass.utils import ClcDynamicPersistentTileScheduler
|
||||
from quack import copy_utils, layout_utils
|
||||
from quack.cute_dsl_utils import ParamsBase
|
||||
|
||||
import sglang.jit_kernel.flash_attn.cute.pipeline as pipeline_custom
|
||||
from sglang.jit_kernel.flash_attn.cute import blackwell_helpers as sm100_utils
|
||||
from sglang.jit_kernel.flash_attn.cute import mma_sm100_desc as sm100_desc
|
||||
from sglang.jit_kernel.flash_attn.cute import utils
|
||||
from sglang.jit_kernel.flash_attn.cute.block_info import BlockInfo
|
||||
from sglang.jit_kernel.flash_attn.cute.block_sparse_utils import (
|
||||
import sglang.kernels.ops.attention.flash_attn.cute.pipeline as pipeline_custom
|
||||
from sglang.kernels.ops.attention.flash_attn.cute import (
|
||||
blackwell_helpers as sm100_utils,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute import mma_sm100_desc as sm100_desc
|
||||
from sglang.kernels.ops.attention.flash_attn.cute import (
|
||||
utils,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.block_info import BlockInfo
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.block_sparse_utils import (
|
||||
get_total_block_count,
|
||||
handle_block_sparse_empty_tile_correction_sm100,
|
||||
produce_block_sparse_loads_sm100,
|
||||
softmax_block_sparse_sm100,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.block_sparsity import BlockSparseTensors
|
||||
from sglang.jit_kernel.flash_attn.cute.cute_dsl_utils import assume_tensor_aligned
|
||||
from sglang.jit_kernel.flash_attn.cute.fa_logging import fa_log, fa_printf
|
||||
from sglang.jit_kernel.flash_attn.cute.mask import AttentionMask
|
||||
from sglang.jit_kernel.flash_attn.cute.named_barrier import NamedBarrierFwdSm100
|
||||
from sglang.jit_kernel.flash_attn.cute.pack_gqa import PackGQA, pack_gqa_layout
|
||||
from sglang.jit_kernel.flash_attn.cute.paged_kv import PagedKVManager
|
||||
from sglang.jit_kernel.flash_attn.cute.seqlen_info import SeqlenInfoQK
|
||||
from sglang.jit_kernel.flash_attn.cute.softmax import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.block_sparsity import (
|
||||
BlockSparseTensors,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.cute_dsl_utils import (
|
||||
assume_tensor_aligned,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.fa_logging import fa_log, fa_printf
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.mask import AttentionMask
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.named_barrier import (
|
||||
NamedBarrierFwdSm100,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.pack_gqa import (
|
||||
PackGQA,
|
||||
pack_gqa_layout,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.paged_kv import PagedKVManager
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.seqlen_info import SeqlenInfoQK
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.softmax import (
|
||||
SoftmaxSm100,
|
||||
apply_score_mod_inner,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.tile_scheduler import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.tile_scheduler import (
|
||||
ClcState,
|
||||
SchedulingMode,
|
||||
SingleTileLPTScheduler,
|
||||
@@ -76,7 +89,7 @@ from sglang.jit_kernel.flash_attn.cute.tile_scheduler import (
|
||||
TileSchedulerArguments,
|
||||
TileSchedulerProtocol,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.utils import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.utils import (
|
||||
AuxData,
|
||||
cvt_bf16x2_ue8m0x2,
|
||||
smid,
|
||||
+3
-1
@@ -8,7 +8,9 @@
|
||||
import cutlass
|
||||
import cutlass.utils as utils_basic
|
||||
|
||||
from sglang.jit_kernel.flash_attn.cute.flash_fwd import FlashAttentionForwardSm80
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.flash_fwd import (
|
||||
FlashAttentionForwardSm80,
|
||||
)
|
||||
|
||||
|
||||
class FlashAttentionForwardSm120(FlashAttentionForwardSm80):
|
||||
+24
-15
@@ -18,33 +18,42 @@ from cutlass.utils import LayoutEnum
|
||||
from quack import copy_utils, layout_utils, sm90_utils
|
||||
from quack.cute_dsl_utils import ParamsBase
|
||||
|
||||
from sglang.jit_kernel.flash_attn.cute import pipeline as pipeline_custom
|
||||
from sglang.jit_kernel.flash_attn.cute import utils
|
||||
from sglang.jit_kernel.flash_attn.cute.block_info import BlockInfo
|
||||
from sglang.jit_kernel.flash_attn.cute.block_sparse_utils import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute import pipeline as pipeline_custom
|
||||
from sglang.kernels.ops.attention.flash_attn.cute import utils
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.block_info import BlockInfo
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.block_sparse_utils import (
|
||||
consume_block_sparse_loads,
|
||||
produce_block_sparse_loads,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.block_sparsity import BlockSparseTensors
|
||||
from sglang.jit_kernel.flash_attn.cute.cute_dsl_utils import assume_tensor_aligned
|
||||
from sglang.jit_kernel.flash_attn.cute.flash_fwd import FlashAttentionForwardBase
|
||||
from sglang.jit_kernel.flash_attn.cute.mask import AttentionMask
|
||||
from sglang.jit_kernel.flash_attn.cute.named_barrier import NamedBarrierFwd
|
||||
from sglang.jit_kernel.flash_attn.cute.pack_gqa import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.block_sparsity import (
|
||||
BlockSparseTensors,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.cute_dsl_utils import (
|
||||
assume_tensor_aligned,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.flash_fwd import (
|
||||
FlashAttentionForwardBase,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.mask import AttentionMask
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.named_barrier import NamedBarrierFwd
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.pack_gqa import (
|
||||
PackGQA,
|
||||
make_packgqa_tiled_tma_atom,
|
||||
pack_gqa_layout,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.paged_kv import PagedKVManager
|
||||
from sglang.jit_kernel.flash_attn.cute.seqlen_info import SeqlenInfoQK
|
||||
from sglang.jit_kernel.flash_attn.cute.softmax import Softmax, apply_score_mod_inner
|
||||
from sglang.jit_kernel.flash_attn.cute.tile_scheduler import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.paged_kv import PagedKVManager
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.seqlen_info import SeqlenInfoQK
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.softmax import (
|
||||
Softmax,
|
||||
apply_score_mod_inner,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.tile_scheduler import (
|
||||
SingleTileLPTScheduler,
|
||||
SingleTileScheduler,
|
||||
SingleTileVarlenScheduler,
|
||||
TileSchedulerArguments,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.utils import AuxData
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.utils import AuxData
|
||||
|
||||
|
||||
class FlashAttentionForwardSm90(FlashAttentionForwardBase):
|
||||
+26
-16
@@ -13,50 +13,60 @@ import torch
|
||||
from cutlass import Float32, Int32
|
||||
from quack.compile_utils import make_fake_tensor as fake_tensor
|
||||
|
||||
from sglang.jit_kernel.flash_attn.cute.cache_utils import get_jit_cache
|
||||
from sglang.jit_kernel.flash_attn.cute.testing import is_fake_mode
|
||||
from sglang.kernels.jit.utils import is_arch_support_pdl
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.cache_utils import get_jit_cache
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.testing import is_fake_mode
|
||||
|
||||
if os.environ.get("CUTE_DSL_PTXAS_PATH", None) is not None:
|
||||
from sglang.jit_kernel.flash_attn.cute import cute_dsl_ptxas # noqa: F401
|
||||
from sglang.kernels.ops.attention.flash_attn.cute import ( # noqa: F401
|
||||
cute_dsl_ptxas,
|
||||
)
|
||||
|
||||
# Patch to dump ptx and then use system ptxas to compile to cubin
|
||||
cute_dsl_ptxas.patch()
|
||||
|
||||
|
||||
from sglang.jit_kernel.flash_attn.cute import fa_logging, utils
|
||||
from sglang.jit_kernel.flash_attn.cute.block_sparsity import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute import fa_logging, utils
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.block_sparsity import (
|
||||
BlockSparseTensorsTorch,
|
||||
get_sparse_q_block_size,
|
||||
normalize_block_sparse_config,
|
||||
to_cute_block_sparse_tensors,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.cu_blocks_kernels import CuSeqlensToBlocksKernel
|
||||
from sglang.jit_kernel.flash_attn.cute.cute_dsl_utils import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.cu_blocks_kernels import (
|
||||
CuSeqlensToBlocksKernel,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.cute_dsl_utils import (
|
||||
get_aux_tensor_metadata,
|
||||
to_cute_aux_tensor,
|
||||
to_cute_tensor,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.flash_fwd import FlashAttentionForwardSm80
|
||||
from sglang.jit_kernel.flash_attn.cute.flash_fwd_combine import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.flash_fwd import (
|
||||
FlashAttentionForwardSm80,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.flash_fwd_combine import (
|
||||
FlashAttentionForwardCombine,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.flash_fwd_mla_sm100 import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.flash_fwd_mla_sm100 import (
|
||||
FlashAttentionMLAForwardSm100,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.flash_fwd_sm90 import FlashAttentionForwardSm90
|
||||
from sglang.jit_kernel.flash_attn.cute.flash_fwd_sm100 import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.flash_fwd_sm90 import (
|
||||
FlashAttentionForwardSm90,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.flash_fwd_sm100 import (
|
||||
DescaleTensors,
|
||||
FlashAttentionForwardSm100,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.flash_fwd_sm120 import FlashAttentionForwardSm120
|
||||
from sglang.jit_kernel.flash_attn.cute.shearing_bias import ShearingBias
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.flash_fwd_sm120 import (
|
||||
FlashAttentionForwardSm120,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.shearing_bias import ShearingBias
|
||||
|
||||
# SM100 head_dim=256 2CTA kernel imports
|
||||
from sglang.jit_kernel.flash_attn.cute.sm100_hd256_2cta_fmha_forward import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.sm100_hd256_2cta_fmha_forward import (
|
||||
BlackwellFusedMultiHeadAttentionForward,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.utils import AuxData
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.utils import AuxData
|
||||
|
||||
|
||||
def _parse_arch_str(arch_str):
|
||||
+4
-4
@@ -10,10 +10,10 @@ from cutlass import Float32, Int32, Uint32, const_expr
|
||||
from cutlass.cutlass_dsl import min as dsl_min
|
||||
from quack import layout_utils
|
||||
|
||||
import sglang.jit_kernel.flash_attn.cute.utils as utils
|
||||
from sglang.jit_kernel.flash_attn.cute.block_info import BlockInfo
|
||||
from sglang.jit_kernel.flash_attn.cute.seqlen_info import SeqlenInfoQK
|
||||
from sglang.jit_kernel.flash_attn.cute.utils import AuxData
|
||||
import sglang.kernels.ops.attention.flash_attn.cute.utils as utils
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.block_info import BlockInfo
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.seqlen_info import SeqlenInfoQK
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.utils import AuxData
|
||||
|
||||
MaskGenFn: TypeAlias = Callable[[int], Uint32]
|
||||
MASK_R2P_CHUNK_SIZE: int = 32
|
||||
+1
-1
@@ -8,7 +8,7 @@ import cutlass.cute as cute
|
||||
from cutlass.cute.nvgpu import cpasync
|
||||
from quack import layout_utils
|
||||
|
||||
import sglang.jit_kernel.flash_attn.cute.utils as utils
|
||||
import sglang.kernels.ops.attention.flash_attn.cute.utils as utils
|
||||
|
||||
|
||||
def pack_gqa_layout(T, qhead_per_kvhead, nheads_kv, head_idx):
|
||||
+1
-1
@@ -9,7 +9,7 @@ from cutlass.cute import FastDivmodDivisor
|
||||
from cutlass.cute.nvgpu import cpasync
|
||||
from quack.cute_dsl_utils import ParamsBase
|
||||
|
||||
from sglang.jit_kernel.flash_attn.cute import utils
|
||||
from sglang.kernels.ops.attention.flash_attn.cute import utils
|
||||
|
||||
|
||||
@dataclass
|
||||
+9
-7
@@ -9,18 +9,20 @@ import cutlass
|
||||
import cutlass.cute as cute
|
||||
from cutlass import Float32, Int32, const_expr
|
||||
|
||||
from sglang.jit_kernel.flash_attn.cute.block_info import BlockInfo
|
||||
from sglang.jit_kernel.flash_attn.cute.copy_utils import tiled_copy_2d
|
||||
from sglang.jit_kernel.flash_attn.cute.cute_dsl_utils import assume_tensor_aligned
|
||||
from sglang.jit_kernel.flash_attn.cute.pack_gqa import pack_gqa_layout
|
||||
from sglang.jit_kernel.flash_attn.cute.seqlen_info import SeqlenInfoQK
|
||||
from sglang.jit_kernel.flash_attn.cute.tile_scheduler import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.block_info import BlockInfo
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.copy_utils import tiled_copy_2d
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.cute_dsl_utils import (
|
||||
assume_tensor_aligned,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.pack_gqa import pack_gqa_layout
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.seqlen_info import SeqlenInfoQK
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.tile_scheduler import (
|
||||
ParamsBase,
|
||||
SingleTileScheduler,
|
||||
SingleTileVarlenScheduler,
|
||||
TileSchedulerArguments,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.utils import get_batch_from_cu_tensor
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.utils import get_batch_from_cu_tensor
|
||||
|
||||
|
||||
class ShearingBias:
|
||||
+13
-11
@@ -13,37 +13,39 @@ import cutlass.utils.blackwell_helpers as sm100_utils
|
||||
from cutlass.cute.typing import Float32, Int32, Int64
|
||||
from cutlass.utils import ClcDynamicPersistentTileScheduler
|
||||
|
||||
from sglang.jit_kernel.flash_attn.cute.flash_fwd_sm100 import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.flash_fwd_sm100 import (
|
||||
_TUNING_CONFIG,
|
||||
DescaleTensors,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.mask import Sm100FusedMask as FusedMask
|
||||
from sglang.jit_kernel.flash_attn.cute.tile_scheduler import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.mask import (
|
||||
Sm100FusedMask as FusedMask,
|
||||
)
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.tile_scheduler import (
|
||||
SM100_TMEM_CAPACITY_COLUMNS,
|
||||
ClcState,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.tile_scheduler import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.tile_scheduler import (
|
||||
Sm100FmhaClcDynamicTileScheduler as FmhaClcDynamicTileScheduler,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.tile_scheduler import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.tile_scheduler import (
|
||||
Sm100FmhaClcDynamicTileSchedulerParams as FmhaClcDynamicTileSchedulerParams,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.tile_scheduler import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.tile_scheduler import (
|
||||
Sm100FmhaStaticTileScheduler as FmhaStaticTileScheduler,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.tile_scheduler import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.tile_scheduler import (
|
||||
Sm100FmhaStaticTileSchedulerParams as FmhaStaticTileSchedulerParams,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.tile_scheduler import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.tile_scheduler import (
|
||||
compute_sm100_fmha_grid as compute_grid,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.tile_scheduler import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.tile_scheduler import (
|
||||
compute_sm100_fmha_grid_clc as compute_grid_clc,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.tile_scheduler import (
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.tile_scheduler import (
|
||||
make_sm100_thread_cooperative_group as make_thread_cooperative_group,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attn.cute.utils import AuxData, ex2_emulation_2
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.utils import AuxData, ex2_emulation_2
|
||||
|
||||
|
||||
class BlackwellFusedMultiHeadAttentionForward:
|
||||
+3
-3
@@ -11,9 +11,9 @@ from cutlass import Boolean, Float32
|
||||
from quack import layout_utils
|
||||
from quack.cute_dsl_utils import ParamsBase
|
||||
|
||||
import sglang.jit_kernel.flash_attn.cute.utils as utils
|
||||
from sglang.jit_kernel.flash_attn.cute.seqlen_info import SeqlenInfoQK
|
||||
from sglang.jit_kernel.flash_attn.cute.utils import AuxData
|
||||
import sglang.kernels.ops.attention.flash_attn.cute.utils as utils
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.seqlen_info import SeqlenInfoQK
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.utils import AuxData
|
||||
|
||||
|
||||
@cute.jit
|
||||
+2
-2
@@ -30,8 +30,8 @@ from cutlass.utils import (
|
||||
from cutlass.utils.hardware_info import HardwareInfo
|
||||
from quack.cute_dsl_utils import ParamsBase
|
||||
|
||||
import sglang.jit_kernel.flash_attn.cute.utils as utils
|
||||
from sglang.jit_kernel.flash_attn.cute.fast_math import clz
|
||||
import sglang.kernels.ops.attention.flash_attn.cute.utils as utils
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.fast_math import clz
|
||||
|
||||
|
||||
class SchedulingMode(IntEnum):
|
||||
+2
-2
@@ -10,8 +10,8 @@ from cutlass import Boolean, Int32, Uint32, const_expr
|
||||
from cutlass.cute.nvgpu import cpasync
|
||||
from quack.cute_dsl_utils import ParamsBase
|
||||
|
||||
from sglang.jit_kernel.flash_attn.cute import utils
|
||||
from sglang.jit_kernel.flash_attn.cute.utils import warp_reduce
|
||||
from sglang.kernels.ops.attention.flash_attn.cute import utils
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.utils import warp_reduce
|
||||
|
||||
|
||||
@dataclass
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user