Improved support for FP4 per_token/block quant

[learning-chip]

PTOAS FP4 VMI lowering support

Needed for https://github.com/learning-chip/TileKernels-vmi/pull/47

Purpose

This change completes the PTOAS half of the logical-to-physical FP4 boundary needed by VMI per_token_cast and per_block_cast. TileLang deliberately emits logical packed-FP4 operations; PTOAS must then choose legal physical parts, layouts, and materialization sequences. The previous lowering treated every 8-bit source as a four-part value and assumed a contiguous result could always be formed directly. That was wrong for narrow FP4 pair carriers and for compact source layouts, causing conversion failures, wrong part selection, or unnecessary scratch traffic.

The current quant audit confirms that these fixes solve the type/layout correctness failures and allow ASC-equivalent physical FP4 sequences to be generated. It does not yet prove that all remaining packed-store, TMA-scale, and FP32-carrier schedules match ASC latency; those require generated VPTO-guided optimization and fresh device measurements.

Base and scope

• Base: main @ 75e4a224d45bb81b7101df97edd1e4a98c0e1b9d.
• Branch: zjw/per_token_block_0903.
• Focused fixes: c5ed27060 (FP4→BF16 layout materialization) and 8b66d922a (narrow FP4 layout and physical-part lowering).

What is delivered

Type-aware FP4 pair lowering

• Adds an explicit f4x2 → bf16x2 cast class instead of treating packed FP4 as a generic 8-bit conversion.
• Keeps width-changing bitcasts strict; there is no global relaxation that could legalize unrelated invalid conversions.
• Supports the BF16 pair carrier used by TileLang reverse conversion, with FP32 widening performed only after the BF16 logical lanes are materialized.

Layout-aware physical part selection

• Physical arity is derived from the source element type, logical lane count, and lane stride.
• Compact FP4 layouts select the legal subset of physical parts (P0, or P0/P2) when the source stride proves that only those parts exist.
• Contiguous and four-part cases still use the full P0..P3 expansion where required; accidental P1 use is no longer possible for a P0/P2 source.
• Preferred-layout tables now include the 128-logical-lane FP4 shape used by H=128 per-block quantization.

Safe materialization and fallback

• Direct E2B lowering is retained only for one-packet-per-part shapes where its physical result is reusable.
• Multi-chunk contiguous broadcasts fall back to generic group-slot materialization instead of forcing an invalid deinterleaved layout.
• Generic deint4ToContiguous materialization remains available when physical widths and types match; direct scratch-free lowering is selected when they do not require a reload.
• Unsupported width/layout relations produce a precise match failure rather than an invalid pto.vmi operation.

Problems solved

These fixes address the failures seen when moving from ASC-equivalent physical references to logical VMI:

• deinterleaved=4 → contiguous no longer fails for legal FP4→BF16 cases;
• 128- and 256-logical-lane values select compatible physical widths;
• compact/lane-stride sources preserve lane order and mask coverage;
• physical UNPK4/vcvt(P0..P3) ordering matches the source layout;
• reverse conversion no longer requires a pair store/barrier/reload when a direct materialization is legal;
• no unsupported pto.vmi is left behind for the covered FP4 shapes.

Validation and performance impact

The unified TileLang branch can now lower the packed FP4 frontend operations needed by the quant kernels. On CANN 9.1/device 0, the dependent BF16→E2M1 per-block H=128 case compiled, passed byte/scale correctness gates, and measured 1.048× ASC. The remaining slow FP4 rows are not unexplained lowering failures: packed pair-store assembly, TMA-column scale materialization, and FP32 carrier width choices still add physical work in specific schedules. Those rows are listed with concrete causes and next steps in the kernel PR and full coverage table.

Follow-up

Use the ASC/CCE sequence and generated VPTO for each remaining slow row as a three-level reference ladder. The next PTOAS optimization should add direct scratch-free packed stores for the legal compact layouts, specialize TMA-column scale strides, and preserve the narrow FP32/E2M1 width choices while keeping the type-aware and layout-aware safeguards above.