[diffusion] model: support VDN-H3 with a hybrid_window_attn_h3 backend (#37903)
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Haocheng Xi <xihc@berkeley.edu> Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Kevin Mi
co-authored by
Claude Fable 5.1
Haocheng Xi
Mick
Claude Fable 5.1
Haocheng Xi
Mick
parent
e91c948057
commit
ff1ce11348
@@ -47,8 +47,22 @@ struct QKNormRopePackKVParams : QKNormRopeParams {
|
||||
uint32_t suffix_tokens;
|
||||
};
|
||||
|
||||
template <bool kPackKV>
|
||||
using QKNormRopeParamsT = std::conditional_t<kPackKV, QKNormRopePackKVParams, QKNormRopeParams>;
|
||||
/// \brief Out-of-place variant: q/k are read (any strides), the normed + roped
|
||||
/// rows are written to q_out/k_out and the inputs stay untouched. Used where a
|
||||
/// consumer still needs the raw projections (VDN-H3's NoPE linear branch).
|
||||
struct QKNormRopeOutOfPlaceParams : QKNormRopeParams {
|
||||
void* __restrict__ q_out_ptr;
|
||||
void* __restrict__ k_out_ptr; // pre-offset by -num_qo_heads * out_head_stride_bytes
|
||||
int64_t q_out_stride_bytes;
|
||||
int64_t k_out_stride_bytes;
|
||||
int64_t out_head_stride_bytes;
|
||||
};
|
||||
|
||||
template <bool kPackKV, bool kOutOfPlace = false>
|
||||
using QKNormRopeParamsT = std::conditional_t<
|
||||
kPackKV,
|
||||
QKNormRopePackKVParams,
|
||||
std::conditional_t<kOutOfPlace, QKNormRopeOutOfPlaceParams, QKNormRopeParams>>;
|
||||
|
||||
constexpr uint32_t kThreadsPerBlock = 256;
|
||||
constexpr uint32_t kWarpsPerBlock = kThreadsPerBlock / device::kWarpThreads;
|
||||
@@ -199,9 +213,11 @@ template <
|
||||
bool kRoundNormBeforeRope,
|
||||
bool kPackKV,
|
||||
bool kCacheHasFullWidth,
|
||||
typename IdType>
|
||||
__global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT<kPackKV> __grid_constant__ params) {
|
||||
typename IdType,
|
||||
bool kOutOfPlace = false>
|
||||
__global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT<kPackKV, kOutOfPlace> __grid_constant__ params) {
|
||||
using namespace device;
|
||||
static_assert(!(kPackKV && kOutOfPlace), "KV packing and out-of-place output are exclusive");
|
||||
|
||||
static_assert(std::is_same_v<DType, fp16_t> || std::is_same_v<DType, bf16_t>);
|
||||
static_assert(kHeadDim <= 256, "Only warp-level fused qknorm+rope is supported");
|
||||
@@ -291,6 +307,13 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT<kPackKV> __grid_c
|
||||
const void* input = load_q ? pointer::offset(q_ptr, token_id * q_stride_bytes, head_id * head_stride_bytes)
|
||||
: pointer::offset(k_ptr, token_id * k_stride_bytes, head_id * head_stride_bytes);
|
||||
void* output = const_cast<void*>(input);
|
||||
if constexpr (kOutOfPlace) {
|
||||
output =
|
||||
load_q ? pointer::offset(
|
||||
params.q_out_ptr, token_id * params.q_out_stride_bytes, head_id * params.out_head_stride_bytes)
|
||||
: pointer::offset(
|
||||
params.k_out_ptr, token_id * params.k_out_stride_bytes, head_id * params.out_head_stride_bytes);
|
||||
}
|
||||
if constexpr (kPackKV) {
|
||||
if (!load_q) {
|
||||
const uint32_t batch_id = token_id / params.suffix_tokens;
|
||||
@@ -446,6 +469,23 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT<kPackKV> __grid_c
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
}
|
||||
|
||||
/// \brief Shared launch tail of the three host runners: pick the index-type
|
||||
/// instantiation, size the persistent grid from the occupancy table, launch.
|
||||
template <auto kKernelI32, auto kKernelI64, typename Params>
|
||||
void launch_qknorm_rope(const Params& params, bool is_int32, uint32_t num_works, DLDevice device, bool use_pdl) {
|
||||
using namespace host;
|
||||
const auto selected_kernel = is_int32 ? kKernelI32 : kKernelI64;
|
||||
const uint32_t kNumSM = runtime::get_sm_count(device.device_id);
|
||||
static const uint32_t kOccupancyTable[2] = {
|
||||
runtime::get_blocks_per_sm(kKernelI32, kThreadsPerBlock),
|
||||
runtime::get_blocks_per_sm(kKernelI64, kThreadsPerBlock),
|
||||
};
|
||||
const auto max_blocks = kOccupancyTable[is_int32 ? 0 : 1] * kNumSM;
|
||||
const auto needed_blocks = div_ceil(num_works, kWarpsPerBlock);
|
||||
const auto num_blocks = std::min(max_blocks, needed_blocks);
|
||||
LaunchKernel(num_blocks, kThreadsPerBlock, device).enable_pdl(use_pdl)(selected_kernel, params);
|
||||
}
|
||||
|
||||
template <
|
||||
int64_t kHeadDim,
|
||||
int64_t kRopeDim,
|
||||
@@ -528,18 +568,106 @@ struct QKNormRopeKernel {
|
||||
.eps = eps,
|
||||
};
|
||||
|
||||
const auto is_int32 = id_type.is_type<int32_t>();
|
||||
const auto selected_kernel = is_int32 ? kernel<int32_t> : kernel<int64_t>;
|
||||
const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id);
|
||||
static const uint32_t kOccupancyTable[2] = {
|
||||
runtime::get_blocks_per_sm(kernel<int32_t>, kThreadsPerBlock),
|
||||
runtime::get_blocks_per_sm(kernel<int64_t>, kThreadsPerBlock),
|
||||
};
|
||||
const auto max_blocks = kOccupancyTable[is_int32 ? 0 : 1] * kNumSM;
|
||||
const auto num_works = (num_qo_heads + num_kv_heads) * num_tokens;
|
||||
const auto needed_blocks = div_ceil(num_works, kWarpsPerBlock);
|
||||
const auto num_blocks = std::min(max_blocks, needed_blocks);
|
||||
LaunchKernel(num_blocks, kThreadsPerBlock, device.unwrap()).enable_pdl(kUsePDL)(selected_kernel, params);
|
||||
launch_qknorm_rope<kernel<int32_t>, kernel<int64_t>>(
|
||||
params, id_type.is_type<int32_t>(), num_works, device.unwrap(), kUsePDL);
|
||||
}
|
||||
};
|
||||
|
||||
template <
|
||||
int64_t kHeadDim,
|
||||
int64_t kRopeDim,
|
||||
bool kIsNeox,
|
||||
bool kUsePDL,
|
||||
typename DType,
|
||||
typename CacheDType,
|
||||
bool kRoundNormBeforeRope,
|
||||
bool kCacheHasFullWidth>
|
||||
struct QKNormRopeOutOfPlaceKernel {
|
||||
static_assert(kHeadDim <= 256, "Only head_dim <= 256 is supported");
|
||||
template <typename IdType>
|
||||
static constexpr auto kernel = fused_qknorm_rope_warp<
|
||||
kHeadDim,
|
||||
kRopeDim,
|
||||
kIsNeox,
|
||||
kUsePDL,
|
||||
DType,
|
||||
CacheDType,
|
||||
kRoundNormBeforeRope,
|
||||
false,
|
||||
kCacheHasFullWidth,
|
||||
IdType,
|
||||
true>;
|
||||
|
||||
/// \brief QK-norm + RoPE from q/k into q_out/k_out; q and k are left untouched.
|
||||
static void
|
||||
run(const tvm::ffi::TensorView q,
|
||||
const tvm::ffi::TensorView k,
|
||||
const tvm::ffi::TensorView q_out,
|
||||
const tvm::ffi::TensorView k_out,
|
||||
const tvm::ffi::TensorView q_weight,
|
||||
const tvm::ffi::TensorView k_weight,
|
||||
const tvm::ffi::TensorView cos_sin_cache,
|
||||
const tvm::ffi::TensorView positions,
|
||||
float eps) {
|
||||
using namespace host;
|
||||
|
||||
auto N = SymbolicSize{"num_tokens"};
|
||||
auto Q = SymbolicSize{"num_qo_heads"};
|
||||
auto K = SymbolicSize{"num_kv_heads"};
|
||||
auto D = SymbolicSize{"head_dim"};
|
||||
auto Dq = SymbolicSize{"q_stride"};
|
||||
auto Dk = SymbolicSize{"k_stride"};
|
||||
auto Dd = SymbolicSize{"head_stride"};
|
||||
auto Dqo = SymbolicSize{"q_out_stride"};
|
||||
auto Dko = SymbolicSize{"k_out_stride"};
|
||||
auto Ddo = SymbolicSize{"out_head_stride"};
|
||||
auto device = SymbolicDevice{};
|
||||
auto id_type = SymbolicDType{};
|
||||
D.set_value(kHeadDim);
|
||||
device.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({N, Q, D}).with_strides({Dq, Dd, 1}).with_dtype<DType>().with_device(device).verify(q);
|
||||
TensorMatcher({N, K, D}).with_strides({Dk, Dd, 1}).with_dtype<DType>().with_device(device).verify(k);
|
||||
TensorMatcher({N, Q, D}).with_strides({Dqo, Ddo, 1}).with_dtype<DType>().with_device(device).verify(q_out);
|
||||
TensorMatcher({N, K, D}).with_strides({Dko, Ddo, 1}).with_dtype<DType>().with_device(device).verify(k_out);
|
||||
TensorMatcher({D}).with_dtype<DType>().with_device(device).verify(q_weight).verify(k_weight);
|
||||
TensorMatcher({-1, kCacheHasFullWidth ? 2 * kRopeDim : kRopeDim})
|
||||
.with_dtype<CacheDType>()
|
||||
.with_device(device)
|
||||
.verify(cos_sin_cache);
|
||||
TensorMatcher({N}).with_dtype<int32_t, int64_t>(id_type).with_device(device).verify(positions);
|
||||
|
||||
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
|
||||
const auto num_qo_heads = static_cast<uint32_t>(Q.unwrap());
|
||||
const auto num_kv_heads = static_cast<uint32_t>(K.unwrap());
|
||||
if (num_tokens == 0 || (num_qo_heads == 0 && num_kv_heads == 0)) return;
|
||||
const auto head_stride_bytes = static_cast<int64_t>(Dd.unwrap() * sizeof(DType));
|
||||
const auto out_head_stride_bytes = static_cast<int64_t>(Ddo.unwrap() * sizeof(DType));
|
||||
|
||||
QKNormRopeOutOfPlaceParams params{};
|
||||
params.q_ptr = q.data_ptr();
|
||||
params.k_ptr = pointer::offset(k.data_ptr(), -static_cast<int64_t>(num_qo_heads) * head_stride_bytes);
|
||||
params.q_weight_ptr = q_weight.data_ptr();
|
||||
params.k_weight_ptr = k_weight.data_ptr();
|
||||
params.cos_sin_cache_ptr = cos_sin_cache.data_ptr();
|
||||
params.positions = positions.data_ptr();
|
||||
params.q_stride_bytes = static_cast<int64_t>(Dq.unwrap() * sizeof(DType));
|
||||
params.k_stride_bytes = static_cast<int64_t>(Dk.unwrap() * sizeof(DType));
|
||||
params.head_stride_bytes = head_stride_bytes;
|
||||
params.num_qo_heads = num_qo_heads;
|
||||
params.num_kv_heads = num_kv_heads;
|
||||
params.num_tokens = num_tokens;
|
||||
params.eps = eps;
|
||||
params.q_out_ptr = q_out.data_ptr();
|
||||
params.k_out_ptr = pointer::offset(k_out.data_ptr(), -static_cast<int64_t>(num_qo_heads) * out_head_stride_bytes);
|
||||
params.q_out_stride_bytes = static_cast<int64_t>(Dqo.unwrap() * sizeof(DType));
|
||||
params.k_out_stride_bytes = static_cast<int64_t>(Dko.unwrap() * sizeof(DType));
|
||||
params.out_head_stride_bytes = out_head_stride_bytes;
|
||||
|
||||
const auto num_works = (num_qo_heads + num_kv_heads) * num_tokens;
|
||||
launch_qknorm_rope<kernel<int32_t>, kernel<int64_t>>(
|
||||
params, id_type.is_type<int32_t>(), num_works, device.unwrap(), kUsePDL);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -655,20 +783,11 @@ struct QKNormRopePackKVKernel {
|
||||
params.prefix_tokens = static_cast<uint32_t>(prefix_tokens);
|
||||
params.suffix_tokens = static_cast<uint32_t>(suffix_tokens);
|
||||
|
||||
const auto is_int32 = id_type.is_type<int32_t>();
|
||||
const auto selected_kernel = is_int32 ? kernel<int32_t> : kernel<int64_t>;
|
||||
const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id);
|
||||
static const uint32_t kOccupancyTable[2] = {
|
||||
runtime::get_blocks_per_sm(kernel<int32_t>, kThreadsPerBlock),
|
||||
runtime::get_blocks_per_sm(kernel<int64_t>, kThreadsPerBlock),
|
||||
};
|
||||
const auto max_blocks = kOccupancyTable[is_int32 ? 0 : 1] * kNumSM;
|
||||
const uint32_t num_prefix_works = static_cast<uint32_t>(batch_size * prefix_tokens) * num_kv_heads;
|
||||
const uint32_t num_works =
|
||||
(num_qo_heads + num_kv_heads) * num_tokens + 2 * num_prefix_works + num_tokens * num_kv_heads;
|
||||
const auto needed_blocks = div_ceil(num_works, kWarpsPerBlock);
|
||||
const auto num_blocks = std::min(max_blocks, needed_blocks);
|
||||
LaunchKernel(num_blocks, kThreadsPerBlock, device.unwrap()).enable_pdl(kUsePDL)(selected_kernel, params);
|
||||
launch_qknorm_rope<kernel<int32_t>, kernel<int64_t>>(
|
||||
params, id_type.is_type<int32_t>(), num_works, device.unwrap(), kUsePDL);
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -0,0 +1,367 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// Fused VDN-H3 delta-rule factors.
|
||||
//
|
||||
// Per (frame, head) the linear branch needs, for M = I + A (128x128 fp32, symmetric positive definite):
|
||||
// transition = diag(alpha) M^-1 [F, H, dk, dk]
|
||||
// injection = B M^-1 [F, H, dv, dk]
|
||||
// The eager path is cholesky + solve_triangular + two GEMMs (~40 launches); this kernel is one
|
||||
// launch: one CTA of 256 threads per matrix, thread (ti, tj) owns rows 8ti.., cols 8tj.. as float2
|
||||
// pairs t[8][4] so the rank-2 updates and the final GEMM run on packed FFMA2 (sm_100+, scalar
|
||||
// fallback elsewhere).
|
||||
//
|
||||
// M^-1 is formed in place by block Gauss-Jordan elimination without pivoting (stable for SPD, the
|
||||
// same class as Cholesky), two pivots per barrier. For the pivot block S = {k, k+1} with
|
||||
// P = M[S,S], R = M[S,:], C = M[:,S], D the rest:
|
||||
// R' = P^-1 R, G = C (raw), D' = D - G R', M'[~S, S] = 0 - G R'[:,S] = -G P^-1
|
||||
// (in-place GJ inverse semantics: the eliminated column receives the inverse column). The next
|
||||
// block's band is prepared in the middle of the current update (software pipelined) and the barrier
|
||||
// sits before the second half of the update so its FMAs overlap the barrier skew and the loads of
|
||||
// the next step. Shared rows are column-swizzled so the float4 reads are bank-conflict free; the
|
||||
// same swizzle is used for the smem copy of M^-1 consumed by the register-tiled GEMM B M^-1.
|
||||
//
|
||||
// Accuracy: the relative error vs fp64 matches the cholesky path; both are dominated by cond(M).
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
namespace sglang {
|
||||
|
||||
namespace vdn_delta_factors {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int kDim = 128; // dk == dv == 128
|
||||
constexpr int kBlockSize = 256; // 16 x 16 tiles of 8 x 8
|
||||
constexpr int kMinBlocksPerSm = 2; // 128 registers per thread
|
||||
constexpr int kXsBytes = kDim * kDim * static_cast<int>(sizeof(float)); // 64 KB dynamic smem
|
||||
constexpr unsigned kFullMask = 0xffffffffu;
|
||||
|
||||
// thread tj's tile columns 8tj..8tj+3 land at 4tj.., 8tj+4..8tj+7 at 64+4tj.., so a quarter warp
|
||||
// reads 8 consecutive 16-byte chunks
|
||||
SGL_DEVICE int swz_lo(int tj) {
|
||||
return 4 * tj;
|
||||
}
|
||||
SGL_DEVICE int swz_hi(int tj) {
|
||||
return 64 + 4 * tj;
|
||||
}
|
||||
|
||||
SGL_DEVICE float rcp_nr(float p) {
|
||||
float r;
|
||||
asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(r) : "f"(p));
|
||||
return fmaf(r, fmaf(-p, r, 1.f), r); // one Newton step: ~0.5 ulp
|
||||
}
|
||||
SGL_DEVICE float4 ld4(const float* p) {
|
||||
return *reinterpret_cast<const float4*>(p);
|
||||
}
|
||||
SGL_DEVICE void st4(float* p, float a, float b, float c, float d) {
|
||||
*reinterpret_cast<float4*>(p) = make_float4(a, b, c, d);
|
||||
}
|
||||
SGL_DEVICE void st4(float* p, float2 a, float2 b) {
|
||||
*reinterpret_cast<float4*>(p) = make_float4(a.x, a.y, b.x, b.y);
|
||||
}
|
||||
SGL_DEVICE float2 f2(float a) {
|
||||
return make_float2(a, a);
|
||||
}
|
||||
// packed fma: (a.x*b.x+c.x, a.y*b.y+c.y); one FFMA2 on sm_100+ (the first operand is a scalar
|
||||
// broadcast in SASS), two FFMA elsewhere. Bitwise identical results either way.
|
||||
SGL_DEVICE float2 fma2(float2 a, float2 b, float2 c) {
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
||||
return __ffma2_rn(a, b, c);
|
||||
#else
|
||||
return make_float2(fmaf(a.x, b.x, c.x), fmaf(a.y, b.y, c.y));
|
||||
#endif
|
||||
}
|
||||
SGL_DEVICE float get(const float4& v, int i) {
|
||||
return i == 0 ? v.x : i == 1 ? v.y : i == 2 ? v.z : v.w;
|
||||
}
|
||||
|
||||
struct Smem {
|
||||
float row[2][2][kDim]; // [buffer][band row m][swizzled col] scaled band rows R'
|
||||
float col[2][kDim * 2]; // [buffer][row*2 + m] multipliers G
|
||||
};
|
||||
|
||||
/// Prepare the band of pivot block {k, k+1}, k = 8*kt1 + R1 (R1 even). Column owners (tj == kt1)
|
||||
/// publish G = C (0 for the band rows) and zero their band columns; row owners (ti == kt1) replace P
|
||||
/// by I, scale R' = P^-1 R and publish it. Executed by all threads (P is broadcast with shuffles).
|
||||
template <int R1, int BUF>
|
||||
SGL_DEVICE void prepare(float2 (&t)[8][4], Smem& sm, int kt1, int ti, int tj, int i0) {
|
||||
constexpr int CP = R1 >> 1;
|
||||
const int src = ((kt1 & 1) << 4) | kt1; // lane of the diagonal tile inside the row-owner warp
|
||||
const float pa = __shfl_sync(kFullMask, t[R1][CP].x, src);
|
||||
const float pb = __shfl_sync(kFullMask, t[R1][CP].y, src);
|
||||
const float pc = __shfl_sync(kFullMask, t[R1 + 1][CP].x, src);
|
||||
const float pd = __shfl_sync(kFullMask, t[R1 + 1][CP].y, src);
|
||||
if (tj == kt1) {
|
||||
const bool diag = (ti == kt1);
|
||||
float* dst = &sm.col[BUF][i0 * 2];
|
||||
float2 g[8];
|
||||
#pragma unroll
|
||||
for (int r = 0; r < 8; ++r) {
|
||||
const bool band = diag && (r == R1 || r == R1 + 1);
|
||||
g[r] = band ? make_float2(0.f, 0.f) : t[r][CP];
|
||||
if (!band) t[r][CP] = make_float2(0.f, 0.f);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int q = 0; q < 4; ++q)
|
||||
st4(dst + 4 * q, g[2 * q], g[2 * q + 1]);
|
||||
}
|
||||
if (ti == kt1) {
|
||||
const float det = fmaf(pa, pd, -pb * pc);
|
||||
const float idet = rcp_nr(det);
|
||||
const float ia = pd * idet, ib = -pb * idet, ic = -pc * idet, id = pa * idet;
|
||||
if (tj == kt1) {
|
||||
t[R1][CP] = make_float2(1.f, 0.f);
|
||||
t[R1 + 1][CP] = make_float2(0.f, 1.f);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int cp = 0; cp < 4; ++cp) {
|
||||
const float2 x = t[R1][cp], y = t[R1 + 1][cp];
|
||||
t[R1][cp] = make_float2(fmaf(ia, x.x, ib * y.x), fmaf(ia, x.y, ib * y.y));
|
||||
t[R1 + 1][cp] = make_float2(fmaf(ic, x.x, id * y.x), fmaf(ic, x.y, id * y.y));
|
||||
}
|
||||
st4(&sm.row[BUF][0][swz_lo(tj)], t[R1][0], t[R1][1]);
|
||||
st4(&sm.row[BUF][0][swz_hi(tj)], t[R1][2], t[R1][3]);
|
||||
st4(&sm.row[BUF][1][swz_lo(tj)], t[R1 + 1][0], t[R1 + 1][1]);
|
||||
st4(&sm.row[BUF][1][swz_hi(tj)], t[R1 + 1][2], t[R1 + 1][3]);
|
||||
}
|
||||
}
|
||||
|
||||
/// Rank-2 update for pivots (8kt + KR, 8kt + KR + 1); the next block's band is prepared in the middle.
|
||||
template <int KR>
|
||||
SGL_DEVICE void step(float2 (&t)[8][4], Smem& sm, int kt, int ti, int tj, int i0) {
|
||||
constexpr int CUR = (KR >> 1) & 1, NXT = CUR ^ 1;
|
||||
constexpr int R1 = (KR + 2) & 7, CP = R1 >> 1; // next block's band rows / column pair
|
||||
float2 rk0[4], rk1[4];
|
||||
float ck0[8], ck1[8]; // negated multipliers (scalar broadcast operands of FFMA2)
|
||||
{
|
||||
const float4 a = ld4(&sm.row[CUR][0][swz_lo(tj)]), b = ld4(&sm.row[CUR][0][swz_hi(tj)]);
|
||||
const float4 c = ld4(&sm.row[CUR][1][swz_lo(tj)]), d = ld4(&sm.row[CUR][1][swz_hi(tj)]);
|
||||
rk0[0] = make_float2(a.x, a.y);
|
||||
rk0[1] = make_float2(a.z, a.w);
|
||||
rk0[2] = make_float2(b.x, b.y);
|
||||
rk0[3] = make_float2(b.z, b.w);
|
||||
rk1[0] = make_float2(c.x, c.y);
|
||||
rk1[1] = make_float2(c.z, c.w);
|
||||
rk1[2] = make_float2(d.x, d.y);
|
||||
rk1[3] = make_float2(d.z, d.w);
|
||||
const float* cp = &sm.col[CUR][i0 * 2];
|
||||
#pragma unroll
|
||||
for (int q = 0; q < 4; ++q) {
|
||||
const float4 v = ld4(cp + 4 * q);
|
||||
ck0[2 * q] = -v.x;
|
||||
ck1[2 * q] = -v.y;
|
||||
ck0[2 * q + 1] = -v.z;
|
||||
ck1[2 * q + 1] = -v.w;
|
||||
}
|
||||
}
|
||||
// part 1: the next band (rows R1, R1+1 fully; column pair CP in the other rows)
|
||||
#pragma unroll
|
||||
for (int cp = 0; cp < 4; ++cp) {
|
||||
t[R1][cp] = fma2(f2(ck1[R1]), rk1[cp], fma2(f2(ck0[R1]), rk0[cp], t[R1][cp]));
|
||||
t[R1 + 1][cp] = fma2(f2(ck1[R1 + 1]), rk1[cp], fma2(f2(ck0[R1 + 1]), rk0[cp], t[R1 + 1][cp]));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int r = 0; r < 8; ++r) {
|
||||
if (r == R1 || r == R1 + 1) continue;
|
||||
t[r][CP] = fma2(f2(ck1[r]), rk1[CP], fma2(f2(ck0[r]), rk0[CP], t[r][CP]));
|
||||
}
|
||||
if (KR != 6 || kt != 15) prepare<R1, NXT>(t, sm, kt + (KR == 6 ? 1 : 0), ti, tj, i0);
|
||||
__syncthreads();
|
||||
// part 2: everything else (registers only; overlaps the barrier skew and the next step's loads)
|
||||
#pragma unroll
|
||||
for (int r = 0; r < 8; ++r) {
|
||||
if (r == R1 || r == R1 + 1) continue;
|
||||
#pragma unroll
|
||||
for (int cp = 0; cp < 4; ++cp) {
|
||||
if (cp == CP) continue;
|
||||
t[r][cp] = fma2(f2(ck1[r]), rk1[cp], fma2(f2(ck0[r]), rk0[cp], t[r][cp]));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* \brief transition = diag(alpha) (I + A)^-1, injection = B (I + A)^-1 for a batch of 128x128 SPD A.
|
||||
*
|
||||
* \param A [N, 128, 128] fp32, symmetric positive semi-definite (I + A is inverted)
|
||||
* \param B [N, 128, 128] fp32
|
||||
* \param alpha [N, 128] fp32 row scales of the transition
|
||||
* \param transition [N, 128, 128] fp32 output
|
||||
* \param injection [N, 128, 128] fp32 output
|
||||
*/
|
||||
__global__ void __launch_bounds__(kBlockSize, kMinBlocksPerSm) vdn_delta_factors_kernel(
|
||||
const float* __restrict__ A,
|
||||
const float* __restrict__ B,
|
||||
const float* __restrict__ alpha,
|
||||
float* __restrict__ transition,
|
||||
float* __restrict__ injection) {
|
||||
extern __shared__ __align__(16) float Xs[]; // [kDim][kDim], column-swizzled copy of (I + A)^-1
|
||||
__shared__ __align__(16) Smem sm;
|
||||
|
||||
const int n = blockIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
const int ti = tid >> 4, tj = tid & 15;
|
||||
const int i0 = ti * 8, j0 = tj * 8;
|
||||
const size_t mat = static_cast<size_t>(n) * kDim * kDim;
|
||||
const float* An = A + mat;
|
||||
const float* Bn = B + mat;
|
||||
|
||||
float2 t[8][4];
|
||||
#pragma unroll
|
||||
for (int r = 0; r < 8; ++r) {
|
||||
const float4 v0 = ld4(An + (i0 + r) * kDim + j0), v1 = ld4(An + (i0 + r) * kDim + j0 + 4);
|
||||
t[r][0] = make_float2(v0.x, v0.y);
|
||||
t[r][1] = make_float2(v0.z, v0.w);
|
||||
t[r][2] = make_float2(v1.x, v1.y);
|
||||
t[r][3] = make_float2(v1.z, v1.w);
|
||||
}
|
||||
if (ti == tj) { // M = I + A
|
||||
#pragma unroll
|
||||
for (int r = 0; r < 8; ++r) {
|
||||
if (r & 1) {
|
||||
t[r][r >> 1].y += 1.f;
|
||||
} else {
|
||||
t[r][r >> 1].x += 1.f;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
prepare<0, 0>(t, sm, 0, ti, tj, i0);
|
||||
__syncthreads();
|
||||
#pragma unroll 1
|
||||
for (int kt = 0; kt < kDim / 8; ++kt) {
|
||||
step<0>(t, sm, kt, ti, tj, i0);
|
||||
step<2>(t, sm, kt, ti, tj, i0);
|
||||
step<4>(t, sm, kt, ti, tj, i0);
|
||||
step<6>(t, sm, kt, ti, tj, i0);
|
||||
}
|
||||
|
||||
// transition = diag(alpha) X, straight from registers
|
||||
{
|
||||
const float* al = alpha + static_cast<size_t>(n) * kDim + i0;
|
||||
const float4 a0 = ld4(al), a1 = ld4(al + 4);
|
||||
const float av[8] = {a0.x, a0.y, a0.z, a0.w, a1.x, a1.y, a1.z, a1.w};
|
||||
float* Tn = transition + mat;
|
||||
#pragma unroll
|
||||
for (int r = 0; r < 8; ++r) {
|
||||
float* row = Tn + (i0 + r) * kDim + j0;
|
||||
st4(row, av[r] * t[r][0].x, av[r] * t[r][0].y, av[r] * t[r][1].x, av[r] * t[r][1].y);
|
||||
st4(row + 4, av[r] * t[r][2].x, av[r] * t[r][2].y, av[r] * t[r][3].x, av[r] * t[r][3].y);
|
||||
}
|
||||
}
|
||||
// stage X in smem (swizzled columns) for the GEMM
|
||||
#pragma unroll
|
||||
for (int r = 0; r < 8; ++r) {
|
||||
float* row = Xs + (i0 + r) * kDim;
|
||||
st4(row + swz_lo(tj), t[r][0], t[r][1]);
|
||||
st4(row + swz_hi(tj), t[r][2], t[r][3]);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// injection = B X (register-tiled GEMM; B rows from global through L1, X from smem)
|
||||
#pragma unroll
|
||||
for (int r = 0; r < 8; ++r)
|
||||
#pragma unroll
|
||||
for (int cp = 0; cp < 4; ++cp)
|
||||
t[r][cp] = make_float2(0.f, 0.f);
|
||||
const float* Bp = Bn + i0 * kDim;
|
||||
#pragma unroll 1
|
||||
for (int k = 0; k < kDim; k += 4) {
|
||||
float4 b[8];
|
||||
#pragma unroll
|
||||
for (int r = 0; r < 8; ++r)
|
||||
b[r] = __ldg(reinterpret_cast<const float4*>(Bp + r * kDim + k));
|
||||
#pragma unroll
|
||||
for (int kk = 0; kk < 4; ++kk) {
|
||||
const float* xr = Xs + (k + kk) * kDim;
|
||||
const float4 x0 = ld4(xr + swz_lo(tj)), x1 = ld4(xr + swz_hi(tj));
|
||||
const float2 xv[4] = {
|
||||
make_float2(x0.x, x0.y), make_float2(x0.z, x0.w), make_float2(x1.x, x1.y), make_float2(x1.z, x1.w)};
|
||||
#pragma unroll
|
||||
for (int r = 0; r < 8; ++r) {
|
||||
const float2 bv = f2(get(b[r], kk));
|
||||
#pragma unroll
|
||||
for (int cp = 0; cp < 4; ++cp)
|
||||
t[r][cp] = fma2(bv, xv[cp], t[r][cp]);
|
||||
}
|
||||
}
|
||||
}
|
||||
{
|
||||
float* Jn = injection + mat;
|
||||
#pragma unroll
|
||||
for (int r = 0; r < 8; ++r) {
|
||||
float* row = Jn + (i0 + r) * kDim + j0;
|
||||
st4(row, t[r][0], t[r][1]);
|
||||
st4(row + 4, t[r][2], t[r][3]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
struct VdnDeltaFactorsKernel {
|
||||
/**
|
||||
* \brief Validate the tensors and launch one CTA per matrix.
|
||||
*
|
||||
* \param transition [N, 128, 128] fp32 output, diag(alpha) (I + A)^-1
|
||||
* \param injection [N, 128, 128] fp32 output, B (I + A)^-1
|
||||
* \param A [N, 128, 128] fp32 SPD statistics (I + A is inverted)
|
||||
* \param B [N, 128, 128] fp32
|
||||
* \param alpha [N, 128] fp32
|
||||
*/
|
||||
static void
|
||||
run(tvm::ffi::TensorView transition,
|
||||
tvm::ffi::TensorView injection,
|
||||
tvm::ffi::TensorView A,
|
||||
tvm::ffi::TensorView B,
|
||||
tvm::ffi::TensorView alpha) {
|
||||
using namespace host;
|
||||
auto N = SymbolicSize{"num_matrices"};
|
||||
auto device = SymbolicDevice{};
|
||||
device.set_options<kDLCUDA>();
|
||||
TensorMatcher({N, kDim, kDim})
|
||||
.with_dtype<fp32_t>()
|
||||
.with_device(device)
|
||||
.verify(transition)
|
||||
.verify(injection)
|
||||
.verify(A)
|
||||
.verify(B);
|
||||
TensorMatcher({N, kDim}).with_dtype<fp32_t>().with_device(device).verify(alpha);
|
||||
const int64_t num = N.unwrap();
|
||||
if (num == 0) return;
|
||||
CHECK_HOST(
|
||||
transition.data_ptr() != A.data_ptr() && transition.data_ptr() != B.data_ptr() &&
|
||||
injection.data_ptr() != A.data_ptr() && injection.data_ptr() != B.data_ptr() &&
|
||||
transition.data_ptr() != injection.data_ptr())
|
||||
<< "vdn_delta_factors outputs must not alias inputs";
|
||||
// every tensor is read or written as float4; a storage offset breaks this
|
||||
const auto aligned16 = [](const void* p) { return reinterpret_cast<uintptr_t>(p) % 16 == 0; };
|
||||
CHECK_HOST(
|
||||
aligned16(A.data_ptr()) && aligned16(B.data_ptr()) && aligned16(alpha.data_ptr()) &&
|
||||
aligned16(transition.data_ptr()) && aligned16(injection.data_ptr()))
|
||||
<< "vdn_delta_factors needs 16-byte aligned tensors";
|
||||
const DLDevice dev = device.unwrap();
|
||||
// 64 KB of dynamic shared memory needs the opt-in, once per device.
|
||||
static bool attr_set[64] = {};
|
||||
const int dev_id = dev.device_id;
|
||||
CHECK_HOST(dev_id >= 0 && dev_id < 64) << "vdn_delta_factors: unexpected device id " << dev_id;
|
||||
if (!attr_set[dev_id]) {
|
||||
CHECK_CUDA(cudaFuncSetAttribute(vdn_delta_factors_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kXsBytes))
|
||||
<< "vdn_delta_factors: cannot reserve " << kXsBytes << " bytes of dynamic shared memory";
|
||||
attr_set[dev_id] = true;
|
||||
}
|
||||
LaunchKernel(static_cast<uint32_t>(num), kBlockSize, dev, kXsBytes)(
|
||||
vdn_delta_factors_kernel,
|
||||
static_cast<const float*>(A.data_ptr()),
|
||||
static_cast<const float*>(B.data_ptr()),
|
||||
static_cast<const float*>(alpha.data_ptr()),
|
||||
static_cast<float*>(transition.data_ptr()),
|
||||
static_cast<float*>(injection.data_ptr()));
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace vdn_delta_factors
|
||||
|
||||
} // namespace sglang
|
||||
@@ -37,6 +37,7 @@ norm/ RMSNorm / LayerNorm / GroupNorm and their fused epilogues
|
||||
modulate/ adaLN modulate, gating, timestep conditioning
|
||||
rope/ rotary embeddings and the QK-norm chains fused into them
|
||||
activation/ SiLU / GLU / GELU fusions
|
||||
quantization/ MXFP8 producers whose scales land in the GEMM's swizzled layout
|
||||
attention/ sparse linear attention, gated delta-net
|
||||
routing/ diffusion-model MoE routing and expert selection
|
||||
layout/ pure data movement: USP/Ulysses relayout, varlen pack, causal pad
|
||||
@@ -140,6 +141,7 @@ tensor copy per residual site.
|
||||
|---|---|---|
|
||||
| `fused_inplace_qknorm_rope` | JIT CUDA | one bf16 rounding step vs split baseline; `round_norm_before_rope=True` makes it exact; supports compact and full-width NeoX/interleaved caches |
|
||||
| `fused_qknorm_rope_pack_kv` | JIT CUDA | as above, also packs prefix K/V |
|
||||
| `fused_qknorm_rope_out_of_place` | JIT CUDA | as above, bit-equal to the in-place kernel; reads strided q/k and writes contiguous copies, inputs untouched (VDN-H3 keeps the raw q/k for its linear branch) |
|
||||
| `try_fused_flux2_qkv_epilogue` | KDA (JIT CUDA) | bit-exact vs the selected BF16 chain | FLUX.2 QK RMSNorm + RoPE + joint QKV packing |
|
||||
| `try_fused_qwen_qkv_epilogue` | JIT CUDA | bit-exact vs the selected BF16 chain | Qwen-Image QK RMSNorm + RoPE + joint QKV writes; SM100+ |
|
||||
| `fused_rope_rotate_half_bitexact` | Triton | bit-exact (elementwise only) |
|
||||
@@ -150,6 +152,22 @@ tensor copy per residual site.
|
||||
| `apply_rotary_embedding` | Triton (+fallbacks) | close; the generic entry point |
|
||||
| `hunyuan_qkv_rope_pack` | Triton | bit-exact; packs QKV and applies RoPE in one pass |
|
||||
|
||||
### MiniMax-H3 / VDN-H3 linear branch
|
||||
|
||||
| Entry point | Backend | Contract |
|
||||
|---|---|---|
|
||||
| `vdn_frame_stats_prep`, `vdn_gather_linear_state` | Triton | bit-exact (same products, fp32 gather) |
|
||||
| `vdn_temporal_conv_act`, `vdn_silu_l2norm`, `vdn_linear_epilogue` | Triton | one rounding at the store, within one bf16 ulp of the eager chain; the model's own inference kernels, mounted unconditionally by the VDN-H3 branch |
|
||||
| `vdn_delta_factors` | JIT CUDA | `(alpha * inv(I + A), B @ inv(I + A))` in one launch; same fp32 accuracy class as the cholesky + solve_triangular chain (cond-dominated); head_dim 128 |
|
||||
|
||||
### MXFP8 producers (online `mxfp8`, cuBLASLt block-scaled GEMM on SM100)
|
||||
|
||||
| Entry point | Backend | Contract |
|
||||
|---|---|---|
|
||||
| `mxfp8_quantize_swizzled` | Triton | bit-exact vs `flashinfer.mxfp8_quantize(x, True)`: e4m3 payload + block-32 E8M0 scales in the `SWIZZLE_32_4_4` layout; weights at load and any bf16 GEMM input |
|
||||
| `silu_mul_mxfp8` | Triton | bit-exact vs eager bf16 `silu(gate) * up` followed by the quantizer above; the fc2 input |
|
||||
| `indexed_scale_shift_mxfp8_` | Triton | bit-exact vs `indexed_scale_shift_bf16_` followed by the quantizer above, optionally keeping the bf16 rows in place; the qkv / fc1 inputs |
|
||||
|
||||
### MoE routing
|
||||
|
||||
| Entry point | Backend | Contract | Applies to |
|
||||
|
||||
@@ -328,6 +328,76 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
|
||||
_CUDA,
|
||||
"Sana-WM bidirectional gated delta-net.",
|
||||
),
|
||||
(
|
||||
"diffusion.fused_qknorm_rope_out_of_place",
|
||||
KernelBackend.JIT,
|
||||
"rope.qknorm_rope_jit:fused_qknorm_rope_out_of_place",
|
||||
_CUDA,
|
||||
"Out-of-place fused QK-norm + RoPE (raw q/k preserved).",
|
||||
),
|
||||
(
|
||||
"diffusion.vdn_delta_factors",
|
||||
KernelBackend.JIT,
|
||||
"attention.vdn_delta_factors_jit:vdn_delta_factors",
|
||||
_CUDA,
|
||||
"VDN-H3 delta rule: fused (I + A)^-1 -> transition / injection (fp32, head_dim 128).",
|
||||
),
|
||||
(
|
||||
"diffusion.vdn_temporal_conv_act",
|
||||
KernelBackend.TRITON,
|
||||
"attention.vdn_linear_branch_triton:vdn_temporal_conv_act",
|
||||
_CUDA,
|
||||
"VDN-H3 linear branch: 5-tap temporal conv + SiLU + L2 norm.",
|
||||
),
|
||||
(
|
||||
"diffusion.vdn_silu_l2norm",
|
||||
KernelBackend.TRITON,
|
||||
"attention.vdn_linear_branch_triton:vdn_silu_l2norm",
|
||||
_CUDA,
|
||||
"VDN-H3 linear branch: SiLU + L2 norm over head_dim.",
|
||||
),
|
||||
(
|
||||
"diffusion.vdn_frame_stats_prep",
|
||||
KernelBackend.TRITON,
|
||||
"attention.vdn_linear_branch_triton:vdn_frame_stats_prep",
|
||||
_CUDA,
|
||||
"VDN-H3 linear branch: frame-statistics GEMM operands in one pass.",
|
||||
),
|
||||
(
|
||||
"diffusion.vdn_gather_linear_state",
|
||||
KernelBackend.TRITON,
|
||||
"attention.vdn_linear_branch_triton:vdn_gather_linear_state",
|
||||
_CUDA,
|
||||
"VDN-H3 linear branch: alpha-bridged boundary gather in one pass.",
|
||||
),
|
||||
(
|
||||
"diffusion.vdn_linear_epilogue",
|
||||
KernelBackend.TRITON,
|
||||
"attention.vdn_linear_branch_triton:vdn_linear_epilogue",
|
||||
_CUDA,
|
||||
"VDN-H3 linear branch: RMSNorm * gate readout epilogue.",
|
||||
),
|
||||
(
|
||||
"diffusion.mxfp8_quantize_swizzled",
|
||||
KernelBackend.TRITON,
|
||||
"quantization.mxfp8_swizzled_triton:mxfp8_quantize_swizzled",
|
||||
_CUDA,
|
||||
"bf16 -> MXFP8 (e4m3, block-32 E8M0 scales in the cuBLASLt swizzled layout).",
|
||||
),
|
||||
(
|
||||
"diffusion.silu_mul_mxfp8",
|
||||
KernelBackend.TRITON,
|
||||
"quantization.mxfp8_swizzled_triton:silu_mul_mxfp8",
|
||||
_CUDA,
|
||||
"SwiGLU + MXFP8 quant for the online mxfp8 fc2 input.",
|
||||
),
|
||||
(
|
||||
"diffusion.indexed_scale_shift_mxfp8_",
|
||||
KernelBackend.TRITON,
|
||||
"quantization.mxfp8_swizzled_triton:indexed_scale_shift_mxfp8_",
|
||||
_CUDA,
|
||||
"Indexed adaLN modulation + MXFP8 quant for the online mxfp8 qkv/fc1 inputs.",
|
||||
),
|
||||
(
|
||||
"diffusion.group_limited_topk",
|
||||
KernelBackend.TRITON,
|
||||
@@ -541,6 +611,24 @@ _EXPORTS: dict[str, str] = {
|
||||
"fused_causal_conv3d_cat_pad_cuda": "sglang.kernels.kda_kernels.causal_conv3d_cat_pad_jit",
|
||||
"fused_causal_conv3d_cat_pad": "layout.causal_conv3d_cat_pad_triton",
|
||||
"pack_qkv_destination_major": "layout.ulysses_qkv_triton",
|
||||
"fused_qknorm_rope_out_of_place": "rope.qknorm_rope_jit",
|
||||
"vdn_delta_factors": "attention.vdn_delta_factors_jit",
|
||||
"can_use_vdn_delta_factors": "attention.vdn_delta_factors_jit",
|
||||
"vdn_temporal_conv_act": "attention.vdn_linear_branch_triton",
|
||||
"can_use_vdn_temporal_conv_act": "attention.vdn_linear_branch_triton",
|
||||
"can_use_vdn_silu_l2norm": "attention.vdn_linear_branch_triton",
|
||||
"can_use_vdn_frame_stats_prep": "attention.vdn_linear_branch_triton",
|
||||
"can_use_vdn_gather_linear_state": "attention.vdn_linear_branch_triton",
|
||||
"can_use_vdn_linear_epilogue": "attention.vdn_linear_branch_triton",
|
||||
"vdn_silu_l2norm": "attention.vdn_linear_branch_triton",
|
||||
"vdn_frame_stats_prep": "attention.vdn_linear_branch_triton",
|
||||
"vdn_gather_linear_state": "attention.vdn_linear_branch_triton",
|
||||
"vdn_linear_epilogue": "attention.vdn_linear_branch_triton",
|
||||
"can_use_mxfp8_swizzled": "quantization.mxfp8_swizzled_triton",
|
||||
"can_use_silu_mul_mxfp8": "quantization.mxfp8_swizzled_triton",
|
||||
"indexed_scale_shift_mxfp8_": "quantization.mxfp8_swizzled_triton",
|
||||
"mxfp8_quantize_swizzled": "quantization.mxfp8_swizzled_triton",
|
||||
"silu_mul_mxfp8": "quantization.mxfp8_swizzled_triton",
|
||||
"can_use_usp_merge_heads": "layout.usp_relayout_jit",
|
||||
"usp_merge_heads": "layout.usp_relayout_jit",
|
||||
"build_inv_indices": "layout.varlen_pack_pad_triton",
|
||||
|
||||
@@ -0,0 +1,102 @@
|
||||
"""Fused VDN-H3 delta-rule factors: (I + A)^-1 folded into the transition and injection.
|
||||
|
||||
One CUDA kernel (block Gauss-Jordan inverse in registers + the two products) replaces the
|
||||
cholesky / solve_triangular / GEMM chain of ``delta_factor_apply`` for the ``vdn_solve`` and
|
||||
``vdn_scaled`` rules. fp32, head_dim 128 only.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.utils import cache_once, load_jit
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
HEAD_DIM = 128
|
||||
_FLOAT4_BYTES = 16
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_vdn_delta_factors_module() -> Module:
|
||||
if torch.cuda.get_device_capability()[0] < 8:
|
||||
raise RuntimeError(
|
||||
"vdn_delta_factors needs SM80 or later (2 x 70 KB shared memory per SM)"
|
||||
)
|
||||
return load_jit(
|
||||
"diffusion_vdn_delta_factors",
|
||||
cuda_files=["diffusion/vdn_delta_factors.cuh"],
|
||||
cuda_wrappers=[
|
||||
("vdn_delta_factors", "vdn_delta_factors::VdnDeltaFactorsKernel::run")
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _aligned(t: torch.Tensor) -> torch.Tensor:
|
||||
# the kernel loads float4; .contiguous() keeps a storage offset, a fresh allocation is aligned
|
||||
return t if t.data_ptr() % _FLOAT4_BYTES == 0 else t.clone()
|
||||
|
||||
|
||||
def _fake_impl(
|
||||
A: torch.Tensor, B: torch.Tensor, alpha: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
del alpha
|
||||
return torch.empty_like(A), torch.empty_like(B)
|
||||
|
||||
|
||||
@register_custom_op(
|
||||
op_name="diffusion_vdn_delta_factors",
|
||||
mutates_args=[],
|
||||
fake_impl=_fake_impl,
|
||||
)
|
||||
def vdn_delta_factors(
|
||||
A: torch.Tensor, B: torch.Tensor, alpha: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""``(alpha[..., :, None] * inv(I + A), B @ inv(I + A))`` for fp32 ``[..., 128, 128]`` SPD ``A``.
|
||||
|
||||
``B`` has the shape of ``A``; ``alpha`` is ``[..., 128]``. Same accuracy as the eager
|
||||
cholesky path (both are dominated by cond(I + A) in fp32).
|
||||
"""
|
||||
A, B, alpha = _aligned(A), _aligned(B), _aligned(alpha)
|
||||
transition = torch.empty_like(A)
|
||||
injection = torch.empty_like(B)
|
||||
module = _jit_vdn_delta_factors_module()
|
||||
module.vdn_delta_factors(
|
||||
transition.view(-1, HEAD_DIM, HEAD_DIM),
|
||||
injection.view(-1, HEAD_DIM, HEAD_DIM),
|
||||
A.view(-1, HEAD_DIM, HEAD_DIM),
|
||||
B.view(-1, HEAD_DIM, HEAD_DIM),
|
||||
alpha.view(-1, HEAD_DIM),
|
||||
)
|
||||
return transition, injection
|
||||
|
||||
|
||||
def can_use_vdn_delta_factors(
|
||||
A: torch.Tensor, B: torch.Tensor, alpha: torch.Tensor
|
||||
) -> bool:
|
||||
return (
|
||||
A.is_cuda
|
||||
and A.dtype is torch.float32
|
||||
and B.dtype is torch.float32
|
||||
and alpha.dtype is torch.float32
|
||||
and A.device == B.device == alpha.device
|
||||
and A.dim() >= 2
|
||||
and A.shape[-1] == HEAD_DIM
|
||||
and A.shape[-2] == HEAD_DIM
|
||||
and B.shape == A.shape
|
||||
and alpha.shape == A.shape[:-1]
|
||||
and A.is_contiguous()
|
||||
and B.is_contiguous()
|
||||
and alpha.is_contiguous()
|
||||
and torch.cuda.get_device_capability(A.device)[0] >= 8
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"can_use_vdn_delta_factors",
|
||||
"vdn_delta_factors",
|
||||
]
|
||||
@@ -0,0 +1,569 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Fused Triton kernels for the VDN-H3 (Video DeltaNet MiniMax-H3) linear branch;
|
||||
each reads its operands once and rounds once at the store.
|
||||
|
||||
vdn_temporal_conv_act 5-tap depthwise temporal conv + SiLU [+ L2 norm]
|
||||
(port of OpenVDN's _tconv_act_kernel)
|
||||
vdn_silu_l2norm SiLU [+ L2 norm] over head_dim, strided input ok
|
||||
vdn_frame_stats_prep the four GEMM operands of the frame statistics
|
||||
(kf16, kf32, kf32 * beta, v * beta) in [F, H, S, d]
|
||||
off one read of k and one of v
|
||||
vdn_gather_linear_state the alpha-bridged boundary gather over the fp32
|
||||
state banks
|
||||
vdn_linear_epilogue RMSNorm(d) * gate with the [F, H, S, d] -> [F*S, H*d]
|
||||
transpose folded into the store
|
||||
|
||||
Contract: vdn_frame_stats_prep and vdn_gather_linear_state are bitwise equal
|
||||
to the eager chains (widening casts, same products, fp32 gather). The three
|
||||
activation kernels round once instead of once per op and sit within one bf16
|
||||
ulp of the eager bf16 chains; OpenVDN ships the same contract, so the branch
|
||||
mounts them unconditionally.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
_BLOCK_T = 16
|
||||
_BLOCK_ROWS = 32
|
||||
|
||||
|
||||
def _pow2_head_dim(head_dim: int) -> bool:
|
||||
return head_dim >= 16 and head_dim & (head_dim - 1) == 0
|
||||
|
||||
|
||||
def _check_head_dim(head_dim: int) -> None:
|
||||
if not _pow2_head_dim(head_dim):
|
||||
raise ValueError(
|
||||
f"head_dim must be a power of two >= 16 (tl.arange), got {head_dim}"
|
||||
)
|
||||
|
||||
|
||||
def _cuda_bf16_rows(t: torch.Tensor) -> bool:
|
||||
return t.is_cuda and t.dtype == torch.bfloat16 and t.stride(-1) == 1
|
||||
|
||||
|
||||
def _i32(t: torch.Tensor) -> torch.Tensor:
|
||||
return t.to(torch.int32).contiguous()
|
||||
|
||||
|
||||
def can_use_vdn_temporal_conv_act(x: torch.Tensor, heads: int, head_dim: int) -> bool:
|
||||
"""x [T, S, heads * head_dim] bf16 on CUDA, power-of-two head_dim."""
|
||||
return (
|
||||
_cuda_bf16_rows(x)
|
||||
and x.ndim == 3
|
||||
and x.shape[-1] == heads * head_dim
|
||||
and _pow2_head_dim(head_dim)
|
||||
and not torch.compiler.is_compiling()
|
||||
)
|
||||
|
||||
|
||||
def can_use_vdn_silu_l2norm(tokens: torch.Tensor) -> bool:
|
||||
"""tokens [N, H, d] bf16 on CUDA (any row/head strides), power-of-two d."""
|
||||
return (
|
||||
_cuda_bf16_rows(tokens)
|
||||
and tokens.ndim == 3
|
||||
and _pow2_head_dim(tokens.shape[-1])
|
||||
and not torch.compiler.is_compiling()
|
||||
)
|
||||
|
||||
|
||||
def can_use_vdn_frame_stats_prep(key: torch.Tensor, value: torch.Tensor) -> bool:
|
||||
"""key/value [F * S, H, d] bf16 on CUDA with matching shapes."""
|
||||
return (
|
||||
_cuda_bf16_rows(key)
|
||||
and _cuda_bf16_rows(value)
|
||||
and key.ndim == 3
|
||||
and key.shape == value.shape
|
||||
and _pow2_head_dim(key.shape[-1])
|
||||
and not torch.compiler.is_compiling()
|
||||
)
|
||||
|
||||
|
||||
def can_use_vdn_gather_linear_state(prefix: torch.Tensor) -> bool:
|
||||
"""prefix/suffix [F, H, dv, dk] fp32 on CUDA, power-of-two dk."""
|
||||
return (
|
||||
prefix.is_cuda
|
||||
and prefix.dtype == torch.float32
|
||||
and prefix.ndim == 4
|
||||
and _pow2_head_dim(prefix.shape[-1])
|
||||
and not torch.compiler.is_compiling()
|
||||
)
|
||||
|
||||
|
||||
def can_use_vdn_linear_epilogue(readout: torch.Tensor) -> bool:
|
||||
"""readout [F, H, S, d] bf16 on CUDA, power-of-two d."""
|
||||
return (
|
||||
_cuda_bf16_rows(readout)
|
||||
and readout.ndim == 4
|
||||
and _pow2_head_dim(readout.shape[-1])
|
||||
and not torch.compiler.is_compiling()
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# temporal conv + SiLU + L2 norm
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _tconv_act_kernel(
|
||||
X,
|
||||
W,
|
||||
OUT,
|
||||
num_frames,
|
||||
tokens_per_frame,
|
||||
channels,
|
||||
BLOCK_T: tl.constexpr,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
L2NORM: tl.constexpr,
|
||||
HEADS: tl.constexpr,
|
||||
FRAME_MAJOR: tl.constexpr,
|
||||
):
|
||||
pid_t = tl.program_id(0)
|
||||
pid_s = tl.program_id(1)
|
||||
pid_h = tl.program_id(2)
|
||||
chan = pid_h * HEAD_DIM + tl.arange(0, HEAD_DIM)
|
||||
rows = pid_t * BLOCK_T + tl.arange(0, BLOCK_T)
|
||||
valid = rows < num_frames
|
||||
|
||||
acc = tl.zeros((BLOCK_T, HEAD_DIM), dtype=tl.float32)
|
||||
for dt in tl.static_range(5):
|
||||
r = rows + dt - 2
|
||||
ok = valid & (r >= 0) & (r < num_frames) # zero padding, both ends
|
||||
v = tl.load(
|
||||
X
|
||||
+ (r[:, None].to(tl.int64) * tokens_per_frame + pid_s) * channels
|
||||
+ chan[None, :],
|
||||
mask=ok[:, None],
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
wd = tl.load(W + chan * 5 + dt).to(tl.float32)
|
||||
acc += v * wd[None, :]
|
||||
|
||||
y = acc * tl.sigmoid(acc) # SiLU
|
||||
if L2NORM:
|
||||
inv = 1.0 / tl.sqrt(tl.maximum(tl.sum(y * y, axis=1), 1e-12))
|
||||
y = y * inv[:, None]
|
||||
if FRAME_MAJOR:
|
||||
# [T, HEADS, S, D]: the readout bmm reads this layout directly
|
||||
dst = (
|
||||
(rows[:, None].to(tl.int64) * HEADS + pid_h) * tokens_per_frame + pid_s
|
||||
) * HEAD_DIM + tl.arange(0, HEAD_DIM)[None, :]
|
||||
else:
|
||||
dst = (rows[:, None].to(tl.int64) * tokens_per_frame + pid_s) * channels + chan[
|
||||
None, :
|
||||
]
|
||||
tl.store(OUT + dst, y.to(OUT.dtype.element_ty), mask=valid[:, None])
|
||||
|
||||
|
||||
def vdn_temporal_conv_act(
|
||||
x: torch.Tensor,
|
||||
w: torch.Tensor,
|
||||
heads: int,
|
||||
head_dim: int,
|
||||
l2norm: bool,
|
||||
frame_major: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""x [T, S, C] bf16 contiguous, w [C, 5] -> [T * S, heads, head_dim], or
|
||||
[T, heads, S, head_dim] with ``frame_major``."""
|
||||
if not x.is_cuda:
|
||||
raise ValueError("vdn_temporal_conv_act is a Triton kernel; x must be on CUDA")
|
||||
_check_head_dim(head_dim)
|
||||
num_frames, tokens_per_frame, channels = x.shape
|
||||
if channels != heads * head_dim:
|
||||
raise ValueError(f"C={channels} != heads*head_dim={heads * head_dim}")
|
||||
if w.shape != (channels, 5):
|
||||
raise ValueError(f"w must be [C, 5], got {tuple(w.shape)}")
|
||||
x = x.contiguous()
|
||||
w = w.contiguous()
|
||||
out = torch.empty_like(x)
|
||||
_tconv_act_kernel[(triton.cdiv(num_frames, _BLOCK_T), tokens_per_frame, heads)](
|
||||
x,
|
||||
w,
|
||||
out,
|
||||
num_frames,
|
||||
tokens_per_frame,
|
||||
channels,
|
||||
BLOCK_T=_BLOCK_T,
|
||||
HEAD_DIM=head_dim,
|
||||
L2NORM=l2norm,
|
||||
HEADS=heads,
|
||||
FRAME_MAJOR=frame_major,
|
||||
num_warps=4,
|
||||
num_stages=2,
|
||||
)
|
||||
if frame_major:
|
||||
return out.view(num_frames, heads, tokens_per_frame, head_dim)
|
||||
return out.view(num_frames * tokens_per_frame, heads, head_dim)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# SiLU + L2 norm on a (possibly strided) [N, H, d] tensor
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _silu_l2norm_kernel(
|
||||
X,
|
||||
OUT,
|
||||
N,
|
||||
stride_n,
|
||||
stride_h,
|
||||
H,
|
||||
tokens_per_frame,
|
||||
BLOCK_N: tl.constexpr,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
L2NORM: tl.constexpr,
|
||||
FRAME_MAJOR: tl.constexpr,
|
||||
):
|
||||
pid_n = tl.program_id(0)
|
||||
pid_h = tl.program_id(1)
|
||||
rows = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
|
||||
valid = rows < N
|
||||
offs = tl.arange(0, HEAD_DIM)
|
||||
x = tl.load(
|
||||
X + rows[:, None].to(tl.int64) * stride_n + pid_h * stride_h + offs[None, :],
|
||||
mask=valid[:, None],
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
y = x * tl.sigmoid(x)
|
||||
if L2NORM:
|
||||
inv = 1.0 / tl.sqrt(tl.maximum(tl.sum(y * y, axis=1), 1e-12))
|
||||
y = y * inv[:, None]
|
||||
if FRAME_MAJOR:
|
||||
# row n = frame * S_ + s -> [F, H, S_, D]
|
||||
frame = rows // tokens_per_frame
|
||||
pos = rows - frame * tokens_per_frame
|
||||
dst = (
|
||||
(frame[:, None].to(tl.int64) * H + pid_h) * tokens_per_frame + pos[:, None]
|
||||
) * HEAD_DIM + offs[None, :]
|
||||
else:
|
||||
dst = (rows[:, None].to(tl.int64) * H + pid_h) * HEAD_DIM + offs[None, :]
|
||||
tl.store(OUT + dst, y.to(OUT.dtype.element_ty), mask=valid[:, None])
|
||||
|
||||
|
||||
def vdn_silu_l2norm(
|
||||
tokens: torch.Tensor, l2norm: bool, per_frame: int | None = None
|
||||
) -> torch.Tensor:
|
||||
"""tokens [N, H, d] (last dim contiguous) -> contiguous [N, H, d], or
|
||||
[N / per_frame, H, per_frame, d] when ``per_frame`` is given."""
|
||||
if not tokens.is_cuda:
|
||||
raise ValueError("vdn_silu_l2norm is a Triton kernel; tokens must be on CUDA")
|
||||
N, H, D = tokens.shape
|
||||
_check_head_dim(D)
|
||||
if tokens.stride(-1) != 1:
|
||||
tokens = tokens.contiguous()
|
||||
frame_major = per_frame is not None
|
||||
if frame_major and (per_frame <= 0 or N % per_frame):
|
||||
raise ValueError(f"per_frame={per_frame} must divide N={N}")
|
||||
shape = (N // per_frame, H, per_frame, D) if frame_major else (N, H, D)
|
||||
out = torch.empty(shape, dtype=tokens.dtype, device=tokens.device)
|
||||
if N == 0:
|
||||
return out
|
||||
_silu_l2norm_kernel[(triton.cdiv(N, _BLOCK_ROWS), H)](
|
||||
tokens,
|
||||
out,
|
||||
N,
|
||||
tokens.stride(0),
|
||||
tokens.stride(1),
|
||||
H,
|
||||
per_frame if frame_major else 1,
|
||||
BLOCK_N=_BLOCK_ROWS,
|
||||
HEAD_DIM=D,
|
||||
L2NORM=l2norm,
|
||||
FRAME_MAJOR=frame_major,
|
||||
num_warps=4,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# frame statistics prologue
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _frame_stats_prep_kernel(
|
||||
K,
|
||||
V,
|
||||
BETA,
|
||||
K16,
|
||||
K32,
|
||||
KB32,
|
||||
VB,
|
||||
tokens_per_frame,
|
||||
H,
|
||||
BLOCK_S: tl.constexpr,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
):
|
||||
pid_s = tl.program_id(0)
|
||||
f = tl.program_id(1)
|
||||
h = tl.program_id(2)
|
||||
s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S)
|
||||
valid = s < tokens_per_frame
|
||||
offs = tl.arange(0, HEAD_DIM)
|
||||
rows = f * tokens_per_frame + s # token rows
|
||||
src = (rows[:, None].to(tl.int64) * H + h) * HEAD_DIM + offs[None, :] # [F*S, H, d]
|
||||
dst = ((f * H + h) * tokens_per_frame + s)[:, None] * HEAD_DIM + offs[
|
||||
None, :
|
||||
] # [F, H, S, d]
|
||||
k = tl.load(K + src, mask=valid[:, None], other=0.0)
|
||||
v = tl.load(V + src, mask=valid[:, None], other=0.0)
|
||||
beta = tl.load(BETA + rows * H + h, mask=valid, other=0.0)
|
||||
k32 = k.to(tl.float32)
|
||||
beta32 = beta.to(tl.float32)
|
||||
tl.store(K16 + dst, k, mask=valid[:, None])
|
||||
tl.store(K32 + dst, k32, mask=valid[:, None])
|
||||
tl.store(KB32 + dst, k32 * beta32[:, None], mask=valid[:, None])
|
||||
vb = (v.to(tl.float32) * beta32[:, None]).to(VB.dtype.element_ty)
|
||||
tl.store(VB + dst, vb, mask=valid[:, None])
|
||||
|
||||
|
||||
def vdn_frame_stats_prep(
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
beta: torch.Tensor,
|
||||
num_frames: int,
|
||||
tokens_per_frame: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""key/value [F*S, H, d] bf16 contiguous, beta [F*S, H] bf16 ->
|
||||
(k16 [F,H,S,d] bf16, k32 fp32, k32*beta fp32, v*beta bf16), all contiguous."""
|
||||
if not key.is_cuda:
|
||||
raise ValueError(
|
||||
"vdn_frame_stats_prep is a Triton kernel; inputs must be on CUDA"
|
||||
)
|
||||
rows, H, D = key.shape
|
||||
_check_head_dim(D)
|
||||
if rows != num_frames * tokens_per_frame:
|
||||
raise ValueError(f"{rows} rows != {num_frames} x {tokens_per_frame}")
|
||||
key = key.contiguous()
|
||||
value = value.contiguous()
|
||||
beta = beta.to(key.dtype).contiguous()
|
||||
shape = (num_frames, H, tokens_per_frame, D)
|
||||
k16 = torch.empty(shape, dtype=key.dtype, device=key.device)
|
||||
k32 = torch.empty(shape, dtype=torch.float32, device=key.device)
|
||||
kb32 = torch.empty(shape, dtype=torch.float32, device=key.device)
|
||||
vb = torch.empty(shape, dtype=value.dtype, device=key.device)
|
||||
_frame_stats_prep_kernel[
|
||||
(triton.cdiv(tokens_per_frame, _BLOCK_ROWS), num_frames, H)
|
||||
](
|
||||
key,
|
||||
value,
|
||||
beta,
|
||||
k16,
|
||||
k32,
|
||||
kb32,
|
||||
vb,
|
||||
tokens_per_frame,
|
||||
H,
|
||||
BLOCK_S=_BLOCK_ROWS,
|
||||
HEAD_DIM=D,
|
||||
num_warps=4,
|
||||
)
|
||||
return k16, k32, kb32, vb
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# readout epilogue: RMSNorm(d) * gate, [F, H, S, d] -> [F*S, H*d]
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _linear_epilogue_kernel(
|
||||
R,
|
||||
W,
|
||||
G,
|
||||
OUT,
|
||||
tokens_per_frame,
|
||||
H,
|
||||
eps,
|
||||
BLOCK_S: tl.constexpr,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
):
|
||||
pid_s = tl.program_id(0)
|
||||
f = tl.program_id(1)
|
||||
h = tl.program_id(2)
|
||||
s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S)
|
||||
valid = s < tokens_per_frame
|
||||
offs = tl.arange(0, HEAD_DIM)
|
||||
src = ((f * H + h) * tokens_per_frame + s)[:, None] * HEAD_DIM + offs[None, :]
|
||||
rows = f * tokens_per_frame + s
|
||||
dst = (rows[:, None].to(tl.int64) * H + h) * HEAD_DIM + offs[None, :]
|
||||
r = tl.load(R + src, mask=valid[:, None], other=0.0).to(tl.float32)
|
||||
ms = tl.sum(r * r, axis=1) / HEAD_DIM
|
||||
w = tl.load(W + offs).to(tl.float32)
|
||||
g = tl.load(G + dst, mask=valid[:, None], other=0.0).to(tl.float32)
|
||||
y = r * (1.0 / tl.sqrt(ms + eps))[:, None] * w[None, :] * g
|
||||
tl.store(OUT + dst, y.to(OUT.dtype.element_ty), mask=valid[:, None])
|
||||
|
||||
|
||||
def vdn_linear_epilogue(
|
||||
readout: torch.Tensor,
|
||||
norm_weight: torch.Tensor,
|
||||
gate: torch.Tensor,
|
||||
eps: float,
|
||||
) -> torch.Tensor:
|
||||
"""readout [F, H, S, d] bf16 contiguous, norm_weight [d], gate [F*S, H, d]
|
||||
-> [F*S, H*d] bf16."""
|
||||
if not readout.is_cuda:
|
||||
raise ValueError(
|
||||
"vdn_linear_epilogue is a Triton kernel; readout must be on CUDA"
|
||||
)
|
||||
F, H, tokens_per_frame, D = readout.shape
|
||||
_check_head_dim(D)
|
||||
readout = readout.contiguous()
|
||||
gate = gate.reshape(F * tokens_per_frame, H, D).to(readout.dtype).contiguous()
|
||||
out = torch.empty(
|
||||
(F * tokens_per_frame, H * D), dtype=readout.dtype, device=readout.device
|
||||
)
|
||||
_linear_epilogue_kernel[(triton.cdiv(tokens_per_frame, _BLOCK_ROWS), F, H)](
|
||||
readout,
|
||||
norm_weight.contiguous(),
|
||||
gate,
|
||||
out,
|
||||
tokens_per_frame,
|
||||
H,
|
||||
float(eps),
|
||||
BLOCK_S=_BLOCK_ROWS,
|
||||
HEAD_DIM=D,
|
||||
num_warps=4,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# boundary gather: prefix[lo-1] * prod alpha + suffix[hi+1] * prod alpha
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _gather_state_kernel(
|
||||
PREFIX,
|
||||
SUFFIX,
|
||||
LOGP, # [F+1, H, dk] fp32 exclusive log-alpha prefix sums
|
||||
TEXT, # [H, dv, dk] fp32 (or PREFIX when HAS_TEXT is False)
|
||||
BEFORE, # [F] int32 prefix row to read (clamped)
|
||||
AFTER, # [F] int32 suffix row to read (clamped)
|
||||
HASB, # [F] int32 0/1
|
||||
HASA, # [F] int32 0/1
|
||||
BRIDGEB, # [F] int32 log-prefix row for the before side
|
||||
BRIDGEA, # [F] int32 log-prefix row for the after side
|
||||
OUT,
|
||||
H,
|
||||
DV,
|
||||
HAS_TEXT: tl.constexpr,
|
||||
BRIDGE: tl.constexpr,
|
||||
BLOCK_V: tl.constexpr,
|
||||
DK: tl.constexpr,
|
||||
):
|
||||
f = tl.program_id(0)
|
||||
h = tl.program_id(1)
|
||||
pid_v = tl.program_id(2)
|
||||
rows = pid_v * BLOCK_V + tl.arange(0, BLOCK_V)
|
||||
cols = tl.arange(0, DK)
|
||||
valid = rows < DV
|
||||
fb = tl.load(BEFORE + f)
|
||||
fa = tl.load(AFTER + f)
|
||||
has_b = tl.load(HASB + f)
|
||||
has_a = tl.load(HASA + f)
|
||||
plane = rows[:, None] * DK + cols[None, :]
|
||||
off_b = ((fb * H + h) * DV) * DK + plane
|
||||
off_a = ((fa * H + h) * DV) * DK + plane
|
||||
sb = tl.load(PREFIX + off_b, mask=valid[:, None], other=0.0)
|
||||
sa = tl.load(SUFFIX + off_a, mask=valid[:, None], other=0.0)
|
||||
if HAS_TEXT:
|
||||
ts = tl.load(TEXT + (h * DV) * DK + plane, mask=valid[:, None], other=0.0)
|
||||
sb = tl.where(has_b != 0, sb, ts)
|
||||
sa = tl.where(has_a != 0, sa, ts)
|
||||
else:
|
||||
sb = tl.where(has_b != 0, sb, 0.0)
|
||||
sa = tl.where(has_a != 0, sa, 0.0)
|
||||
if BRIDGE:
|
||||
bb = tl.load(BRIDGEB + f)
|
||||
ba = tl.load(BRIDGEA + f)
|
||||
lp_t1 = tl.load(LOGP + ((f + 1) * H + h) * DK + cols)
|
||||
lp_t = tl.load(LOGP + (f * H + h) * DK + cols)
|
||||
lp_bb = tl.load(LOGP + (bb * H + h) * DK + cols)
|
||||
lp_ba = tl.load(LOGP + (ba * H + h) * DK + cols)
|
||||
sb = sb * tl.exp(lp_t1 - lp_bb)[None, :]
|
||||
sa = sa * tl.exp(lp_ba - lp_t)[None, :]
|
||||
out = sb + sa
|
||||
tl.store(
|
||||
OUT + ((f * H + h) * DV) * DK + plane,
|
||||
out.to(OUT.dtype.element_ty),
|
||||
mask=valid[:, None],
|
||||
)
|
||||
|
||||
|
||||
def vdn_gather_linear_state(
|
||||
prefix: torch.Tensor,
|
||||
suffix: torch.Tensor,
|
||||
alpha: torch.Tensor,
|
||||
text_state: torch.Tensor | None,
|
||||
*,
|
||||
before_idx: torch.Tensor,
|
||||
after_idx: torch.Tensor,
|
||||
has_before: torch.Tensor,
|
||||
has_after: torch.Tensor,
|
||||
bridge_before: torch.Tensor,
|
||||
bridge_after: torch.Tensor,
|
||||
bridge: bool,
|
||||
out_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""The boundary gather of the linear branch as one kernel over the
|
||||
[F, H, dv, dk] fp32 state banks."""
|
||||
if not prefix.is_cuda:
|
||||
raise ValueError(
|
||||
"vdn_gather_linear_state is a Triton kernel; inputs must be on CUDA"
|
||||
)
|
||||
F, H, DV, DK = prefix.shape
|
||||
_check_head_dim(DK)
|
||||
prefix = prefix.contiguous()
|
||||
suffix = suffix.contiguous()
|
||||
if bridge:
|
||||
log_alpha = torch.log(alpha.float().clamp_min(1e-12))
|
||||
logp = torch.cat(
|
||||
[torch.zeros_like(log_alpha[:1]), log_alpha.cumsum(0)]
|
||||
).contiguous()
|
||||
else:
|
||||
logp = prefix # unused
|
||||
text = text_state.float().contiguous() if text_state is not None else prefix
|
||||
out = torch.empty(prefix.shape, dtype=out_dtype, device=prefix.device)
|
||||
_gather_state_kernel[(F, H, triton.cdiv(DV, _BLOCK_ROWS))](
|
||||
prefix,
|
||||
suffix,
|
||||
logp,
|
||||
text,
|
||||
_i32(before_idx),
|
||||
_i32(after_idx),
|
||||
_i32(has_before),
|
||||
_i32(has_after),
|
||||
_i32(bridge_before),
|
||||
_i32(bridge_after),
|
||||
out,
|
||||
H,
|
||||
DV,
|
||||
HAS_TEXT=text_state is not None,
|
||||
BRIDGE=bridge,
|
||||
BLOCK_V=_BLOCK_ROWS,
|
||||
DK=DK,
|
||||
num_warps=4,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
__all__ = [
|
||||
"can_use_vdn_frame_stats_prep",
|
||||
"can_use_vdn_gather_linear_state",
|
||||
"can_use_vdn_linear_epilogue",
|
||||
"can_use_vdn_silu_l2norm",
|
||||
"can_use_vdn_temporal_conv_act",
|
||||
"vdn_frame_stats_prep",
|
||||
"vdn_gather_linear_state",
|
||||
"vdn_linear_epilogue",
|
||||
"vdn_silu_l2norm",
|
||||
"vdn_temporal_conv_act",
|
||||
]
|
||||
@@ -0,0 +1 @@
|
||||
"""Block-scaled (MXFP8) activation quantizers with GEMM-ready scale layouts."""
|
||||
@@ -0,0 +1,318 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""MXFP8 producers: e4m3 payload plus one E8M0 scale per 32 elements along K,
|
||||
the scales in the cuBLASLt ``SWIZZLE_32_4_4`` layout that
|
||||
``torch.nn.functional.scaled_mm(..., BlockWise1x32)`` consumes on SM100.
|
||||
|
||||
Scale ``(r, c)`` of the ``[rows, K/32]`` scale matrix lives at byte
|
||||
``((r // 128) * ceil(K/32 / 4) + c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + c % 4``,
|
||||
rows padded to 128 and scale columns to 4, padding zero. Exponent
|
||||
``e = ceil(log2(amax / 448))`` exactly from the float bits; ``q = e4m3(x * 2**-e)``;
|
||||
scale byte ``e + 127``. Every producer quantizes the bf16-rounded value the
|
||||
unfused bf16 kernel stores, so each is byte-exact against that kernel followed
|
||||
by ``mxfp8_quantize_swizzled`` (itself byte-exact vs ``flashinfer.mxfp8_quantize``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.kernels.ops.diffusion.common.numerics import round_bf16_to_fp32
|
||||
|
||||
_E4M3 = torch.float8_e4m3fn
|
||||
|
||||
|
||||
def _scale_numel(rows: int, k: int) -> int:
|
||||
n_groups = k // 32
|
||||
return -(-rows // 128) * 128 * (-(-n_groups // 4) * 4)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _mx_e8m0_from_amax(amax):
|
||||
bits = amax.to(tl.int32, bitcast=True)
|
||||
e0 = ((bits >> 23) & 0xFF) - 135
|
||||
# amax / 2**e0 lies in [256, 512); bump e0 when it exceeds 448 = 1.75 * 2**8
|
||||
thr = (((e0 + 135) << 23) | 0x600000).to(tl.float32, bitcast=True)
|
||||
e = e0 + (amax > thr).to(tl.int32)
|
||||
e = tl.maximum(e, -127)
|
||||
inv = ((127 - e) << 23).to(tl.float32, bitcast=True)
|
||||
return e + 127, inv
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _mx_scale_offsets(r, c, n_col_blocks):
|
||||
tile = (r // 128) * n_col_blocks + (c // 4)
|
||||
return tile * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + (c % 4)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _mxfp8_quant_kernel(
|
||||
x_ptr,
|
||||
q_ptr,
|
||||
s_ptr,
|
||||
rows,
|
||||
k,
|
||||
n_groups,
|
||||
n_col_blocks,
|
||||
stride_x,
|
||||
BLOCK_R: tl.constexpr,
|
||||
G: tl.constexpr,
|
||||
):
|
||||
pid_r = tl.program_id(0)
|
||||
pid_g = tl.program_id(1)
|
||||
r = pid_r * BLOCK_R + tl.arange(0, BLOCK_R)
|
||||
g = pid_g * G + tl.arange(0, G)
|
||||
c = pid_g * (G * 32) + tl.arange(0, G * 32)
|
||||
rmask = r < rows
|
||||
mask = rmask[:, None] & (c < k)[None, :]
|
||||
x = tl.load(
|
||||
x_ptr + r[:, None].to(tl.int64) * stride_x + c[None, :], mask=mask, other=0.0
|
||||
).to(tl.float32)
|
||||
x3 = tl.reshape(x, [BLOCK_R, G, 32])
|
||||
amax = tl.max(tl.abs(x3), axis=2)
|
||||
sbyte, inv = _mx_e8m0_from_amax(amax)
|
||||
q = tl.reshape(x3 * inv[:, :, None], [BLOCK_R, G * 32])
|
||||
tl.store(
|
||||
q_ptr + r[:, None].to(tl.int64) * k + c[None, :],
|
||||
q.to(tl.float8e4nv),
|
||||
mask=mask,
|
||||
)
|
||||
smask = rmask[:, None] & (g < n_groups)[None, :]
|
||||
tl.store(
|
||||
s_ptr + _mx_scale_offsets(r[:, None], g[None, :], n_col_blocks),
|
||||
sbyte.to(tl.uint8),
|
||||
mask=smask,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _silu_mul_mxfp8_kernel(
|
||||
x_ptr,
|
||||
q_ptr,
|
||||
s_ptr,
|
||||
rows,
|
||||
hidden,
|
||||
n_groups,
|
||||
n_col_blocks,
|
||||
stride_row,
|
||||
BLOCK_R: tl.constexpr,
|
||||
G: tl.constexpr,
|
||||
):
|
||||
pid_r = tl.program_id(0)
|
||||
pid_c = tl.program_id(1)
|
||||
r = pid_r * BLOCK_R + tl.arange(0, BLOCK_R)
|
||||
g = pid_c * G + tl.arange(0, G)
|
||||
c = pid_c * (G * 32) + tl.arange(0, G * 32)
|
||||
rmask = r < rows
|
||||
mask = rmask[:, None] & (c < hidden)[None, :]
|
||||
base = x_ptr + r[:, None].to(tl.int64) * stride_row
|
||||
gate = tl.load(base + c[None, :], mask=mask, other=0.0).to(tl.float32)
|
||||
up = tl.load(base + hidden + c[None, :], mask=mask, other=0.0).to(tl.float32)
|
||||
act = (gate * tl.sigmoid(gate)).to(tl.bfloat16).to(tl.float32)
|
||||
prod = (act * up).to(tl.bfloat16).to(tl.float32)
|
||||
p3 = tl.reshape(prod, [BLOCK_R, G, 32])
|
||||
amax = tl.max(tl.abs(p3), axis=2)
|
||||
sbyte, inv = _mx_e8m0_from_amax(amax)
|
||||
q = tl.reshape(p3 * inv[:, :, None], [BLOCK_R, G * 32])
|
||||
tl.store(
|
||||
q_ptr + r[:, None].to(tl.int64) * hidden + c[None, :],
|
||||
q.to(tl.float8e4nv),
|
||||
mask=mask,
|
||||
)
|
||||
smask = rmask[:, None] & (g < n_groups)[None, :]
|
||||
tl.store(
|
||||
s_ptr + _mx_scale_offsets(r[:, None], g[None, :], n_col_blocks),
|
||||
sbyte.to(tl.uint8),
|
||||
mask=smask,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _indexed_scale_shift_mxfp8_kernel(
|
||||
x_ptr,
|
||||
q_ptr,
|
||||
s_ptr,
|
||||
shift_ptr,
|
||||
scale_ptr,
|
||||
indices_ptr,
|
||||
hidden_size,
|
||||
n_groups,
|
||||
n_col_blocks,
|
||||
stride_x_row,
|
||||
stride_shift_row,
|
||||
stride_scale_row,
|
||||
stride_indices,
|
||||
STORE_BF16: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
):
|
||||
row = tl.program_id(0)
|
||||
columns = tl.arange(0, BLOCK_N)
|
||||
mask = columns < hidden_size
|
||||
index = tl.load(indices_ptr + row * stride_indices)
|
||||
xrow = x_ptr + row.to(tl.int64) * stride_x_row
|
||||
x = tl.load(xrow + columns, mask=mask, other=0.0).to(tl.float32)
|
||||
shift = tl.load(
|
||||
shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0
|
||||
).to(tl.float32)
|
||||
scale = tl.load(
|
||||
scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0
|
||||
).to(tl.float32)
|
||||
# the rounding points of _indexed_scale_shift_bf16_kernel
|
||||
one_plus_scale = round_bf16_to_fp32(1.0 + scale)
|
||||
scaled = round_bf16_to_fp32(x * one_plus_scale)
|
||||
out = round_bf16_to_fp32(scaled + shift)
|
||||
if STORE_BF16:
|
||||
tl.store(xrow + columns, out, mask=mask)
|
||||
v3 = tl.reshape(out, [BLOCK_N // 32, 32])
|
||||
amax = tl.max(tl.abs(v3), axis=1)
|
||||
sbyte, inv = _mx_e8m0_from_amax(amax)
|
||||
q = tl.reshape(v3 * inv[:, None], [BLOCK_N])
|
||||
tl.store(
|
||||
q_ptr + row.to(tl.int64) * hidden_size + columns,
|
||||
q.to(tl.float8e4nv),
|
||||
mask=mask,
|
||||
)
|
||||
g = tl.arange(0, BLOCK_N // 32)
|
||||
tl.store(
|
||||
s_ptr + _mx_scale_offsets(row, g, n_col_blocks),
|
||||
sbyte.to(tl.uint8),
|
||||
mask=g < n_groups,
|
||||
)
|
||||
|
||||
|
||||
def can_use_mxfp8_swizzled(x: torch.Tensor) -> bool:
|
||||
"""Row-major bf16 CUDA 2D tensor with K % 32 == 0, outside torch.compile."""
|
||||
return (
|
||||
x.is_cuda
|
||||
and x.ndim == 2
|
||||
and x.dtype == torch.bfloat16
|
||||
and x.stride(-1) == 1
|
||||
and x.shape[-1] % 32 == 0
|
||||
and not torch.compiler.is_compiling()
|
||||
)
|
||||
|
||||
|
||||
def _alloc(
|
||||
rows: int, k: int, device: torch.device
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
q = torch.empty(rows, k, dtype=_E4M3, device=device)
|
||||
s = torch.zeros(_scale_numel(rows, k), dtype=torch.uint8, device=device)
|
||||
return q, s
|
||||
|
||||
|
||||
def mxfp8_quantize_swizzled(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""bf16 ``[rows, k]`` -> ``(fp8 [rows, k], swizzled e8m0 scale bytes)``."""
|
||||
if not can_use_mxfp8_swizzled(x):
|
||||
raise ValueError("expected a row-major bf16 CUDA [rows, k] tensor, k % 32 == 0")
|
||||
rows, k = x.shape
|
||||
q, s = _alloc(rows, k, x.device)
|
||||
if rows == 0:
|
||||
return q, s
|
||||
n_groups = k // 32
|
||||
n_col_blocks = -(-n_groups // 4)
|
||||
block_r, g = 32, 8
|
||||
grid = (triton.cdiv(rows, block_r), triton.cdiv(n_groups, g))
|
||||
with torch.get_device_module().device(x.device):
|
||||
_mxfp8_quant_kernel[grid](
|
||||
x,
|
||||
q,
|
||||
s,
|
||||
rows,
|
||||
k,
|
||||
n_groups,
|
||||
n_col_blocks,
|
||||
x.stride(0),
|
||||
BLOCK_R=block_r,
|
||||
G=g,
|
||||
num_warps=4,
|
||||
)
|
||||
return q, s
|
||||
|
||||
|
||||
def can_use_silu_mul_mxfp8(hidden: torch.Tensor) -> bool:
|
||||
return can_use_mxfp8_swizzled(hidden) and (hidden.shape[-1] // 2) % 32 == 0
|
||||
|
||||
|
||||
def silu_mul_mxfp8(hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""``hidden [rows, 2n]`` bf16 (gate | up) -> quantized ``silu(gate) * up``."""
|
||||
if not can_use_silu_mul_mxfp8(hidden):
|
||||
raise ValueError(
|
||||
"expected a row-major bf16 CUDA [rows, 2 * n] tensor, n % 32 == 0"
|
||||
)
|
||||
rows, twice = hidden.shape
|
||||
n = twice // 2
|
||||
q, s = _alloc(rows, n, hidden.device)
|
||||
if rows == 0:
|
||||
return q, s
|
||||
n_groups = n // 32
|
||||
n_col_blocks = -(-n_groups // 4)
|
||||
block_r, g = 16, 8
|
||||
grid = (triton.cdiv(rows, block_r), triton.cdiv(n_groups, g))
|
||||
with torch.get_device_module().device(hidden.device):
|
||||
_silu_mul_mxfp8_kernel[grid](
|
||||
hidden,
|
||||
q,
|
||||
s,
|
||||
rows,
|
||||
n,
|
||||
n_groups,
|
||||
n_col_blocks,
|
||||
hidden.stride(0),
|
||||
BLOCK_R=block_r,
|
||||
G=g,
|
||||
num_warps=4,
|
||||
)
|
||||
return q, s
|
||||
|
||||
|
||||
def indexed_scale_shift_mxfp8_(
|
||||
x: torch.Tensor,
|
||||
shift: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
indices: torch.Tensor,
|
||||
*,
|
||||
keep_bf16: bool,
|
||||
) -> tuple[torch.Tensor | None, torch.Tensor, torch.Tensor]:
|
||||
"""``x * (1 + scale[idx]) + shift[idx]`` -> ``(x | None, fp8, scales)``; with
|
||||
``keep_bf16`` the bf16 result is also written into ``x`` and returned."""
|
||||
if not can_use_mxfp8_swizzled(x):
|
||||
raise ValueError(
|
||||
"expected a row-major bf16 CUDA [rows, hidden] tensor, hidden % 32 == 0"
|
||||
)
|
||||
rows, hidden_size = x.shape
|
||||
q, s = _alloc(rows, hidden_size, x.device)
|
||||
if rows == 0:
|
||||
return (x if keep_bf16 else None), q, s
|
||||
n_groups = hidden_size // 32
|
||||
n_col_blocks = -(-n_groups // 4)
|
||||
block_n = triton.next_power_of_2(hidden_size)
|
||||
with torch.get_device_module().device(x.device):
|
||||
_indexed_scale_shift_mxfp8_kernel[(rows,)](
|
||||
x,
|
||||
q,
|
||||
s,
|
||||
shift,
|
||||
scale,
|
||||
indices,
|
||||
hidden_size,
|
||||
n_groups,
|
||||
n_col_blocks,
|
||||
x.stride(0),
|
||||
shift.stride(0),
|
||||
scale.stride(0),
|
||||
indices.stride(0),
|
||||
STORE_BF16=keep_bf16,
|
||||
BLOCK_N=block_n,
|
||||
num_warps=8,
|
||||
)
|
||||
return (x if keep_bf16 else None), q, s
|
||||
|
||||
|
||||
__all__ = [
|
||||
"can_use_mxfp8_swizzled",
|
||||
"can_use_silu_mul_mxfp8",
|
||||
"indexed_scale_shift_mxfp8_",
|
||||
"mxfp8_quantize_swizzled",
|
||||
"silu_mul_mxfp8",
|
||||
]
|
||||
@@ -33,6 +33,7 @@ def _jit_qknorm_rope_module(
|
||||
round_norm_before_rope: bool,
|
||||
pack_kv: bool = False,
|
||||
cache_has_full_width: bool = False,
|
||||
out_of_place: bool = False,
|
||||
) -> Module:
|
||||
args = make_cpp_args(
|
||||
head_dim,
|
||||
@@ -44,8 +45,12 @@ def _jit_qknorm_rope_module(
|
||||
round_norm_before_rope,
|
||||
cache_has_full_width,
|
||||
)
|
||||
op_name = "qknorm_rope_pack_kv" if pack_kv else "qknorm_rope"
|
||||
kernel_name = "QKNormRopePackKVKernel" if pack_kv else "QKNormRopeKernel"
|
||||
if pack_kv:
|
||||
op_name, kernel_name = "qknorm_rope_pack_kv", "QKNormRopePackKVKernel"
|
||||
elif out_of_place:
|
||||
op_name, kernel_name = "qknorm_rope_out_of_place", "QKNormRopeOutOfPlaceKernel"
|
||||
else:
|
||||
op_name, kernel_name = "qknorm_rope", "QKNormRopeKernel"
|
||||
return load_jit(
|
||||
op_name,
|
||||
*args,
|
||||
@@ -182,6 +187,46 @@ def fused_inplace_qknorm_rope(
|
||||
module.qknorm_rope(q, k, q_weight, k_weight, cos_sin_cache, positions, eps)
|
||||
|
||||
|
||||
@register_custom_op(mutates_args=["q_out", "k_out"])
|
||||
def fused_qknorm_rope_out_of_place(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_out: torch.Tensor,
|
||||
k_out: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
*,
|
||||
is_neox: bool,
|
||||
eps: float = 1e-6,
|
||||
head_dim: int = 0,
|
||||
rope_dim: int = 0,
|
||||
round_norm_before_rope: bool = False,
|
||||
cache_has_full_width: bool = False,
|
||||
) -> None:
|
||||
"""QK-norm + RoPE from ``q``/``k`` (any strides) into ``q_out``/``k_out``;
|
||||
the inputs are left untouched. Same arithmetic as the in-place kernel."""
|
||||
head_dim = head_dim or q.size(-1)
|
||||
if not rope_dim:
|
||||
cache_width = cos_sin_cache.size(-1)
|
||||
rope_dim = cache_width // 2 if cache_has_full_width else cache_width
|
||||
module = _jit_qknorm_rope_module(
|
||||
head_dim,
|
||||
rope_dim,
|
||||
is_neox,
|
||||
q.dtype,
|
||||
cos_sin_cache.dtype,
|
||||
round_norm_before_rope,
|
||||
False,
|
||||
cache_has_full_width,
|
||||
True,
|
||||
)
|
||||
module.qknorm_rope_out_of_place(
|
||||
q, k, q_out, k_out, q_weight, k_weight, cos_sin_cache, positions, eps
|
||||
)
|
||||
|
||||
|
||||
@register_custom_op(mutates_args=["q", "packed_kv"])
|
||||
def fused_qknorm_rope_pack_kv(
|
||||
q: torch.Tensor,
|
||||
|
||||
+27
@@ -452,6 +452,33 @@ MODELS = {
|
||||
],
|
||||
"force_eager": True,
|
||||
},
|
||||
# OpenVDN paper workload: 1344x768, 14.375 s (latent_t 102), t2va, 9 grid points = 8 NFE
|
||||
"vdn-h3": {
|
||||
"path": "OpenVDN/vdn-minimax-h3",
|
||||
"prompt": (
|
||||
"A curious raccoon peers through a vibrant field of yellow "
|
||||
"sunflowers, its eyes wide with interest."
|
||||
),
|
||||
"seed": 1000,
|
||||
"config_overrides": {
|
||||
"task": "t2va",
|
||||
"conditions": [],
|
||||
"target": {
|
||||
"short_edge": 768,
|
||||
"aspect_ratio": "16:9",
|
||||
"duration_seconds": 14.375,
|
||||
},
|
||||
"num_inference_steps": 9,
|
||||
},
|
||||
"extra_args": [
|
||||
"--num-gpus=8",
|
||||
"--quantization=fp8",
|
||||
"--performance-mode=speed",
|
||||
"--enable-torch-compile=false",
|
||||
"--warmup-steps=2",
|
||||
],
|
||||
"force_eager": True,
|
||||
},
|
||||
# Source-tracked extras from current registry / GPU test coverage.
|
||||
"longcat-image": {
|
||||
"path": "meituan-longcat/LongCat-Image",
|
||||
|
||||
@@ -9,7 +9,7 @@ SGLang diffusion features an end-to-end unified pipeline for accelerating diffus
|
||||
## Key Features
|
||||
|
||||
SGLang Diffusion has the following features:
|
||||
- Broad model support: Wan, FastWan, FLUX, Qwen-Image, LongCat-Image, Z-Image, Ideogram 4, Krea-2, Cosmos3, LTX-2/LTX-2.3/LTX-2.5, MiniMax-H3, FastH3, LingBot Video MoE, LingBot World, SANA-Video/SANA-WM, JoyEcho, MOVA, GLM-Image, ERNIE-Image, Hunyuan3D, and more
|
||||
- Broad model support: Wan, FastWan, FLUX, Qwen-Image, LongCat-Image, Z-Image, Ideogram 4, Krea-2, Cosmos3, LTX-2/LTX-2.3/LTX-2.5, MiniMax-H3, FastH3, VDN-H3, LingBot Video MoE, LingBot World, SANA-Video/SANA-WM, JoyEcho, MOVA, GLM-Image, ERNIE-Image, Hunyuan3D, and more
|
||||
- Fast inference speed: empowered by optimized `sgl-kernel` kernels, scheduler/runtime improvements, caching acceleration, and native diffusion hot-path optimizations
|
||||
- Ease of use: OpenAI-compatible api, CLI, and python sdk support
|
||||
- Multi-platform support:
|
||||
|
||||
@@ -2,6 +2,9 @@
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
|
||||
from sglang.multimodal_gen.configs.models.dits.minimax_h3_vdn import (
|
||||
VDNHybridAttentionArchConfig,
|
||||
)
|
||||
|
||||
MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT = 64
|
||||
MINIMAX_H3_ADALN_MODALITY_NUM = 3
|
||||
@@ -48,6 +51,8 @@ class MiniMaxH3DiTArchConfig(DiTArchConfig):
|
||||
),
|
||||
r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$": r"blocks.\1.attn.out_proj.\2",
|
||||
r"^transformer_blocks\.(\d+)\.attn\.to_gate_compress\.(.*)$": r"blocks.\1.attn.to_gate_compress.\2",
|
||||
# VDN-H3 hybrid attention module (see minimax_h3_vdn_attention)
|
||||
r"^transformer_blocks\.(\d+)\.attn\.(linear_attention|softmax_gate|to_out_linear)\.(.*)$": r"blocks.\1.attn.hybrid.\2.\3",
|
||||
r"^transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$": r"blocks.\1.attn.q_norm.\2",
|
||||
r"^transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$": r"blocks.\1.attn.k_norm.\2",
|
||||
r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.(.*)$": r"blocks.\1.mlp.fc1.\2",
|
||||
@@ -102,6 +107,8 @@ class MiniMaxH3DiTArchConfig(DiTArchConfig):
|
||||
checkpoint_uses_diffusers_layout: bool = False
|
||||
adaln_affine_input_dim: int | None = None
|
||||
has_gate_compress: bool = False
|
||||
# VDN-H3: None for the dense model; set from transformer/config.json
|
||||
hybrid_attention: VDNHybridAttentionArchConfig | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
@@ -110,6 +117,10 @@ class MiniMaxH3DiTArchConfig(DiTArchConfig):
|
||||
if len(self.patch_size) != 3:
|
||||
raise ValueError(f"patch_size must have 3 values, got {self.patch_size}.")
|
||||
self.num_channels_latents = self.latents_dim
|
||||
if isinstance(self.hybrid_attention, dict):
|
||||
self.hybrid_attention = VDNHybridAttentionArchConfig.from_transform_config(
|
||||
self.hybrid_attention
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -133,6 +144,11 @@ class MiniMaxH3DiTConfig(DiTConfig):
|
||||
model_dict["adaln_affine_input_dim"] = source_model_dict["time_embed_dim"]
|
||||
model_dict["time_embed_dim"] = source_model_dict["adaln_rank"]
|
||||
model_dict["adaln_curve_grid"] = source_model_dict["time_table_size"]
|
||||
hybrid = model_dict.get("hybrid_attention")
|
||||
if isinstance(hybrid, dict):
|
||||
model_dict["hybrid_attention"] = (
|
||||
VDNHybridAttentionArchConfig.from_transform_config(hybrid)
|
||||
)
|
||||
super().update_model_arch(model_dict)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""VDN-H3 hybrid attention architecture config (window softmax + linear branch)."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import msgspec
|
||||
|
||||
VDN_H3_DELTA_RULES = ("vdn_solve", "sana_scaled", "vdn_scaled")
|
||||
VDN_H3_BRIDGE_MODES = ("alpha", "none")
|
||||
VDN_H3_ANCHOR_FRAME_MODES = ("none", "columns", "rows", "both")
|
||||
VDN_H3_SHORT_CONV_TARGETS = ("q", "k", "v")
|
||||
|
||||
|
||||
class VDNHybridAttentionArchConfig(msgspec.Struct):
|
||||
"""VDN-H3 hybrid attention (window softmax + frame-wise linear branch); the
|
||||
resolved ``hybrid_attention`` transform config the overlay copies into
|
||||
``transformer/config.json``. A dense checkpoint has none."""
|
||||
|
||||
# frame t is in chunk t // chunk and attends chunks [c - radius, c + radius];
|
||||
# chunk 0 means a centered frame window
|
||||
chunk: int = 5
|
||||
radius: int = 1
|
||||
# "both": frames 0 and F-1 dense as rows and columns, so the branch skips them
|
||||
anchor_frames: str = "both"
|
||||
enable_softmax_gate: bool = True
|
||||
delta_rule: str = "vdn_solve"
|
||||
linear_head_dim: int = 128
|
||||
bridge: str = "alpha"
|
||||
a_fp32: bool = True
|
||||
enable_text_state: bool = True
|
||||
short_conv: tuple[str, ...] = ("k", "v")
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.delta_rule not in VDN_H3_DELTA_RULES:
|
||||
raise ValueError(
|
||||
f"hybrid_attention.delta_rule={self.delta_rule!r}; expected one of "
|
||||
f"{VDN_H3_DELTA_RULES}"
|
||||
)
|
||||
if self.bridge not in VDN_H3_BRIDGE_MODES:
|
||||
raise ValueError(
|
||||
f"hybrid_attention.bridge={self.bridge!r}; expected one of "
|
||||
f"{VDN_H3_BRIDGE_MODES}"
|
||||
)
|
||||
if self.anchor_frames not in VDN_H3_ANCHOR_FRAME_MODES:
|
||||
raise ValueError(
|
||||
f"hybrid_attention.anchor_frames={self.anchor_frames!r}; expected "
|
||||
f"one of {VDN_H3_ANCHOR_FRAME_MODES}"
|
||||
)
|
||||
if any(t not in VDN_H3_SHORT_CONV_TARGETS for t in self.short_conv) or len(
|
||||
set(self.short_conv)
|
||||
) != len(self.short_conv):
|
||||
raise ValueError(
|
||||
f"hybrid_attention.short_conv={self.short_conv!r}; expected a "
|
||||
f"distinct subset of {VDN_H3_SHORT_CONV_TARGETS}"
|
||||
)
|
||||
if self.chunk < 0 or self.radius < 0:
|
||||
raise ValueError("hybrid_attention.chunk and radius must be >= 0")
|
||||
if self.linear_head_dim <= 0:
|
||||
raise ValueError("hybrid_attention.linear_head_dim must be positive")
|
||||
|
||||
@classmethod
|
||||
def from_transform_config(
|
||||
cls, config: dict[str, Any]
|
||||
) -> "VDNHybridAttentionArchConfig":
|
||||
"""Build from VDN's nested v2 transform config."""
|
||||
soft = dict(config.get("softmax_attention", {}))
|
||||
lin = dict(config.get("linear_attention", {}))
|
||||
short_conv = lin.get("short_conv", {"targets": []})
|
||||
targets = (
|
||||
short_conv.get("targets", [])
|
||||
if isinstance(short_conv, dict)
|
||||
else list(short_conv or [])
|
||||
)
|
||||
return cls(
|
||||
chunk=int(soft.get("chunk", 0)),
|
||||
radius=int(soft["radius"]),
|
||||
anchor_frames=str(config.get("anchor_frames", "none")),
|
||||
enable_softmax_gate=bool(config.get("enable_softmax_gate", True)),
|
||||
delta_rule=str(lin.get("delta_rule", "vdn_solve")),
|
||||
linear_head_dim=int(lin["linear_head_dim"]),
|
||||
bridge=str(lin.get("bridge", "alpha")),
|
||||
a_fp32=bool(lin.get("a_fp32", True)),
|
||||
enable_text_state=bool(lin.get("enable_text_state", False)),
|
||||
short_conv=tuple(targets),
|
||||
)
|
||||
|
||||
def window_bounds(self, num_frames: int) -> list[tuple[int, int]]:
|
||||
"""Per-frame inclusive softmax-window bounds [lo, hi], unclamped."""
|
||||
if self.chunk <= 0:
|
||||
return [(t - self.radius, t + self.radius) for t in range(num_frames)]
|
||||
return [
|
||||
(
|
||||
((t // self.chunk) - self.radius) * self.chunk,
|
||||
((t // self.chunk) + self.radius + 1) * self.chunk - 1,
|
||||
)
|
||||
for t in range(num_frames)
|
||||
]
|
||||
|
||||
def full_cover(self, num_frames: int) -> bool:
|
||||
"""True when every frame's window already spans the whole clip, i.e.
|
||||
the softmax branch IS dense attention and the linear branch is off."""
|
||||
return all(
|
||||
lo <= 0 and hi >= num_frames - 1
|
||||
for lo, hi in self.window_bounds(num_frames)
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["VDNHybridAttentionArchConfig"]
|
||||
@@ -0,0 +1,76 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""VDN-H3 pipeline config: the MiniMax-H3 deployment envelope for the hybrid
|
||||
attention checkpoint."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import (
|
||||
MiniMaxH3PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import (
|
||||
AttentionBackendEnum,
|
||||
current_platform,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class VDNH3PipelineConfig(MiniMaxH3PipelineConfig):
|
||||
"""VDN-H3 (hybrid window-softmax + Video Delta linear attention, 8-NFE DMD2
|
||||
distill; t2va and fl2va on one fl2va partition): the deployment envelope;
|
||||
the arch config comes from the materialized ``transformer/config.json``."""
|
||||
|
||||
def validate_quality_deployment(self, server_args) -> None:
|
||||
raise ValueError(
|
||||
'quality="high" is audited only for the base MiniMax-H3 50-step '
|
||||
"4xH200 deployment; the VDN-H3 8-step hybrid checkpoint has no "
|
||||
'audited high-quality deployment. Use quality="lossless".'
|
||||
)
|
||||
|
||||
def validate_server_args(self, server_args) -> None:
|
||||
if server_args.model_variant is not None:
|
||||
raise ValueError(
|
||||
"VDN-H3 ships one weight partition (fl2va, serving t2va and "
|
||||
"fl2va); --model-variant does not apply. Ref2VA was not trained; "
|
||||
"use MiniMaxAI/MiniMax-H3 --model-variant ref2va."
|
||||
)
|
||||
quantization = (server_args.quantization or "").lower()
|
||||
if quantization in ("none", "bf16"):
|
||||
server_args.quantization = None
|
||||
elif (
|
||||
current_platform.is_blackwell() or current_platform.is_sm120()
|
||||
) and quantization in ("", "fp8"):
|
||||
# the block-scaled GEMM exists on SM100+ only; before that fp8 stays per-channel
|
||||
server_args.quantization = "mxfp8"
|
||||
# an unset backend would resolve to the platform default (dense FA)
|
||||
if server_args.attention_backend is None and not (
|
||||
server_args.component_attention_backends or {}
|
||||
).get("transformer"):
|
||||
server_args.attention_backend = "hybrid_window_attn_h3"
|
||||
selected_backend = self.resolve_transformer_attention_backend(server_args)
|
||||
if selected_backend is not AttentionBackendEnum.HYBRID_WINDOW_ATTN_H3:
|
||||
# the base-H3+LoRA equivalence smoke runs plain attention on purpose
|
||||
config = server_args.attention_backend_config or {}
|
||||
if not bool(config.get("vdn_h3_dense_smoke", False)):
|
||||
raise ValueError(
|
||||
"VDN-H3 requires --attention-backend hybrid_window_attn_h3 for "
|
||||
f"the transformer (got {selected_backend}); a dense backend "
|
||||
"would skip the linear branch and the softmax gates and "
|
||||
"produce the wrong model, not a slower one. Pass "
|
||||
"--attention-backend-config '{\"vdn_h3_dense_smoke\": true}' "
|
||||
"only for the base-H3+LoRA equivalence smoke."
|
||||
)
|
||||
if int(server_args.ring_degree or 1) > 1:
|
||||
raise ValueError(
|
||||
"VDN-H3 does not support --ring-degree > 1; use Ulysses sequence "
|
||||
"parallelism."
|
||||
)
|
||||
if server_args.enable_torch_compile or server_args.enable_breakable_cuda_graph:
|
||||
# BCG keeps one pool per captured segment and exhausts 183 GB at 104k rows
|
||||
raise ValueError(
|
||||
"VDN-H3 hybrid attention is not validated under torch.compile or "
|
||||
"the breakable CUDA graph yet; disable them."
|
||||
)
|
||||
super().validate_server_args(server_args)
|
||||
|
||||
|
||||
__all__ = ["VDNH3PipelineConfig"]
|
||||
@@ -0,0 +1,36 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""VDN-H3 sampling params: the 8-NFE grid on the MiniMax-H3 request surface."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from sglang.multimodal_gen.configs.sample.minimax_h3 import MiniMaxH3SamplingParams
|
||||
|
||||
|
||||
@dataclass
|
||||
class VDNH3SamplingParams(MiniMaxH3SamplingParams):
|
||||
"""VDN-H3: nine sigma grid points, i.e. the eight distilled DiT forwards
|
||||
(VDN counts NFEs; SGLang counts sigma grid points). The turbo adapter is
|
||||
only valid at 8 NFE with video shift 12 / audio shift 3 (the defaults)."""
|
||||
|
||||
num_inference_steps: int = 9
|
||||
|
||||
def _validate(self) -> None:
|
||||
super()._validate()
|
||||
if self.num_inference_steps != 9:
|
||||
raise ValueError(
|
||||
"VDN-H3 is distilled for exactly nine sigma grid points (eight DiT "
|
||||
f"forwards); got num_inference_steps={self.num_inference_steps}. "
|
||||
"Use MiniMaxAI/MiniMax-H3 for other schedules."
|
||||
)
|
||||
if self.task is not None and self.task.strip().lower() not in (
|
||||
"t2va",
|
||||
"fl2va",
|
||||
):
|
||||
raise ValueError(
|
||||
"VDN-H3 serves t2va and fl2va; ref2va was not trained (got "
|
||||
f"task={self.task!r}). Use MiniMaxAI/MiniMax-H3 --model-variant "
|
||||
"ref2va for that task."
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["VDNH3SamplingParams"]
|
||||
@@ -84,6 +84,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
|
||||
LTX23PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.ltx_2_5 import LTX25PipelineConfig
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3_vdn import (
|
||||
VDNH3PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.mova import (
|
||||
MOVA360PConfig,
|
||||
MOVA720PConfig,
|
||||
@@ -170,6 +173,7 @@ from sglang.multimodal_gen.configs.sample.minimax_h3 import (
|
||||
FastH3SamplingParams,
|
||||
MiniMaxH3SamplingParams,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sample.minimax_h3_vdn import VDNH3SamplingParams
|
||||
from sglang.multimodal_gen.configs.sample.mova import (
|
||||
MOVA_360P_SamplingParams,
|
||||
MOVA_720P_SamplingParams,
|
||||
@@ -350,6 +354,7 @@ KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS: Dict[str, str] = {
|
||||
"minimaxai/minimax-h3": "MiniMaxH3Pipeline",
|
||||
"minimax/minimax-h3": "MiniMaxH3Pipeline",
|
||||
"fastvideo/fastvideo-fasth3-4-step-preview-v1-vsa-datafree": "FastH3Pipeline",
|
||||
"openvdn/vdn-minimax-h3": "VDNH3Pipeline",
|
||||
"lerobot/pi05": "Pi05Pipeline",
|
||||
"pi05": "Pi05Pipeline",
|
||||
"pi0.5": "Pi05Pipeline",
|
||||
@@ -1016,6 +1021,7 @@ def _register_configs():
|
||||
model_detectors=[
|
||||
lambda model_id: (
|
||||
"minimaxh3" in model_id.lower().replace("-", "").replace("_", "")
|
||||
and "vdn" not in model_id.lower()
|
||||
)
|
||||
],
|
||||
)
|
||||
@@ -1038,6 +1044,19 @@ def _register_configs():
|
||||
)
|
||||
],
|
||||
)
|
||||
register_configs(
|
||||
sampling_param_cls=VDNH3SamplingParams,
|
||||
pipeline_config_cls=VDNH3PipelineConfig,
|
||||
hf_model_paths=[
|
||||
"OpenVDN/vdn-minimax-h3",
|
||||
],
|
||||
model_detectors=[
|
||||
lambda model_id: (
|
||||
"vdn" in model_id.lower()
|
||||
and "minimaxh3" in model_id.lower().replace("-", "").replace("_", "")
|
||||
)
|
||||
],
|
||||
)
|
||||
# FLUX
|
||||
register_configs(
|
||||
sampling_param_cls=FluxSamplingParams,
|
||||
|
||||
+521
@@ -0,0 +1,521 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""VDN-H3 window-softmax backend on the MiniMax-H3 packed layout.
|
||||
|
||||
An exact softmax over a chunk-aligned frame window: frame t belongs to chunk
|
||||
t // chunk and attends to chunks [c - radius, c + radius]; frames 0 and F-1
|
||||
are dense anchors; text and audio rows are dense both ways; padding rows sit
|
||||
outside every mask; a per-(token, head) sigmoid gate scales the output. The
|
||||
linear branch (``minimax_h3_vdn.py``) covers the window's complement. The
|
||||
metadata is request-static and installed once per request through
|
||||
``set_forward_context``. The window runs as a union of dense varlen
|
||||
FlashAttention calls: the dense-query rows against all keys, then per-chunk
|
||||
gathered [globals | window | anchors] K/V; same math as a masked kernel up to
|
||||
bf16 reduction order.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import msgspec
|
||||
import torch
|
||||
|
||||
from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func
|
||||
from sglang.multimodal_gen.configs.models.dits.minimax_h3_vdn import (
|
||||
VDNHybridAttentionArchConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends import (
|
||||
flash_attn as _flash_attn_backend,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
|
||||
AttentionBackend,
|
||||
AttentionImpl,
|
||||
AttentionMetadata,
|
||||
AttentionMetadataBuilder,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.dits.minimax_h3_vdn import VDNH3Layout
|
||||
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
|
||||
|
||||
_DIT_BLOCK_PREFIX = re.compile(r"^blocks\.(\d+)\.")
|
||||
|
||||
|
||||
class HybridWindowAttentionH3Backend(AttentionBackend):
|
||||
accept_output_buffer: bool = False
|
||||
|
||||
@staticmethod
|
||||
def get_supported_head_sizes() -> list[int]:
|
||||
return [64, 128]
|
||||
|
||||
@staticmethod
|
||||
def get_enum() -> AttentionBackendEnum:
|
||||
return AttentionBackendEnum.HYBRID_WINDOW_ATTN_H3
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls() -> type[HybridWindowAttentionH3Impl]:
|
||||
return HybridWindowAttentionH3Impl
|
||||
|
||||
@staticmethod
|
||||
def get_metadata_cls() -> type[HybridWindowAttentionH3Metadata]:
|
||||
return HybridWindowAttentionH3Metadata
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type[HybridWindowAttentionH3MetadataBuilder]:
|
||||
return HybridWindowAttentionH3MetadataBuilder
|
||||
|
||||
|
||||
def window_mask_frames(
|
||||
hybrid: VDNHybridAttentionArchConfig, num_frames: int
|
||||
) -> tuple[list[tuple[int, int]], set[int], set[int]]:
|
||||
"""(clamped per-frame window bounds, dense-ROW frames, dense-COLUMN frames)."""
|
||||
bounds = [
|
||||
(max(lo, 0), min(hi, num_frames - 1))
|
||||
for lo, hi in hybrid.window_bounds(num_frames)
|
||||
]
|
||||
anchors = {0, num_frames - 1} if hybrid.anchor_frames != "none" else set()
|
||||
dense_rows = anchors if hybrid.anchor_frames in ("rows", "both") else set()
|
||||
dense_cols = anchors if hybrid.anchor_frames in ("columns", "both") else set()
|
||||
return bounds, dense_rows, dense_cols
|
||||
|
||||
|
||||
def window_mask_reference(
|
||||
hybrid: VDNHybridAttentionArchConfig, layout: VDNH3Layout, device: torch.device
|
||||
) -> torch.Tensor:
|
||||
"""Dense boolean [used, used] mask of the softmax branch, for tests."""
|
||||
used = layout.used
|
||||
keep = torch.ones(used, used, dtype=torch.bool, device=device)
|
||||
vs, ve = layout.video_start, layout.video_end
|
||||
bounds, dense_rows, dense_cols = window_mask_frames(hybrid, layout.num_frames)
|
||||
tpf = layout.tokens_per_frame
|
||||
rows = torch.arange(vs, ve, device=device)
|
||||
frame_of = (rows - vs) // tpf
|
||||
qf = frame_of[:, None]
|
||||
kf = frame_of[None, :]
|
||||
lo = torch.tensor([b[0] for b in bounds], device=device)[qf]
|
||||
hi = torch.tensor([b[1] for b in bounds], device=device)[qf]
|
||||
inside = (kf >= lo) & (kf <= hi)
|
||||
for f in dense_rows:
|
||||
inside |= qf == f
|
||||
for f in dense_cols:
|
||||
inside |= kf == f
|
||||
keep[vs:ve, vs:ve] = inside
|
||||
return keep
|
||||
|
||||
|
||||
def _merge_ranges(ranges: list[tuple[int, int]]) -> list[tuple[int, int]]:
|
||||
out: list[tuple[int, int]] = []
|
||||
for a, b in sorted(ranges):
|
||||
if out and out[-1][1] >= a:
|
||||
out[-1] = (out[-1][0], max(out[-1][1], b))
|
||||
else:
|
||||
out.append((a, b))
|
||||
return out
|
||||
|
||||
|
||||
def _cat_ranges(ranges: list[tuple[int, int]], *, device: torch.device) -> torch.Tensor:
|
||||
if not ranges:
|
||||
return torch.empty(0, dtype=torch.long, device=device)
|
||||
return torch.cat(
|
||||
[torch.arange(a, b, device=device, dtype=torch.long) for a, b in ranges]
|
||||
)
|
||||
|
||||
|
||||
def _chunk_groups(
|
||||
raw_bounds: list[tuple[int, int]], dense_rows: set[int]
|
||||
) -> list[list[int]]:
|
||||
# consecutive window frames with identical bounds share one varlen segment
|
||||
groups: list[list[int]] = []
|
||||
for f in range(len(raw_bounds)):
|
||||
if f in dense_rows:
|
||||
continue
|
||||
if (
|
||||
groups
|
||||
and raw_bounds[groups[-1][-1]] == raw_bounds[f]
|
||||
and groups[-1][-1] == f - 1
|
||||
):
|
||||
groups[-1].append(f)
|
||||
else:
|
||||
groups.append([f])
|
||||
return groups
|
||||
|
||||
|
||||
class _ChunkGroup(msgspec.Struct, frozen=True):
|
||||
frames: list[int]
|
||||
query_rows: torch.Tensor
|
||||
kv_rows: torch.Tensor
|
||||
|
||||
|
||||
class _WindowPass(msgspec.Struct, frozen=True):
|
||||
query_rows: torch.Tensor
|
||||
query_slice: tuple[int, int] | None # set when the query rows are contiguous
|
||||
kv_rows: torch.Tensor
|
||||
cu_q: torch.Tensor
|
||||
cu_k: torch.Tensor
|
||||
max_q: int
|
||||
max_k: int
|
||||
|
||||
|
||||
def _window_pass(
|
||||
layout: VDNH3Layout, groups: list[_ChunkGroup], device: torch.device
|
||||
) -> _WindowPass:
|
||||
query_lens = [int(group.query_rows.numel()) for group in groups]
|
||||
kv_lens = [int(group.kv_rows.numel()) for group in groups]
|
||||
frames = [frame for group in groups for frame in group.frames]
|
||||
contiguous = frames == list(range(frames[0], frames[0] + len(frames)))
|
||||
zero = torch.zeros(1, dtype=torch.long)
|
||||
return _WindowPass(
|
||||
query_rows=torch.cat([group.query_rows for group in groups]),
|
||||
query_slice=(
|
||||
(layout.frame_rows(frames[0])[0], layout.frame_rows(frames[-1])[1])
|
||||
if contiguous
|
||||
else None
|
||||
),
|
||||
kv_rows=torch.cat([group.kv_rows for group in groups]),
|
||||
cu_q=torch.cat([zero, torch.tensor(query_lens).cumsum(0)]).to(
|
||||
device, torch.int32
|
||||
),
|
||||
cu_k=torch.cat([zero, torch.tensor(kv_lens).cumsum(0)]).to(device, torch.int32),
|
||||
max_q=max(query_lens),
|
||||
max_k=max(kv_lens),
|
||||
)
|
||||
|
||||
|
||||
def _window_passes(
|
||||
layout: VDNH3Layout,
|
||||
groups: list[_ChunkGroup],
|
||||
max_gather_rows: int,
|
||||
device: torch.device,
|
||||
) -> list[_WindowPass]:
|
||||
# one varlen call per pass; consecutive chunk groups fill up to max_gather_rows
|
||||
passes: list[_WindowPass] = []
|
||||
current: list[_ChunkGroup] = []
|
||||
current_rows = 0
|
||||
for group in groups:
|
||||
rows = int(group.kv_rows.numel())
|
||||
if current and current_rows + rows > max_gather_rows:
|
||||
passes.append(_window_pass(layout, current, device))
|
||||
current, current_rows = [], 0
|
||||
current.append(group)
|
||||
current_rows += rows
|
||||
if current:
|
||||
passes.append(_window_pass(layout, current, device))
|
||||
return passes
|
||||
|
||||
|
||||
class _DecomposedPlan:
|
||||
"""Query-row groups with identical kept key sets, as dense varlen calls:
|
||||
the dense-q rows against all ``used`` keys, then each chunk of frames
|
||||
against its gathered [globals | window | anchors] keys."""
|
||||
|
||||
__slots__ = ("dense_q", "dense_cu_q", "dense_cu_k", "passes")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
layout: VDNH3Layout,
|
||||
hybrid: VDNHybridAttentionArchConfig,
|
||||
device: torch.device,
|
||||
max_gather_rows: int = 200_000,
|
||||
) -> None:
|
||||
used, num_frames = layout.used, layout.num_frames
|
||||
bounds, dense_rows, dense_cols = window_mask_frames(hybrid, num_frames)
|
||||
rows = functools.partial(_cat_ranges, device=device)
|
||||
dense_ranges = _merge_ranges(
|
||||
layout.global_ranges + [layout.frame_rows(f) for f in sorted(dense_rows)]
|
||||
)
|
||||
self.dense_q = rows(dense_ranges)
|
||||
# built once: a tensor from a Python list costs a pageable H2D copy + sync
|
||||
self.dense_cu_q = torch.tensor(
|
||||
[0, int(self.dense_q.numel())], dtype=torch.int32, device=device
|
||||
)
|
||||
self.dense_cu_k = torch.tensor([0, used], dtype=torch.int32, device=device)
|
||||
groups = []
|
||||
for frames in _chunk_groups(hybrid.window_bounds(num_frames), dense_rows):
|
||||
lo, hi = bounds[frames[0]]
|
||||
kv_frames = sorted(set(range(lo, hi + 1)) | dense_cols)
|
||||
groups.append(
|
||||
_ChunkGroup(
|
||||
frames=frames,
|
||||
query_rows=rows(
|
||||
_merge_ranges([layout.frame_rows(f) for f in frames])
|
||||
),
|
||||
kv_rows=rows(
|
||||
_merge_ranges(
|
||||
layout.global_ranges
|
||||
+ [layout.frame_rows(f) for f in kv_frames]
|
||||
)
|
||||
),
|
||||
)
|
||||
)
|
||||
self.passes = _window_passes(layout, groups, max_gather_rows, device)
|
||||
window_rows = sum(int(p.query_rows.numel()) for p in self.passes)
|
||||
covered = int(self.dense_q.numel()) + window_rows
|
||||
if covered != used:
|
||||
raise ValueError(
|
||||
f"window decomposition covers {covered} of {used} packed rows"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class HybridWindowAttentionH3Metadata(AttentionMetadata):
|
||||
layout: VDNH3Layout
|
||||
# radius >= F: the window IS dense attention and the linear branch is off
|
||||
full_cover: bool
|
||||
decomposed: _DecomposedPlan | None = None
|
||||
# (cos_sin [seq_len, rope_dim] bf16, positions [seq_len]) under Ulysses, else None
|
||||
rope_cache_full: tuple[torch.Tensor, torch.Tensor] | None = None
|
||||
|
||||
|
||||
class HybridWindowAttentionH3MetadataBuilder(AttentionMetadataBuilder):
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def prepare(self) -> None:
|
||||
pass
|
||||
|
||||
def build( # type: ignore[override]
|
||||
self,
|
||||
*,
|
||||
layout: VDNH3Layout,
|
||||
hybrid: VDNHybridAttentionArchConfig,
|
||||
device: torch.device,
|
||||
rope_cache_full: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
current_timestep: int = 0,
|
||||
max_gather_rows: int = 200_000,
|
||||
**kwargs: dict[str, Any],
|
||||
) -> HybridWindowAttentionH3Metadata:
|
||||
full_cover = hybrid.full_cover(layout.num_frames)
|
||||
decomposed = None
|
||||
if not full_cover:
|
||||
decomposed = _DecomposedPlan(
|
||||
layout, hybrid, device, max_gather_rows=max_gather_rows
|
||||
)
|
||||
return HybridWindowAttentionH3Metadata(
|
||||
current_timestep=current_timestep,
|
||||
layout=layout,
|
||||
full_cover=full_cover,
|
||||
decomposed=decomposed,
|
||||
rope_cache_full=rope_cache_full,
|
||||
)
|
||||
|
||||
|
||||
def _fa_varlen(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
*,
|
||||
cu_q: torch.Tensor,
|
||||
cu_k: torch.Tensor,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
scale: float,
|
||||
out: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
attn_out = flash_attn_varlen_func(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seqlens_q=cu_q,
|
||||
cu_seqlens_k=cu_k,
|
||||
max_seqlen_q=max_q,
|
||||
max_seqlen_k=max_k,
|
||||
softmax_scale=scale,
|
||||
causal=False,
|
||||
ver=_flash_attn_backend.fa_ver,
|
||||
out=out,
|
||||
)
|
||||
attn_out = attn_out[0] if isinstance(attn_out, tuple) else attn_out
|
||||
if out is not None and attn_out.data_ptr() != out.data_ptr():
|
||||
out.copy_(attn_out)
|
||||
return out
|
||||
return attn_out
|
||||
|
||||
|
||||
class HybridWindowAttentionH3Impl(AttentionImpl):
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
causal: bool,
|
||||
softmax_scale: float,
|
||||
num_kv_heads: int | None = None,
|
||||
prefix: str = "",
|
||||
**extra_impl_args,
|
||||
) -> None:
|
||||
self.num_heads = num_heads
|
||||
self.head_size = head_size
|
||||
self.softmax_scale = softmax_scale
|
||||
self.prefix = prefix
|
||||
match = _DIT_BLOCK_PREFIX.match(prefix)
|
||||
self.layer_idx = int(match.group(1)) if match else None
|
||||
# non-DiT callers (the token refiner) resolve this backend too: dense FA
|
||||
self._dense_fallback = _flash_attn_backend.FlashAttentionImpl(
|
||||
num_heads=num_heads,
|
||||
head_size=head_size,
|
||||
causal=causal,
|
||||
softmax_scale=softmax_scale,
|
||||
num_kv_heads=num_kv_heads,
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata: AttentionMetadata,
|
||||
) -> torch.Tensor:
|
||||
"""Dense FlashAttention for the non-DiT layers this backend reaches."""
|
||||
return self._dense_fallback.forward(query, key, value, attn_metadata)
|
||||
|
||||
def dense_varlen(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
*,
|
||||
cu_seqlens: torch.Tensor,
|
||||
max_seqlen: int,
|
||||
cu_seqlens_host: tuple[int, ...] | None = None,
|
||||
) -> torch.Tensor:
|
||||
return self._dense_fallback.forward_varlen(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
cu_seqlens_host=cu_seqlens_host,
|
||||
)
|
||||
|
||||
def forward_varlen(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
*,
|
||||
cu_seqlens: torch.Tensor,
|
||||
max_seqlen: int,
|
||||
cu_seqlens_host: tuple[int, ...] | None = None,
|
||||
attn_metadata: HybridWindowAttentionH3Metadata | None = None,
|
||||
softmax_gate: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""query/key/value: [T, H, D] packed rows (post-norm, post-RoPE) ->
|
||||
[T, H, D]; ``softmax_gate`` [T, H] scales the output per (row, head).
|
||||
Rows at and past ``used`` (padding) are zero."""
|
||||
if self.layer_idx is not None and attn_metadata is None:
|
||||
raise RuntimeError(
|
||||
"hybrid_window_attn_h3 needs per-request attention metadata "
|
||||
"from the MiniMax-H3 denoising stage; none was set in the "
|
||||
"forward context."
|
||||
)
|
||||
if self.layer_idx is None:
|
||||
return self.dense_varlen(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
cu_seqlens_host=cu_seqlens_host,
|
||||
)
|
||||
|
||||
meta = attn_metadata
|
||||
layout = meta.layout
|
||||
bounds = (
|
||||
cu_seqlens_host
|
||||
if cu_seqlens_host is not None
|
||||
else tuple(int(item) for item in cu_seqlens.tolist())
|
||||
)
|
||||
used = int(bounds[1])
|
||||
if used != layout.used or query.shape[0] != layout.seq_len:
|
||||
raise ValueError(
|
||||
f"hybrid_window_attn_h3 metadata was built for used={layout.used} "
|
||||
f"of seq_len={layout.seq_len} rows, got used={used} of "
|
||||
f"{query.shape[0]}. The request metadata and the packed layout "
|
||||
"have diverged."
|
||||
)
|
||||
|
||||
if meta.full_cover:
|
||||
out = self.dense_varlen(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
cu_seqlens_host=cu_seqlens_host,
|
||||
)
|
||||
else:
|
||||
out = self._decomposed(query, key, value, meta.decomposed, used)
|
||||
|
||||
if softmax_gate is not None:
|
||||
out.mul_(softmax_gate.to(out.dtype).unsqueeze(-1))
|
||||
if used < out.shape[0]:
|
||||
out[used:].zero_()
|
||||
return out
|
||||
|
||||
def _decomposed(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
plan: _DecomposedPlan,
|
||||
used: int,
|
||||
) -> torch.Tensor:
|
||||
out = torch.empty_like(query)
|
||||
key_used = key[:used]
|
||||
value_used = value[:used]
|
||||
if not key_used.is_contiguous():
|
||||
key_used = key_used.contiguous()
|
||||
if not value_used.is_contiguous():
|
||||
value_used = value_used.contiguous()
|
||||
if plan.dense_q.numel():
|
||||
out[plan.dense_q] = _fa_varlen(
|
||||
torch.index_select(query, 0, plan.dense_q),
|
||||
key_used,
|
||||
value_used,
|
||||
cu_q=plan.dense_cu_q,
|
||||
cu_k=plan.dense_cu_k,
|
||||
max_q=int(plan.dense_q.numel()),
|
||||
max_k=used,
|
||||
scale=self.softmax_scale,
|
||||
)
|
||||
for window in plan.passes:
|
||||
# index_select on contiguous copies takes the vectorized gather kernel
|
||||
keys = torch.index_select(key_used, 0, window.kv_rows)
|
||||
values = torch.index_select(value_used, 0, window.kv_rows)
|
||||
if window.query_slice is not None:
|
||||
start, stop = window.query_slice
|
||||
_fa_varlen(
|
||||
query[start:stop],
|
||||
keys,
|
||||
values,
|
||||
cu_q=window.cu_q,
|
||||
cu_k=window.cu_k,
|
||||
max_q=window.max_q,
|
||||
max_k=window.max_k,
|
||||
scale=self.softmax_scale,
|
||||
out=out[start:stop],
|
||||
)
|
||||
else:
|
||||
out[window.query_rows] = _fa_varlen(
|
||||
torch.index_select(query, 0, window.query_rows),
|
||||
keys,
|
||||
values,
|
||||
cu_q=window.cu_q,
|
||||
cu_k=window.cu_k,
|
||||
max_q=window.max_q,
|
||||
max_k=window.max_k,
|
||||
scale=self.softmax_scale,
|
||||
)
|
||||
del keys, values
|
||||
return out
|
||||
|
||||
|
||||
__all__ = [
|
||||
"HybridWindowAttentionH3Backend",
|
||||
"HybridWindowAttentionH3Impl",
|
||||
"HybridWindowAttentionH3Metadata",
|
||||
"HybridWindowAttentionH3MetadataBuilder",
|
||||
"window_mask_frames",
|
||||
"window_mask_reference",
|
||||
]
|
||||
@@ -91,6 +91,11 @@ def adjust_scalar_to_fused_array(
|
||||
class LinearMethodBase(QuantizeMethodBase):
|
||||
"""Base class for different (maybe quantized) linear methods."""
|
||||
|
||||
def accepts_mxfp8_input(self, layer: torch.nn.Module) -> bool:
|
||||
"""Whether ``apply`` takes a prequantized ``(e4m3 input, swizzled E8M0
|
||||
block scales)`` tuple for this layer."""
|
||||
return False
|
||||
|
||||
@abstractmethod
|
||||
def create_weights(
|
||||
self,
|
||||
|
||||
@@ -76,6 +76,12 @@ class MXFP8Config(SRTFp8Config, QuantizationConfig):
|
||||
return UnquantizedLinearMethod()
|
||||
if current_platform.is_npu():
|
||||
return NPUMXFP8LinearMethod(self)
|
||||
if not self.is_checkpoint_fp8_serialized:
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.mxfp8_online import (
|
||||
MXFP8OnlineLinearMethod,
|
||||
)
|
||||
|
||||
return MXFP8OnlineLinearMethod(self)
|
||||
return SRTFp8LinearMethod(self)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,104 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Online MXFP8 for diffusion linears: ``--quantization mxfp8`` on a bf16 checkpoint.
|
||||
|
||||
Weights quantize at load to e4m3 with one E8M0 scale per 32 elements along K
|
||||
(scales in the cuBLASLt 128x4 swizzled layout, ``mxfp8_quantize_swizzled``);
|
||||
activations take the same block quant per call unless the producer hands over
|
||||
a prequantized ``(fp8, swizzled scales)`` tuple; the GEMM is cuBLASLt's
|
||||
block-scaled ``torch.nn.functional.scaled_mm``. Layers with K not a multiple
|
||||
of 32, N not a multiple of 16, non-bf16 params, or pre-Blackwell GPUs keep
|
||||
the per-channel fp8 path.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
from torch.nn import Module
|
||||
from torch.nn.functional import ScalingType, SwizzleType, scaled_mm
|
||||
|
||||
from sglang.kernels.ops.diffusion import (
|
||||
can_use_mxfp8_swizzled,
|
||||
mxfp8_quantize_swizzled,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8LinearMethod
|
||||
|
||||
_E8M0 = torch.float8_e8m0fnu
|
||||
|
||||
|
||||
def mxfp8_scaled_mm(
|
||||
a: torch.Tensor,
|
||||
a_scale: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
weight_scale: torch.Tensor,
|
||||
*,
|
||||
bias: Optional[torch.Tensor],
|
||||
output_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
return scaled_mm(
|
||||
a,
|
||||
weight.t(),
|
||||
scale_a=a_scale,
|
||||
scale_recipe_a=ScalingType.BlockWise1x32,
|
||||
swizzle_a=SwizzleType.SWIZZLE_32_4_4,
|
||||
scale_b=weight_scale,
|
||||
scale_recipe_b=ScalingType.BlockWise1x32,
|
||||
swizzle_b=SwizzleType.SWIZZLE_32_4_4,
|
||||
bias=bias,
|
||||
output_dtype=output_dtype,
|
||||
)
|
||||
|
||||
|
||||
class MXFP8OnlineLinearMethod(Fp8LinearMethod):
|
||||
def __init__(self, quant_config) -> None:
|
||||
super().__init__(quant_config)
|
||||
# SRTFp8Config(use_mxfp8) sets a block size; the fallback path loads bf16
|
||||
self.block_quant = False
|
||||
|
||||
def process_weights_after_loading(self, layer: Module) -> None:
|
||||
layer.mxfp8 = (
|
||||
not self.use_marlin
|
||||
and can_use_mxfp8_swizzled(layer.weight)
|
||||
and torch.cuda.get_device_capability(layer.weight.device)[0] >= 10
|
||||
and layer.weight.shape[0] % 16 == 0
|
||||
)
|
||||
if not layer.mxfp8:
|
||||
super().process_weights_after_loading(layer)
|
||||
return
|
||||
qweight, scale = mxfp8_quantize_swizzled(layer.weight.data.contiguous())
|
||||
layer.weight = torch.nn.Parameter(qweight, requires_grad=False)
|
||||
layer.weight_scale = torch.nn.Parameter(scale.view(_E8M0), requires_grad=False)
|
||||
layer.input_scale = None
|
||||
|
||||
def accepts_mxfp8_input(self, layer: Module) -> bool:
|
||||
return bool(layer.mxfp8)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
if not layer.mxfp8:
|
||||
return super().apply(layer, x, bias)
|
||||
if isinstance(x, tuple):
|
||||
a, a_scale = x
|
||||
a_scale = a_scale.view(_E8M0)
|
||||
lead, output_dtype = a.shape[:-1], torch.bfloat16
|
||||
else:
|
||||
lead, output_dtype = x.shape[:-1], x.dtype
|
||||
a, a_scale = mxfp8_quantize_swizzled(
|
||||
x.reshape(-1, x.shape[-1]).contiguous()
|
||||
)
|
||||
a_scale = a_scale.view(_E8M0)
|
||||
out = mxfp8_scaled_mm(
|
||||
a.reshape(-1, a.shape[-1]),
|
||||
a_scale,
|
||||
layer.weight,
|
||||
layer.weight_scale,
|
||||
bias=bias,
|
||||
output_dtype=output_dtype,
|
||||
)
|
||||
return out.view(*lead, out.shape[-1])
|
||||
|
||||
|
||||
__all__ = ["MXFP8OnlineLinearMethod", "mxfp8_scaled_mm"]
|
||||
@@ -22,10 +22,14 @@ from sglang.kernels.ops.activation.activation import (
|
||||
)
|
||||
from sglang.kernels.ops.diffusion import (
|
||||
can_use_fused_inplace_qknorm_rope,
|
||||
can_use_mxfp8_swizzled,
|
||||
can_use_silu_mul_mxfp8,
|
||||
fused_inplace_qknorm_rope,
|
||||
indexed_gate_bf16,
|
||||
indexed_gate_bf16_,
|
||||
indexed_scale_shift_bf16_,
|
||||
indexed_scale_shift_mxfp8_,
|
||||
silu_mul_mxfp8,
|
||||
)
|
||||
from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm
|
||||
from sglang.multimodal_gen import envs
|
||||
@@ -79,6 +83,9 @@ from sglang.multimodal_gen.runtime.models.dits.minimax_h3_adaln_cache import (
|
||||
from sglang.multimodal_gen.runtime.models.dits.minimax_h3_adaln_cache import (
|
||||
native_adaln_weight_files,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.dits.minimax_h3_vdn_attention import (
|
||||
MiniMaxH3VDNHybridAttention,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.parameter import BlockQuantScaleParameter
|
||||
from sglang.multimodal_gen.runtime.platforms import (
|
||||
AttentionBackendEnum,
|
||||
@@ -362,6 +369,12 @@ def _rotate_half(x: torch.Tensor) -> torch.Tensor:
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
|
||||
|
||||
def _accepts_mxfp8_input(linear: nn.Module) -> bool:
|
||||
return linear.quant_method is not None and linear.quant_method.accepts_mxfp8_input(
|
||||
linear
|
||||
)
|
||||
|
||||
|
||||
def _modulate_scale_shift(
|
||||
x: torch.Tensor,
|
||||
shift: torch.Tensor,
|
||||
@@ -827,6 +840,10 @@ class MiniMaxH3Attention(nn.Module):
|
||||
quant_config=None,
|
||||
prefix=f"{prefix}.to_gate_compress",
|
||||
)
|
||||
# VDN-H3: None for the dense model and for the token refiner
|
||||
self.hybrid = MiniMaxH3VDNHybridAttention.build(
|
||||
arch, quant_config, prefix=prefix, local_heads=self.num_heads
|
||||
)
|
||||
|
||||
def _set_attention_backend(self, backend) -> None:
|
||||
if (
|
||||
@@ -1032,9 +1049,12 @@ class MiniMaxH3Attention(nn.Module):
|
||||
subblock_sparse_query_block_mask: torch.Tensor | None = None,
|
||||
ulysses_active: bool = False,
|
||||
ring_active: bool = False,
|
||||
x_prequant: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""x: [T, hidden] packed thd rows -> [T, hidden].
|
||||
|
||||
``x_prequant``: ``x`` already quantized for ``qkv_proj`` as ``(fp8, scales)``.
|
||||
|
||||
Operation order: fused qkv projection -> per-head q/k RMSNorm -> RoPE
|
||||
on q/k -> variable-length non-causal flash attention -> output projection.
|
||||
|
||||
@@ -1054,11 +1074,29 @@ class MiniMaxH3Attention(nn.Module):
|
||||
)
|
||||
|
||||
total = x.shape[0]
|
||||
qkv, _ = self.qkv_proj(x)
|
||||
qkv, _ = self.qkv_proj(x if x_prequant is None else x_prequant)
|
||||
q, k, v = qkv.split(self.local_inner_dim, dim=-1)
|
||||
q = q.view(total, self.num_heads, self.head_dim)
|
||||
k = k.view(total, self.num_heads, self.head_dim)
|
||||
v = v.view(total, self.num_heads, self.head_dim)
|
||||
if (
|
||||
self.hybrid is not None
|
||||
and self._attention_backend_enum
|
||||
is AttentionBackendEnum.HYBRID_WINDOW_ATTN_H3
|
||||
):
|
||||
return self.hybrid(
|
||||
self,
|
||||
x,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
rope_cache=rope_cache,
|
||||
cu_seqlens=cu_seqlens,
|
||||
cu_seqlens_host=cu_seqlens_host,
|
||||
max_seqlen=max_seqlen,
|
||||
ulysses_active=ulysses_active,
|
||||
ring_active=ring_active,
|
||||
)
|
||||
if rope_cache is None:
|
||||
q, k = _apply_qk_norm(
|
||||
q,
|
||||
@@ -1160,7 +1198,7 @@ class MiniMaxH3MLP(nn.Module):
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if x.device.type == "mps":
|
||||
if not isinstance(x, tuple) and x.device.type == "mps":
|
||||
out = torch.empty_like(x)
|
||||
for start in range(0, x.shape[0], _MPS_MLP_TOKEN_CHUNK_SIZE):
|
||||
stop = min(start + _MPS_MLP_TOKEN_CHUNK_SIZE, x.shape[0])
|
||||
@@ -1173,6 +1211,9 @@ class MiniMaxH3MLP(nn.Module):
|
||||
torch.mps.empty_cache()
|
||||
return out
|
||||
hidden, _ = self.fc1(x)
|
||||
if _accepts_mxfp8_input(self.fc2) and can_use_silu_mul_mxfp8(hidden):
|
||||
out, _ = self.fc2(silu_mul_mxfp8(hidden))
|
||||
return out
|
||||
hidden = _silu_mul(hidden, reuse_input=self.reuse_fc1_activation)
|
||||
out, _ = self.fc2(hidden)
|
||||
return out
|
||||
@@ -1384,11 +1425,20 @@ class MiniMaxH3DiTBlock(nn.Module):
|
||||
# a block-local buffer.
|
||||
residual = x
|
||||
h = self.norm1(x)
|
||||
h = _modulate_scale_shift(
|
||||
h, shift_msa, scale_msa, combined_indices, dtype=_BF16_DTYPE
|
||||
)
|
||||
h_prequant = None
|
||||
if _accepts_mxfp8_input(self.attn.qkv_proj) and can_use_mxfp8_swizzled(h):
|
||||
# the bf16 modulated rows stay in h for the VDN branch projections
|
||||
h, h_fp8, h_scales = indexed_scale_shift_mxfp8_(
|
||||
h, shift_msa, scale_msa, combined_indices, keep_bf16=True
|
||||
)
|
||||
h_prequant = (h_fp8, h_scales)
|
||||
else:
|
||||
h = _modulate_scale_shift(
|
||||
h, shift_msa, scale_msa, combined_indices, dtype=_BF16_DTYPE
|
||||
)
|
||||
h = self.attn(
|
||||
h,
|
||||
x_prequant=h_prequant,
|
||||
rope_cache=rope_cache,
|
||||
cu_seqlens=cu_seqlens,
|
||||
cu_seqlens_host=cu_seqlens_host,
|
||||
@@ -1408,10 +1458,16 @@ class MiniMaxH3DiTBlock(nn.Module):
|
||||
|
||||
residual = x
|
||||
h = self.norm2(x)
|
||||
h = _modulate_scale_shift(
|
||||
h, shift_mlp, scale_mlp, combined_indices, dtype=_BF16_DTYPE
|
||||
)
|
||||
h = self.mlp(h)
|
||||
if _accepts_mxfp8_input(self.mlp.fc1) and can_use_mxfp8_swizzled(h):
|
||||
_, h_fp8, h_scales = indexed_scale_shift_mxfp8_(
|
||||
h, shift_mlp, scale_mlp, combined_indices, keep_bf16=False
|
||||
)
|
||||
h = self.mlp((h_fp8, h_scales))
|
||||
else:
|
||||
h = _modulate_scale_shift(
|
||||
h, shift_mlp, scale_mlp, combined_indices, dtype=_BF16_DTYPE
|
||||
)
|
||||
h = self.mlp(h)
|
||||
# `residual` is block-local here (see above), so this stays in-place
|
||||
# even while Cache-DiT is attached.
|
||||
return _modulate_gate(
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,647 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""VDN-H3 hybrid attention inside a MiniMax-H3 DiT block: per-head softmax gate,
|
||||
Video Delta linear branch and its output projection, the Ulysses exchange that
|
||||
shards both branches by head, and the request-static attention metadata.
|
||||
``MiniMaxH3Attention`` owns one instance as ``hybrid`` and hands it raw q/k/v."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any, Callable, Mapping
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from sglang.kernels.ops.diffusion import (
|
||||
fused_qknorm_rope_out_of_place,
|
||||
usp_merge_heads,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.models.dits.minimax_h3 import MiniMaxH3DiTArchConfig
|
||||
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
|
||||
get_ring_ctx,
|
||||
get_ulysses_ctx,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.hybrid_window_attn_h3 import (
|
||||
HybridWindowAttentionH3Metadata,
|
||||
HybridWindowAttentionH3MetadataBuilder,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.linear import RowParallelLinear
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
|
||||
QuantizationConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.usp import _a2a_staging_buffer
|
||||
from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
|
||||
from sglang.multimodal_gen.runtime.models.dits.minimax_h3_vdn import (
|
||||
MiniMaxH3VDNLinearBranch,
|
||||
VDNSoftmaxGate,
|
||||
vdn_h3_layout_from_packed,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
|
||||
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
|
||||
eager_on_graph,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import MiniMaxH3Attention
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
_FP32_DTYPE = torch.float32
|
||||
_BF16_DTYPE = torch.bfloat16
|
||||
|
||||
# one side stream per process and device: under Ulysses the linear readout and
|
||||
# its a2a run on it while FA4 holds the current stream
|
||||
_LINEAR_STREAMS: dict[int, torch.cuda.Stream] = {}
|
||||
|
||||
|
||||
def _linear_branch_stream(device: torch.device) -> torch.cuda.Stream:
|
||||
index = device.index if device.index is not None else torch.cuda.current_device()
|
||||
stream = _LINEAR_STREAMS.get(index)
|
||||
if stream is None:
|
||||
# high priority: the readout's small kernels fill the SM slots FA4's
|
||||
# persistent CTAs leave at their tails
|
||||
stream = torch.cuda.Stream(device=device, priority=-1)
|
||||
_LINEAR_STREAMS[index] = stream
|
||||
return stream
|
||||
|
||||
|
||||
class MiniMaxH3VDNHybridAttention(nn.Module):
|
||||
"""out = to_out(gate_sm * window_softmax(q, k, v)) + to_out_linear(branch(q, k, v))."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
arch: MiniMaxH3DiTArchConfig,
|
||||
quant_config: QuantizationConfig | None,
|
||||
*,
|
||||
prefix: str,
|
||||
local_heads: int,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
hybrid = arch.hybrid_attention
|
||||
self.softmax_gate: VDNSoftmaxGate | None = None
|
||||
if hybrid.enable_softmax_gate:
|
||||
self.softmax_gate = VDNSoftmaxGate(
|
||||
arch.hidden_size,
|
||||
arch.num_attention_heads,
|
||||
prefix=f"{prefix}.softmax_gate",
|
||||
)
|
||||
self.linear_attention = MiniMaxH3VDNLinearBranch(
|
||||
arch, hybrid, local_heads=local_heads, prefix=f"{prefix}.linear_attention"
|
||||
)
|
||||
self.to_out_linear = RowParallelLinear(
|
||||
arch.num_attention_heads * arch.attention_head_dim,
|
||||
arch.hidden_size,
|
||||
bias=False,
|
||||
input_is_parallel=True,
|
||||
params_dtype=_BF16_DTYPE,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.to_out_linear",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def build(
|
||||
cls,
|
||||
arch: MiniMaxH3DiTArchConfig,
|
||||
quant_config: QuantizationConfig | None,
|
||||
*,
|
||||
prefix: str,
|
||||
local_heads: int,
|
||||
) -> MiniMaxH3VDNHybridAttention | None:
|
||||
"""None for the dense model and the token refiner (VDN converts the DiT blocks only)."""
|
||||
if arch.hybrid_attention is None or not prefix.startswith("blocks."):
|
||||
return None
|
||||
return cls(
|
||||
arch, quant_config, prefix=f"{prefix}.hybrid", local_heads=local_heads
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
attention: MiniMaxH3Attention,
|
||||
x: torch.Tensor,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
*,
|
||||
rope_cache: tuple[torch.Tensor, torch.Tensor] | None,
|
||||
cu_seqlens: torch.Tensor,
|
||||
cu_seqlens_host: tuple[int, ...] | None,
|
||||
max_seqlen: int,
|
||||
ulysses_active: bool,
|
||||
ring_active: bool,
|
||||
) -> torch.Tensor:
|
||||
"""out = to_out(gate_sm * window_softmax) + to_out_linear(branch); gates and
|
||||
beta are row-local and computed before the core exchanges rows for heads."""
|
||||
if ring_active:
|
||||
raise NotImplementedError(
|
||||
"VDN-H3 hybrid attention does not support ring parallelism"
|
||||
)
|
||||
total = x.shape[0]
|
||||
softmax_gate = self.softmax_gate(x) if self.softmax_gate is not None else None
|
||||
beta = self.linear_attention.beta(x)
|
||||
gate_hidden, _ = self.linear_attention.output_gate.down(x)
|
||||
attention_core = (
|
||||
_hybrid_attention_core_bcg
|
||||
if attention.bcg_breakpoint
|
||||
else _minimax_h3_hybrid_attention_core_impl
|
||||
)
|
||||
softmax_out, linear_out = attention_core(
|
||||
attention,
|
||||
x,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
softmax_gate,
|
||||
beta,
|
||||
gate_hidden,
|
||||
rope_cache=rope_cache,
|
||||
cu_seqlens=cu_seqlens,
|
||||
cu_seqlens_host=cu_seqlens_host,
|
||||
max_seqlen=max_seqlen,
|
||||
ulysses_active=ulysses_active,
|
||||
)
|
||||
out, _ = attention.out_proj(softmax_out.reshape(total, -1))
|
||||
if linear_out is not None:
|
||||
linear_proj, _ = self.to_out_linear(linear_out)
|
||||
if linear_proj.shape[0] == total:
|
||||
out.add_(linear_proj)
|
||||
else:
|
||||
# single-rank path: the readout covers the video rows only
|
||||
layout = get_forward_context().attn_metadata.layout
|
||||
out[layout.video_start : layout.video_end].add_(linear_proj)
|
||||
return out
|
||||
|
||||
|
||||
def _vdn_frame_partial_sums(
|
||||
x: torch.Tensor,
|
||||
*,
|
||||
row_start: int,
|
||||
video_start: int,
|
||||
video_end: int,
|
||||
num_frames: int,
|
||||
tokens_per_frame: int,
|
||||
) -> torch.Tensor:
|
||||
# fp32 [F, hidden] sums of this rank's video rows; whole frames as one reduction
|
||||
hidden = x.shape[-1]
|
||||
sums = torch.zeros(num_frames, hidden, dtype=_FP32_DTYPE, device=x.device)
|
||||
lo = max(row_start, video_start)
|
||||
hi = min(row_start + x.shape[0], video_end)
|
||||
if lo >= hi:
|
||||
return sums
|
||||
rows = x[lo - row_start : hi - row_start]
|
||||
first_frame, offset = divmod(lo - video_start, tokens_per_frame)
|
||||
lead = (tokens_per_frame - offset) % tokens_per_frame
|
||||
lead = min(lead, rows.shape[0])
|
||||
if lead:
|
||||
sums[first_frame] += rows[:lead].sum(0, dtype=_FP32_DTYPE)
|
||||
first_frame += 1
|
||||
full = (rows.shape[0] - lead) // tokens_per_frame
|
||||
if full:
|
||||
sums[first_frame : first_frame + full] = (
|
||||
rows[lead : lead + full * tokens_per_frame]
|
||||
.view(full, tokens_per_frame, hidden)
|
||||
.sum(1, dtype=_FP32_DTYPE)
|
||||
)
|
||||
tail = lead + full * tokens_per_frame
|
||||
if tail < rows.shape[0]:
|
||||
sums[first_frame + full] += rows[tail:].sum(0, dtype=_FP32_DTYPE)
|
||||
return sums
|
||||
|
||||
|
||||
def _vdn_a2a_rows_to_heads(
|
||||
field: torch.Tensor,
|
||||
*,
|
||||
ulysses_ws: int,
|
||||
role: str,
|
||||
process_group: torch.distributed.ProcessGroup,
|
||||
) -> tuple[torch.distributed.Work, torch.Tensor]:
|
||||
# [L, H, d] row shard -> contiguous [S, H / ws, d] of this rank's heads
|
||||
rows, total_heads, head_dim = field.shape
|
||||
local_heads = total_heads // ulysses_ws
|
||||
send = _a2a_staging_buffer(
|
||||
role + "_send",
|
||||
(ulysses_ws, rows, local_heads, head_dim),
|
||||
field.dtype,
|
||||
field.device,
|
||||
)
|
||||
send.copy_(field.view(rows, ulysses_ws, local_heads, head_dim).permute(1, 0, 2, 3))
|
||||
recv = _a2a_staging_buffer(
|
||||
role + "_recv",
|
||||
(ulysses_ws * rows, local_heads, head_dim),
|
||||
field.dtype,
|
||||
field.device,
|
||||
)
|
||||
work = torch.distributed.all_to_all_single(
|
||||
recv, send, group=process_group, async_op=True
|
||||
)
|
||||
return work, recv
|
||||
|
||||
|
||||
def _vdn_a2a_heads_to_rows(
|
||||
out: torch.Tensor,
|
||||
*,
|
||||
ulysses_ws: int,
|
||||
role: str,
|
||||
process_group: torch.distributed.ProcessGroup,
|
||||
) -> tuple[torch.distributed.Work, torch.Tensor]:
|
||||
# [S, H / ws, d] -> [ws, L, H / ws, d] source-rank major; _vdn_merge_heads after wait
|
||||
seq_len, local_heads, head_dim = out.shape
|
||||
rows = seq_len // ulysses_ws
|
||||
recv = _a2a_staging_buffer(
|
||||
role + "_recv", (ulysses_ws, rows, local_heads, head_dim), out.dtype, out.device
|
||||
)
|
||||
work = torch.distributed.all_to_all_single(
|
||||
recv, out.contiguous(), group=process_group, async_op=True
|
||||
)
|
||||
return work, recv
|
||||
|
||||
|
||||
def _vdn_merge_heads(recv: torch.Tensor) -> torch.Tensor:
|
||||
# [ws, L, h, d] -> [L, ws * h, d]; rank-major heads are the global head order
|
||||
ulysses_ws, rows, local_heads, head_dim = recv.shape
|
||||
merged = usp_merge_heads(recv.view(ulysses_ws, rows, 1, local_heads, head_dim))
|
||||
return merged.reshape(rows, ulysses_ws * local_heads, head_dim)
|
||||
|
||||
|
||||
def _vdn_window_softmax(
|
||||
attention: MiniMaxH3Attention,
|
||||
meta: HybridWindowAttentionH3Metadata,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
*,
|
||||
softmax_gate: torch.Tensor | None,
|
||||
rope_cache: tuple[torch.Tensor, torch.Tensor],
|
||||
cu_seqlens: torch.Tensor,
|
||||
cu_seqlens_host: tuple[int, ...] | None,
|
||||
max_seqlen: int,
|
||||
) -> torch.Tensor:
|
||||
# the branch keeps reading the raw q/k, so norm + RoPE write copies
|
||||
cos_sin_cache, positions = rope_cache
|
||||
if attention._use_fused_qknorm_rope and not torch.compiler.is_compiling():
|
||||
q_sm = torch.empty(q.shape, dtype=q.dtype, device=q.device)
|
||||
k_sm = torch.empty(k.shape, dtype=k.dtype, device=k.device)
|
||||
fused_qknorm_rope_out_of_place(
|
||||
q,
|
||||
k,
|
||||
q_sm,
|
||||
k_sm,
|
||||
attention.q_norm.weight,
|
||||
attention.k_norm.weight,
|
||||
cos_sin_cache,
|
||||
positions,
|
||||
is_neox=True,
|
||||
eps=attention.q_norm.eps,
|
||||
head_dim=attention.head_dim,
|
||||
rope_dim=cos_sin_cache.shape[-1],
|
||||
round_norm_before_rope=True,
|
||||
)
|
||||
else:
|
||||
from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import (
|
||||
_apply_qk_norm,
|
||||
_apply_rope_qk,
|
||||
)
|
||||
|
||||
q_sm, k_sm = _apply_qk_norm(
|
||||
q.clone(), k.clone(), attention.q_norm, attention.k_norm, attention.head_dim
|
||||
)
|
||||
q_sm, k_sm = _apply_rope_qk(q_sm, k_sm, cos_sin_cache, positions)
|
||||
return attention._attention_impl.forward_varlen(
|
||||
q_sm,
|
||||
k_sm,
|
||||
v,
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
cu_seqlens_host=cu_seqlens_host,
|
||||
attn_metadata=meta,
|
||||
softmax_gate=softmax_gate,
|
||||
)
|
||||
|
||||
|
||||
def _vdn_linear_readout(
|
||||
attention: MiniMaxH3Attention,
|
||||
meta: HybridWindowAttentionH3Metadata,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
*,
|
||||
beta: torch.Tensor,
|
||||
linear_gate: torch.Tensor,
|
||||
frame_mean: torch.Tensor,
|
||||
head_range: slice | None,
|
||||
) -> torch.Tensor:
|
||||
# video rows in, [V, h * d] out; text rows seed the state
|
||||
layout = meta.layout
|
||||
video = slice(layout.video_start, layout.video_end)
|
||||
text = slice(0, layout.text_len)
|
||||
return attention.hybrid.linear_attention(
|
||||
q_raw=q[video],
|
||||
k_raw=k[video],
|
||||
v_raw=v[video],
|
||||
beta=beta[video],
|
||||
gate=linear_gate[video],
|
||||
frame_mean=frame_mean,
|
||||
layout=layout,
|
||||
text_k_raw=k[text],
|
||||
text_v_raw=v[text],
|
||||
text_beta=beta[text],
|
||||
heads=head_range,
|
||||
)
|
||||
|
||||
|
||||
def _minimax_h3_hybrid_attention_core_impl(
|
||||
attention: MiniMaxH3Attention,
|
||||
x: torch.Tensor,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
softmax_gate: torch.Tensor | None,
|
||||
beta: torch.Tensor,
|
||||
gate_hidden: torch.Tensor,
|
||||
*,
|
||||
rope_cache: tuple[torch.Tensor, torch.Tensor] | None,
|
||||
cu_seqlens: torch.Tensor,
|
||||
cu_seqlens_host: tuple[int, ...] | None,
|
||||
max_seqlen: int,
|
||||
ulysses_active: bool,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
meta = get_forward_context().attn_metadata
|
||||
if not isinstance(meta, HybridWindowAttentionH3Metadata):
|
||||
raise RuntimeError(
|
||||
"VDN-H3 hybrid attention needs HybridWindowAttentionH3Metadata in the "
|
||||
"forward context; the MiniMax-H3 denoising stage installs it per request "
|
||||
f"(got {type(meta).__name__})."
|
||||
)
|
||||
layout = meta.layout
|
||||
softmax = functools.partial(
|
||||
_vdn_window_softmax,
|
||||
attention,
|
||||
meta,
|
||||
cu_seqlens=cu_seqlens,
|
||||
cu_seqlens_host=cu_seqlens_host,
|
||||
max_seqlen=max_seqlen,
|
||||
)
|
||||
if not ulysses_active:
|
||||
if rope_cache is None:
|
||||
raise RuntimeError("VDN-H3 hybrid attention requires the RoPE cache")
|
||||
softmax_out = softmax(q, k, v, softmax_gate=softmax_gate, rope_cache=rope_cache)
|
||||
if meta.full_cover:
|
||||
return softmax_out, None
|
||||
frame_mean = (
|
||||
x[layout.video_start : layout.video_end]
|
||||
.view(layout.num_frames, layout.tokens_per_frame, x.shape[-1])
|
||||
.mean(dim=1, dtype=_FP32_DTYPE)
|
||||
)
|
||||
readout = _vdn_linear_readout(
|
||||
attention,
|
||||
meta,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
beta=beta,
|
||||
linear_gate=attention.hybrid.linear_attention.output_gate.up_gate(
|
||||
gate_hidden
|
||||
),
|
||||
frame_mean=frame_mean,
|
||||
head_range=None,
|
||||
)
|
||||
return softmax_out, readout
|
||||
|
||||
return _vdn_ulysses_hybrid_core(
|
||||
attention,
|
||||
meta,
|
||||
x,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
softmax_gate=softmax_gate,
|
||||
beta=beta,
|
||||
gate_hidden=gate_hidden,
|
||||
softmax=softmax,
|
||||
)
|
||||
|
||||
|
||||
def _vdn_ulysses_hybrid_core(
|
||||
attention: MiniMaxH3Attention,
|
||||
meta: HybridWindowAttentionH3Metadata,
|
||||
x: torch.Tensor,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
*,
|
||||
softmax_gate: torch.Tensor | None,
|
||||
beta: torch.Tensor,
|
||||
gate_hidden: torch.Tensor,
|
||||
softmax: Callable[..., torch.Tensor],
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_group
|
||||
|
||||
layout = meta.layout
|
||||
if meta.rope_cache_full is None:
|
||||
raise RuntimeError("VDN-H3 under Ulysses needs the full-sequence RoPE cache")
|
||||
sp_group = get_sp_group()
|
||||
process_group = sp_group.ulysses_group
|
||||
ulysses_ws, ulysses_rank = get_ulysses_ctx()
|
||||
local_rows, head_dim = x.shape[0], q.shape[2]
|
||||
local_heads = q.shape[1] // ulysses_ws
|
||||
seq_len = local_rows * ulysses_ws
|
||||
|
||||
# the q/k/v exchange is in flight while the frame sums and the gate hidden go out
|
||||
inflight = [
|
||||
_vdn_a2a_rows_to_heads(
|
||||
field,
|
||||
ulysses_ws=ulysses_ws,
|
||||
role=f"vdn_{name}",
|
||||
process_group=process_group,
|
||||
)
|
||||
for name, field in (("q", q), ("k", k), ("v", v))
|
||||
]
|
||||
frame_sums = _vdn_frame_partial_sums(
|
||||
x,
|
||||
row_start=ulysses_rank * local_rows,
|
||||
video_start=layout.video_start,
|
||||
video_end=layout.video_end,
|
||||
num_frames=layout.num_frames,
|
||||
tokens_per_frame=layout.tokens_per_frame,
|
||||
)
|
||||
frame_work = torch.distributed.all_reduce(
|
||||
frame_sums, group=sp_group.device_group, async_op=True
|
||||
)
|
||||
# the per-head scalars (beta, softmax gate) ride one more async field
|
||||
scalars = [beta] if softmax_gate is None else [beta, softmax_gate]
|
||||
inflight.append(
|
||||
_vdn_a2a_rows_to_heads(
|
||||
torch.stack(scalars, dim=-1),
|
||||
ulysses_ws=ulysses_ws,
|
||||
role="vdn_scalars",
|
||||
process_group=process_group,
|
||||
)
|
||||
)
|
||||
head_range = slice(ulysses_rank * local_heads, (ulysses_rank + 1) * local_heads)
|
||||
linear_gate = attention.hybrid.linear_attention.output_gate.up_gate(
|
||||
sp_group.all_gather(gate_hidden.contiguous(), dim=0), heads=head_range
|
||||
)
|
||||
|
||||
for work, _ in inflight:
|
||||
work.wait()
|
||||
q, k, v, scalars = (recv for _, recv in inflight)
|
||||
beta = scalars[..., 0]
|
||||
if softmax_gate is not None:
|
||||
softmax_gate = scalars[..., 1]
|
||||
softmax = functools.partial(
|
||||
softmax, q, k, v, softmax_gate=softmax_gate, rope_cache=meta.rope_cache_full
|
||||
)
|
||||
if meta.full_cover:
|
||||
softmax_out = softmax()
|
||||
frame_work.wait()
|
||||
return _vdn_return_to_rows(
|
||||
softmax_out, None, ulysses_ws=ulysses_ws, process_group=process_group
|
||||
)
|
||||
|
||||
def linear_branch() -> torch.Tensor:
|
||||
frame_work.wait()
|
||||
readout = _vdn_linear_readout(
|
||||
attention,
|
||||
meta,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
beta=beta,
|
||||
linear_gate=linear_gate,
|
||||
frame_mean=frame_sums / layout.tokens_per_frame,
|
||||
head_range=head_range,
|
||||
)
|
||||
# rows go back to their owners: pad the non-video rows with zeros
|
||||
linear_out = q.new_zeros(seq_len, local_heads, head_dim)
|
||||
linear_out[layout.video_start : layout.video_end] = readout.view(
|
||||
-1, local_heads, head_dim
|
||||
)
|
||||
return linear_out
|
||||
|
||||
a2a_back = functools.partial(
|
||||
_vdn_a2a_heads_to_rows, ulysses_ws=ulysses_ws, process_group=process_group
|
||||
)
|
||||
# the linear readout and its a2a run on the side stream while FA4 (issued
|
||||
# first) holds the current one; both read the exchanged q/k/v, frame sums and
|
||||
# gate, so the side stream waits for them and the current stream joins before the merge
|
||||
main_stream = torch.cuda.current_stream(q.device)
|
||||
side_stream = _linear_branch_stream(q.device)
|
||||
side_stream.wait_stream(main_stream)
|
||||
softmax_out = softmax()
|
||||
softmax_work, softmax_recv = a2a_back(softmax_out, role="vdn_out0")
|
||||
with torch.cuda.stream(side_stream):
|
||||
linear_work, linear_recv = a2a_back(linear_branch(), role="vdn_out1")
|
||||
main_stream.wait_stream(side_stream)
|
||||
softmax_work.wait()
|
||||
linear_work.wait()
|
||||
merged_softmax = _vdn_merge_heads(softmax_recv)
|
||||
merged_linear = _vdn_merge_heads(linear_recv)
|
||||
return merged_softmax, merged_linear.reshape(merged_linear.shape[0], -1)
|
||||
|
||||
|
||||
def _vdn_return_to_rows(
|
||||
softmax_out: torch.Tensor,
|
||||
linear_out: torch.Tensor | None,
|
||||
*,
|
||||
ulysses_ws: int,
|
||||
process_group: torch.distributed.ProcessGroup,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
# [S, h, d] per branch -> ([L, H, d], [L, H * d] or None), both trips in flight together
|
||||
branch_outputs = [out for out in (softmax_out, linear_out) if out is not None]
|
||||
inflight = [
|
||||
_vdn_a2a_heads_to_rows(
|
||||
out, ulysses_ws=ulysses_ws, role=f"vdn_out{i}", process_group=process_group
|
||||
)
|
||||
for i, out in enumerate(branch_outputs)
|
||||
]
|
||||
merged = []
|
||||
for work, recv in inflight:
|
||||
work.wait()
|
||||
merged.append(_vdn_merge_heads(recv))
|
||||
linear_rows = (
|
||||
merged[1].reshape(merged[1].shape[0], -1) if linear_out is not None else None
|
||||
)
|
||||
return merged[0], linear_rows
|
||||
|
||||
|
||||
_hybrid_attention_core_bcg = eager_on_graph(True)(
|
||||
_minimax_h3_hybrid_attention_core_impl
|
||||
)
|
||||
|
||||
|
||||
def prepare_hybrid_attention_metadata(
|
||||
*,
|
||||
model,
|
||||
packed: Mapping[str, torch.Tensor],
|
||||
latent_shape: tuple[int, int, int],
|
||||
server_args,
|
||||
device: torch.device,
|
||||
) -> Callable[[int], Any] | None:
|
||||
"""Request-static metadata (window plan, packed layout, full-sequence RoPE
|
||||
cache under Ulysses) for every step and block; None for other backends."""
|
||||
model._resolve_attention_backend_once()
|
||||
if (
|
||||
model._resolved_attention_backend
|
||||
is not AttentionBackendEnum.HYBRID_WINDOW_ATTN_H3
|
||||
):
|
||||
return None
|
||||
hybrid = model.arch.hybrid_attention
|
||||
if hybrid is None:
|
||||
raise ValueError(
|
||||
"--attention-backend hybrid_window_attn_h3 needs a VDN-H3 checkpoint "
|
||||
"(transformer/config.json with hybrid_attention); this checkpoint has "
|
||||
"no linear branch. Use --attention-backend fa for MiniMax-H3."
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import (
|
||||
_rope_cos_sin_cache,
|
||||
)
|
||||
|
||||
config = server_args.attention_backend_config or {}
|
||||
max_gather_rows = int(config.get("vdn_max_gather_rows", 200_000))
|
||||
latent_t, latent_h, latent_w = latent_shape
|
||||
layout = vdn_h3_layout_from_packed(
|
||||
packed, latent_t=latent_t, latent_h=latent_h, latent_w=latent_w
|
||||
)
|
||||
rope_cache_full = None
|
||||
ulysses_ws, _ = get_ulysses_ctx()
|
||||
ring_ws, _ = get_ring_ctx()
|
||||
if ring_ws > 1:
|
||||
raise ValueError("VDN-H3 does not support ring parallelism")
|
||||
if ulysses_ws > 1:
|
||||
# QK-norm + RoPE run on the head shard after the all-to-all: full-sequence cache
|
||||
with torch.inference_mode():
|
||||
img_position_ids = (
|
||||
packed["img_position_ids"][None].to(torch.float32).to(device)
|
||||
)
|
||||
rope_freqs = model.rope(img_position_ids)
|
||||
rope_cache_full = (
|
||||
_rope_cos_sin_cache(rope_freqs, dtype=torch.bfloat16),
|
||||
torch.arange(layout.seq_len, device=device, dtype=torch.long),
|
||||
)
|
||||
metadata = HybridWindowAttentionH3MetadataBuilder().build(
|
||||
layout=layout,
|
||||
hybrid=hybrid,
|
||||
device=device,
|
||||
rope_cache_full=rope_cache_full,
|
||||
max_gather_rows=max_gather_rows,
|
||||
)
|
||||
logger.info(
|
||||
"VDN-H3 hybrid attention: frames=%d tokens/frame=%d text=%d "
|
||||
"used=%d/%d chunk=%d radius=%d anchors=%s full_cover=%s",
|
||||
layout.num_frames,
|
||||
layout.tokens_per_frame,
|
||||
layout.text_len,
|
||||
layout.used,
|
||||
layout.seq_len,
|
||||
hybrid.chunk,
|
||||
hybrid.radius,
|
||||
hybrid.anchor_frames,
|
||||
metadata.full_cover,
|
||||
)
|
||||
|
||||
def build(step_index: int):
|
||||
return metadata
|
||||
|
||||
return build
|
||||
|
||||
|
||||
__all__ = ["MiniMaxH3VDNHybridAttention", "prepare_hybrid_attention_metadata"]
|
||||
@@ -0,0 +1,16 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from sglang.multimodal_gen.runtime.pipelines.minimax_h3_pipeline import (
|
||||
MiniMaxH3Pipeline,
|
||||
)
|
||||
|
||||
|
||||
class VDNH3Pipeline(MiniMaxH3Pipeline):
|
||||
"""VDN-H3 on the MiniMax-H3 pipeline: the model overlay materializes a
|
||||
base-H3 layout (LoRAs prefused, linear branch attached), so only the DiT
|
||||
blocks' attention and its backend differ."""
|
||||
|
||||
pipeline_name = "VDNH3Pipeline"
|
||||
default_model_subfolder = None
|
||||
|
||||
|
||||
EntryClass = [VDNH3Pipeline]
|
||||
+12
@@ -749,6 +749,18 @@ class MiniMaxH3DenoisingStage(DenoisingStage):
|
||||
server_args=server_args,
|
||||
device=device,
|
||||
)
|
||||
if build_vsa_h3_step_metadata is None:
|
||||
from sglang.multimodal_gen.runtime.models.dits.minimax_h3_vdn_attention import (
|
||||
prepare_hybrid_attention_metadata,
|
||||
)
|
||||
|
||||
build_vsa_h3_step_metadata = prepare_hybrid_attention_metadata(
|
||||
model=model,
|
||||
packed=packed,
|
||||
latent_shape=(ctx.latent_t, ctx.latent_h, ctx.latent_w),
|
||||
server_args=server_args,
|
||||
device=device,
|
||||
)
|
||||
positive = MiniMaxH3DenoiseBranch(
|
||||
packed=packed,
|
||||
text_embeddings=emb["hidden_states"],
|
||||
|
||||
@@ -290,6 +290,52 @@ class _VideoSparseAttentionH3BackendResolver(_CudaAttentionBackendResolver):
|
||||
) from e
|
||||
|
||||
|
||||
class _HybridWindowAttentionH3BackendResolver(_CudaAttentionBackendResolver):
|
||||
backend = AttentionBackendEnum.HYBRID_WINDOW_ATTN_H3
|
||||
|
||||
# the window rides FlashAttention varlen: FA4 on SM100 / SM103 / SM120, FA3 on
|
||||
# SM90; SM80 / SM86 / SM89 run FA3's Sm80 mainloop (FA2-class throughput)
|
||||
supported_capabilities = {
|
||||
(8, 0),
|
||||
(8, 6),
|
||||
(8, 9),
|
||||
(9, 0),
|
||||
(10, 0),
|
||||
(10, 3),
|
||||
(12, 0),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def resolve(cls, platform) -> str:
|
||||
capability = platform.get_device_capability()
|
||||
capability_tuple = (
|
||||
(capability.major, capability.minor) if capability is not None else None
|
||||
)
|
||||
if capability_tuple not in cls.supported_capabilities:
|
||||
found = capability.as_version_str() if capability else "unknown"
|
||||
raise ValueError(
|
||||
"hybrid_window_attn_h3 (VDN-H3) needs compute capability 8.0 / "
|
||||
"8.6 / 8.9 (Ampere, Ada), 9.0 (Hopper), 10.0 (B200 / GB200), "
|
||||
"10.3 (B300 / GB300) or 12.0 (RTX PRO 6000 Blackwell); this "
|
||||
f"device reports {found}."
|
||||
)
|
||||
if not platform._prepare_flash_attention_for_blackwell():
|
||||
raise RuntimeError(
|
||||
"hybrid_window_attn_h3 requires FlashAttention for its dense legs"
|
||||
)
|
||||
try:
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.hybrid_window_attn_h3 import ( # noqa: F401
|
||||
HybridWindowAttentionH3Backend,
|
||||
)
|
||||
|
||||
return "sglang.multimodal_gen.runtime.layers.attention.backends.hybrid_window_attn_h3.HybridWindowAttentionH3Backend"
|
||||
except Exception as e:
|
||||
logger.error("Failed to import hybrid_window_attn_h3 backend: %s", str(e))
|
||||
raise ImportError(
|
||||
"hybrid_window_attn_h3 needs FlashAttention and Triton."
|
||||
) from e
|
||||
|
||||
|
||||
class _CubeSparseAttentionBackendResolver(_CudaAttentionBackendResolver):
|
||||
backend = AttentionBackendEnum.CUBE_SPARSE_ATTN
|
||||
|
||||
@@ -485,6 +531,7 @@ _CUDA_ATTENTION_BACKEND_RESOLVERS = {
|
||||
_SpargeAttentionBackendResolver,
|
||||
_VideoSparseAttentionBackendResolver,
|
||||
_VideoSparseAttentionH3BackendResolver,
|
||||
_HybridWindowAttentionH3BackendResolver,
|
||||
_CubeSparseAttentionBackendResolver,
|
||||
_SparseVideoGen2AttentionBackendResolver,
|
||||
_SolAttnBackendResolver,
|
||||
@@ -652,7 +699,9 @@ class CudaPlatformBase(Platform):
|
||||
|
||||
@classmethod
|
||||
def _prepare_flash_attention_for_blackwell(cls) -> bool:
|
||||
if not cls.is_blackwell():
|
||||
# the FA4 CuTe package ships an sm120 forward kernel; the default FA backend
|
||||
# still resolves to SDPA on SM120 before reaching this
|
||||
if not (cls.is_blackwell() or cls.is_sm120()):
|
||||
return True
|
||||
|
||||
try:
|
||||
|
||||
@@ -36,6 +36,7 @@ class AttentionBackendEnum(enum.Enum):
|
||||
SPARGE_ATTN = enum.auto()
|
||||
VIDEO_SPARSE_ATTN = enum.auto()
|
||||
VIDEO_SPARSE_ATTN_H3 = enum.auto()
|
||||
HYBRID_WINDOW_ATTN_H3 = enum.auto()
|
||||
SPARSE_VIDEO_GEN_2_ATTN = enum.auto()
|
||||
VMOBA_ATTN = enum.auto()
|
||||
AITER = enum.auto()
|
||||
@@ -59,6 +60,7 @@ class AttentionBackendEnum(enum.Enum):
|
||||
AttentionBackendEnum.SLIDING_TILE_ATTN,
|
||||
AttentionBackendEnum.VIDEO_SPARSE_ATTN,
|
||||
AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3,
|
||||
AttentionBackendEnum.HYBRID_WINDOW_ATTN_H3,
|
||||
AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN,
|
||||
AttentionBackendEnum.VMOBA_ATTN,
|
||||
AttentionBackendEnum.SLA_ATTN,
|
||||
|
||||
@@ -46,6 +46,10 @@ BUILTIN_MODEL_OVERLAY_REGISTRY: dict[str, dict[str, Any]] = {
|
||||
"overlay_repo_id": "kevin-mi/FastH3-4step-Preview-overlay",
|
||||
"overlay_revision": "f769cb8001dae335089de7b250364335bc7cb183",
|
||||
},
|
||||
"OpenVDN/vdn-minimax-h3": {
|
||||
"overlay_repo_id": "kevin-mi/VDN-H3-overlay",
|
||||
"overlay_revision": "0ad315a05b914c4003af4d26152d288c2506a609",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -688,6 +688,47 @@ MINIMAX_H3_FOUR_GPU_H100_CASES = [
|
||||
run_models_api_check=False,
|
||||
run_t2v_input_reference_check=False,
|
||||
),
|
||||
DiffusionTestCase(
|
||||
"vdn_h3_t2va_4gpu_h100",
|
||||
DiffusionServerArgs(
|
||||
model_path="OpenVDN/vdn-minimax-h3",
|
||||
modality="video",
|
||||
num_gpus=4,
|
||||
extras=[
|
||||
"--attention-backend",
|
||||
"hybrid_window_attn_h3",
|
||||
"--enable-torch-compile",
|
||||
"false",
|
||||
],
|
||||
),
|
||||
DiffusionSamplingParams(
|
||||
prompt=(
|
||||
"A curious raccoon peers through a vibrant field of yellow "
|
||||
"sunflowers, its eyes wide with interest."
|
||||
),
|
||||
output_size="1344x768",
|
||||
seconds=5,
|
||||
output_format="mp4",
|
||||
expect_audio_output=True,
|
||||
num_outputs_per_prompt=1,
|
||||
extras={
|
||||
"task": "t2va",
|
||||
"conditions": [],
|
||||
"target": {
|
||||
"short_edge": 768,
|
||||
"aspect_ratio": "16:9",
|
||||
"duration_seconds": 5.0,
|
||||
},
|
||||
"num_inference_steps": 9,
|
||||
"seed": 42,
|
||||
},
|
||||
),
|
||||
run_perf_check=False,
|
||||
run_consistency_check=False,
|
||||
run_component_accuracy_check=False,
|
||||
run_models_api_check=False,
|
||||
run_t2v_input_reference_check=False,
|
||||
),
|
||||
DiffusionTestCase(
|
||||
"fasth3_t2va_vsa_4gpu_h100",
|
||||
DiffusionServerArgs(
|
||||
@@ -1428,6 +1469,7 @@ STANDALONE_FILES = {
|
||||
"../single_test_file/test_dp_serving_2_gpu.py",
|
||||
"../single_test_file/test_pynccl_a2a_capture_2_gpu.py",
|
||||
"../single_test_file/test_usp_replicated_parity_2_gpu.py",
|
||||
"../single_test_file/test_vdn_ulysses_exchange_2_gpu.py",
|
||||
],
|
||||
}
|
||||
|
||||
@@ -1473,6 +1515,8 @@ STANDALONE_FILE_EST_TIMES = {
|
||||
"../single_test_file/test_pynccl_a2a_capture_2_gpu.py": 180.0,
|
||||
# two SDPA parity checks on 128+6 rows
|
||||
"../single_test_file/test_usp_replicated_parity_2_gpu.py": 180.0,
|
||||
# no model load; two small all-to-alls
|
||||
"../single_test_file/test_vdn_ulysses_exchange_2_gpu.py": 60.0,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""VDN-H3 Ulysses exchange: the field-major row->head all-to-all and its
|
||||
inverse plus the head merge, checked against plain slicing on 2 GPUs."""
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from sglang.test.test_utils import run_distributed_test
|
||||
|
||||
|
||||
def _check(rank: int) -> None:
|
||||
from sglang.multimodal_gen.runtime.models.dits.minimax_h3_vdn_attention import (
|
||||
_vdn_a2a_heads_to_rows,
|
||||
_vdn_a2a_rows_to_heads,
|
||||
_vdn_merge_heads,
|
||||
)
|
||||
|
||||
world = dist.get_world_size()
|
||||
device = torch.device("cuda", rank)
|
||||
heads, head_dim, local_rows = 6, 32, 24
|
||||
local_heads = heads // world
|
||||
seq = local_rows * world
|
||||
# every rank builds the same global tensors, so a shard is checked by slicing
|
||||
g = torch.Generator(device="cpu").manual_seed(3)
|
||||
qkv = torch.randn(seq, 3 * heads * head_dim, generator=g).to(device, torch.bfloat16)
|
||||
q = qkv.view(seq, 3, heads, head_dim)[:, 0] # strided, as the split qkv is
|
||||
scalars = torch.randn(seq, heads, 2, generator=g).to(device, torch.bfloat16)
|
||||
rows = slice(rank * local_rows, (rank + 1) * local_rows)
|
||||
mine = slice(rank * local_heads, (rank + 1) * local_heads)
|
||||
with torch.inference_mode():
|
||||
for name, field in (("q", q), ("scalars", scalars)):
|
||||
work, recv = _vdn_a2a_rows_to_heads(
|
||||
field[rows], ulysses_ws=world, role=name, process_group=dist.group.WORLD
|
||||
)
|
||||
work.wait()
|
||||
assert recv.is_contiguous()
|
||||
assert torch.equal(recv, field[:, mine])
|
||||
# inverse: this rank's heads for every row -> row shard, every head
|
||||
work, back = _vdn_a2a_heads_to_rows(
|
||||
q[:, mine].contiguous() * 2,
|
||||
ulysses_ws=world,
|
||||
role="out",
|
||||
process_group=dist.group.WORLD,
|
||||
)
|
||||
work.wait()
|
||||
assert torch.equal(_vdn_merge_heads(back), q[rows] * 2)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
|
||||
def test_exchange_round_trip_two_ranks() -> None:
|
||||
run_distributed_test(_check, world_size=2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_exchange_round_trip_two_ranks()
|
||||
@@ -107,6 +107,51 @@ class TestCudaAttentionBackendSelection(unittest.TestCase):
|
||||
|
||||
fake_flash_attn.set_fa_ver.assert_called_once_with(4)
|
||||
|
||||
def test_hybrid_window_h3_on_sm120_uses_fa4(self):
|
||||
FakeCudaPlatform.is_sm120_device = True
|
||||
FakeCudaPlatform.device_capability = DeviceCapability(12, 0)
|
||||
fa_module = "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn"
|
||||
fake_flash_attn = ModuleType(fa_module)
|
||||
fake_flash_attn.set_fa_ver = Mock()
|
||||
hybrid_module = "sglang.multimodal_gen.runtime.layers.attention.backends.hybrid_window_attn_h3"
|
||||
fake_hybrid = ModuleType(hybrid_module)
|
||||
fake_hybrid.HybridWindowAttentionH3Backend = object
|
||||
|
||||
with patch.dict(
|
||||
"sys.modules", {fa_module: fake_flash_attn, hybrid_module: fake_hybrid}
|
||||
):
|
||||
self.assertEqual(
|
||||
self.resolve(AttentionBackendEnum.HYBRID_WINDOW_ATTN_H3),
|
||||
f"{hybrid_module}.HybridWindowAttentionH3Backend",
|
||||
)
|
||||
|
||||
fake_flash_attn.set_fa_ver.assert_called_once_with(4)
|
||||
|
||||
def test_hybrid_window_h3_on_ampere_keeps_fa3(self):
|
||||
FakeCudaPlatform.device_capability = DeviceCapability(8, 0)
|
||||
fa_module = "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn"
|
||||
fake_flash_attn = ModuleType(fa_module)
|
||||
fake_flash_attn.set_fa_ver = Mock()
|
||||
hybrid_module = "sglang.multimodal_gen.runtime.layers.attention.backends.hybrid_window_attn_h3"
|
||||
fake_hybrid = ModuleType(hybrid_module)
|
||||
fake_hybrid.HybridWindowAttentionH3Backend = object
|
||||
|
||||
with patch.dict(
|
||||
"sys.modules", {fa_module: fake_flash_attn, hybrid_module: fake_hybrid}
|
||||
):
|
||||
self.assertEqual(
|
||||
self.resolve(AttentionBackendEnum.HYBRID_WINDOW_ATTN_H3),
|
||||
f"{hybrid_module}.HybridWindowAttentionH3Backend",
|
||||
)
|
||||
|
||||
fake_flash_attn.set_fa_ver.assert_not_called()
|
||||
|
||||
def test_hybrid_window_h3_rejects_unsupported_capability(self):
|
||||
FakeCudaPlatform.device_capability = DeviceCapability(7, 5)
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "12.0"):
|
||||
self.resolve(AttentionBackendEnum.HYBRID_WINDOW_ATTN_H3)
|
||||
|
||||
def test_default_backend_uses_torch_sdpa_on_sm120(self):
|
||||
FakeCudaPlatform.is_sm120_device = True
|
||||
|
||||
|
||||
@@ -0,0 +1,264 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""hybrid_window_attn_h3 must reproduce a masked dense softmax with exactly the VDN
|
||||
mask on a ragged packed layout, and dense attention once the window covers the clip."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.configs.models.dits.minimax_h3_vdn import (
|
||||
VDNHybridAttentionArchConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.hybrid_window_attn_h3 import (
|
||||
HybridWindowAttentionH3Impl,
|
||||
HybridWindowAttentionH3MetadataBuilder,
|
||||
window_mask_frames,
|
||||
window_mask_reference,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.dits.minimax_h3_vdn import VDNH3Layout
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
|
||||
# torch.cuda.is_available() is also True under ROCm, but the backend lives on the
|
||||
# CUDA platform only (RocmPlatform rejects hybrid_window_attn_h3 outright and has
|
||||
# no _prepare_flash_attention_for_blackwell), so gate on the platform itself.
|
||||
requires_cuda = pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="hybrid_window_attn_h3 kernels need NVIDIA CUDA",
|
||||
)
|
||||
|
||||
# ragged on purpose: 70 and 100 are not tile multiples, 12 frames is not a chunk multiple
|
||||
TEXT_LEN = 70
|
||||
AUDIO_ROWS = 100
|
||||
NUM_FRAMES = 12
|
||||
FRAME_H, FRAME_W = 6, 8
|
||||
TOKENS_PER_FRAME = FRAME_H * FRAME_W
|
||||
HEADS = 4
|
||||
HEAD_DIM = 128
|
||||
|
||||
|
||||
def _layout() -> VDNH3Layout:
|
||||
video_start = TEXT_LEN + AUDIO_ROWS
|
||||
used = video_start + NUM_FRAMES * TOKENS_PER_FRAME
|
||||
seq_len = (used + 63) // 64 * 64
|
||||
return VDNH3Layout(
|
||||
seq_len=seq_len,
|
||||
used=used,
|
||||
text_len=TEXT_LEN,
|
||||
video_start=video_start,
|
||||
num_frames=NUM_FRAMES,
|
||||
tokens_per_frame=TOKENS_PER_FRAME,
|
||||
frame_height=FRAME_H,
|
||||
frame_width=FRAME_W,
|
||||
)
|
||||
|
||||
|
||||
def _hybrid(**overrides) -> VDNHybridAttentionArchConfig:
|
||||
kwargs = dict(chunk=5, radius=1, anchor_frames="both")
|
||||
kwargs.update(overrides)
|
||||
return VDNHybridAttentionArchConfig(**kwargs)
|
||||
|
||||
|
||||
def _qkv(device, seed: int = 7):
|
||||
layout = _layout()
|
||||
generator = torch.Generator(device="cpu").manual_seed(seed)
|
||||
tensors = [
|
||||
torch.randn(
|
||||
(layout.seq_len, HEADS, HEAD_DIM), generator=generator, dtype=torch.float32
|
||||
).to(device=device, dtype=torch.bfloat16)
|
||||
for _ in range(3)
|
||||
]
|
||||
return layout, tensors
|
||||
|
||||
|
||||
def _masked_reference(q, k, v, mask: torch.Tensor, used: int) -> torch.Tensor:
|
||||
qf = q[:used].float().permute(1, 0, 2)
|
||||
kf = k[:used].float().permute(1, 0, 2)
|
||||
vf = v[:used].float().permute(1, 0, 2)
|
||||
scores = qf @ kf.transpose(-2, -1) / math.sqrt(HEAD_DIM)
|
||||
scores = scores.masked_fill(~mask[None], float("-inf"))
|
||||
return (torch.softmax(scores, dim=-1) @ vf).permute(1, 0, 2)
|
||||
|
||||
|
||||
def _prepare_flash_attention() -> None:
|
||||
# the platform resolver runs this before the first forward; a direct impl must too
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
|
||||
current_platform._prepare_flash_attention_for_blackwell()
|
||||
|
||||
|
||||
def _impl() -> HybridWindowAttentionH3Impl:
|
||||
_prepare_flash_attention()
|
||||
impl = HybridWindowAttentionH3Impl(
|
||||
num_heads=HEADS,
|
||||
head_size=HEAD_DIM,
|
||||
causal=False,
|
||||
softmax_scale=HEAD_DIM**-0.5,
|
||||
num_kv_heads=HEADS,
|
||||
prefix="blocks.3.attn",
|
||||
)
|
||||
assert impl.layer_idx == 3
|
||||
return impl
|
||||
|
||||
|
||||
def _run(impl, meta, layout, q, k, v, gate=None):
|
||||
cu = torch.tensor(
|
||||
[0, layout.used, layout.seq_len], dtype=torch.int32, device=q.device
|
||||
)
|
||||
return impl.forward_varlen(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seqlens=cu,
|
||||
max_seqlen=layout.used,
|
||||
cu_seqlens_host=(0, layout.used, layout.seq_len),
|
||||
attn_metadata=meta,
|
||||
softmax_gate=gate,
|
||||
)
|
||||
|
||||
|
||||
def test_window_bounds_and_anchor_frames() -> None:
|
||||
hybrid = _hybrid()
|
||||
bounds, dense_rows, dense_cols = window_mask_frames(hybrid, NUM_FRAMES)
|
||||
# chunk 5, radius 1: frame 7 (chunk 1) sees chunks 0..2 = frames 0..14 -> clamped 11
|
||||
assert bounds[7] == (0, 11)
|
||||
assert bounds[0] == (0, 9)
|
||||
assert bounds[11] == (5, 11)
|
||||
assert dense_rows == {0, NUM_FRAMES - 1} and dense_cols == {0, NUM_FRAMES - 1}
|
||||
assert not hybrid.full_cover(NUM_FRAMES)
|
||||
assert _hybrid(radius=NUM_FRAMES).full_cover(NUM_FRAMES)
|
||||
# 102 = 20 * 5 + 2: the last chunk is short but still a whole chunk
|
||||
raw = hybrid.window_bounds(102)
|
||||
assert raw[100] == raw[101] == (95, 109)
|
||||
assert raw[99] == (90, 104)
|
||||
|
||||
|
||||
def test_mask_reference_partition_is_exact() -> None:
|
||||
"""Every (video q, video k) pair is in the softmax window or in the
|
||||
linear branch's complement exactly once; anchors are absent from the
|
||||
branch."""
|
||||
layout = _layout()
|
||||
hybrid = _hybrid()
|
||||
mask = window_mask_reference(hybrid, layout, torch.device("cpu"))
|
||||
# globals dense both ways
|
||||
assert mask[: layout.video_start].all() and mask[:, : layout.video_start].all()
|
||||
assert mask[layout.video_end : layout.used].all()
|
||||
bounds = hybrid.window_bounds(NUM_FRAMES)
|
||||
vs, tpf = layout.video_start, TOKENS_PER_FRAME
|
||||
for qf in range(NUM_FRAMES):
|
||||
row = mask[vs + qf * tpf, vs : layout.video_end].view(NUM_FRAMES, tpf)
|
||||
frame_kept = row.all(dim=1)
|
||||
assert (frame_kept == row.any(dim=1)).all() # whole frames
|
||||
if qf in (0, NUM_FRAMES - 1):
|
||||
assert frame_kept.all()
|
||||
continue
|
||||
lo, hi = max(bounds[qf][0], 0), min(bounds[qf][1], NUM_FRAMES - 1)
|
||||
expected_softmax = {f for f in range(lo, hi + 1)} | {0, NUM_FRAMES - 1}
|
||||
# the branch covers the inner frames 1..F-2 outside the window
|
||||
expected_linear = {f for f in range(1, NUM_FRAMES - 1) if f < lo or f > hi}
|
||||
kept = {f for f in range(NUM_FRAMES) if frame_kept[f]}
|
||||
assert kept == expected_softmax
|
||||
assert kept.isdisjoint(expected_linear)
|
||||
assert kept | expected_linear == set(range(NUM_FRAMES))
|
||||
|
||||
|
||||
@requires_cuda
|
||||
def test_window_matches_masked_dense() -> None:
|
||||
device = torch.device("cuda")
|
||||
layout, (q, k, v) = _qkv(device)
|
||||
hybrid = _hybrid()
|
||||
meta = HybridWindowAttentionH3MetadataBuilder().build(
|
||||
layout=layout, hybrid=hybrid, device=device
|
||||
)
|
||||
assert not meta.full_cover
|
||||
out = _run(_impl(), meta, layout, q, k, v)
|
||||
mask = window_mask_reference(hybrid, layout, device)
|
||||
reference = _masked_reference(q, k, v, mask, layout.used)
|
||||
diff = (out[: layout.used].float() - reference).abs().max().item()
|
||||
assert diff < 2e-2, f"window vs masked dense max diff {diff}"
|
||||
assert torch.all(out[layout.used :] == 0)
|
||||
|
||||
|
||||
@requires_cuda
|
||||
def test_full_cover_matches_dense() -> None:
|
||||
device = torch.device("cuda")
|
||||
layout, (q, k, v) = _qkv(device, seed=11)
|
||||
hybrid = _hybrid(radius=NUM_FRAMES)
|
||||
meta = HybridWindowAttentionH3MetadataBuilder().build(
|
||||
layout=layout, hybrid=hybrid, device=device
|
||||
)
|
||||
assert meta.full_cover
|
||||
out = _run(_impl(), meta, layout, q, k, v)
|
||||
full = torch.ones(layout.used, layout.used, dtype=torch.bool, device=device)
|
||||
reference = _masked_reference(q, k, v, full, layout.used)
|
||||
diff = (out[: layout.used].float() - reference).abs().max().item()
|
||||
assert diff < 2e-2, f"full cover vs dense max diff {diff}"
|
||||
|
||||
|
||||
@requires_cuda
|
||||
def test_decomposed_passes_are_arithmetic_neutral() -> None:
|
||||
"""Bounding the gathered K/V rows per pass splits the window into several
|
||||
varlen calls without changing any query's kept set."""
|
||||
device = torch.device("cuda")
|
||||
layout, (q, k, v) = _qkv(device, seed=9)
|
||||
one = HybridWindowAttentionH3MetadataBuilder().build(
|
||||
layout=layout, hybrid=_hybrid(), device=device
|
||||
)
|
||||
many = HybridWindowAttentionH3MetadataBuilder().build(
|
||||
layout=layout, hybrid=_hybrid(), device=device, max_gather_rows=1
|
||||
)
|
||||
assert len(one.decomposed.passes) == 1 and len(many.decomposed.passes) > 1
|
||||
impl = _impl()
|
||||
a = _run(impl, one, layout, q, k, v).clone()
|
||||
b = _run(impl, many, layout, q, k, v)
|
||||
assert torch.equal(a, b)
|
||||
|
||||
|
||||
@requires_cuda
|
||||
def test_dense_fallback_off_the_dit_blocks() -> None:
|
||||
"""The token refiner resolves the same backend but runs plain dense FA."""
|
||||
device = torch.device("cuda")
|
||||
layout, (q, k, v) = _qkv(device, seed=5)
|
||||
_prepare_flash_attention()
|
||||
impl = HybridWindowAttentionH3Impl(
|
||||
num_heads=HEADS,
|
||||
head_size=HEAD_DIM,
|
||||
causal=False,
|
||||
softmax_scale=HEAD_DIM**-0.5,
|
||||
num_kv_heads=HEADS,
|
||||
prefix="token_refiner.blocks.0.attn",
|
||||
)
|
||||
assert impl.layer_idx is None
|
||||
out = _run(impl, None, layout, q, k, v)
|
||||
full = torch.ones(layout.used, layout.used, dtype=torch.bool, device=device)
|
||||
reference = _masked_reference(q, k, v, full, layout.used)
|
||||
assert (out[: layout.used].float() - reference).abs().max().item() < 2e-2
|
||||
|
||||
|
||||
@requires_cuda
|
||||
def test_metadata_layout_mismatch_is_rejected() -> None:
|
||||
device = torch.device("cuda")
|
||||
layout, (q, k, v) = _qkv(device)
|
||||
meta = HybridWindowAttentionH3MetadataBuilder().build(
|
||||
layout=layout, hybrid=_hybrid(), device=device
|
||||
)
|
||||
cu = torch.tensor(
|
||||
[0, layout.used - 48, layout.seq_len], dtype=torch.int32, device=device
|
||||
)
|
||||
with pytest.raises(ValueError, match="diverged"):
|
||||
_impl().forward_varlen(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seqlens=cu,
|
||||
max_seqlen=layout.used - 48,
|
||||
cu_seqlens_host=(0, layout.used - 48, layout.seq_len),
|
||||
attn_metadata=meta,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -238,7 +238,9 @@ def test_cache_dit_preservation_only_makes_first_gate_out_of_place():
|
||||
block.norm1 = torch.nn.Identity()
|
||||
block.norm2 = torch.nn.Identity()
|
||||
block.attn = _KwargIdentity()
|
||||
block.attn.qkv_proj = SimpleNamespace(quant_method=UnquantizedLinearMethod())
|
||||
block.mlp = torch.nn.Identity()
|
||||
block.mlp.fc1 = SimpleNamespace(quant_method=UnquantizedLinearMethod())
|
||||
gate_modes = []
|
||||
|
||||
def fake_gate(residual, _gate, _other, _indices, *, dtype, allow_inplace=True):
|
||||
|
||||
@@ -0,0 +1,703 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""VDN-H3 (hybrid window-softmax + Video Delta linear attention MiniMax-H3):
|
||||
registration, admission, and the linear branch's arithmetic against
|
||||
step-by-step references (no weights, CPU + small CUDA shapes)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import re
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.configs.models.dits.minimax_h3 import (
|
||||
MiniMaxH3DiTArchConfig,
|
||||
MiniMaxH3DiTConfig,
|
||||
VDNHybridAttentionArchConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import (
|
||||
MiniMaxH3PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3_vdn import (
|
||||
VDNH3PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sample.minimax_h3_vdn import VDNH3SamplingParams
|
||||
from sglang.multimodal_gen.registry import (
|
||||
get_model_info,
|
||||
get_non_diffusers_pipeline_name,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.dits.minimax_h3_vdn import (
|
||||
TEXT_STATE_SCALE,
|
||||
MiniMaxH3VDNLinearBranch,
|
||||
VDNH3Layout,
|
||||
delta_factor_apply,
|
||||
frame_statistics,
|
||||
gather_linear_state,
|
||||
run_scans,
|
||||
vdn_h3_layout_from_packed,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import (
|
||||
AttentionBackendEnum,
|
||||
current_platform,
|
||||
)
|
||||
|
||||
VDN_MODEL_ID = "OpenVDN/vdn-minimax-h3"
|
||||
requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
|
||||
# admission resolves hybrid_window_attn_h3, which only the CUDA platform registers
|
||||
requires_cuda_backend = pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="hybrid_window_attn_h3 admission needs NVIDIA CUDA",
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# registration and admission
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_registry_resolves_vdn_h3_configs() -> None:
|
||||
info = get_model_info(VDN_MODEL_ID)
|
||||
assert info.sampling_param_cls is VDNH3SamplingParams
|
||||
assert info.pipeline_config_cls is VDNH3PipelineConfig
|
||||
assert get_non_diffusers_pipeline_name(VDN_MODEL_ID) == "VDNH3Pipeline"
|
||||
assert get_non_diffusers_pipeline_name("/models/OpenVDN/vdn-minimax-h3") == (
|
||||
"VDNH3Pipeline"
|
||||
)
|
||||
# the base H3 detector must not swallow the VDN id (it contains "minimax-h3")
|
||||
base = get_model_info("MiniMaxAI/MiniMax-H3")
|
||||
assert base.pipeline_config_cls is MiniMaxH3PipelineConfig
|
||||
|
||||
|
||||
def test_vdn_h3_sampling_defaults_and_rejections() -> None:
|
||||
params = VDNH3SamplingParams(prompt="p")
|
||||
assert params.num_inference_steps == 9 # 8 NFE
|
||||
with pytest.raises(ValueError, match="exactly nine sigma grid points"):
|
||||
VDNH3SamplingParams(prompt="p", num_inference_steps=8)
|
||||
fl2va = VDNH3SamplingParams(
|
||||
prompt="p",
|
||||
task="fl2va",
|
||||
conditions=[
|
||||
{"type": "image", "uri": "x.png", "role": "keyframe", "frame_index": 0}
|
||||
],
|
||||
target={"short_edge": 768, "aspect_ratio": "auto", "duration_seconds": 5.0},
|
||||
)
|
||||
assert fl2va.task == "fl2va"
|
||||
with pytest.raises(ValueError, match="ref2va was not trained"):
|
||||
VDNH3SamplingParams(
|
||||
prompt="p",
|
||||
task="ref2va",
|
||||
conditions=[{"type": "image", "uri": "x.png", "role": "reference"}],
|
||||
target={"short_edge": 768, "aspect_ratio": "auto", "duration_seconds": 5.0},
|
||||
)
|
||||
|
||||
|
||||
def _server_args(**overrides) -> SimpleNamespace:
|
||||
args = dict(
|
||||
model_variant=None,
|
||||
attention_backend=None,
|
||||
component_attention_backends={},
|
||||
attention_backend_config=None,
|
||||
ring_degree=1,
|
||||
enable_torch_compile=False,
|
||||
enable_breakable_cuda_graph=False,
|
||||
quantization=None,
|
||||
)
|
||||
args.update(overrides)
|
||||
ns = SimpleNamespace(**args)
|
||||
ns.resolve_component_attention_backend = lambda name: (
|
||||
(
|
||||
AttentionBackendEnum[str(ns.component_attention_backends[name]).upper()]
|
||||
if name in ns.component_attention_backends
|
||||
else None
|
||||
),
|
||||
None,
|
||||
)
|
||||
return ns
|
||||
|
||||
|
||||
@requires_cuda_backend
|
||||
def test_vdn_h3_pipeline_config_rejections() -> None:
|
||||
config = VDNH3PipelineConfig()
|
||||
with pytest.raises(ValueError, match="--model-variant does not apply"):
|
||||
config.validate_server_args(_server_args(model_variant="ref2va"))
|
||||
with pytest.raises(
|
||||
ValueError, match="requires --attention-backend hybrid_window_attn_h3"
|
||||
):
|
||||
config.validate_server_args(_server_args(attention_backend="fa"))
|
||||
with pytest.raises(ValueError, match="ring-degree"):
|
||||
config.validate_server_args(_server_args(ring_degree=2))
|
||||
with pytest.raises(ValueError, match="torch.compile"):
|
||||
config.validate_server_args(_server_args(enable_torch_compile=True))
|
||||
with pytest.raises(ValueError, match="breakable CUDA graph"):
|
||||
config.validate_server_args(_server_args(enable_breakable_cuda_graph=True))
|
||||
with pytest.raises(ValueError, match="no.*audited high-quality deployment"):
|
||||
config.validate_quality_deployment(server_args=None)
|
||||
args = _server_args()
|
||||
config.validate_server_args(args)
|
||||
assert args.attention_backend == "hybrid_window_attn_h3"
|
||||
|
||||
|
||||
@requires_cuda_backend
|
||||
def test_vdn_h3_quantization_defaults_to_mxfp8_on_blackwell(monkeypatch) -> None:
|
||||
"""Online MXFP8 is the default on SM100+ (SM120 included) and what `fp8`
|
||||
maps to there; `bf16` opts out; before SM100 the block-scaled GEMM does
|
||||
not exist, so an unset flag stays bf16 and `fp8` stays the per-channel
|
||||
path."""
|
||||
from sglang.multimodal_gen.configs.pipeline_configs import minimax_h3_vdn as module
|
||||
|
||||
def resolved(
|
||||
quantization: str | None, blackwell: bool, sm120: bool = False
|
||||
) -> str | None:
|
||||
monkeypatch.setattr(module.current_platform, "is_blackwell", lambda: blackwell)
|
||||
monkeypatch.setattr(module.current_platform, "is_sm120", lambda: sm120)
|
||||
args = _server_args(quantization=quantization)
|
||||
VDNH3PipelineConfig().validate_server_args(args)
|
||||
return args.quantization
|
||||
|
||||
assert resolved(None, True) == "mxfp8"
|
||||
assert resolved("fp8", True) == "mxfp8"
|
||||
assert resolved("bf16", True) is None
|
||||
assert resolved(None, False) is None
|
||||
assert resolved("fp8", False) == "fp8"
|
||||
assert resolved(None, False, sm120=True) == "mxfp8"
|
||||
assert resolved("fp8", False, sm120=True) == "mxfp8"
|
||||
assert resolved("bf16", False, sm120=True) is None
|
||||
|
||||
|
||||
def test_hybrid_arch_config_from_transform_config_and_mapping() -> None:
|
||||
transform = {
|
||||
"anchor_frames": "both",
|
||||
"enable_softmax_gate": True,
|
||||
"linear_attention": {
|
||||
"a_fp32": True,
|
||||
"bridge": "alpha",
|
||||
"delta_rule": "vdn_solve",
|
||||
"enable_text_state": True,
|
||||
"linear_head_dim": 128,
|
||||
"short_conv": {"targets": ["k", "v"]},
|
||||
},
|
||||
"softmax_attention": {"chunk": 5, "radius": 1},
|
||||
}
|
||||
dit = MiniMaxH3DiTConfig()
|
||||
dit.update_model_arch({"hybrid_attention": transform, "num_layers": 2})
|
||||
hybrid = dit.arch_config.hybrid_attention
|
||||
assert isinstance(hybrid, VDNHybridAttentionArchConfig)
|
||||
assert hybrid.short_conv == ("k", "v") and hybrid.chunk == 5
|
||||
assert MiniMaxH3DiTConfig().arch_config.hybrid_attention is None
|
||||
with pytest.raises(ValueError, match="delta_rule"):
|
||||
VDNHybridAttentionArchConfig(delta_rule="bogus")
|
||||
|
||||
mapping = MiniMaxH3DiTArchConfig().param_names_mapping
|
||||
for source, expected in (
|
||||
(
|
||||
"transformer_blocks.7.attn.linear_attention.alpha.A_log",
|
||||
"blocks.7.attn.hybrid.linear_attention.alpha.A_log",
|
||||
),
|
||||
(
|
||||
"transformer_blocks.7.attn.softmax_gate.up.bias",
|
||||
"blocks.7.attn.hybrid.softmax_gate.up.bias",
|
||||
),
|
||||
(
|
||||
"transformer_blocks.7.attn.to_out_linear.weight",
|
||||
"blocks.7.attn.hybrid.to_out_linear.weight",
|
||||
),
|
||||
):
|
||||
targets = [
|
||||
re.sub(pattern, target if isinstance(target, str) else target[0], source)
|
||||
for pattern, target in mapping.items()
|
||||
if re.match(pattern, source)
|
||||
]
|
||||
assert targets == [expected], (source, targets)
|
||||
|
||||
|
||||
def test_layout_from_packed_t2va() -> None:
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.packed_sequence import (
|
||||
minimax_h3_packed_sequence,
|
||||
)
|
||||
|
||||
packed = minimax_h3_packed_sequence(
|
||||
text_len=70,
|
||||
latent_t=12,
|
||||
latent_h=12,
|
||||
latent_w=16,
|
||||
audio_t=50,
|
||||
include_keyframe_cond=False,
|
||||
)
|
||||
layout = vdn_h3_layout_from_packed(packed, latent_t=12, latent_h=12, latent_w=16)
|
||||
assert layout.text_len == 70
|
||||
assert layout.video_start == 70 + 100
|
||||
assert layout.tokens_per_frame == 48 and layout.frame_size == (6, 8)
|
||||
assert layout.used == 70 + 100 + 12 * 48
|
||||
assert layout.seq_len == int(packed["seq_len"]) and layout.seq_len % 64 == 0
|
||||
assert layout.global_ranges == [(0, 170)]
|
||||
|
||||
|
||||
def test_layout_from_packed_fl2va_keeps_keyframe_rows_global() -> None:
|
||||
"""Keyframe rows must land in the global ranges, not in the video span
|
||||
the linear branch scans."""
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.packed_sequence import (
|
||||
minimax_h3_packed_sequence,
|
||||
)
|
||||
|
||||
packed = minimax_h3_packed_sequence(
|
||||
text_len=70,
|
||||
latent_t=12,
|
||||
latent_h=12,
|
||||
latent_w=16,
|
||||
audio_t=50,
|
||||
include_keyframe_cond=True,
|
||||
keyframe_frame_indices=[0, -1],
|
||||
frame_count=45,
|
||||
)
|
||||
layout = vdn_h3_layout_from_packed(packed, latent_t=12, latent_h=12, latent_w=16)
|
||||
assert layout.text_len == 70
|
||||
assert layout.video_start == 70 + 2 * 48 + 100
|
||||
assert layout.used == layout.video_end == layout.video_start + 12 * 48
|
||||
assert layout.global_ranges == [(0, layout.video_start)]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# the linear branch arithmetic
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
FRAMES, HEADS, TOKENS, HEAD_DIM = 6, 2, 16, 32
|
||||
|
||||
|
||||
def _random_stats(device, seed=0):
|
||||
g = torch.Generator(device="cpu").manual_seed(seed)
|
||||
k = torch.nn.functional.normalize(
|
||||
torch.randn(FRAMES, HEADS, TOKENS, HEAD_DIM, generator=g), dim=-1
|
||||
).to(device)
|
||||
v = torch.randn(FRAMES, HEADS, TOKENS, HEAD_DIM, generator=g).to(device)
|
||||
beta = torch.rand(FRAMES, HEADS, TOKENS, generator=g).to(device)
|
||||
alpha = torch.rand(FRAMES, HEADS, HEAD_DIM, generator=g).to(device) * 0.5 + 0.5
|
||||
return k, v, beta, alpha
|
||||
|
||||
|
||||
def test_frame_statistics_and_delta_rule_match_dense_algebra() -> None:
|
||||
k, v, beta, alpha = _random_stats("cpu")
|
||||
A, B = frame_statistics(k, v, beta, a_fp32=True)
|
||||
A_ref = torch.einsum("fhsk,fhs,fhsl->fhkl", k, beta, k)
|
||||
B_ref = torch.einsum("fhsv,fhs,fhsk->fhvk", v, beta, k)
|
||||
assert torch.allclose(A, A_ref, atol=1e-4) and torch.allclose(B, B_ref, atol=1e-4)
|
||||
transition, injection = delta_factor_apply(
|
||||
"vdn_solve", alpha, A, B, tokens_per_frame=TOKENS
|
||||
)
|
||||
inv = torch.linalg.inv(torch.eye(HEAD_DIM) + A)
|
||||
assert torch.allclose(transition, alpha.unsqueeze(-1) * inv, atol=1e-4)
|
||||
assert torch.allclose(injection, B @ inv, atol=1e-4)
|
||||
|
||||
|
||||
def test_scans_match_step_reference_and_text_seed() -> None:
|
||||
k, v, beta, alpha = _random_stats("cpu", seed=1)
|
||||
A, B = frame_statistics(k, v, beta, a_fp32=True)
|
||||
transition, injection = delta_factor_apply(
|
||||
"vdn_solve", alpha, A, B, tokens_per_frame=TOKENS
|
||||
)
|
||||
text_state = torch.randn(HEADS, HEAD_DIM, HEAD_DIM)
|
||||
prefix, suffix = run_scans(transition, injection, text_state)
|
||||
state = text_state.clone()
|
||||
for f in range(FRAMES):
|
||||
state = state @ transition[f] + injection[f]
|
||||
assert torch.allclose(prefix[f], state, atol=1e-4)
|
||||
state = text_state.clone()
|
||||
for f in range(FRAMES - 1, -1, -1):
|
||||
state = state @ transition[f] + injection[f]
|
||||
assert torch.allclose(suffix[f], state, atol=1e-4)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("world", [1, 2, 5, 10])
|
||||
def test_frame_partial_sums_match_index_add(world: int) -> None:
|
||||
"""The Ulysses frame-mean partial sums (reshape-sum over whole frames plus
|
||||
two edge rows sums, deterministic) equal the index_add formulation."""
|
||||
from sglang.multimodal_gen.runtime.models.dits.minimax_h3_vdn_attention import (
|
||||
_vdn_frame_partial_sums,
|
||||
)
|
||||
|
||||
frames, tpf, hidden = 12, 50, 64
|
||||
video_start = 20
|
||||
video_end = video_start + frames * tpf
|
||||
seq = video_end + 30 # 650 rows, divisible by every ``world`` above
|
||||
g = torch.Generator(device="cpu").manual_seed(0)
|
||||
x = torch.randn(seq, hidden, generator=g).to(torch.bfloat16)
|
||||
local = seq // world
|
||||
total = torch.zeros(frames, hidden)
|
||||
for rank in range(world):
|
||||
total += _vdn_frame_partial_sums(
|
||||
x[rank * local : (rank + 1) * local],
|
||||
row_start=rank * local,
|
||||
video_start=video_start,
|
||||
video_end=video_end,
|
||||
num_frames=frames,
|
||||
tokens_per_frame=tpf,
|
||||
)
|
||||
ref = x[video_start:video_end].float().view(frames, tpf, hidden).sum(1)
|
||||
torch.testing.assert_close(total, ref, rtol=1e-5, atol=1e-3)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
@pytest.mark.parametrize("anchor_frames", ["both", "none"])
|
||||
@pytest.mark.parametrize("reference", ["frame_chain_scans", "eager_kernels"])
|
||||
def test_branch_forward_matches_reference(
|
||||
anchor_frames: str, reference: str, monkeypatch
|
||||
) -> None:
|
||||
"""The shipped branch (boundary scans, fused Triton kernels) against the
|
||||
same module with the plain frame-chain scans, and against the eager
|
||||
kernel chain; both across the anchor-frame shift of the chunk grid."""
|
||||
from sglang.multimodal_gen.runtime.models.dits import minimax_h3_vdn as module
|
||||
|
||||
device = torch.device("cuda")
|
||||
hidden, heads, head_dim = 64, 4, 32
|
||||
num_frames, fh, fw = 13, 4, 6
|
||||
tpf = fh * fw
|
||||
hybrid = VDNHybridAttentionArchConfig(
|
||||
chunk=3, radius=1, anchor_frames=anchor_frames, linear_head_dim=head_dim
|
||||
)
|
||||
branch = _branch(hybrid, heads, hidden, head_dim, seed=0).to(device)
|
||||
layout = VDNH3Layout(
|
||||
seq_len=64 * 6,
|
||||
used=10 + num_frames * tpf,
|
||||
text_len=10,
|
||||
video_start=10,
|
||||
num_frames=num_frames,
|
||||
tokens_per_frame=tpf,
|
||||
frame_height=fh,
|
||||
frame_width=fw,
|
||||
)
|
||||
g = torch.Generator(device="cpu").manual_seed(1)
|
||||
V = num_frames * tpf
|
||||
q, k, v = (
|
||||
torch.randn(V, heads, head_dim, generator=g).to(device, torch.bfloat16)
|
||||
for _ in range(3)
|
||||
)
|
||||
tk, tv = (
|
||||
torch.randn(10, heads, head_dim, generator=g).to(device, torch.bfloat16)
|
||||
for _ in range(2)
|
||||
)
|
||||
x = torch.randn(V, hidden, generator=g).to(device, torch.bfloat16)
|
||||
tx = torch.randn(10, hidden, generator=g).to(device, torch.bfloat16)
|
||||
kwargs = dict(
|
||||
q_raw=q,
|
||||
k_raw=k,
|
||||
v_raw=v,
|
||||
beta=branch.beta(x),
|
||||
gate=branch.output_gate(x),
|
||||
frame_mean=x.view(num_frames, tpf, hidden).mean(1, dtype=torch.float32),
|
||||
layout=layout,
|
||||
text_k_raw=tk,
|
||||
text_v_raw=tv,
|
||||
text_beta=branch.beta(tx),
|
||||
)
|
||||
fast = branch(**kwargs)
|
||||
if reference == "frame_chain_scans":
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"run_boundary_scans",
|
||||
lambda t, i, ts, *, chunk, frame_offset=0: run_scans(t, i, ts),
|
||||
)
|
||||
else:
|
||||
branch.fused_kernels = False
|
||||
expected = branch(**kwargs)
|
||||
assert expected.abs().sum() > 0
|
||||
# fused kernels skip the eager chain's bf16 roundings; the scans re-associate fp32
|
||||
tolerance = {"frame_chain_scans": 5e-3, "eager_kernels": 2e-2}[reference]
|
||||
rel = (fast.float() - expected.float()).norm() / expected.float().norm()
|
||||
assert rel < tolerance, rel
|
||||
|
||||
|
||||
def test_gather_is_the_exact_window_complement() -> None:
|
||||
"""With alpha = 1 and one-hot frame indicators, the gathered state must
|
||||
be exactly the indicator of the frames outside the window."""
|
||||
num_frames = 9
|
||||
hybrid = VDNHybridAttentionArchConfig(chunk=3, radius=1, anchor_frames="none")
|
||||
bounds = hybrid.window_bounds(num_frames)
|
||||
# injection[f] = one-hot(f) laid along dv; transition = identity
|
||||
eye = torch.eye(num_frames)
|
||||
injection = (
|
||||
eye.view(num_frames, 1, num_frames, 1)
|
||||
.expand(num_frames, 1, num_frames, 1)
|
||||
.clone()
|
||||
)
|
||||
transition = torch.eye(1).view(1, 1, 1, 1).expand(num_frames, 1, 1, 1).clone()
|
||||
prefix, suffix = run_scans(transition, injection, None)
|
||||
alpha = torch.ones(num_frames, 1, 1)
|
||||
gathered = gather_linear_state(
|
||||
prefix,
|
||||
suffix,
|
||||
alpha,
|
||||
bounds,
|
||||
bridge="alpha",
|
||||
text_state=None,
|
||||
out_dtype=torch.float32,
|
||||
)
|
||||
for t in range(num_frames):
|
||||
lo, hi = max(bounds[t][0], 0), min(bounds[t][1], num_frames - 1)
|
||||
expected = torch.tensor(
|
||||
[1.0 if (f < lo or f > hi) else 0.0 for f in range(num_frames)]
|
||||
)
|
||||
assert torch.equal(gathered[t, 0, :, 0], expected), (t, gathered[t, 0, :, 0])
|
||||
|
||||
|
||||
def test_gather_text_state_decays_over_skipped_frames() -> None:
|
||||
"""A clip-end frame reads the text state decayed by prod alpha over
|
||||
exactly the frames between the boundary and t (VDN's bridge indices)."""
|
||||
num_frames = 4
|
||||
bounds = [
|
||||
(t, t) for t in range(num_frames)
|
||||
] # radius 0: complement = everything else
|
||||
prefix = torch.zeros(num_frames, 1, 1, 1)
|
||||
suffix = torch.zeros(num_frames, 1, 1, 1)
|
||||
alpha = torch.tensor([0.5, 0.25, 0.5, 0.5]).view(num_frames, 1, 1)
|
||||
text_state = torch.ones(1, 1, 1)
|
||||
out = gather_linear_state(
|
||||
prefix,
|
||||
suffix,
|
||||
alpha,
|
||||
bounds,
|
||||
bridge="alpha",
|
||||
text_state=text_state,
|
||||
out_dtype=torch.float32,
|
||||
).view(num_frames)
|
||||
# frame 0 reads the text state through alpha[0]; frame 3 through alpha[3]
|
||||
assert math.isclose(out[0].item(), 0.5, rel_tol=1e-6)
|
||||
assert math.isclose(out[3].item(), 0.5, rel_tol=1e-6)
|
||||
assert out[1].item() == 0.0, "both neighbours in range and zero"
|
||||
|
||||
|
||||
def test_temporal_shift_features_match_conv1d() -> None:
|
||||
from sglang.multimodal_gen.runtime.models.dits.minimax_h3_vdn import _temporal_shift
|
||||
|
||||
x = torch.randn(7, 5, 6) # [F, S, C]
|
||||
w = torch.randn(6, 5)
|
||||
got = _temporal_shift(x, w)
|
||||
ref = (
|
||||
torch.nn.functional.conv1d(
|
||||
x.permute(1, 2, 0).reshape(5, 6, 7), w.view(6, 1, 5), padding=2, groups=6
|
||||
)
|
||||
.reshape(5, 6, 7)
|
||||
.permute(2, 0, 1)
|
||||
)
|
||||
assert torch.allclose(got, ref, atol=1e-5)
|
||||
|
||||
|
||||
def _branch(
|
||||
hybrid: VDNHybridAttentionArchConfig,
|
||||
heads: int,
|
||||
hidden: int,
|
||||
head_dim: int,
|
||||
seed: int,
|
||||
):
|
||||
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
|
||||
maybe_init_distributed_environment_and_model_parallel,
|
||||
model_parallel_is_initialized,
|
||||
)
|
||||
from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import (
|
||||
ensure_distributed_env_defaults,
|
||||
)
|
||||
|
||||
if not model_parallel_is_initialized():
|
||||
ensure_distributed_env_defaults()
|
||||
maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1)
|
||||
arch = MiniMaxH3DiTArchConfig(
|
||||
num_attention_heads=heads, attention_head_dim=head_dim, hidden_size=hidden
|
||||
)
|
||||
branch = MiniMaxH3VDNLinearBranch(arch, hybrid, local_heads=heads)
|
||||
g = torch.Generator(device="cpu").manual_seed(seed)
|
||||
with torch.no_grad():
|
||||
for name, p in branch.named_parameters():
|
||||
if name.endswith("A_log"):
|
||||
p.copy_(torch.log(torch.empty_like(p).uniform_(1, 16, generator=g)))
|
||||
elif name.endswith("dt_bias"):
|
||||
p.copy_(torch.randn(p.shape, generator=g) - 3)
|
||||
elif name.endswith("norm.weight"):
|
||||
p.copy_(torch.ones_like(p))
|
||||
else:
|
||||
p.copy_(
|
||||
torch.randn(p.shape, generator=g, dtype=torch.float32).to(p.dtype)
|
||||
* 0.1
|
||||
)
|
||||
return branch
|
||||
|
||||
|
||||
@requires_cuda
|
||||
def test_branch_head_slice_equals_full_run() -> None:
|
||||
"""The Ulysses contract: the branch is per-head independent given
|
||||
(beta, gate, alpha), so a head-sliced run equals the full run's slice."""
|
||||
device = torch.device("cuda")
|
||||
hidden, heads, head_dim = 64, 4, 32
|
||||
num_frames, fh, fw = 7, 4, 6
|
||||
tpf = fh * fw
|
||||
hybrid = VDNHybridAttentionArchConfig(
|
||||
chunk=2, radius=1, anchor_frames="both", linear_head_dim=head_dim
|
||||
)
|
||||
branch = _branch(hybrid, heads, hidden, head_dim, seed=0).to(device)
|
||||
layout = VDNH3Layout(
|
||||
seq_len=64 * 4,
|
||||
used=10 + num_frames * tpf,
|
||||
text_len=10,
|
||||
video_start=10,
|
||||
num_frames=num_frames,
|
||||
tokens_per_frame=tpf,
|
||||
frame_height=fh,
|
||||
frame_width=fw,
|
||||
)
|
||||
g = torch.Generator(device="cpu").manual_seed(1)
|
||||
V = num_frames * tpf
|
||||
q, k, v = (
|
||||
torch.randn(V, heads, head_dim, generator=g).to(device, torch.bfloat16)
|
||||
for _ in range(3)
|
||||
)
|
||||
tk, tv = (
|
||||
torch.randn(10, heads, head_dim, generator=g).to(device, torch.bfloat16)
|
||||
for _ in range(2)
|
||||
)
|
||||
x = torch.randn(V, hidden, generator=g).to(device, torch.bfloat16)
|
||||
tx = torch.randn(10, hidden, generator=g).to(device, torch.bfloat16)
|
||||
beta, gate = branch.beta(x), branch.output_gate(x)
|
||||
tbeta = branch.beta(tx)
|
||||
frame_mean = x.view(num_frames, tpf, hidden).mean(1, dtype=torch.float32)
|
||||
full = branch(
|
||||
q_raw=q,
|
||||
k_raw=k,
|
||||
v_raw=v,
|
||||
beta=beta,
|
||||
gate=gate,
|
||||
frame_mean=frame_mean,
|
||||
layout=layout,
|
||||
text_k_raw=tk,
|
||||
text_v_raw=tv,
|
||||
text_beta=tbeta,
|
||||
).view(V, heads, head_dim)
|
||||
# anchors read zero
|
||||
assert torch.all(full[:tpf] == 0) and torch.all(full[-tpf:] == 0)
|
||||
assert full[tpf:-tpf].abs().sum() > 0
|
||||
|
||||
# the same module on a head range of the full sequence
|
||||
hs = slice(1, 3)
|
||||
part = branch(
|
||||
q_raw=q[:, hs],
|
||||
k_raw=k[:, hs],
|
||||
v_raw=v[:, hs],
|
||||
beta=beta[:, hs],
|
||||
gate=gate[:, hs],
|
||||
frame_mean=frame_mean,
|
||||
layout=layout,
|
||||
text_k_raw=tk[:, hs],
|
||||
text_v_raw=tv[:, hs],
|
||||
text_beta=tbeta[:, hs],
|
||||
heads=hs,
|
||||
).view(V, 2, head_dim)
|
||||
diff = (part.float() - full[:, hs].float()).abs().max().item()
|
||||
assert diff < 2e-2, f"head slice vs full run max diff {diff}"
|
||||
|
||||
|
||||
@requires_cuda
|
||||
def test_branch_matches_eager_reference_algorithm() -> None:
|
||||
"""The module against a from-scratch spelling of VDN's _readout (no
|
||||
skip_ends), including the text state seed."""
|
||||
device = torch.device("cuda")
|
||||
hidden, heads, head_dim = 48, 2, 32
|
||||
num_frames, fh, fw = 5, 3, 4
|
||||
tpf = fh * fw
|
||||
hybrid = VDNHybridAttentionArchConfig(
|
||||
chunk=0, radius=1, anchor_frames="none", linear_head_dim=head_dim, short_conv=()
|
||||
)
|
||||
branch = _branch(hybrid, heads, hidden, head_dim, seed=2).to(device)
|
||||
layout = VDNH3Layout(
|
||||
seq_len=256,
|
||||
used=8 + num_frames * tpf,
|
||||
text_len=8,
|
||||
video_start=8,
|
||||
num_frames=num_frames,
|
||||
tokens_per_frame=tpf,
|
||||
frame_height=fh,
|
||||
frame_width=fw,
|
||||
)
|
||||
g = torch.Generator(device="cpu").manual_seed(3)
|
||||
V = num_frames * tpf
|
||||
q, k, v = (
|
||||
torch.randn(V, heads, head_dim, generator=g).to(device, torch.bfloat16)
|
||||
for _ in range(3)
|
||||
)
|
||||
tk, tv = (
|
||||
torch.randn(8, heads, head_dim, generator=g).to(device, torch.bfloat16)
|
||||
for _ in range(2)
|
||||
)
|
||||
x = torch.randn(V, hidden, generator=g).to(device, torch.bfloat16)
|
||||
tx = torch.randn(8, hidden, generator=g).to(device, torch.bfloat16)
|
||||
beta, gate, tbeta = branch.beta(x), branch.output_gate(x), branch.beta(tx)
|
||||
frame_mean = x.view(num_frames, tpf, hidden).mean(1, dtype=torch.float32)
|
||||
got = branch(
|
||||
q_raw=q,
|
||||
k_raw=k,
|
||||
v_raw=v,
|
||||
beta=beta,
|
||||
gate=gate,
|
||||
frame_mean=frame_mean,
|
||||
layout=layout,
|
||||
text_k_raw=tk,
|
||||
text_v_raw=tv,
|
||||
text_beta=tbeta,
|
||||
)
|
||||
|
||||
# reference in fp32
|
||||
def feat(t, l2):
|
||||
y = torch.nn.functional.silu(t.float())
|
||||
return torch.nn.functional.normalize(y, dim=-1, eps=1e-6) if l2 else y
|
||||
|
||||
qf, kf, vf = feat(q, True), feat(k, True), feat(v, False)
|
||||
bounds = hybrid.window_bounds(num_frames)
|
||||
kb = kf.view(num_frames, tpf, heads, head_dim).permute(0, 2, 1, 3)
|
||||
vb = vf.view(num_frames, tpf, heads, head_dim).permute(0, 2, 1, 3)
|
||||
bb = beta.float().view(num_frames, tpf, heads).permute(0, 2, 1)
|
||||
A = torch.einsum("fhsk,fhs,fhsl->fhkl", kb, bb, kb)
|
||||
B = torch.einsum("fhsv,fhs,fhsk->fhvk", vb, bb, kb)
|
||||
inv = torch.linalg.inv(torch.eye(head_dim, device=device) + A)
|
||||
alpha = branch.alpha(frame_mean)
|
||||
trans = alpha.unsqueeze(-1) * inv
|
||||
inj = B @ inv
|
||||
# text state
|
||||
tkf, tvf = feat(tk, True), feat(tv, False)
|
||||
tkb = tkf.view(1, 8, heads, head_dim).permute(0, 2, 1, 3)
|
||||
tvb = tvf.view(1, 8, heads, head_dim).permute(0, 2, 1, 3)
|
||||
tbb = tbeta.float().view(1, 8, heads).permute(0, 2, 1)
|
||||
tA = torch.einsum("fhsk,fhs,fhsl->fhkl", tkb, tbb, tkb)[0]
|
||||
tB = torch.einsum("fhsv,fhs,fhsk->fhvk", tvb, tbb, tkb)[0]
|
||||
text_state = TEXT_STATE_SCALE * (
|
||||
tB @ torch.linalg.inv(torch.eye(head_dim, device=device) + tA)
|
||||
)
|
||||
prefix, suffix = [], [None] * num_frames
|
||||
s = text_state.clone()
|
||||
for f in range(num_frames):
|
||||
s = s @ trans[f] + inj[f]
|
||||
prefix.append(s)
|
||||
s = text_state.clone()
|
||||
for f in range(num_frames - 1, -1, -1):
|
||||
s = s @ trans[f] + inj[f]
|
||||
suffix[f] = s
|
||||
outs = []
|
||||
for t in range(num_frames):
|
||||
lo, hi = bounds[t]
|
||||
left = prefix[lo - 1] if lo - 1 >= 0 else text_state
|
||||
right = suffix[hi + 1] if hi + 1 < num_frames else text_state
|
||||
a_before = torch.prod(alpha[max(lo, 0) : t + 1], dim=0)
|
||||
a_after = torch.prod(alpha[t : min(hi, num_frames - 1) + 1], dim=0)
|
||||
state = left * a_before.unsqueeze(1) + right * a_after.unsqueeze(1)
|
||||
qt = qf.view(num_frames, tpf, heads, head_dim)[t] # [S, H, d]
|
||||
ro = torch.einsum("shk,hvk->shv", qt, state)
|
||||
ms = ro.pow(2).mean(-1, keepdim=True)
|
||||
ro = ro * torch.rsqrt(ms + branch.norm.eps) * branch.norm.weight.float()
|
||||
outs.append(ro)
|
||||
ref = (torch.cat(outs) * gate.float()).reshape(V, heads * head_dim)
|
||||
diff = (got.float() - ref).abs().max().item()
|
||||
scale = ref.abs().max().item()
|
||||
assert diff < 3e-2 * max(scale, 1.0), (
|
||||
f"branch vs reference max diff {diff} (scale {scale})"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -0,0 +1,115 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Online MXFP8 GEMM path (``MXFP8Config()`` on a bf16 checkpoint): load-time
|
||||
block quant, the prequantized (e4m3, swizzled scales) input from the fused
|
||||
SwiGLU kernel, and the per-layer fallback to the per-channel fp8 path."""
|
||||
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
|
||||
|
||||
def _init_parallel() -> None:
|
||||
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
|
||||
maybe_init_distributed_environment_and_model_parallel,
|
||||
model_parallel_is_initialized,
|
||||
)
|
||||
from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import (
|
||||
ensure_distributed_env_defaults,
|
||||
)
|
||||
|
||||
if not model_parallel_is_initialized():
|
||||
ensure_distributed_env_defaults()
|
||||
maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1)
|
||||
|
||||
|
||||
def _layer(
|
||||
in_f: int, out_f: int, bias: bool, params_dtype: torch.dtype = torch.bfloat16
|
||||
):
|
||||
from sglang.multimodal_gen.runtime.layers.linear import RowParallelLinear
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.mxfp8 import MXFP8Config
|
||||
|
||||
return RowParallelLinear(
|
||||
in_f,
|
||||
out_f,
|
||||
bias=bias,
|
||||
params_dtype=params_dtype,
|
||||
quant_config=MXFP8Config(),
|
||||
prefix="mlp.fc2",
|
||||
).to("cuda")
|
||||
|
||||
|
||||
def test_mxfp8_linear_matches_bf16_and_accepts_prequantized() -> None:
|
||||
if torch.cuda.get_device_capability()[0] < 10:
|
||||
pytest.skip("cuBLASLt MXFP8 block scaling requires Blackwell or newer")
|
||||
_init_parallel()
|
||||
from sglang.kernels.ops.diffusion import silu_mul_mxfp8
|
||||
|
||||
in_f, out_f, rows = 512, 384, 200
|
||||
layer = _layer(in_f, out_f, bias=False)
|
||||
g = torch.Generator(device="cpu").manual_seed(1)
|
||||
weight = (torch.randn(out_f, in_f, generator=g) * 0.02).to("cuda", torch.bfloat16)
|
||||
with torch.no_grad():
|
||||
layer.weight.copy_(weight)
|
||||
layer.quant_method.process_weights_after_loading(layer)
|
||||
assert layer.mxfp8 and layer.quant_method.accepts_mxfp8_input(layer)
|
||||
assert layer.weight.dtype == torch.float8_e4m3fn
|
||||
assert layer.weight_scale.dtype == torch.float8_e8m0fnu
|
||||
|
||||
x = torch.randn(rows, in_f, generator=g).to("cuda", torch.bfloat16)
|
||||
out, _ = layer(x)
|
||||
ref = x.float() @ weight.float().t()
|
||||
rel = ((out.float() - ref).norm() / ref.norm()).item()
|
||||
assert rel < 0.05, rel
|
||||
|
||||
hidden = torch.randn(rows, 2 * in_f, generator=g).to("cuda", torch.bfloat16)
|
||||
act = torch.nn.functional.silu(hidden[:, :in_f]) * hidden[:, in_f:]
|
||||
out_tensor, _ = layer(act)
|
||||
out_tuple, _ = layer(silu_mul_mxfp8(hidden))
|
||||
assert torch.equal(out_tensor, out_tuple)
|
||||
|
||||
|
||||
def _assert_aligned_layer_falls_back(params_dtype: torch.dtype) -> None:
|
||||
_init_parallel()
|
||||
layer = _layer(512, 384, bias=False, params_dtype=params_dtype)
|
||||
g = torch.Generator(device="cpu").manual_seed(1)
|
||||
weight = (torch.randn(384, 512, generator=g) * 0.02).to("cuda", params_dtype)
|
||||
with torch.no_grad():
|
||||
layer.weight.copy_(weight)
|
||||
layer.quant_method.process_weights_after_loading(layer)
|
||||
assert not layer.mxfp8 and not layer.quant_method.accepts_mxfp8_input(layer)
|
||||
x = torch.randn(200, 512, generator=g).to("cuda", params_dtype)
|
||||
out, _ = layer(x)
|
||||
ref = x.float() @ weight.float().t()
|
||||
rel = ((out.float() - ref).norm() / ref.norm()).item()
|
||||
assert rel < 0.05, rel
|
||||
|
||||
|
||||
def test_pre_blackwell_aligned_layer_falls_back_to_channelwise() -> None:
|
||||
if torch.cuda.get_device_capability()[0] >= 10:
|
||||
pytest.skip("requires a pre-Blackwell GPU")
|
||||
_assert_aligned_layer_falls_back(torch.bfloat16)
|
||||
|
||||
|
||||
def test_fp16_layer_falls_back_to_channelwise() -> None:
|
||||
"""The swizzled quantizer takes bf16 only; an fp16 layer must not fail at load."""
|
||||
_assert_aligned_layer_falls_back(torch.float16)
|
||||
|
||||
|
||||
def test_unaligned_layer_falls_back_to_channelwise() -> None:
|
||||
"""The block-scaled GEMM needs K % 32 == 0; such a layer keeps the
|
||||
per-channel fp8 path and still answers a forward. K is 16 rather than a
|
||||
smaller odd size because the fallback's scaled GEMM still wants K % 16 == 0
|
||||
(ROCm rejects anything else outright)."""
|
||||
_init_parallel()
|
||||
layer = _layer(16, 128, bias=True)
|
||||
layer.quant_method.process_weights_after_loading(layer)
|
||||
assert not layer.mxfp8 and not layer.quant_method.accepts_mxfp8_input(layer)
|
||||
out, _ = layer(torch.randn(4, 16, device="cuda", dtype=torch.bfloat16))
|
||||
assert out.shape == (4, 128)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
Reference in New Issue
Block a user