mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-11-02 08:33:19 +08:00
T.copy_ue4m3_scale_tile and T.ue4m3_scale_tile_source_coords were NVFP4-specific helpers on the general T.* surface. Remove both and move the abstraction down a level: - The example now declares the packed scale source with its honest "tile rows" shape ([n_mn_blocks * n_k_blocks * block_MN, block_K // 64] uint32): the blockscaled_chunk_kmajor layout stores each (mn_block, k_block) tile contiguously, so with this view every tile is a plain rectangular slice and staging is a two-line T.Parallel loop with row-offset addressing only. The host obtains the view zero-copy via swizzle_blockscaled_chunk_kmajor_scale_words(...).reshape(-1, 4). - tilelang.quantize.nvfp4 gains blockscaled_chunk_kmajor_tile_source_coords, the hand-written offset formula for staging schedules that keep the 2-D [MN, K/64] view (the maint benchmark's producer warp group uses it). Compiler-internal staging was prototyped first (plain T.copy of the slice) and rejected for now: the raw-pointer bulk-TMA path this lowers to on some shapes never arrives on its mbarrier (deadlock, tracked upstream), and the synchronous fallback stalls the producer warps (-8.5% at 2048^3). The example comment documents this so the loop can become a plain T.copy once fixed. Validation on RTX 5090: 123 tests pass; example outputs bitwise-identical to the previous macro implementation on 5 shapes including M/N tail tiles (257x384x512, 130x128x256, 128x136x256); example perf at parity (100.1-100.4% over 2048^3/4096^3/8192^3/8192x8192x1536, interleaved A/B, generated CUDA identical modulo a loop variable name); maint WS benchmark kernel source byte-identical with --verify passing and perf unchanged (641.9 / 1310.5 / 1149.5 TF at 2048^3 / 8192^3 / 8192x8192x1536). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>