biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
2,095 Commits 1 Branch 17 Tags
e0f60b9f052d2aa1fdd5ff2f7df922a00f26e76b
Commit Graph

19 Commits

Author SHA1 Message Date
biondizzle
7b3f6cb13c Fix fused router: use run_nvfp4_fused_router wrapper, correct CuTe tensor API 
- kernel wrapper converts torch tensors to CuTe tensors with mark_layout_dynamic
- test uses the wrapper instead of calling kernel.run() directly
- mat_b/scale_b are now torch tensors (converted inside wrapper)
 2026-06-01 09:19:48 +00:00
biondizzle
483e759d53 Fix: use tensor.mark_layout_dynamic() method (not cute.mark_layout_dynamic) 2026-06-01 09:16:33 +00:00
biondizzle
2412745b21 Test fix: slice NVFP4 logits to actual expert count (GEMM padding) 2026-06-01 09:15:06 +00:00
biondizzle
9b86b2b414 Test: fix fused router test - proper NVFP4 quantization and CuTe tensor setup 
- Use quantize_to_nvfp4 for weight quantization
- Use quantize_activation_nvfp4 with computed global_scale
- Get mat_b and scale_b from Nvfp4Linear after finalize_weights
- Compare against both BF16 reference and NVFP4 GEMM reference
 2026-06-01 08:56:20 +00:00
biondizzle
b94f8d4ed8 Test: fused router kernel vs BF16 reference path 
- BF16 GEMM + activation_topk as reference
- NVFP4 GEMM + fused router epilogue as test target
- Proper NVFP4 quantization and CuTe tensor creation
- Cosine similarity and topk_ids matching validation
 2026-06-01 08:54:24 +00:00
biondizzle
2433700a69 Fused router kernel: rewrite epilogue with proper CuTeDSL constructs 
- Replace Python lists with individual scalar variables (s0..s5, i0..i5, a0..a5)
- Replace min-heap sift-down with fully unrolled sorted insertion
  (descending order, no dynamic indexing, no while loops)
- Replace raw SMEM pointer arithmetic with CuTeDSL SMEM tensors
  (s_merge_s, s_merge_i, s_merge_a)
- Replace cute.where with cute.math.fmax
- Fix expert index calculation: col + tile_n_offset + subtile_idx * epi_n
- Top-6 accumulates across all N-tiles (for E=384 with 3 tiles of 128)
- Add iter_acc_early_release for overlapping accumulator
- Rewrite test to compare fused kernel vs 2-kernel reference path
- Remove stale memory doc
 2026-06-01 08:49:39 +00:00
biondizzle
0553117af6 Simplify fused router test: compare fused vs 2-kernel NVFP4 path 2026-06-01 07:10:55 +00:00
biondizzle
44a0e59808 Fix fused router test: use quantize_weight_to_nvfp4 (correct function name) 2026-06-01 07:08:56 +00:00
biondizzle
940f37fb6c NVFP4 fused router kernel: full rewrite with proper block-scaled GEMM setup 
Major fixes:
- Added tiled_mma_sfb creation (always CtaGroup.ONE, rounded N)
- Added mma_tiler_sfb, cta_tile_shape_mnk_sfb, cluster_layout_sfb_vmnk
- Use blockscaled_utils.make_smem_layout_sfa/sfb (with sf_vec_size)
  instead of sm100_utils (which doesn't support block-scaled SF layouts)
- Proper TMEM column accounting for SFA + SFB + accumulator
- Fixed make_blockscaled_trivial_tiled_mma argument order
  (a_dtype, b_dtype, a_major, b_major, sf_dtype, sf_vec_size, cta_group, mma_inst_shape)
- Fixed SFB TMA atom to use tiled_mma_sfb and cluster_layout_sfb_vmnk
- Fixed SFB partition_SFB to use tiled_mma_sfb.get_slice
- Fixed SFB global tile partitioning to use mma_tiler_sfb
- Fixed mainloop_s2t_copy_and_partition to use TMEM fragments
  (make_fragment_SFA/SFB) as the tSF parameter
- Updated run_nvfp4_fused_router wrapper to accept processed weight
  tensors from Nvfp4Linear._mat_b and _scale_b
- Updated test to properly build Nvfp4Linear and use processed weights

The old code was a rough sketch that never worked — it was missing
the entire tiled_mma_sfb infrastructure, used wrong SMEM layout
functions, and had broken TMA atom setup for scale factors.
 2026-06-01 07:08:12 +00:00
biondizzle
e6803b450d rewrite: simplified fused router test (reference + import check) 2026-06-01 06:53:17 +00:00
biondizzle
262cec262d fix: add shape assertions to fused router test 2026-06-01 06:51:47 +00:00
biondizzle
db07d17a62 fix: set activation global scale in fused router test 2026-06-01 06:50:41 +00:00
biondizzle
2abb4a19d9 fix: set gs and ws2 fields for Nvfp4Linear in fused router test 2026-06-01 06:49:43 +00:00
biondizzle
61c04f7152 fix: Nvfp4Linear field is sf not scale_b 2026-06-01 06:48:39 +00:00
biondizzle
982f245c67 fix: use correct Nvfp4Linear field names (fp4, scale_b, gsb) 2026-06-01 06:47:15 +00:00
biondizzle
16af96380f fix: use internal fields for Nvfp4Linear weight setup in test 2026-06-01 06:46:05 +00:00
biondizzle
7f1f224c78 fix: quantize_weight_to_nvfp4 returns 3 values, not 4 2026-06-01 06:43:53 +00:00
biondizzle
27fd847dd0 fix: correct quantize function name in fused router test 2026-06-01 06:41:54 +00:00
biondizzle
0873d65253 test: add fused router kernel test 
Compares NVFP4 fused CuTeDSL kernel against reference
(Nvfp4Linear + activation_topk) for correctness.
 2026-06-01 06:40:46 +00:00